From 1cb786a8466e185653828e20dd51db5295a4798c Mon Sep 17 00:00:00 2001 From: Punisheroot <44579963+Punisheroot@users.noreply.github.com> Date: Tue, 4 Aug 2026 00:58:35 +0200 Subject: [PATCH] perf: use Vec in ArrowBytesMap Replace BufferBuilder with Vec for accumulating values in ArrowBytesMap, converting the Vec into an Arrow Buffer without copying when materializing the output array. Add Criterion benchmarks covering short unique values, long unique values, and long values with low cardinality. --- datafusion/physical-expr-common/Cargo.toml | 4 + .../benches/arrow_bytes_map.rs | 82 +++++++++++++++++++ .../physical-expr-common/src/binary_map.rs | 22 ++--- 3 files changed, 97 insertions(+), 11 deletions(-) create mode 100644 datafusion/physical-expr-common/benches/arrow_bytes_map.rs diff --git a/datafusion/physical-expr-common/Cargo.toml b/datafusion/physical-expr-common/Cargo.toml index d1ee7feb29db1..903f5a6a901ac 100644 --- a/datafusion/physical-expr-common/Cargo.toml +++ b/datafusion/physical-expr-common/Cargo.toml @@ -65,3 +65,7 @@ rand = { workspace = true } [[bench]] harness = false name = "compare_nested" + +[[bench]] +harness = false +name = "arrow_bytes_map" diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs new file mode 100644 index 0000000000000..7c8cdc3b4c50e --- /dev/null +++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, StringArray}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +fn make_short_strings(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS).map(|index| format!("{:04x}", index % cardinality)); + Arc::new(StringArray::from_iter_values(values)) +} + +fn make_long_strings(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS).map(|index| { + let value = (index % cardinality) as u32; + format!( + "{value:08x}{:08x}{:08x}{:08x}", + value.wrapping_mul(17), + value.wrapping_mul(31), + value.wrapping_mul(127) + ) + }); + Arc::new(StringArray::from_iter_values(values)) +} + +fn bench_arrow_bytes_map(c: &mut Criterion) { + let cases = [ + // Exercises inline entry storage while still growing the output buffer. + ("short_unique", make_short_strings(NUM_ROWS)), + // Exercises repeated buffer growth and out-of-line entry storage. + ("long_unique", make_long_strings(NUM_ROWS)), + // Fits the distinct values in the initial buffer and repeats comparisons. + ("long_low_cardinality", make_long_strings(128)), + ]; + + let mut group = c.benchmark_group("arrow_bytes_map"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + + for (name, values) in cases { + group.bench_function(name, |b| { + b.iter(|| { + let mut map = ArrowBytesMap::::new(OutputType::Utf8); + let mut next_payload = 0; + map.insert_if_new( + &values, + |_| { + let payload = next_payload; + next_payload += 1; + payload + }, + |payload| { + black_box(payload); + }, + ); + black_box(map.into_state()) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_arrow_bytes_map); +criterion_main!(benches); diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ad184d6500d56..44ca35c7f8708 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -19,12 +19,12 @@ //! StringArray / LargeStringArray / BinaryArray / LargeBinaryArray. use arrow::array::{ - Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray, - NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, GenericBinaryArray, GenericStringArray, NullBufferBuilder, + OffsetSizeTrait, cast::AsArray, types::{ByteArrayType, GenericBinaryType, GenericStringType}, }; -use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; @@ -218,8 +218,8 @@ where map: hashbrown::hash_table::HashTable>, /// Total size of the map in bytes map_size: usize, - /// In progress arrow `Buffer` containing all values - buffer: BufferBuilder, + /// In progress buffer containing all values + buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used /// directly to create the final `GenericBinaryArray`. The `i`th string is /// stored in the range `offsets[i]..offsets[i+1]` in `buffer`. Null values @@ -248,7 +248,7 @@ where output_type, map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), map_size: 0, - buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY), + buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), hashes_buffer: vec![], @@ -405,7 +405,7 @@ where // Put the small values into buffer and offsets so it appears // the output array, but store the actual bytes inline for // comparison - self.buffer.append_slice(value); + self.buffer.extend_from_slice(value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); let new_header = Entry { @@ -433,7 +433,7 @@ where // Need to compare the bytes in the buffer // SAFETY: buffer is only appended to, and we correctly inserted values and offsets let existing_value = - unsafe { self.buffer.as_slice().get_unchecked(header.range()) }; + unsafe { self.buffer.get_unchecked(header.range()) }; value == existing_value }); @@ -446,7 +446,7 @@ where // appears the output array, and store that offset // so the bytes can be compared if needed let offset = self.buffer.len(); // offset of start for data - self.buffer.append_slice(value); + self.buffer.extend_from_slice(value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); @@ -488,7 +488,7 @@ where map: _, map_size: _, offsets, - mut buffer, + buffer, random_state: _, hashes_buffer: _, null, @@ -502,7 +502,7 @@ where // SAFETY: the offsets were constructed correctly in `insert_if_new` -- // monotonically increasing, overflows were checked. let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; - let values = buffer.finish(); + let values = Buffer::from_vec(buffer); match output_type { OutputType::Binary => {