This is an automated email from the ASF dual-hosted git repository.
andygrove 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 1feef7aba8 perf: compute spark_size list lengths with Arrow length
kernel (#5233)
1feef7aba8 is described below
commit 1feef7aba891385349234b4fb8d204c3f3f371ac
Author: ChenChen Lai <[email protected]>
AuthorDate: Thu Aug 6 02:01:50 2026 +0800
perf: compute spark_size list lengths with Arrow length kernel (#5233)
---
.../contributor-guide/optimizing_expressions.md | 3 +-
native/spark-expr/benches/array_size.rs | 75 ++++--
native/spark-expr/src/array_funcs/size.rs | 256 ++++++++++++++++-----
3 files changed, 257 insertions(+), 77 deletions(-)
diff --git a/docs/source/contributor-guide/optimizing_expressions.md
b/docs/source/contributor-guide/optimizing_expressions.md
index a000f91398..ecdb8a7263 100644
--- a/docs/source/contributor-guide/optimizing_expressions.md
+++ b/docs/source/contributor-guide/optimizing_expressions.md
@@ -149,7 +149,8 @@ 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`: compute list lengths from offsets, zero
allocation
|
+| **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
|
+| **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
|
| **`memcpy` from a precomputed buffer** |
Char-by-char `push` into a scratch `String`
| `spark_lpad`: pad from a precomputed repeating pad
buffer, write directly into the builder
|
diff --git a/native/spark-expr/benches/array_size.rs
b/native/spark-expr/benches/array_size.rs
index f6d3559da7..afb340c06e 100644
--- a/native/spark-expr/benches/array_size.rs
+++ b/native/spark-expr/benches/array_size.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{ArrayRef, Int32Array, ListArray};
+use arrow::array::{ArrayRef, Int32Array, LargeListArray, ListArray};
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{DataType, Field};
use criterion::{criterion_group, criterion_main, Criterion};
@@ -24,9 +24,9 @@ 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 {
+/// Build a `ListArray` of `rows` lists, each with `elems_per_row` Int32
elements.
+/// When `with_nulls` is true every 10th row is null.
+fn create_list_array(rows: usize, elems_per_row: usize, with_nulls: bool) ->
ArrayRef {
let total = rows * elems_per_row;
let values = Int32Array::from((0..total as i32).collect::<Vec<i32>>());
@@ -36,9 +36,33 @@ fn create_list_array(rows: usize, elems_per_row: usize) ->
ArrayRef {
offsets.push((i * elems_per_row) as i32);
}
- let nulls = NullBuffer::from((0..rows).map(|i| i % 10 !=
0).collect::<Vec<bool>>());
+ let nulls =
+ with_nulls.then(|| 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),
+ nulls,
+ ))
+}
+
+/// Build a `LargeListArray` (i64 offsets) of `rows` lists with `elems_per_row`
+/// Int32 elements. Every 10th row is null. LargeList exercises the extra
+/// Int64->Int32 cast on top of the length kernel.
+fn create_large_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(0i64);
+ for i in 1..=rows {
+ offsets.push((i * elems_per_row) as i64);
+ }
+
+ 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(LargeListArray::new(
field,
OffsetBuffer::new(offsets.into()),
Arc::new(values),
@@ -49,17 +73,38 @@ fn create_list_array(rows: usize, elems_per_row: usize) ->
ArrayRef {
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 mut bench = |name: &str, arr: &ArrayRef| {
+ let args = vec![ColumnarValue::Array(Arc::clone(arr))];
+ c.bench_function(name, |b| {
+ b.iter(|| black_box(spark_size(black_box(&args)).unwrap()))
+ });
+ };
+
+ // 10%-null shapes: match the pre-existing coverage.
+ bench(
+ "spark_size: list of short arrays",
+ &create_list_array(rows, 5, true),
+ );
+ bench(
+ "spark_size: list of long arrays",
+ &create_list_array(rows, 64, true),
+ );
+
+ // No-null shape: `CometSize.convert` wraps size() in a `CASE WHEN
isnotnull(child)`
+ // that filters null rows out before the THEN branch runs, so in a real
Comet plan
+ // spark_size_list_like only ever sees a null-free array. This shape
measures that
+ // path.
+ bench(
+ "spark_size: list, no nulls",
+ &create_list_array(rows, 5, false),
+ );
- 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()))
- });
+ // LargeList: exposes the Int64 length -> Int32 cast (extra allocation) on
top of
+ // the length kernel.
+ bench(
+ "spark_size: LargeList (10% null)",
+ &create_large_list_array(rows, 5),
+ );
}
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 5ad4f1670a..0d555806b7 100644
--- a/native/spark-expr/src/array_funcs/size.rs
+++ b/native/spark-expr/src/array_funcs/size.rs
@@ -15,8 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::builder::Int32Builder;
-use arrow::array::{Array, ArrayRef, GenericListArray, Int32Array,
OffsetSizeTrait};
+use arrow::array::{Array, ArrayRef, Int32Array};
+use arrow::compute::kernels::length::length;
+use arrow::compute::{cast_with_options, CastOptions};
use arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_err, DataFusionError, Result as
DataFusionResult, ScalarValue};
use datafusion::logical_expr::{
@@ -91,99 +92,111 @@ impl ScalarUDFImpl for SparkSizeFunc {
}
fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
- let mut builder = Int32Array::builder(array.len());
-
match array.data_type() {
- DataType::List(_) => {
- let list_array = array
- .as_any()
- .downcast_ref::<arrow::array::ListArray>()
- .ok_or_else(|| DataFusionError::Internal("Expected
ListArray".to_string()))?;
- 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()))?;
- append_list_sizes(&mut builder, list_array);
- }
- DataType::FixedSizeList(_, size) => {
- let fixed_list_array = array
- .as_any()
- .downcast_ref::<arrow::array::FixedSizeListArray>()
- .ok_or_else(|| {
- DataFusionError::Internal("Expected
FixedSizeListArray".to_string())
- })?;
-
- for i in 0..fixed_list_array.len() {
- if fixed_list_array.is_null(i) {
- builder.append_value(-1); // Spark behavior: return -1 for
null
- } else {
- builder.append_value(*size);
- }
- }
+ // List / LargeList / FixedSizeList: reuse Arrow's vectorized length
kernel.
+ 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 {
- let map_len = map_array.value_length(i);
- builder.append_value(map_len);
+ builder.append_value(map_array.value_length(i));
}
}
+ Ok(Arc::new(builder.finish()))
}
_ => {
- return exec_err!(
+ exec_err!(
"size function only supports arrays and maps, got: {:?}",
array.data_type()
- );
+ )
}
}
-
- 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);
+/// 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() {
+ DataType::Int32 => lengths,
+ // Unsafe cast: overflow must error, not become null then get
rewritten to -1.
+ DataType::Int64 => cast_with_options(
+ lengths.as_ref(),
+ &DataType::Int32,
+ &CastOptions {
+ safe: false,
+ ..Default::default()
+ },
+ )?,
+ other => {
+ return exec_err!("unexpected type from length kernel: {other:?}");
}
+ };
+
+ // 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.
+ if array.null_count() == 0 {
+ return Ok(lengths);
+ }
+
+ let int_lengths = lengths
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| DataFusionError::Internal("Expected Int32Array from
length".to_string()))?;
+
+ // `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();
+ let Some(nulls) = nulls else {
+ return Ok(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)))
}
fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue,
DataFusionError> {
match scalar {
+ // ScalarValue::{List,LargeList,FixedSizeList,Map} each wrap an array
with
+ // exactly one row; read the row's element count from the offset buffer
+ // (matches the array path, avoids `value(0)` slicing).
ScalarValue::List(array) => {
- // ScalarValue::List contains a ListArray with exactly one row.
- // We need the length of that row's contents, not the row count.
if array.is_null(0) {
Ok(ScalarValue::Int32(Some(-1))) // Spark behavior: return -1
for null
} else {
- let len = array.value(0).len() as i32;
- Ok(ScalarValue::Int32(Some(len)))
+ Ok(ScalarValue::Int32(Some(array.value_length(0))))
}
}
ScalarValue::LargeList(array) => {
if array.is_null(0) {
Ok(ScalarValue::Int32(Some(-1)))
} else {
- let len = array.value(0).len() as i32;
+ // Spark arrays are capped near Integer.MAX_VALUE; overflow
shouldn't
+ // happen in practice but must error rather than silently wrap.
+ let len = i32::try_from(array.value_length(0)).map_err(|_| {
+ DataFusionError::Execution("size(): list length exceeds
i32::MAX".to_string())
+ })?;
Ok(ScalarValue::Int32(Some(len)))
}
}
@@ -191,16 +204,14 @@ fn spark_size_scalar(scalar: &ScalarValue) ->
Result<ScalarValue, DataFusionErro
if array.is_null(0) {
Ok(ScalarValue::Int32(Some(-1)))
} else {
- let len = array.value(0).len() as i32;
- Ok(ScalarValue::Int32(Some(len)))
+ Ok(ScalarValue::Int32(Some(array.value_length())))
}
}
ScalarValue::Map(array) => {
if array.is_null(0) {
Ok(ScalarValue::Int32(Some(-1)))
} else {
- let len = array.value_length(0);
- Ok(ScalarValue::Int32(Some(len)))
+ Ok(ScalarValue::Int32(Some(array.value_length(0))))
}
}
ScalarValue::Null => {
@@ -254,6 +265,29 @@ mod tests {
assert_eq!(result.value(3), 0); // [] has 0 elements
}
+ #[test]
+ fn test_spark_size_array_no_nulls() {
+ // Fast path (the shape Comet actually runs after the CASE-WHEN
filter):
+ // input has no nulls, so spark_size_list_like returns the length
kernel
+ // output directly without touching the values buffer.
+ let value_data = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0, 3, 5, 5,
6].into());
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let list_array =
+ ListArray::try_new(field, value_offsets, Arc::new(value_data),
None).unwrap();
+
+ let array_ref: ArrayRef = Arc::new(list_array);
+ let result = spark_size_array(&array_ref).unwrap();
+ let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+
+ // Expected: [3, 2, 0, 1]; no null buffer on the output.
+ assert_eq!(result.null_count(), 0);
+ assert_eq!(result.value(0), 3);
+ assert_eq!(result.value(1), 2);
+ assert_eq!(result.value(2), 0);
+ assert_eq!(result.value(3), 1);
+ }
+
#[test]
fn test_spark_size_scalar() {
// Test non-null list with 3 elements
@@ -484,4 +518,104 @@ mod tests {
assert_eq!(result.value(2), -1); // null returns -1
assert_eq!(result.value(3), 0); // [] has 0 elements
}
+
+ #[test]
+ fn test_spark_size_sliced_list_array() {
+ // Slicing is the classic trap for buffer-level ops: the null buffer,
values
+ // buffer, and offsets all logically start at `array.offset()`, not 0.
Pin the
+ // behavior so a future edit that touches buffers directly cannot
regress it.
+ //
+ // Full array: [[1, 2], [3], [4, 5, 6], null, [7]]
+ let value_data = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7]);
+ let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0, 2, 3, 6,
6, 7].into());
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+
+ 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 list_array = ListArray::try_new(
+ field,
+ value_offsets,
+ Arc::new(value_data),
+ null_buffer.finish(),
+ )
+ .unwrap();
+
+ // Skip the first two rows: sliced view is [[4, 5, 6], null, [7]].
+ let sliced: ArrayRef = Arc::new(list_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_large_list() {
+ use arrow::array::LargeListArray;
+
+ // Non-null LargeList: [10, 20, 30, 40] → 4
+ let values = Int32Array::from(vec![10, 20, 30, 40]);
+ let offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, 4].into());
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let large_list_array =
+ LargeListArray::try_new(Arc::clone(&field), offsets,
Arc::new(values), None).unwrap();
+ let scalar = ScalarValue::LargeList(Arc::new(large_list_array));
+ assert_eq!(
+ spark_size_scalar(&scalar).unwrap(),
+ ScalarValue::Int32(Some(4))
+ );
+
+ // Null LargeList row → -1
+ let empty_values = Int32Array::from(Vec::<i32>::new());
+ let null_offsets = arrow::buffer::OffsetBuffer::new(vec![0i64,
0].into());
+ let mut null_buffer = NullBufferBuilder::new(1);
+ null_buffer.append(false);
+ let null_large_list = LargeListArray::try_new(
+ field,
+ null_offsets,
+ Arc::new(empty_values),
+ null_buffer.finish(),
+ )
+ .unwrap();
+ let scalar = ScalarValue::LargeList(Arc::new(null_large_list));
+ assert_eq!(
+ spark_size_scalar(&scalar).unwrap(),
+ ScalarValue::Int32(Some(-1))
+ );
+ }
+
+ #[test]
+ fn test_spark_size_scalar_fixed_size_list() {
+ use arrow::array::FixedSizeListArray;
+
+ // Non-null FixedSizeList of size 4 → 4. `value_length()` on FSL takes
no index,
+ // unlike `value_length(i)` on List/LargeList — pin the correct API is
called.
+ let values = Int32Array::from(vec![10, 20, 30, 40]);
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let fsl_array = FixedSizeListArray::new(Arc::clone(&field), 4,
Arc::new(values), None);
+ let scalar = ScalarValue::FixedSizeList(Arc::new(fsl_array));
+ assert_eq!(
+ spark_size_scalar(&scalar).unwrap(),
+ ScalarValue::Int32(Some(4))
+ );
+
+ // Null FSL row → -1
+ let null_values = Int32Array::from(vec![0, 0, 0, 0]);
+ let mut null_buffer = NullBufferBuilder::new(1);
+ null_buffer.append(false);
+ let null_fsl_array =
+ FixedSizeListArray::new(field, 4, Arc::new(null_values),
null_buffer.finish());
+ let scalar = ScalarValue::FixedSizeList(Arc::new(null_fsl_array));
+ assert_eq!(
+ spark_size_scalar(&scalar).unwrap(),
+ ScalarValue::Int32(Some(-1))
+ );
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
