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::>()); @@ -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::>()); + let nulls = + with_nulls.then(|| NullBuffer::from((0..rows).map(|i| i % 10 != 0).collect::>())); 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::>()); + + 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::>()); + 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 { - let mut builder = Int32Array::builder(array.len()); - match array.data_type() { - DataType::List(_) => { - let list_array = array - .as_any() - .downcast_ref::() - .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::() - .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::() - .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::() .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( - builder: &mut Int32Builder, - list_array: &GenericListArray, -) { - 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 { + 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::() + .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 { 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 { 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::().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::().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::::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)) + ); + } }