Every PPL query whose plan includes a late-materialization fetch phase
(any `sort | head K | fields <non-sort-key>` shape) paid a fixed ~5s in
the LATE_MATERIALIZATION stage regardless of rows, columns, or bytes
scanned, with a "timed out waiting for reduce teardown" WARN per query.
Root cause is a teardown cycle in DatafusionReduceSink.closeImpl:
close() waited on reduceDone, the drain was parked in stream_next
waiting for a pipeline-breaking SortExec/TopK to emit, the TopK was
waiting for input EOF, and the input senders were only closed after
close() returned. The 5s await timeout was the only thing breaking the
cycle. Cancellation could not break it either, for two Rust-side
reasons:
- QUERY_REGISTRY was keyed one-tracker-per-context_id, but an LM query
opens TWO reduce sinks (built at graph-build time) that both register
under the same ctx.taskId(). The second insert overwrote the first
tracker and the first Drop removed the shared key, so cancel_query
found no entry (silent no-op) while the sibling stream still ran.
- stream_next re-fetched its token from the registry on every call; a
stale registry returned None, degrading cancellable_or to a bare
uncancellable await.
Fixes:
- Rust: QUERY_REGISTRY holds Vec<Arc<QueryTracker>> per id. Register
appends; Drop removes only its own tracker (ptr_eq) and drops the key
when empty; cancel_query cancels every tracker under the id and logs
registry misses instead of silently no-oping.
- Rust: stream_next uses its handle's own cancellation token via new
QueryTrackingContext::cancellation_token(), immune to registry churn.
- Java: closeImpl's REDUCING branch signals end-of-input first —
tryClose() on each input sender (dropping a sender closes its mpsc,
which the native plan reads as EOF, letting the TopK emit and the
drain finish naturally with rows preserved). tryClose, not close(): a
feeder parked in send() holds the sender read lock, and a blocking
close would deadlock (covered by the existing
testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock). If tryClose
fails the producer is live, so fall back to cancel; if EOF does not
finish the drain within 5s, WARN + cancel so close() can never hang.
Cancel-first teardown was not just slow but wrong once cancellation
worked: it aborted the TopK before EOF and projections mixing sort keys
with fetched columns returned 0 rows.
Measured on a 749M-doc 5-shard composite index: 5.15s -> ~0.19s per
query (~27x), all shapes row-correct, zero teardown WARNs.
Tests:
- test_sibling_stream_drop_does_not_orphan_survivor (Rust)
- test_cancel_query_cancels_all_siblings (Rust)
- testCloseWhileReducingSignalsEofAndPreservesRows (Java) — fails on
the old code both ways: 5s close stall without the fix, 0 rows with
cancel-first teardown.
Description
Every PPL query whose plan includes a late-materialization fetch phase (any
... | sort <col> | head K | fields <non-sort-key-col>shape) pays a fixed ~5s in theLATE_MATERIALIZATIONstage — invariant to row count, column count, and bytes scanned — and logs one WARN per query:Profile of an affected query (50 rows out of 749M docs; scan itself is 60–90ms):
Root cause: a teardown cycle that only the 5s timeout could break
DatafusionReduceSink.closeImpl(REDUCING branch) firedcancel_queryand then waited onreduceDone.await(5, SECONDS). That wait sat inside a cycle:close()waits onreduceDone,reduceDone.countDown()runs only after the drain loop exits,stream_nextwaiting for a pipeline-breakingSortExec/TopKto emit,close()returned.The cancel could not break the cycle either, because of two Rust-side defects:
QUERY_REGISTRYwasDashMap<i64, Arc<QueryTracker>>— one tracker percontext_id. But a query with an LM stage opens two coordinator reduce sinks (both constructed at graph-build time, before any stage runs), and both register under the samectx.taskId(). The secondinsertsilently overwrote the first tracker, and the first stream to close removed the shared key inDrop— socancel_queryfound no entry (silent no-op) while the sibling stream was still running. Verified live:cancel_query: NO REGISTRY ENTRY for context_id=943 (no-op). live ids=[]one line afterfireCancelQuery: taskId=943.stream_nextre-fetched its cancellation token from the registry on every call. Once the entry was gone, it gotNone, andcancellable_or(None, ...)degrades to a bare, structurally uncancellable await.jstack during the stall confirms it: the drain thread is RUNNABLE inside the FFM downcall with
cpu=0.43ms— waiting, not working.Why cancel-first teardown was also wrong (not just slow)
With the Rust fixes in place, cancellation actually works — and cancel-first teardown then aborts the TopK before EOF, so projections mixing sort keys with fetched columns returned 0 rows. The close-path semantics are "no more input is coming", not "abort"; EOF is the correct signal.
The fix
Rust (
query_tracker.rs,api.rs):QUERY_REGISTRYbecomesDashMap<i64, Vec<Arc<QueryTracker>>>: register appends;Dropremoves only its own tracker (Arc::ptr_eqretain) and drops the key when empty;cancel_querycancels every tracker under the id, and logs registry misses instead of silently no-oping.stream_nextuses its handle's own token via newQueryTrackingContext::cancellation_token()— immune to registry churn.Java (
DatafusionReduceSink,DatafusionPartitionSender):closeImpl's REDUCING branch signals EOF first:tryClose()each input sender (dropping a sender closes its mpsc; the native plan reads that as end-of-input, the TopK emits, the drain finishes naturally, rows preserved).tryClose()(new), notclose(): a feeder parked insend()holds the sender read lock; a blocking close would deadlock — the exact scenario covered by the existingtestCloseWhileFeederParkedOnFullChannelDoesNotDeadlock, which caught the first version of this patch.tryClosefails (live producer mid-send) → fall back to cancel (dropped rows are correct semantics; the consumer already stopped listening). If EOF doesn't finish the drain within 5s → WARN + cancel, soclose()can never hang.Measured impact
On a 749M-doc, 5-shard composite (parquet-primary) index,
where <lead-sort-key>=<v> | sort - <second-key> | head 50 | fields <2 non-sort-key cols>:head Kvariants, no-sort, and aggregation shapes all row-correct after the fix.Related Issues
Related to the observability gap in #22601 (the LM stage emits no plan/metrics, which is why this required jstack + DEBUG logs to diagnose).
Check List
test_sibling_stream_drop_does_not_orphan_survivor,test_cancel_query_cancels_all_siblings(new), 23/23query_trackertests pass.testCloseWhileReducingSignalsEofAndPreservesRows(new) — fails on the old code both ways (5s close stall pre-fix; 0 rows with cancel-first teardown). 12/12DatafusionReduceSinkTestspass.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.