You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "SubhamSinghal (via GitHub)" <gi...@apache.org> on 2026/04/30 15:01:03 UTC
[PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
SubhamSinghal opened a new pull request, #21962:
URL: https://github.com/apache/datafusion/pull/21962
## Which issue does this PR close?
Related to the TODO at `materializing_stream.rs:283` (from [#17429](https://github.com/apache/datafusion/pull/17429#discussion_r2324056495)): spilled `BufferedBatch` join key arrays are not tracked in memory reservation.
## Rationale for this change
When a `BufferedBatch` is spilled to disk in Sort Merge Join, only the `RecordBatch` data is written to the IPC file. The `join_arrays` (evaluated join key columns) remain in memory because the merge-scan comparator needs them to detect key group boundaries.
Before this fix, these in-memory `join_arrays` were **invisible to the memory pool**:
allocate_reservation():
try_grow(size_estimation) → FAILS (pool full)
spill batch to disk
→ join_arrays still in memory, but reservation was never grown
→ pool thinks 0 bytes are used for this batch
free_reservation():
if InMemory → shrink(size_estimation)
if Spilled → no-op ← correct (nothing was grown), but join_arrays are invisible
With many spilled batches for a skewed key (e.g., millions of rows sharing the same join key), the untracked `join_arrays` memory accumulates. The memory pool cannot account for this when making spill decisions for concurrent operators.
## What changes are included in this PR?
**Memory accounting fix** (`materializing_stream.rs`):
- Add `reserved_amount` field to `BufferedBatch` — tracks how much memory was **actually reserved** in the pool for this batch
- Add `join_arrays_mem()` helper — computes total memory of join key arrays
- `allocate_reservation()`: after spilling, calls `try_grow(join_arrays_mem)` to track the remaining in-memory data. If the pool is too tight for even that, `reserved_amount` stays 0
(best-effort, safe)
- `free_reservation()`: shrinks by `reserved_amount` instead of checking `InMemory` variant. Invariant: only shrink by what was actually grown — no underflow risk
| Scenario | `try_grow` | `reserved_amount` | `try_shrink` | Safe? |
|----------|-----------|-------------------|-------------|-------|
| InMemory | Ok(size_estimation) | size_estimation | size_estimation | Yes |
| Spilled, tracked | Ok(join_arrays_mem) | join_arrays_mem | join_arrays_mem | Yes |
| Spilled, pool tight | Err | 0 | 0 (no-op) | Yes |
**Tests** (`tests.rs`):
- `spill_many_batches_same_key` — 10+5 batches all sharing key=1, verifies correctness under heavy spilling
- `spill_string_join_keys` — Utf8 join keys to exercise larger `join_arrays` footprint
- `spill_mixed_keys_some_match` — multiple distinct keys with partial matching, tests Full outer join NULL rows from spilled batches
- `spill_join_arrays_memory_accounting` — verifies memory pool is fully released after join completes (`memory_pool.reserved() == 0`) and `peak_mem_used > 0`
## Are these changes tested?
Yes. Four new tests added covering heavy spilling with same-key batches, string join keys, mixed keys with partial matching, and memory pool accounting verification.
## Are there any user-facing changes?
No.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#discussion_r3193919876
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -978,6 +996,20 @@ impl MaterializingSortMergeJoinStream {
.unwrap(); // Operation only return None if no batches are spilled, here we ensure that at least one batch is spilled
buffered_batch.batch = BufferedBatchState::Spilled(spill_file);
+
+ // Track remaining in-memory data (join key arrays) that
+ // stay in memory even after the batch is spilled. This is
+ // much smaller than the full batch, so try_grow should
+ // usually succeed. If it fails, reserved_amount stays 0 -
+ // best-effort tracking, free_reservation will safely be a no-op.
+ let join_arrays_mem = buffered_batch.join_arrays_mem();
+ if self.reservation.try_grow(join_arrays_mem).is_ok() {
Review Comment:
I think the spill path can still leave retained join key arrays invisible to the memory pool.
Right now, if the full batch `try_grow(size_estimation)` fails because the pool is full, and the follow-up `try_grow(join_arrays_mem)` also fails, we spill the IPC batch but still push `buffered_batch` with `reserved_amount = 0`.
At that point the operator is still holding the retained `join_arrays`, but the pool is no longer aware of them when making spill decisions for other operators. This seems like the same invariant violation we were trying to avoid.
I think this can still happen with concurrent reservations or when the memory limit is below a single join-array allocation, and in those cases many skewed spilled batches could accumulate untracked memory.
Can we make the retained in-memory portion accounted deterministically here? For example, by growing or resizing the reservation after the physical memory is retained, or by returning an error instead of continuing untracked.
It would also be great to add a regression test that covers the no-headroom path where `try_grow(join_arrays_mem)` fails, since the current test only exercises the successful reservation case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#issuecomment-4353914929
Thanks @comphead! The `MergeAsPartialUDF` wrapper is a clever way to get `PartialMerge` semantics without reaching into DataFusion's aggregation logic.
### Streaming shim asymmetry
I checked the Spark source at `v3.5.8`, `v4.0.2`, and `v4.1.1`. `StreamSourceAwareSparkPlan` doesn't exist before 4.x, so needing a different 3.x approach is unavoidable. The shapes of the two checks do differ in a way worth flagging though.
On 3.x the shim does `plan.logicalLink.exists(_.isStreaming)`. That only looks at the logical link of the root physical node. `LogicalPlan.isStreaming` does recurse through logical children, but if the root physical node has `logicalLink = None` (a rule built a fresh node without copying the link) the check silently returns `false`.
On 4.x the shim does `plan.exists { case p: StreamSourceAwareSparkPlan => p.getStream.isDefined }`. `SparkPlan.exists` walks the whole physical tree, and the leaf check does not depend on logical link preservation.
Would `plan.exists(_.logicalLink.exists(_.isStreaming))` on 3.x be a safer analogue, or is there a reason you only want to inspect the root? Totally fine if this is out of scope for the PartialMerge work, just wanted to surface it.
### `sum_int.rs` merge_batch rewrite
The prior `merge_batch` only read row 0 of the state array, which looks like a latent bug whenever a merge batch had more than one row. Switching to `update_batch(states)` fixes that. Do you think it's worth adding a small direct test that exercises `merge_batch` with a multi-row state array independent of the `PartialMerge` plumbing, so the fix is captured even if the `MergeAsPartial` path changes later?
### `collect_set` in the description
The description lists `FIRST` / `LAST` / `collect_set` as falling back, but the serde only filters `First` and `Last`. Is `collect_set` covered entirely by the streaming plan fallback? If so, a line of comment near the `First`/`Last` filter explaining the split would help future readers.
### `distinct` test lost its aggregate count assertion
The previous version used `checkSparkAnswerAndNumOfAggregates` to assert how many native aggregate stages ran. The new version drops that. Since this PR is specifically about enabling more stages to stay native, would it be worth keeping a mode count assertion in one or two of the cases so a future regression that silently falls back one stage still shows up? Same thread as the existing review comment.
### Test name mismatch
`partialMerge - distinct + non-distinct with first() FILTER (Expand pattern)` doesn't use `FIRST` or `FILTER` in the queries it runs. Was that the intended coverage, or did the queries drift from the original intent?
### Leftover configs
In `partialMerge - cnt distinct + sum`, are `spark.comet.cast.allowIncompatible` and `spark.comet.expression.Cast.allowIncompatible` actually needed? Looks like they might be leftover from local debugging.
### AVG(DISTINCT) coverage
The serde relaxation also admits `SUM(DISTINCT x)` and `AVG(DISTINCT x)` for single column. The `distinct` test covers `COUNT` and `SUM` variants but I didn't see `AVG(DISTINCT)` exercised. Could we add one case, ideally mixed with a non distinct aggregate like `AVG(DISTINCT a) + SUM(b)`, so the Expand plan is exercised for a non count distinct?
### `planner.rs` mode literals
The match on `agg.mode` uses `0`, `1`, `2` directly. Would it be cleaner to match on the generated `AggregateMode` enum so the proto and native sides can't desync silently?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "SubhamSinghal (via GitHub)" <gi...@apache.org>.
SubhamSinghal commented on PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#issuecomment-4355059187
@mbutrovich Thanks for reviewing this PR. I have removed redundant tests. I have added UT which would always hit spill path, with this `peak_memory` would be 0 in main and non-zero in this branch.
Added `set_max` call in spill path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#discussion_r3198707608
##########
datafusion/physical-plan/src/joins/sort_merge_join/tests.rs:
##########
@@ -2487,6 +2487,223 @@ async fn overallocation_multi_batch_spill() -> Result<()> {
Ok(())
}
+/// Verifies that `peak_mem_used` reflects join_arrays memory on the spill path.
+///
+/// Uses a memory limit smaller than a single batch's `size_estimation` so that
+/// every batch spills — the `Ok` arm of `allocate_reservation` is never hit.
+/// Before the fix, `peak_mem_used` would stay 0 because `set_max` was only
+/// called in the `Ok` arm. After the fix, the spill path calls
+/// `grow(join_arrays_mem)` + `set_max`, so `peak_mem_used > 0`.
+#[tokio::test]
+async fn spill_join_arrays_memory_accounting() -> Result<()> {
+ use arrow::array::Array;
+
+ let left_batch = build_table_i32(
+ ("a1", &vec![0, 1]),
+ ("b1", &vec![1, 1]),
+ ("c1", &vec![4, 5]),
+ );
+ let size_estimation = left_batch.get_array_memory_size()
+ + Int32Array::from(vec![1, 1]).get_array_memory_size()
+ + 2usize.next_power_of_two() * size_of::<usize>()
+ + size_of::<std::ops::Range<usize>>()
+ + size_of::<usize>();
+ let join_arrays_mem = Int32Array::from(vec![1, 1]).get_array_memory_size();
+
+ // Memory limit: too small for a full batch, large enough for join_arrays.
+ // Every batch hits the Err arm → spills → grow(join_arrays_mem).
+ let memory_limit = (size_estimation + join_arrays_mem) / 2;
+ assert!(
+ memory_limit < size_estimation && memory_limit > join_arrays_mem,
+ "limit {memory_limit} must be between join_arrays_mem {join_arrays_mem} \
+ and size_estimation {size_estimation}"
+ );
+
+ let left_batches: Vec<RecordBatch> = (0..4)
+ .map(|i| {
+ build_table_i32(
+ ("a1", &vec![i * 2, i * 2 + 1]),
+ ("b1", &vec![1, 1]),
+ ("c1", &vec![100 + i, 101 + i]),
+ )
+ })
+ .collect();
+ let left = build_table_from_batches(left_batches);
+
+ let right_batches: Vec<RecordBatch> = (0..2)
+ .map(|i| {
+ build_table_i32(
+ ("a2", &vec![i * 2, i * 2 + 1]),
+ ("b2", &vec![1, 1]),
+ ("c2", &vec![200 + i, 201 + i]),
+ )
+ })
+ .collect();
+ let right = build_table_from_batches(right_batches);
+
+ let on = vec![(
+ Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+ )];
+ let sort_options = vec![SortOptions::default(); on.len()];
+
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(memory_limit, 1.0)
+ .with_disk_manager_builder(
+ DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
+ )
+ .build_arc()?;
+
+ let session_config = SessionConfig::default().with_batch_size(50);
+ let task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(session_config)
+ .with_runtime(Arc::clone(&runtime)),
+ );
+
+ let join = join_with_options(
+ Arc::clone(&left),
+ Arc::clone(&right),
+ on.clone(),
+ Inner,
+ sort_options,
+ NullEquality::NullEqualsNothing,
+ )?;
+
+ let stream = join.execute(0, task_ctx)?;
+ let result = common::collect(stream).await.unwrap();
+
+ assert!(!result.is_empty(), "Expected non-empty join result");
+
+ let metrics = join.metrics().unwrap();
+ assert!(
+ metrics.spill_count().unwrap() > 0,
+ "Expected spilling to occur"
+ );
+
+ // Before the fix, peak_mem_used was 0 here because set_max was only
+ // called in the Ok arm of allocate_reservation, which is never reached
+ // when every batch spills. After the fix, the spill path calls
+ // grow(join_arrays_mem) + set_max unconditionally.
+ let peak_mem = metrics
+ .sum_by_name("peak_mem_used")
+ .map(|m| m.as_usize())
+ .unwrap_or(0);
Review Comment:
The new tests look good and `peak_mem > 0` definitely catches the old missing-accounting path.
As a possible follow-up improvement, it might be worth tightening this to something like `peak_mem >= join_arrays_mem` if that value is available in scope. That would also help catch partial accounting regressions. Non-blocking though.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -235,6 +235,14 @@ pub(super) struct BufferedBatch {
pub null_joined: Vec<usize>,
/// Size estimation used for reserving / releasing memory
pub size_estimation: usize,
+ /// Actual amount tracked in the memory reservation for this batch.
+ ///
+ /// - `InMemory`: equals `size_estimation` (full batch + join_arrays + metadata)
+ /// - `Spilled`: equals join_arrays memory if `try_grow` succeeded after spill, else 0
Review Comment:
Nice catch fixing the accounting path here 👍
One small thing: the `reserved_amount` field doc still says spilled batches only track join-array memory if `try_grow` succeeds, otherwise `0`. Since the implementation now uses unconditional `grow(join_arrays_mem)`, I think the doc comment should be updated to match the new behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich merged PR #21962:
URL: https://github.com/apache/datafusion/pull/21962
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#issuecomment-4390861669
Thanks for turning this around quickly, the force-grow version is easier to reason about than the conditional one.
A few things I wanted to ask about before signing off:
**grow() can push over the pool limit**
`reservation.grow()` is infallible, so in the no-headroom case (which `spill_join_arrays_no_headroom` exercises) we'll end up with `pool.reserved() > memory_limit`. I think that's the right call since the memory is physically there, but it does mean the next operator to call `try_grow` sees a pool that's already over limit and fails. Could we add a one-line comment at the `grow()` site noting this is intentional? Future readers running into a "why is the pool over limit" question will appreciate the pointer.
**Recomputing join_arrays memory**
`BufferedBatch::new` already sums `get_array_memory_size()` over `join_arrays` as part of `size_estimation` (line 262-265). `join_arrays_mem()` walks the same arrays again on the spill path. Cheap, but would you be open to caching the sum as a field during construction? One less place for the two size calculations to drift if someone touches the estimation formula later.
**try_shrink vs shrink on free**
Given the invariant that we only shrink by what we grew, would `shrink()` read more symmetrically with the force-grow on the alloc side? The `?` on `try_shrink(reserved_amount)` is propagating a `Result` that shouldn't be reachable. Not a blocker, just noticed the asymmetry.
**peak_mem assertion strength**
The two new tests check `peak_mem > 0`, which catches the "never tracked at all" regression nicely. Would `peak_mem >= join_arrays_mem` be worth tightening to? That would also catch a future bug where we track something but not the right amount. Happy to leave as-is if you think the current form is enough.
Perf-wise I looked at this against #20729 and I don't think the concerns there apply here. The extra `grow()` only fires on the spill branch, and spilling already pays for disk I/O, so one more pool call on a custom backend is in the noise. Non-spill workloads see the same pool-call count as before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "SubhamSinghal (via GitHub)" <gi...@apache.org>.
SubhamSinghal commented on PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#issuecomment-4390372416
Thanks for the review @kosiew
I've updated the implementation to use unconditional grow() instead:
**Rationale:** The join key arrays are physically in memory. Not tracking them gives the pool a stale view and could let concurrent operators over-allocate. Since grow() is infallible (returns ()), the accounting is now deterministic — no conditional path.
I've also added a regression test (spill_join_arrays_no_headroom) that configures the pool smaller than join_arrays_mem,
verifying:
- The join completes correctly
- peak_mem_used > 0 (force-grow always tracks)
- runtime.memory_pool.reserved() == 0 after completion (grow/shrink balanced)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: track join_arrays memory in reservation after SMJ spill [datafusion]
Posted by "SubhamSinghal (via GitHub)" <gi...@apache.org>.
SubhamSinghal commented on PR #21962:
URL: https://github.com/apache/datafusion/pull/21962#issuecomment-4396928119
@mbutrovich @kosiew I have resolved following comments:
1. comment around usage of grow() instead of try_grow()
2. Recomputing join_arrays memory -- added join_arrays_mem field which holds cached size
3. try_shrink vs shrink on free -- replaced try_shrink with shrink
4. peak_mem assertion strength -- replaced test assertion with peak_mem >= join_arrays_mem
5. fixed reserved_amount doc which still refer try_grow()
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org