Repository navigation
Sequence try grow - #6948
Sequence try grow#6948nat-openai wants to merge 8 commits into
Conversation
88bd7e5 to
f147d0d
Compare
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed all 10 changed files and eight commits from base e555dba6f5c8c62ceca8c9c6e311a51b8fbfeb70 to head d4016cd2af4ba0d16f5cf5fc8ca87d6d0b0cb444. The PR remains non-draft. The snapshot and live discussion contain no existing reviews or comments. One introduced P2 issue was reproduced.
Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-memory-pr, and review-comet-ffi-pr.
Summary
- Prior state and problem: Integral
sequenceallocated its output without consulting the task pool. Fallible allocation alone did not enforce the task’s memory budget. - Design approach:
SparkSequencereceives the session’s existing pool through the scalar-function factory. Columns reserve construction scratch first and values second. Scalars reserve their complete batch together. - Correctness: Spark’s
Sequenceimplementation was checked across 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0. Null propagation, default steps, integral boundaries, and overflow behavior remain consistent. However, the new admission failure can abort a sort that previously completed by spilling. - Compatibility analysis: Compared the affected paths with
branch-1.1. Construction-pool enforcement and resource-error reporting are new behavior. Temporal fallback and version-specific Spark error conversion remain unchanged. The reproduced loss of spill progress is a regression from both the base and the released allocation policy. - Key design decisions: Invocation-scoped registration avoids reducing fair-pool shares while expressions are idle.
Volatileprevents native folding from bypassing scalar-batch admission. Preserving the outer session pool retains its tracking wrappers. - Implementation sketch: The shared kernel validates and caches lengths, constructs values and offsets under a reservation, and releases that reservation on return or failure. The factory signature change also updates modulo’s plumbing without changing its arithmetic.
- Performance: Staged admission avoids repeating row validation on successful column invocations. The pooled benchmark includes consumer tracking and scalar expansion. The author reports preliminary overhead of roughly 5–7% for some short cases and 12.5% for singleton rows. Those measurements were not independently verified, and production JNI overhead remains unmeasured here.
- Design: Separating construction accounting from downstream retention is understandable and documented. The unresolved problem is progress under memory pressure: downstream operators cannot reclaim their buffered data before an upstream sequence refusal terminates execution.
- Abstraction & complexity: A pool-aware UDF and session-aware factory are justified by the runtime dependency. The allocation helpers and cold validation path have distinct responsibilities. No additional abstraction issue met the P1/P2 reporting bar.
- Behavioral changes worth calling out: Pool refusals now fail sequence evaluation, including when each batch fits individually but a downstream sort occupies the remaining budget. Output arrays and FFI exports deliberately outlive the construction reservation. On-heap execution remains unbounded.
- Suggested improvements: Make admission cooperate with downstream spilling or protected construction headroom, and retain an integration regression test that actually spills. The current test documenting refusal behind a buffering consumer does not preserve the previously successful behavior.
Exact-head CI: the label check passed. Comet CI and CodeQL concluded action_required, so neither provides a validation verdict.
Validation: 25 sequence unit tests, eight UDF registration tests, and 25 modulo tests passed against locked DataFusion 55.2.0 and Arrow 59.3.0. A disposable integration test compared the exact base kernel with the head factory under identical bounded pools and reproduced the finding at two limits. The checkout is clean after removing that test. The native core planner test, JVM tests, Spark SQL matrix, and production JNI benchmarks were not run. The reproduction uses DataFusion’s bounded pool rather than Spark’s JNI memory bridge.
| }; | ||
| if scalar_length.is_none() { | ||
| if let Some(reservation) = &reservation { | ||
| reservation.try_grow(value_bytes::<T>(total)?)?; |
There was a problem hiding this comment.
[P2] Preserve downstream spilling before propagating this refusal. When sequence feeds a native sort, the sorter can retain enough earlier batches that this values reservation fails. SortExec polls its input before reaching its own reservation/spill logic, so propagating the error here aborts a query that could release memory by spilling. With 64 batches of 64 rows, 64-element BIGINT sequences, a 220,000-byte pool, and 64 KiB of sort spill reservation, the exact base kernel returns all 4,096 rows with 47 spills. The head fails on the third batch with ResourcesExhausted for Sequence, before any spill. Each sequence batch fits comfortably in the pool by itself. Please coordinate admission with reclamation or protect construction headroom before downstream buffering consumes it, and add a spilling integration regression test.
Evidence: Ran a disposable Rust integration test using locked DataFusion 55.2.0 / Arrow 59.3.0: cargo test -p datafusion-comet-spark-expr --locked --test review_sequence_admission review_sequence_before_spill -- --nocapture. It compiled sequence.rs from exact base e555dba alongside the head factory. Fixture: one MemTable partition containing 64 batches of 64 rows; for batch n and row i, BIGINT id=4095-64*n-i, a=n%2, b=a+63. Query: SELECT id, spark_sequence(a, b) AS vals FROM t ORDER BY vals, id. Configuration: target_partitions=1, batch_size=64, sort_spill_reservation_bytes=65536, GreedyMemoryPool=220000. Base: Ok(4096), spill_count=47. Head: only 128 input rows consumed, spill_count=0, refusal for 32.0 KB with 772 bytes of sequence scratch reserved and 19.1 KB available. At pool limit 480000, base again returns 4096 rows with 15 spills while head fails before its first spill. Assertions passed for both comparisons. Reproducer saved at /tmp/comet-6948-review_sequence_admission.rs; output at /tmp/comet-6948-repro-final.log. The disposable checkout test was removed.
There was a problem hiding this comment.
With 64 batches of 64 rows, 64-element BIGINT sequences, a 220,000-byte pool, and 64 KiB of sort spill reservation, the exact base kernel returns all 4,096 rows with 47 spills.
This does feel a little bit like a contrived example. On the main branch, the sequence memory reservations are untracked, and with such small sequences, there's probably ample system headroom for them to succeed. Combining that with a very small memory pool (for tracked memory) ensures that we can demonstrate spilling in the sort operator while hitting the limit when we try to reserve space in sequence.
That's not to say I'm discounting this point in general though, and I'm considering those proposals.
Please coordinate admission with reclamation
IIUC, this would mean adding pretty novel cross-native-operator memory coordination? If so, I image we need to think about the design of that pretty carefully.
or protect construction headroom before downstream buffering consumes it
Codex points me toward reserve_memory_for_merge in sort.rs as as a precedent for this sort of thing. I'm trying to understand how this buffer could be sized such that it provides enough space for any computable input that would fit in memory without starving other operators... guarantees seem hard to come by and it might remain possible to come up with adversarial counter examples (large sequences; small pools) that break any assumptions.
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @nat-openai. A few things before the code itself:
check-pr-titlewill fail onSequence try grow, since it needs a conventional prefix such asperf: reserve memory before building sequence output.- The branch is 169 commits behind
mainand a trial merge conflicts inplanner.rs,array_funcs/mod.rs,sequence.rs,comet_scalar_funcs.rs,modulo_expr.rsandCometArrayExpressionSuite.scala. Your draft #6939 ("Sequence invocation ceiling") takes a different approach to the same part of #5712 and conflicts with this one. Which of the two should reviewers look at? CometFairMemoryPoolonmainnow treats spillable and non-spillable consumers differently, so the spill-before-refusal repro from the earlier review may need re-checking under it.- The description reports 5 to 7 percent overhead on short shapes and 12.5 percent on singleton rows, measured with
GreedyMemoryPool. Eachtry_growagainst the Spark-backed pool is a JNI call, so I expect the real cost to be higher, but I have not measured it.
| "spark_sequence" => { | ||
| make_comet_scalar_udf!("spark_sequence", spark_sequence, data_type) | ||
| // Preserve the exact outer pool, including per-plan tracking and task-shared policy. | ||
| Ok(Arc::new(ScalarUDF::new_from_impl(SparkSequence::new( | ||
| data_type, | ||
| Arc::clone(&session_state.runtime_env().memory_pool), | ||
| )))) | ||
| } |
There was a problem hiding this comment.
Since this branch was cut, #7171 added a separate spark_sequence_recomputed arm here for Spark 4.0.5, 4.1.4 and 4.2.1, built with make_comet_scalar_udf!. A rebase that keeps that arm as is leaves those versions outside the pool. Could the overflow mode become a field of SparkSequence, so both names go through the reservation, with a test of the recomputed name under a bounded pool?
| let mut lengths = sequence_vec::<usize>(num_rows).map_err(scratch_error)?; | ||
| let total = if let Some(len) = scalar_length { | ||
| lengths.resize(num_rows, len); | ||
| scalar_total.unwrap() |
There was a problem hiding this comment.
For all-scalar inputs scalar_length is constant, but this still allocates and reserves a num_rows usize length cache (64 KiB at 8192 rows). Could the scalar path pass cache_lengths = false? The test that pins "four cached lengths" would change with it.
| fn sequence_vec<T>(capacity: usize) -> Result<Vec<T>> { | ||
| let mut values = Vec::new(); | ||
| values.try_reserve_exact(capacity).map_err(|error| { | ||
| DataFusionError::ResourcesExhausted(format!("Sequence allocation failed: {error}")) | ||
| })?; | ||
| Ok(values) |
There was a problem hiding this comment.
On main a failed try_reserve_exact becomes SequenceBatchTooLarge, whose message tells the user to lower spark.comet.batchSize, and its doc in native/common/src/error.rs says it covers allocator refusal. Here allocator and pool refusals become a plain ResourcesExhausted with no hint. Could they keep the existing variant, or could its doc and message be updated?
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
The PR backlog is huge and the maintainers don't have the bandwidth to review them all, so here is an automated review. Feel free to ignore it.
Thanks for tackling the accounting gap for integral sequence. The staged reservation, the cleanup on error and unwind, and the Volatile signature that stops constant folding from bypassing admission all read well. My main concern is the one already raised in the thread on sequence.rs around line 351 and I think it blocks the merge as written.
Against branch-1.2, the released kernel allocates the values buffer with try_reserve_exact and never consults the pool. A query where a native sort or aggregate sits above sequence therefore completes today by spilling. With this change, reserve_construction and the later reservation.try_grow(value_bytes::<T>(total)?) run while the downstream operator is still holding its buffered batches, and the refusal propagates as a query error before that operator gets a chance to spill. The PR's own test existing_buffered_reservations_can_refuse_a_sequence_before_spilling pins exactly this, and the new section in memory_management.md says it plainly. So a user who upgrades can see a failing query that used to succeed, with every individual batch fitting in the pool.
The earlier discussion left open whether this only shows up with a GreedyMemoryPool. I read CometFairMemoryPool::try_grow on current main to check. A consumer that can spill is limited to (task_share - unspillable) / num_spill, so with one sort in the task it may hold essentially the whole task share. The Sequence consumer is non-spillable, so it is only checked against pool_size and then Spark's own grant. When the sort is close to its limit, that leaves no headroom for the next batch, and the request is refused. The sort's merge reservation does not help because it is held by the sort and is already counted as unspillable. So the default fair_unified pool has the same failure, just at a different threshold.
Could the reservation be made recoverable instead of fatal? One direction is to hold a small construction headroom up front, the way reserve_memory_for_merge does in the sorter, so a normal batch is never refused while a downstream operator is still able to spill. Another is to apply the hard try_grow refusal only to requests far larger than a normal batch, which is the case the PR is really guarding against. Whichever policy is chosen, an integration test where sequence feeds a sort that spills, and which also passes on branch-1.2, would show the regression is gone. The current Spark test with a tiny memoryPool.fraction only shows refusal and recovery with a smaller batch.
The branch also conflicts with main now. main has a separate spark_sequence_recomputed arm in comet_scalar_funcs.rs, and as comphead noted, a plain rebase would leave those Spark versions outside the pool. Folding the overflow mode into SparkSequence as a field would cover both names and let the bounded-pool tests cover the recomputed variant too.
Which issue does this PR close?
Addresses part 2 of #5712: memory accounting for integral
sequenceconstruction.Rationale for this change
Integral
sequencecan allocate a large Arrow values buffer without consulting Comet’s memory pool. Fallible allocation handles allocator refusals, but does not enforce the task’s memory budget.This change reserves construction memory through the task’s existing pool before allocating buffers. The reservation lasts only for the invocation; downstream buffering operators remain responsible for accounting for the output they retain.
What changes are included in this PR?
SessionState, preserving the session’s exact memory pool and tracking wrappers.try_growreservations for column inputs: reserve offsets, validity buffers and length scratch first, then reserve values after validating and caching each row’s length once.The check follows the existing pool policy and introduces no new memory setting. A refusal can occur before a downstream native operator reaches its spill step; this does not add native reclamation or guarantee a process-memory ceiling. On-heap mode retains its existing unbounded pool policy.
How are these changes tested?
sequence_pooledCriterion group exercises the UDF with a tracked, bounded pool across short, long, nullable and scalar inputs.Preliminary diagnostics retain approximately 5–7% overhead for the two-element column/null cases and 12.5% for default singleton rows; performance parity is not claimed. Validation against the pinned DataFusion 55.2, JVM test execution and production JNI performance measurements remain pending.