viirya commented on code in PR #24684:
URL: https://github.com/apache/datafusion/pull/24684#discussion_r3878880381
##########
datafusion/spark/src/function/map/utils.rs:
##########
@@ -225,31 +246,48 @@ fn map_deduplicate_keys(
);
}
keys_mask_builder.append_value(true);
+ values_mask_builder.append_value(true);
key_to_output_idx.insert(key, value_indices.len());
value_indices.push(abs_value_idx);
new_last_offset += 1;
}
} else {
// The result entry is NULL — no keys/values emitted. Still pad the
- // mask so it stays aligned with `flat_keys`.
+ // 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);
Review 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.
##########
datafusion/spark/src/function/map/utils.rs:
##########
@@ -161,28 +161,38 @@ fn map_deduplicate_keys(
values_nulls: Option<&NullBuffer>,
last_value_wins: bool,
) -> Result<(ArrayRef, ArrayRef, OffsetBuffer<i32>)> {
+ const MIN_RETAINED_LOOKUP_CAPACITY: usize = 16;
Review 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.
##########
datafusion/spark/src/function/map/utils.rs:
##########
@@ -161,28 +161,38 @@ fn map_deduplicate_keys(
values_nulls: Option<&NullBuffer>,
last_value_wins: bool,
) -> Result<(ArrayRef, ArrayRef, OffsetBuffer<i32>)> {
+ const MIN_RETAINED_LOOKUP_CAPACITY: usize = 16;
+ const MAX_RETAINED_LOOKUP_CAPACITY_RATIO: usize = 4;
+
let offsets_len = keys_offsets.len();
let mut new_offsets = Vec::with_capacity(offsets_len);
let mut cur_keys_offset = keys_offsets
.first()
.map(|offset| *offset as usize)
.unwrap_or(0);
- let mut cur_values_offset = values_offsets
+ let values_start_offset = values_offsets
.first()
.map(|offset| *offset as usize)
.unwrap_or(0);
+ let mut cur_values_offset = values_start_offset;
let mut new_last_offset = 0;
new_offsets.push(new_last_offset);
// Mirror Spark's `ArrayBasedMapBuilder`: the first occurrence of a key
// fixes its position in the output; under LAST_WIN a later duplicate
- // overwrites that slot's value. `keys_mask` selects the first-seen keys,
- // `value_indices` records the source index in `flat_values` to materialize
- // for each output slot (updated in place on overwrite).
+ // overwrites that slot's value. `keys_mask` selects the first-seen keys.
+ // Keep a matching `values_mask` for the common case where no overwrite
+ // occurs, so Arrow's all-true filter path can share the original value
+ // buffers. Use `take` for value reordering after a LAST_WIN overwrite or
+ // for narrowed large child offsets that cannot be used directly by
`slice`.
let mut keys_mask_builder = BooleanBuilder::new();
+ let mut values_mask_builder = BooleanBuilder::new();
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);
Review 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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
