sunchao commented on code in PR #7106:
URL: https://github.com/apache/datafusion-comet/pull/7106#discussion_r4361263148
##########
native/spark-expr/src/conversion_funcs/cast.rs:
##########
@@ -876,7 +881,75 @@ fn cast_binary_to_string<O: OffsetSizeTrait>(
let num_rows = input.len();
let offsets = input.value_offsets();
- let value_bytes = offsets[num_rows].as_usize() - offsets[0].as_usize();
+ let start = offsets[0].as_usize();
+ let end = offsets[num_rows].as_usize();
+ // Bytes of the non-null rows. A null row can still span bytes, for
example after `nullif`.
+ let live_bytes = match input.nulls().filter(|nulls| nulls.null_count() >
0) {
+ Some(nulls) => nulls
+ .valid_indices()
+ .map(|i| offsets[i + 1].as_usize() - offsets[i].as_usize())
+ .sum(),
+ None => end - start,
+ };
+
+ // Check the bytes the rows span once. If they are valid UTF-8, reuse the
input's value buffer,
+ // or copy them when they fill less than half of its allocation; otherwise
copy the rows in the
+ // valid prefix and decode the rest. Skip this when null rows span more
than half of those
+ // bytes: decoding the live rows one at a time then reads less and returns
a compact buffer
+ // instead of keeping the null rows' bytes alive.
+ let mut valid_prefix = "";
+ if matches!(
+ spark_cast_options.binary_output_style,
+ None | Some(BinaryOutputStyle::Utf8)
+ ) && 2 * live_bytes >= end - start
+ {
+ let data = input.value_data();
+ match std::str::from_utf8(&data[start..end]) {
+ Ok(valid) => {
+ // An offset splits a code point only if it points at a
continuation byte.
+ if offsets.iter().all(|offset| {
+ let offset = offset.as_usize();
+ offset == end || !matches!(data[offset], 0x80..=0xBF)
+ }) {
+ // Reuse a buffer only when the rows fill at least half of
its allocation. A
+ // slice of a large array shares the whole allocation, and
memory accounting
+ // such as a sort's charges all of it to every batch that
holds it, so the
+ // values are copied then, without decoding them again.
+ let fills = |buffer: &Buffer, used: usize| 2 * used >=
buffer.capacity();
+ let offset_buffer = input.offsets().inner().inner();
+ let offsets = if start == 0 && fills(offset_buffer,
offset_buffer.len()) {
+ input.offsets().clone()
+ } else {
+ // Rebases a slice's offsets onto the narrowed value
buffer.
+ OffsetBuffer::from_lengths(input.offsets().lengths())
+ };
+ let values = if fills(input.values(), end - start) {
+ input.values().slice_with_length(start, end - start)
+ } else {
+ Buffer::from_slice_ref(&data[start..end])
+ };
+ // SAFETY: the value buffer is valid UTF-8 and every
offset falls on a char
+ // boundary, which is what `GenericStringArray::try_new`
checks.
+ let strings = unsafe {
+ GenericStringArray::<O>::new_unchecked(
+ offsets,
+ values,
+ input.nulls().cloned(),
Review Comment:
[P2] [P2] Compact oversized validity buffers alongside values and offsets
The fast path unconditionally clones `input.nulls()`, retaining the parent
bitmap even when the other buffers are compacted. This affects nullable binary
aggregate keys cast before sorting, such as `SELECT CAST(k AS STRING) AS s FROM
t GROUP BY id, k SORT BY s` with `spark.comet.batchSize=128`. In a reproduction
with 262,144 groups, six-byte valid values and one null per 128 rows, each
output batch retained a 32,768-byte bitmap instead of the base’s 64-byte
bitmap. The sorter charges that allocation for every batch: with a 64 MiB pool,
the base had zero spills while the head had two and spilled every output row.
Compacting only the validity bitmap restored zero spills with identical
results. Could the same capacity-versus-visible-size policy also cover this
buffer? This extends the earlier value-buffer fix to a separate allocation that
still causes avoidable spilling.
Evidence: Ran `LD_LIBRARY_PATH=/usr/lib/jvm/java-21-openjdk-amd64/lib/server
/tmp/pr7106-1023-review-validity`. Source:
`/tmp/pr7106-d297-this-pass-validity-confirm.rs`. Output:
`/tmp/pr7106-1023-review-validity.log`. The bounded DataFusion 55.2.0
AggregateExec → ProjectionExec → SortExec probe calls the compiled exact-head
Comet `spark_cast`. Its baseline helper was verified byte-for-byte against base
`39380cca813bc0b861353de2cb53af960dd2a87f`. All variants returned identical
262,144-row sorted results and passed Arrow validation. Validity capacity per
128-row batch was base=64 bytes, head=32,768 bytes, compacted=64 bytes. Spill
counts were 0, 2 and 0 respectively. Arrow 59.3.0 `BooleanBuffer::slice` shares
the backing buffer, DataFusion’s `MaterializedAggregateOutput::next_batch`
emits shared slices, and the sorter reserves each batch’s retained allocation
size. This measures actual spill behavior, not end-to-end Spark runtime.
--
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]
