From 6968c91a823de76d171869f01d6dfb0cf3fb4088 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Fri, 7 Aug 2026 23:27:58 +0800 Subject: [PATCH] perf: build spark_size LargeList lengths from i64 offsets (6.6x faster) --- native/spark-expr/benches/array_size.rs | 25 ++- native/spark-expr/src/array_funcs/size.rs | 181 +++++++++++++++++++--- 2 files changed, 177 insertions(+), 29 deletions(-) diff --git a/native/spark-expr/benches/array_size.rs b/native/spark-expr/benches/array_size.rs index afb340c06e..217d9cde37 100644 --- a/native/spark-expr/benches/array_size.rs +++ b/native/spark-expr/benches/array_size.rs @@ -48,9 +48,8 @@ fn create_list_array(rows: usize, elems_per_row: usize, with_nulls: bool) -> Arr } /// 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 { +/// Int32 elements. When `with_nulls` is true every 10th row is null. +fn create_large_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::>()); @@ -60,13 +59,14 @@ fn create_large_list_array(rows: usize, elems_per_row: usize) -> ArrayRef { offsets.push((i * elems_per_row) as i64); } - 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(LargeListArray::new( field, OffsetBuffer::new(offsets.into()), Arc::new(values), - Some(nulls), + nulls, )) } @@ -99,11 +99,20 @@ fn criterion_benchmark(c: &mut Criterion) { &create_list_array(rows, 5, false), ); - // LargeList: exposes the Int64 length -> Int32 cast (extra allocation) on top of - // the length kernel. + // LargeList: builds Int32 lengths from i64 offsets (see spark_size_large_list_from_offsets). + // Shapes mirror List coverage (short/long + 10% null, short + no nulls). bench( "spark_size: LargeList (10% null)", - &create_large_list_array(rows, 5), + &create_large_list_array(rows, 5, true), + ); + bench( + "spark_size: LargeList of long arrays", + &create_large_list_array(rows, 64, true), + ); + // No-null LargeList: production path after CometSize's isnotnull filter. + bench( + "spark_size: LargeList, no nulls", + &create_large_list_array(rows, 5, false), ); } diff --git a/native/spark-expr/src/array_funcs/size.rs b/native/spark-expr/src/array_funcs/size.rs index 0d555806b7..b8a7cbe9ac 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::{Array, ArrayRef, Int32Array, LargeListArray}; +use arrow::buffer::NullBuffer; use arrow::compute::kernels::length::length; use arrow::compute::{cast_with_options, CastOptions}; use arrow::datatypes::{DataType, Field}; @@ -25,6 +26,21 @@ use datafusion::logical_expr::{ }; use std::sync::Arc; +/// Shared by the LargeList array path and `ScalarValue::LargeList` overflow guard. +const SIZE_OVERFLOW_MSG: &str = "size(): list length exceeds i32::MAX"; + +/// Rewrite null slots to `-1` (Spark `size` of null). Shared by List and LargeList paths. +fn ints_with_nulls_as_neg_one(mut values: Vec, nulls: Option<&NullBuffer>) -> Int32Array { + if let Some(nulls) = nulls { + if nulls.null_count() > 0 { + for i in (!nulls.inner()).set_indices() { + values[i] = -1; + } + } + } + Int32Array::from(values) +} + /// Spark size() function that returns the size of arrays or maps. /// Returns -1 for null inputs (Spark behavior differs from standard SQL). pub fn spark_size(args: &[ColumnarValue]) -> Result { @@ -93,10 +109,11 @@ impl ScalarUDFImpl for SparkSizeFunc { fn spark_size_array(array: &ArrayRef) -> Result { match array.data_type() { - // List / LargeList / FixedSizeList: reuse Arrow's vectorized length kernel. - DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(..) => { - spark_size_list_like(array) - } + // LargeList: build Int32 lengths directly from i64 offsets (avoids + // length→Int64→cast→Int32). + DataType::LargeList(_) => spark_size_large_list_from_offsets(array), + // List / FixedSizeList: reuse Arrow's vectorized length kernel. + DataType::List(_) | 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 @@ -123,9 +140,9 @@ fn spark_size_array(array: &ArrayRef) -> Result { } } -/// 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. +/// Compute Spark `size()` for List / FixedSizeList via Arrow's `length` kernel, +/// then rewrite null inputs to `-1` (Spark's legacy/compatible size-of-null +/// behavior for this UDF). /// /// Patches the values buffer in place rather than using `zip`; `zip` goes /// through `MutableArrayData` and roughly doubled runtime on the `array_size` @@ -161,19 +178,75 @@ fn spark_size_list_like(array: &ArrayRef) -> Result { .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. + // `into_parts` + shared null→-1 rewrite: O(n) values copy, then O(null_count) + // writes via inverted validity. 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))); + Ok(Arc::new(ints_with_nulls_as_neg_one( + values.to_vec(), + nulls.as_ref(), + ))) +} + +/// Compute Spark `size()` for LargeList by subtracting adjacent i64 offsets into +/// Int32 lengths. Avoids Arrow's `length` kernel (which returns Int64 for +/// LargeList) and the subsequent Int32 cast allocation. +/// +/// When the full offset span fits in `i32`, every per-row length does too, so the +/// hot path uses an unchecked `as i32` + `collect` (vectorization-friendly). +/// Otherwise fall back to per-row `i32::try_from` with [`SIZE_OVERFLOW_MSG`]. +/// +/// Null inputs become `-1`, matching `spark_size_list_like`. +fn spark_size_large_list_from_offsets(array: &ArrayRef) -> Result { + let list = array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Expected LargeListArray".to_string()))?; + + let offsets = list.offsets(); + // Offsets are monotonically non-decreasing, so no per-row length can exceed + // the full span. If the span fits in i32, every row length does too. + // `OffsetBuffer` always has at least one element. + let range = *offsets.last().unwrap() - *offsets.first().unwrap(); + let values = if range > i32::MAX as i64 { + spark_size_large_list_lengths_checked(offsets, list.nulls())? + } else { + offsets.windows(2).map(|w| (w[1] - w[0]) as i32).collect() }; - let mut values = values.to_vec(); - for i in (!nulls.inner()).set_indices() { - values[i] = -1; + + // 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. + if list.null_count() == 0 { + return Ok(Arc::new(Int32Array::from(values))); } - Ok(Arc::new(Int32Array::from(values))) + + // `null_count() > 0` implies a null buffer is present. + let nulls = list.nulls().unwrap(); + Ok(Arc::new(ints_with_nulls_as_neg_one(values, Some(nulls)))) +} + +/// Per-row checked Int32 conversion used when the overall offset span exceeds +/// `i32::MAX` (so an individual row might overflow even though most do not). +/// +/// Null rows skip `try_from`: Arrow does not require a null row's offset delta to +/// be zero, and the caller rewrites those slots to `-1` anyway (same as the +/// unchecked fast path). +fn spark_size_large_list_lengths_checked( + offsets: &arrow::buffer::OffsetBuffer, + nulls: Option<&NullBuffer>, +) -> Result, DataFusionError> { + let mut values = Vec::with_capacity(offsets.len() - 1); + for (i, w) in offsets.windows(2).enumerate() { + if nulls.is_some_and(|n| n.is_null(i)) { + values.push(0); // overwritten to -1 by the caller + continue; + } + let len = i32::try_from(w[1] - w[0]) + .map_err(|_| DataFusionError::Execution(SIZE_OVERFLOW_MSG.to_string()))?; + values.push(len); + } + Ok(values) } fn spark_size_scalar(scalar: &ScalarValue) -> Result { @@ -194,9 +267,8 @@ fn spark_size_scalar(scalar: &ScalarValue) -> Result Result().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_large_list_length_overflow() { + // A non-null row whose length is i32::MAX + 1 must error (same contract + // as the old safe: false Int64→Int32 cast / scalar try_from). Call the + // checked helper directly with a valid OffsetBuffer — no invalid ArrayData. + let offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, (i32::MAX as i64) + 1].into()); + let err = spark_size_large_list_lengths_checked(&offsets, None).unwrap_err(); + assert!( + err.to_string().contains(SIZE_OVERFLOW_MSG), + "unexpected error: {err}" + ); + } + + #[test] + fn test_spark_size_large_list_checked_null_row_skips_overflow() { + // Null row offset delta may exceed i32::MAX; checked path must not error + // (matches the fast path, which rewrites null slots to -1 afterward). + let offsets = arrow::buffer::OffsetBuffer::new( + vec![0i64, (i32::MAX as i64) + 1, (i32::MAX as i64) + 1].into(), + ); + let nulls = NullBuffer::from(vec![false, true]); // row 0 null, row 1 empty + let values = spark_size_large_list_lengths_checked(&offsets, Some(&nulls)).unwrap(); + assert_eq!(values, vec![0, 0]); + let result = ints_with_nulls_as_neg_one(values, Some(&nulls)); + assert_eq!(result.value(0), -1); + assert_eq!(result.value(1), 0); + } + #[test] fn test_spark_size_scalar_large_list() { use arrow::array::LargeListArray;