This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 2b35fb8eec perf: vectorize Map in spark_size via offset buffer (#5395)
2b35fb8eec is described below
commit 2b35fb8eec0fccff56f705073de572b34f0fe3dc
Author: ChenChen Lai <[email protected]>
AuthorDate: Fri Aug 21 00:23:04 2026 +0800
perf: vectorize Map in spark_size via offset buffer (#5395)
* perf: vectorize Map in spark_size via offset buffer
* check style
* fix Prettier
---
.../contributor-guide/optimizing_expressions.md | 2 +-
native/spark-expr/benches/array_size.rs | 63 ++++++-
native/spark-expr/src/array_funcs/size.rs | 186 +++++++++++++++++----
3 files changed, 218 insertions(+), 33 deletions(-)
diff --git a/docs/source/contributor-guide/optimizing_expressions.md
b/docs/source/contributor-guide/optimizing_expressions.md
index ecdb8a7263..4b11dd0487 100644
--- a/docs/source/contributor-guide/optimizing_expressions.md
+++ b/docs/source/contributor-guide/optimizing_expressions.md
@@ -149,7 +149,7 @@ for the lightest one that fits.
| **Preallocate builders to known size** |
Repeated buffer growth/reallocation
| `spark_unhex`: preallocate `BinaryBuilder` to the known
output length
|
| **Compile-time lookup tables** |
Per-element range matches / branching
| `spark_unhex`: 256-entry hex table instead of per-digit
range match
|
| **Cache compiled regex** (thread-local, keyed by the constant arg) |
`Regex::new()` per row
| `parse_url` QUERY-with-key (50x): the key is constant
across a batch
|
-| **Read from the offset buffer directly** |
`list_array.value(i)` allocating a sliced `ArrayRef` per row
| `spark_size` Map path and scalar path: `value_length` /
offsets instead of materializing a sliced array
|
+| **Read from the offset buffer directly** |
`list_array.value(i)` allocating a sliced `ArrayRef` per row
| `spark_size` scalar path: `value_length` instead of
materializing a sliced array; Map array path (~6-13x):
`MapArray::offsets().windows(2)` diff, `length` kernel doesn't support Map |
| **Reuse Arrow `length` kernel + cheap null rewrite** | Per-row
`value_length` / builder loop, or `zip` via `MutableArrayData`
| `spark_size` (~12x on no-null List): `length` for
List/LargeList/FixedSizeList, then `into_parts` + `set_indices` to patch null
slots to `-1` |
| **Typed scans over flat values buffers + hash probe** | A
per-element Arrow `eq`/compute kernel that allocates per call
| `spark_arrays_overlap` (up to 18x): scan buffers
directly, hash probe for large lists
|
| **ASCII / byte-offset fast path** |
`chars().count()` and per-char UTF-8 decoding
| `substring` (up to 10x), `spark_lpad` (2x): slice by
byte offset when input is ASCII
|
diff --git a/native/spark-expr/benches/array_size.rs
b/native/spark-expr/benches/array_size.rs
index afb340c06e..c82468f462 100644
--- a/native/spark-expr/benches/array_size.rs
+++ b/native/spark-expr/benches/array_size.rs
@@ -15,9 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{ArrayRef, Int32Array, LargeListArray, ListArray};
+use arrow::array::{ArrayRef, Int32Array, LargeListArray, ListArray, MapArray,
StructArray};
use arrow::buffer::{NullBuffer, OffsetBuffer};
-use arrow::datatypes::{DataType, Field};
+use arrow::datatypes::{DataType, Field, Fields};
use criterion::{criterion_group, criterion_main, Criterion};
use datafusion::physical_plan::ColumnarValue;
use datafusion_comet_spark_expr::spark_size;
@@ -70,6 +70,52 @@ fn create_large_list_array(rows: usize, elems_per_row:
usize) -> ArrayRef {
))
}
+/// Build a `MapArray` of `rows` maps, each with `entries_per_row` Int32βInt32
+/// entries. When `with_nulls` is true every 10th row is null. Map entry counts
+/// come from the offset buffer, so per-row diff is O(1) already; this bench
+/// isolates per-row dispatch + builder overhead vs. one-shot offset
differencing.
+fn create_map_array(rows: usize, entries_per_row: usize, with_nulls: bool) ->
ArrayRef {
+ let total = rows * entries_per_row;
+ let keys = Int32Array::from((0..total as i32).collect::<Vec<i32>>());
+ let values = Int32Array::from((0..total as i32).collect::<Vec<i32>>());
+
+ let mut offsets = Vec::with_capacity(rows + 1);
+ offsets.push(0i32);
+ for i in 1..=rows {
+ offsets.push((i * entries_per_row) as i32);
+ }
+
+ let key_field = Arc::new(Field::new("key", DataType::Int32, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+ let entries = StructArray::new(
+ Fields::from(vec![key_field, value_field]),
+ vec![Arc::new(keys), Arc::new(values)],
+ None,
+ );
+
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(Fields::from(vec![
+ Field::new("key", DataType::Int32, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
+
+ let nulls =
+ with_nulls.then(|| NullBuffer::from((0..rows).map(|i| i % 10 !=
0).collect::<Vec<bool>>()));
+ Arc::new(
+ MapArray::try_new(
+ map_field,
+ OffsetBuffer::new(offsets.into()),
+ entries,
+ nulls,
+ false,
+ )
+ .unwrap(),
+ )
+}
+
fn criterion_benchmark(c: &mut Criterion) {
let rows = 8192;
@@ -105,6 +151,19 @@ fn criterion_benchmark(c: &mut Criterion) {
"spark_size: LargeList (10% null)",
&create_large_list_array(rows, 5),
);
+
+ // Map shapes: no length-kernel support for MapArray, so this path
differences
+ // offsets directly. Cover the two shapes the issue calls for β with and
without
+ // nulls β since the null rewrite is the only branch that touches the
values
+ // buffer.
+ bench(
+ "spark_size: map (10% null)",
+ &create_map_array(rows, 5, true),
+ );
+ bench(
+ "spark_size: map, no nulls",
+ &create_map_array(rows, 5, false),
+ );
}
criterion_group!(benches, criterion_benchmark);
diff --git a/native/spark-expr/src/array_funcs/size.rs
b/native/spark-expr/src/array_funcs/size.rs
index 0d555806b7..2311da5fa9 100644
--- a/native/spark-expr/src/array_funcs/size.rs
+++ b/native/spark-expr/src/array_funcs/size.rs
@@ -97,23 +97,9 @@ fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef,
DataFusionError> {
DataType::List(_) | DataType::LargeList(_) |
DataType::FixedSizeList(..) => {
spark_size_list_like(array)
}
- // Map is not supported by the length kernel; keep the offset-based
path.
- DataType::Map(_, _) => {
- let map_array = array
- .as_any()
- .downcast_ref::<arrow::array::MapArray>()
- .ok_or_else(|| DataFusionError::Internal("Expected
MapArray".to_string()))?;
-
- let mut builder = Int32Array::builder(map_array.len());
- for i in 0..map_array.len() {
- if map_array.is_null(i) {
- builder.append_value(-1); // Spark behavior: return -1 for
null
- } else {
- builder.append_value(map_array.value_length(i));
- }
- }
- Ok(Arc::new(builder.finish()))
- }
+ // Map is not supported by the length kernel; compute row lengths from
+ // the entry-offset buffer directly.
+ DataType::Map(_, _) => spark_size_map(array),
_ => {
exec_err!(
"size function only supports arrays and maps, got: {:?}",
@@ -126,10 +112,6 @@ fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef,
DataFusionError> {
/// Compute Spark `size()` for list-like arrays via Arrow's `length` kernel,
then
/// rewrite null inputs to `-1` (Spark's legacy/compatible size-of-null
behavior
/// for this UDF). LargeList lengths are Int64 and are cast to Int32.
-///
-/// Patches the values buffer in place rather than using `zip`; `zip` goes
-/// through `MutableArrayData` and roughly doubled runtime on the `array_size`
-/// criterion shapes.
fn spark_size_list_like(array: &ArrayRef) -> Result<ArrayRef, DataFusionError>
{
let lengths = length(array.as_ref())?;
let lengths = match lengths.data_type() {
@@ -151,7 +133,8 @@ fn spark_size_list_like(array: &ArrayRef) ->
Result<ArrayRef, DataFusionError> {
// Fast path for the production shape: `CometSize.convert` wraps size() in
a
// `CASE WHEN isnotnull(child)` that filters null rows out before the THEN
// branch runs, so this function only ever sees a null-free array in a real
- // Comet plan. Return the length kernel output as-is.
+ // Comet plan. Return the length kernel output as-is; skip the downcast and
+ // `Int32Array::clone` that `rewrite_nulls_to_minus_one` would otherwise
pay.
if array.null_count() == 0 {
return Ok(lengths);
}
@@ -159,21 +142,48 @@ fn spark_size_list_like(array: &ArrayRef) ->
Result<ArrayRef, DataFusionError> {
let int_lengths = lengths
.as_any()
.downcast_ref::<Int32Array>()
- .ok_or_else(|| DataFusionError::Internal("Expected Int32Array from
length".to_string()))?;
+ .ok_or_else(|| DataFusionError::Internal("Expected Int32Array from
length".to_string()))?
+ .clone();
+ Ok(rewrite_nulls_to_minus_one(int_lengths))
+}
- // `set_indices()` on the inverted validity visits only null slots
- // (O(null_count)). We still `to_vec()` the values buffer (O(n)) so we can
- // write `-1` into those slots; `into_parts` avoids an extra values-buffer
- // clone beyond that copy. Prefer this over scanning every validity bit.
- let (_, values, nulls) = int_lengths.clone().into_parts();
+/// Compute Spark `size()` for `MapArray` by differencing the entry-offset
+/// buffer. `MapArray::offsets()` returns the sliced offsets, so `windows(2)`
+/// stays correct on sliced inputs. Applies the same null β `-1` rewrite as the
+/// list path.
+fn spark_size_map(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
+ let map_array = array
+ .as_any()
+ .downcast_ref::<arrow::array::MapArray>()
+ .ok_or_else(|| DataFusionError::Internal("Expected
MapArray".to_string()))?;
+
+ let offsets = map_array.offsets();
+ let values: Vec<i32> = offsets.windows(2).map(|w| w[1] - w[0]).collect();
+ let lengths = Int32Array::new(values.into(), map_array.nulls().cloned());
+ Ok(rewrite_nulls_to_minus_one(lengths))
+}
+
+/// Rewrite null slots to `-1` (Spark size-of-null semantics) and drop the null
+/// buffer. Returns the input untouched when there are no nulls.
+///
+/// Patches the values buffer in place rather than using `zip`; `zip` goes
+/// through `MutableArrayData` and roughly doubled runtime on the `array_size`
+/// criterion shapes. `set_indices()` on the inverted validity visits only null
+/// slots (O(null_count)); `into_parts` avoids an extra values-buffer clone
+/// beyond the `to_vec()` copy needed to write into the null slots.
+fn rewrite_nulls_to_minus_one(lengths: Int32Array) -> ArrayRef {
+ if lengths.null_count() == 0 {
+ return Arc::new(lengths);
+ }
+ let (_, values, nulls) = lengths.into_parts();
let Some(nulls) = nulls else {
- return Ok(Arc::new(Int32Array::new(values, None)));
+ return Arc::new(Int32Array::new(values, None));
};
let mut values = values.to_vec();
for i in (!nulls.inner()).set_indices() {
values[i] = -1;
}
- Ok(Arc::new(Int32Array::from(values)))
+ Arc::new(Int32Array::from(values))
}
fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue,
DataFusionError> {
@@ -375,6 +385,122 @@ mod tests {
assert_eq!(result.value(3), -1); // null map returns -1
}
+ #[test]
+ fn test_spark_size_map_no_nulls() {
+ use arrow::array::{Int32Array, MapArray, StringArray};
+
+ // Fast path for Map: no null buffer, so rewrite_nulls_to_minus_one
returns
+ // the offset-diff Int32Array as-is without touching the values buffer.
+ let keys = StringArray::from(vec![Some("a"), Some("b"), Some("c"),
Some("d")]);
+ let values = Int32Array::from(vec![Some(1), Some(2), Some(3),
Some(4)]);
+ let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 2, 3,
3, 4].into());
+
+ let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+
+ let entries = arrow::array::StructArray::new(
+ arrow::datatypes::Fields::from(vec![key_field, value_field]),
+ vec![Arc::new(keys), Arc::new(values)],
+ None,
+ );
+
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(arrow::datatypes::Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
+
+ let map_array = MapArray::try_new(map_field, entry_offsets, entries,
None, false).unwrap();
+
+ let array_ref: ArrayRef = Arc::new(map_array);
+ let result = spark_size_array(&array_ref).unwrap();
+ let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+
+ // No nulls on input; no null buffer on output.
+ assert_eq!(result.null_count(), 0);
+ assert_eq!(result.value(0), 2);
+ assert_eq!(result.value(1), 1);
+ assert_eq!(result.value(2), 0);
+ assert_eq!(result.value(3), 1);
+ }
+
+ #[test]
+ fn test_spark_size_sliced_map_array() {
+ use arrow::array::{Int32Array, MapArray, StringArray};
+
+ // Slicing pins that `MapArray::offsets()` returns sliced-aware
offsets so
+ // `windows(2)` differences the correct rows. Same trap as the
sliced-List
+ // test above.
+ //
+ // Full map: [{"a":1,"b":2}, {"c":3}, {"d":4,"e":5,"f":6}, null,
{"g":7}]
+ let keys = StringArray::from(vec![
+ Some("a"),
+ Some("b"),
+ Some("c"),
+ Some("d"),
+ Some("e"),
+ Some("f"),
+ Some("g"),
+ ]);
+ let values = Int32Array::from(vec![
+ Some(1),
+ Some(2),
+ Some(3),
+ Some(4),
+ Some(5),
+ Some(6),
+ Some(7),
+ ]);
+ let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 2, 3,
6, 6, 7].into());
+
+ let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+
+ let entries = arrow::array::StructArray::new(
+ arrow::datatypes::Fields::from(vec![key_field, value_field]),
+ vec![Arc::new(keys), Arc::new(values)],
+ None,
+ );
+
+ let mut null_buffer = NullBufferBuilder::new(5);
+ null_buffer.append(true);
+ null_buffer.append(true);
+ null_buffer.append(true);
+ null_buffer.append(false);
+ null_buffer.append(true);
+
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(arrow::datatypes::Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
+
+ let map_array = MapArray::try_new(
+ map_field,
+ entry_offsets,
+ entries,
+ null_buffer.finish(),
+ false,
+ )
+ .unwrap();
+
+ // Skip the first two rows: sliced view is [{d,e,f}, null, {g}].
+ let sliced: ArrayRef = Arc::new(map_array.slice(2, 3));
+ let result = spark_size_array(&sliced).unwrap();
+ let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+
+ assert_eq!(result.len(), 3);
+ assert_eq!(result.value(0), 3);
+ assert_eq!(result.value(1), -1);
+ assert_eq!(result.value(2), 1);
+ }
+
#[test]
fn test_spark_size_scalar_map() {
use arrow::array::{Int32Array, MapArray, StringArray};
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
