This is an automated email from the ASF dual-hosted git repository.
mbutrovich 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 edb7b65c98 perf: optimize spark_size in spark-expr (#4877)
edb7b65c98 is described below
commit edb7b65c983d73063bb5060e097df3711ccaf122
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 10 10:07:57 2026 -0600
perf: optimize spark_size in spark-expr (#4877)
---
native/spark-expr/Cargo.toml | 4 ++
native/spark-expr/benches/array_size.rs | 66 +++++++++++++++++++++++++++++++
native/spark-expr/src/array_funcs/size.rs | 40 ++++++++++---------
3 files changed, 91 insertions(+), 19 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 800fe3ecb1..42a7bcb97b 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -123,3 +123,7 @@ harness = false
[[bench]]
name = "to_time"
harness = false
+
+[[bench]]
+name = "array_size"
+harness = false
diff --git a/native/spark-expr/benches/array_size.rs
b/native/spark-expr/benches/array_size.rs
new file mode 100644
index 0000000000..f6d3559da7
--- /dev/null
+++ b/native/spark-expr/benches/array_size.rs
@@ -0,0 +1,66 @@
+// 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, Int32Array, ListArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::spark_size;
+use std::hint::black_box;
+use std::sync::Arc;
+
+/// Build a `ListArray` of `rows` lists, each with `elems_per_row` Int32
elements,
+/// with every 10th row null.
+fn create_list_array(rows: usize, elems_per_row: usize) -> ArrayRef {
+ let total = rows * elems_per_row;
+ 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 * elems_per_row) as i32);
+ }
+
+ let nulls = NullBuffer::from((0..rows).map(|i| i % 10 !=
0).collect::<Vec<bool>>());
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ Arc::new(ListArray::new(
+ field,
+ OffsetBuffer::new(offsets.into()),
+ Arc::new(values),
+ Some(nulls),
+ ))
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let rows = 8192;
+
+ let short_lists = create_list_array(rows, 5);
+ c.bench_function("spark_size: list of short arrays", |b| {
+ let args = vec![ColumnarValue::Array(Arc::clone(&short_lists))];
+ b.iter(|| black_box(spark_size(black_box(&args)).unwrap()))
+ });
+
+ let long_lists = create_list_array(rows, 64);
+ c.bench_function("spark_size: list of long arrays", |b| {
+ let args = vec![ColumnarValue::Array(Arc::clone(&long_lists))];
+ b.iter(|| black_box(spark_size(black_box(&args)).unwrap()))
+ });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/src/array_funcs/size.rs
b/native/spark-expr/src/array_funcs/size.rs
index f206b299d6..9bed27b33b 100644
--- a/native/spark-expr/src/array_funcs/size.rs
+++ b/native/spark-expr/src/array_funcs/size.rs
@@ -15,7 +15,8 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{Array, ArrayRef, Int32Array};
+use arrow::array::builder::Int32Builder;
+use arrow::array::{Array, ArrayRef, GenericListArray, Int32Array,
OffsetSizeTrait};
use arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_err, DataFusionError, Result as
DataFusionResult, ScalarValue};
use datafusion::logical_expr::{
@@ -98,30 +99,14 @@ fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef,
DataFusionError> {
.as_any()
.downcast_ref::<arrow::array::ListArray>()
.ok_or_else(|| DataFusionError::Internal("Expected
ListArray".to_string()))?;
-
- for i in 0..list_array.len() {
- if list_array.is_null(i) {
- builder.append_value(-1); // Spark behavior: return -1 for
null
- } else {
- let list_len = list_array.value(i).len() as i32;
- builder.append_value(list_len);
- }
- }
+ append_list_sizes(&mut builder, list_array);
}
DataType::LargeList(_) => {
let list_array = array
.as_any()
.downcast_ref::<arrow::array::LargeListArray>()
.ok_or_else(|| DataFusionError::Internal("Expected
LargeListArray".to_string()))?;
-
- for i in 0..list_array.len() {
- if list_array.is_null(i) {
- builder.append_value(-1); // Spark behavior: return -1 for
null
- } else {
- let list_len = list_array.value(i).len() as i32;
- builder.append_value(list_len);
- }
- }
+ append_list_sizes(&mut builder, list_array);
}
DataType::FixedSizeList(_, size) => {
let fixed_list_array = array
@@ -165,6 +150,23 @@ fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef,
DataFusionError> {
Ok(Arc::new(builder.finish()))
}
+/// Append the element count of each list row to `builder`, using `-1` for null
+/// rows (Spark's behavior). `value_length` reads the row's element count from
the
+/// offset buffer, avoiding the per-row allocation that `value(i).len()` would
incur
+/// from materializing a sliced array.
+fn append_list_sizes<O: OffsetSizeTrait>(
+ builder: &mut Int32Builder,
+ list_array: &GenericListArray<O>,
+) {
+ for i in 0..list_array.len() {
+ if list_array.is_null(i) {
+ builder.append_value(-1); // Spark behavior: return -1 for null
+ } else {
+ builder.append_value(list_array.value_length(i).as_usize() as i32);
+ }
+ }
+}
+
fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue,
DataFusionError> {
match scalar {
ScalarValue::List(array) => {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
