perf(spark): reduce map deduplication allocations - #24684
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24684 +/- ##
==========================================
+ Coverage 81.45% 81.59% +0.13%
==========================================
Files 1118 1123 +5
Lines 399685 407010 +7325
Branches 399685 407010 +7325
==========================================
+ Hits 325576 332107 +6531
- Misses 55103 55460 +357
- Partials 19006 19443 +437 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Align the value filter with the starting list offset while keeping the zero-offset path free of extra slice allocations. Add regression coverage for wide-to-small rows, LAST_WIN ordering, sliced values, and null rows with unequal child lengths.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed head f89176ef1b67482e400dd5966a11fe74e345995c against base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0; no actionable P1/P2 regressions found. Checked value selection for independently sliced inputs, null rows, both duplicate-key policies, and the lookup-table clear/shrink path.
Validation on the review devbox: all 285 datafusion-spark unit tests passed (including 14 focused map tests), all three Spark-map SQL regression files passed, and an exact-base/head helper harness passed 12,288 comparisons across 12 value types. The scratch harness tests the Rust helpers, not a live Spark/Comet integration. No performance benchmark or full workspace test rerun was performed. All 38 current CI checks passed.
pingz-oai
left a comment
pingz-oai
left a comment
There was a problem hiding this comment.
Reviewed exact head f89176ef1b67482e400dd5966a11fe74e345995c against base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0. One reproducible P2 regression is documented inline.
| let flat_values = if values_start_offset == 0 { | ||
| Cow::Borrowed(flat_values) | ||
| } else { | ||
| Cow::Owned(flat_values.slice(values_start_offset, values_mask.len())) |
There was a problem hiding this comment.
[P2] Preserve large child offsets in the new filter path
A valid map_from_arrays input can now panic when its LargeList values start at child offset 2_147_483_648 and no duplicate overwrite occurs. A key row [7] and a LargeList<Null> value row with offsets [2_147_483_648, 2_147_483_649] previously produce {7: NULL}; this slice instead fails its bounds assertion. Both input arrays pass Arrow's validate_full(), and a NullArray child makes this reproducible without allocating a huge buffer.
The existing get_list_offsets conversion narrows the offsets to negative i32 values. The base's take still works because Arrow 59.2.0 reinterprets its Int32 indices as UInt32, recovering the intended child index. This new path instead passes the sign-extended usize (18446744071562067968 on 64-bit builds) directly to slice.
Minimal helper reproduction using the caller's actual offset conversion, inside the existing utils test module:
use arrow::array::{LargeListArray, NullArray};
let start = 1_i64 << 31;
let keys: ArrayRef = Arc::new(Int32Array::from(vec![7]));
let values: ArrayRef = Arc::new(LargeListArray::new(
Arc::new(Field::new("item", DataType::Null, true)),
OffsetBuffer::new(vec![start, start + 1].into()),
Arc::new(NullArray::new(start as usize + 1)),
None,
));
values.to_data().validate_full().unwrap();
map_from_keys_values_offsets_nulls(
&keys,
get_list_values(&values).unwrap(),
&[0, 1],
&get_list_offsets(&values).unwrap(),
None,
None,
false,
).unwrap();A focused comparison of the base (4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0) and this head's helper bodies, using the locked Arrow dependencies, confirms base success and head panic under both EXCEPTION and LAST_WIN. Preserve the wide offset when slicing, or retain the previous take path for narrowed negative offsets.
pingz-oai
left a comment
pingz-oai
left a comment
There was a problem hiding this comment.
Reviewed exact head dd83e70097aa9d3b68e3590379e60bca3929c8f3 against base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0. One reproducible P2 regression is documented inline.
| // masks used by filter so they stay aligned with their flat arrays. | ||
| keys_mask_builder.append_n(num_keys_entries, false); | ||
| if !needs_value_take { | ||
| values_mask_builder.append_n(num_values_entries, false); |
There was a problem hiding this comment.
[P2] Avoid allocating a bitmap for large skipped value spans
When map_from_arrays receives a NULL typed key-list row, its values row may be much wider and must be ignored. This new append allocates a bit for every ignored value, even for a List<Null> whose child has no payload buffer. For example, the valid two-row input below uses only 92 bytes of Arrow buffers. Base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0 returns [NULL, {7: NULL}], but this head attempts a 134,217,728-byte allocation here. Under the same 64 MiB process address-space limit, base succeeds and head aborts with memory allocation of 134217728 bytes failed, under both EXCEPTION and LAST_WIN. A small valid Arrow input can therefore terminate a memory-limited query worker.
Reproduction inside the existing utils test module (both input arrays pass validate_full()):
use arrow::array::{ListArray, NullArray};
let n = 1_i32 << 30;
let keys: ArrayRef = Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Int32, false)),
OffsetBuffer::new(vec![0, 0, 1].into()),
Arc::new(Int32Array::from(vec![7])),
Some(NullBuffer::from(vec![false, true])),
));
let values: ArrayRef = Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Null, true)),
OffsetBuffer::new(vec![0, n, n + 1].into()),
Arc::new(NullArray::new(n as usize + 1)),
None,
));
keys.to_data().validate_full().unwrap();
values.to_data().validate_full().unwrap();
map_from_keys_values_offsets_nulls(
get_list_values(&keys).unwrap(),
get_list_values(&values).unwrap(),
&get_list_offsets(&keys).unwrap(),
&get_list_offsets(&values).unwrap(),
keys.nulls(), values.nulls(), false,
).unwrap();The exact base/head helper comparison used the locked Arrow 59.2.0 dependencies. Previously the skipped row added no value indices, so take needed space only for the one surviving entry. All offsets here are non-negative, so the new negative-offset fallback does not help. Please retain take for disproportionately large skipped value spans, or otherwise avoid materializing their all-false bitmap.
There was a problem hiding this comment.
Rechecked dd83e70097aa9d3b68e3590379e60bca3929c8f3 against base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0 using Arrow 59.2.0. This P2 still reproduces in a bounded probe: with a skipped outer values span of 1 << 20, the largest allocation grows from 224 B on base to 262,144 B on head under both duplicate-key policies, although both return the same two-row result with one retained value. No OOM experiment was needed.
The harness compiles the exact helper bodies with only Int32-key ScalarValue and error-plumbing adapters; it is not a full DataFusion/Spark integration run. The new mask still scales with ignored input rather than retained output.
[Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.]
pingz-oai
left a comment
pingz-oai
left a comment
There was a problem hiding this comment.
Reviewed exact head dd83e70097aa9d3b68e3590379e60bca3929c8f3 against base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0. One additional reproducible P2 regression is documented inline. The existing finding was revalidated and is not duplicated.
| } else { | ||
| Cow::Owned(flat_values.slice(values_start_offset, values_mask.len())) | ||
| }; | ||
| filter(flat_values.as_ref(), &values_mask)? |
There was a problem hiding this comment.
[P2] Skip null nested-list children in the filtered values path
When a map value is itself a NULL list with a nonempty child span, this new filter() call can materialize that ignored span. This is separate from the outer values_mask_builder.append_n issue: both outer offsets below are [0, 1, 2], both masks contain just two bits ([true, false]), and the retained first map value is NULL. No duplicate or negative offset selects the take fallback.
With n = 1 << 30, both inputs pass Arrow's validate_full() and together use only 172 bytes of buffers. Under the same 64 MiB process address-space limit, base 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0 returns [{7: NULL}, NULL], while this head aborts with memory allocation of 134217728 bytes failed. I reproduced this under both EXCEPTION and LAST_WIN. A small valid map_from_arrays batch can therefore terminate a memory-limited query worker.
Reproduction inside the existing utils test module:
use arrow::array::{ListArray, NullArray};
let n = 1_i32 << 30;
let offsets = OffsetBuffer::new(vec![0_i32, 1, 2].into());
let keys: ArrayRef = Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Int32, false)),
offsets.clone(),
Arc::new(Int32Array::from(vec![7, 8])),
Some(NullBuffer::from(vec![true, false])),
));
let inner: ArrayRef = Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Null, true)),
OffsetBuffer::new(vec![0, n, n + 1].into()),
Arc::new(NullArray::new(n as usize + 1)),
Some(NullBuffer::from(vec![false, true])),
));
let values: ArrayRef = Arc::new(ListArray::new(
Arc::new(Field::new("item", inner.data_type().clone(), true)),
offsets,
inner,
None,
));
keys.to_data().validate_full().unwrap();
values.to_data().validate_full().unwrap();
map_from_keys_values_offsets_nulls(
get_list_values(&keys).unwrap(),
get_list_values(&values).unwrap(),
&get_list_offsets(&keys).unwrap(),
&get_list_offsets(&values).unwrap(),
keys.nulls(), values.nulls(), false, // also reproduces with true
).unwrap();In locked Arrow 59.2.0, filtering List uses MutableArrayData: the selected NULL list's offsets still extend its full child span, and inherited null handling allocates a child validity bitmap. The previous take_list instead iterates output_nulls.valid_indices() and skips this span. The independently compiled base/head helper probes used the exact helper bodies, with a standalone adapter for the unchanged Int32-key and error plumbing. With a smaller n = 65_536, both complete, but the head allocates an 8,192-byte bitmap and retains 65,536 unused child elements where the base retains zero.
Please retain take for this partial-filter case, or otherwise avoid materializing children of NULL nested values. A guard based only on large skipped outer-row spans would miss this input.
There was a problem hiding this comment.
Rechecked the nested-null case on dd83e70097aa9d3b68e3590379e60bca3929c8f3 with pinned Arrow 59.2.0. It still reproduces independently of the skipped-outer-span finding: for inner offsets [0, 65536, 65537], inner validity [false, true], outer offsets [0, 1, 2], and key-row validity [true, false], base retains zero inner child elements (largest allocation 224 B), while head retains 65,536 unused elements (largest allocation 8,192 B). Both duplicate-key policies show the same result.
This bounded test uses the exact base/head helper bodies with Int32-key and error-plumbing adapters. Arrow filter takes the MutableArrayData path for the partial List selection, while take skips the null child span. The negative-offset guard and Cow slice do not address this case. No huge allocation or full-engine run was performed.
[Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.]
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
Thanks for pushing this — the two follow-up commits clearly landed real fixes. I re-measured all three of @pingz-oai's reports on the current head (dd83e70097) against base 4fcaa01c with a counting global allocator, tracking the peak single allocation, and two of them are genuinely resolved:
| scenario | base | head dd83e70 |
|---|---|---|
NULL nested value with a large child span (:281) |
9,680 B | 332 B |
LargeList values near the i32 offset limit (:279) |
panicked | fixed, with a regression test |
The Cow/slice rework and the needs_value_take guard for narrowed offsets both do their job, and the :281 path actually comes out better than base.
The third one is still there, though. Details inline.
One thing worth raising at the PR level rather than inline: this is a performance change without any before/after numbers. Your notes mention local runs were blocked by a mirror missing blake3 1.8.7 — that constraint should be gone now (arrow 59.2.0 resolves normally again), so a quick benchmark on the two paths you're optimising would make the gain concrete. It matters more than usual here because the :259 case shows the change is not a strict improvement across input shapes.
| // masks used by filter so they stay aligned with their flat arrays. | ||
| keys_mask_builder.append_n(num_keys_entries, false); | ||
| if !needs_value_take { | ||
| values_mask_builder.append_n(num_values_entries, false); |
There was a problem hiding this comment.
The large-skipped-span allocation from the :259 thread still reproduces on this head. Same shape @pingz-oai described — a NULL key row whose values row spans a large offset range, with no negative offset and no LAST_WIN overwrite:
base 4fcaa01c |
head dd83e70 |
|
|---|---|---|
| peak single allocation | 332 B | 268,435,456 B |
Reproduced under both EXCEPTION and LAST_WIN, from an input of a few hundred bytes of Arrow buffers (flat_values as a NullArray, values_offsets = [0, 1<<30, (1<<30)+1], keys null buffer [false, true]).
The if !needs_value_take guard added in dd83e700 doesn't cover this case: needs_value_take only becomes true once a negative offset or an overwrite has been seen, and this input has neither, so the append_n still allocates one bit per ignored value. The guard narrowed the trigger surface rather than removing it.
Since the ignored span contributes nothing to the output, could the mask be built over only the retained value range instead of the full flat length — or the whole filter path skipped when a row's value span is disproportionate to what it contributes? I don't want to prescribe the shape; you know this code's constraints better. But as it stands a small valid map_from_arrays batch can still allocate hundreds of MiB where base allocated hundreds of bytes.
| values_nulls: Option<&NullBuffer>, | ||
| last_value_wins: bool, | ||
| ) -> Result<(ArrayRef, ArrayRef, OffsetBuffer<i32>)> { | ||
| const MIN_RETAINED_LOOKUP_CAPACITY: usize = 16; |
There was a problem hiding this comment.
MIN_RETAINED_LOOKUP_CAPACITY = 16 and MAX_RETAINED_LOOKUP_CAPACITY_RATIO = 4 are the tuning knobs for the whole second optimisation, but there's nothing recording where they came from. A sentence on what they're trading off (and whether they were measured or chosen as round numbers) would help whoever revisits this.
Related, on the comparison itself: HashMap::capacity() reports the load-factor-adjusted capacity rather than the raw bucket count, so capacity() > target * 4 is a looser test than it reads as. Worth a note that the heuristic only needs to be approximately right, if that's the intent.
| let mut value_indices: Vec<i32> = Vec::new(); | ||
| // LargeList offsets can narrow to negative i32 values. Keep take's index | ||
| // handling for those offsets instead of sign-extending them for slice. | ||
| let mut needs_value_take = values_offsets.first().is_some_and(|offset| *offset < 0); |
There was a problem hiding this comment.
This flag can flip mid-loop, after values_mask_builder has already been appended to for earlier rows. That's safe today because the mask is then discarded in favour of take, but the safety rests on an invariant that isn't written down: the values mask is only ever read when needs_value_take is false. Once the flag is set, the mask is deliberately left inconsistent (NULL rows stop padding it).
A comment stating that would protect the next change here — someone adding another consumer of values_mask would otherwise have no signal that it can be a partial mask.
|
run benchmark map |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing dev/codex/datafusion-map-dedup-allocations (dd83e70) to 4fcaa01 (merge-base) diff Run configurationrun benchmark mapResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing dev/codex/datafusion-map-dedup-allocations (dd83e70) to 4fcaa01 (merge-base) diff Run configurationrun benchmark mapCPU Details (lscpu)Details
Resource Usagemap — base (merge-base)
map — branch
File an issue against this benchmark runner |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @ywskycn the PR looks consistent and makes sense, one thing came from automated review
P1 — Keys path assumes keys_offsets[0] == 0; sliced inputs produce a corrupt map (pre-existing, not a regression, but directly adjacent to this PR)
Description: This PR makes the values path offset-correct (slice flat_values by
values_start_offset, utils.rs:276-281) but the keys path is still
filter(&flat_keys, &keys_mask) (utils.rs:268, an unchanged context line). keys_mask
is built positionally starting at the first processed entry, so bit i maps to
flat_keys[i], not flat_keys[keys_offsets[0] + i]. When keys_offsets[0] != 0 the
keys are selected from the wrong region while the values (now) come from the correct
region, so keys and values disagree.
Reason: Wrong query results (a corrupt Map) for any sliced list input. get_list_values
returns the full child and get_list_offsets returns non-rebased offsets, so a sliced
list column yields offsets[0] > 0. make_scalar_function (datafusion/functions/src/utils.rs:110-135)
does not rebase offsets, and multiple operators emit offset>0 batches downstream
(hash-repartition reordered_batch.slice(start,len), repartition/mod.rs:1273;
LIMIT ... OFFSET batch.slice(self.skip, ..), limit.rs:663). map_from_entries is the
most exposed caller because keys and values share one non-zero offset.
Evidence: Reproduced against the applied PR (probe replicating map_from_arrays_inner
with a list sliced to rows [1..3], keys_offsets=[2,5,7], full child len 7):
- result keys =
[10,11,20,21,22](wrong — read from offset 0) - result values =
[c,d,e,f,g](correct — sliced region, fixed by this PR)
This is pre-existing (values used absolute-indextakebefore, which was already
offset-correct; keys were already wrong). Since this PR is specifically hardening offset
handling for this function and adds a non-zero-values-offset test, fixing keys
symmetrically here is cheap and on-theme: sliceflat_keysthe same way, or keeptake
for keys. At minimum add a failing test and file a follow-up.
ziting-openai
left a comment
ziting-openai
left a comment
There was a problem hiding this comment.
Reviewed current head dd83e70097aa9d3b68e3590379e60bca3929c8f3: no PR-introduced P1 and two significant P2 allocation regressions remain. Fresh bounded Arrow 59.2.0 evidence is in the existing skipped outer-span and null nested-list threads; no duplicate finding was opened. The narrowed-large-offset regression is fixed. The sliced-key issue noted elsewhere is unchanged baseline behavior.
Approving with those two P2s still open; this does not mark them fixed. Validation: 22 standalone extracted-helper tests passed (6 base, 14 head, 2 bounded comparative probes), with Int32-key/error adapters and pinned Arrow 59.2.0. Full DataFusion, SQL, Spark/Comet integration, and performance benchmarks were not run.
[Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.]
ziting-openai
left a comment
ziting-openai
left a comment
There was a problem hiding this comment.
Reviewed exact head b0b359d1ea69beafc1c33a6db5be59cb1fe8ddf0 against 4fcaa01c721ba18c10a5ac65446aa48ddf68e9b0. No introduced P1 or significant P2 findings remain.
The existing skipped outer-span and null nested-list allocation findings are fixed. Fresh bounded Arrow 59.2.0 comparisons under both duplicate policies restore the largest allocation to 224 B in each case; the nested case retains zero hidden child elements. Distinct values still share buffers, and skipped/overwritten values use take.
Validation: all 291 actual datafusion-spark unit tests passed, including 20 map tests, and all three Spark map SQL regression files passed. A separate extracted-helper harness passed 41 tests, including 7,680 base/head comparisons across 15 value types. Only that standalone harness substitutes Int32-key ScalarValue/error plumbing. Full workspace/platform tests and timing benchmarks were not run locally. At preflight, 31 CI checks had succeeded and seven were still running.
[Posted by Codex on behalf of ziting-openai using the spark-pr-review-memo skill.]
8 participants
Which issue does this PR close?
Rationale for this change
map_from_entries/map_from_arrayscan do unnecessary work when deduplicating map keys:LAST_WINoverwrite still materialize values throughtake(), copying value buffers unnecessarily.This adds avoidable CPU, allocation, and memory overhead for common and skewed map workloads.
What changes are included in this PR?
filter()when noLAST_WINoverwrite occurs.take()after an actual duplicate key overwrite changes value ordering.Are these changes tested?
cargo fmt --all --checkblake3 1.8.7; CI should run the full suite.Are there any user-facing changes?
No semantic or API changes. This only reduces map construction overhead.