perf: build spark_size LargeList lengths from i64 offsets

docs/source/contributor-guide/expression-audits/collection_funcs.md

@@ -48,6 +48,6 @@
 - Spark 4.0.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8.
 - Spark 4.1.1 (audited 2026-05-27): byte-for-byte identical to 3.5.8.
 - Both `ArrayType` and `MapType` inputs are `Compatible` and run natively; every other child type is `Unsupported`.
-- Performance (tuned 2026-07-10, PR [#4877](https://github.com/apache/datafusion-comet/pull/4877)): compute list row sizes from the offset buffer instead of allocating a sliced `ArrayRef` per row via `list_array.value(i)`, removing one heap allocation per row. ~94% faster. Benchmark: `benches/array_size.rs`.
+- Performance (tuned 2026-08-22, PRs [#4877](https://github.com/apache/datafusion-comet/pull/4877), [#5233](https://github.com/apache/datafusion-comet/pull/5233), [#5395](https://github.com/apache/datafusion-comet/pull/5395), [#5300](https://github.com/apache/datafusion-comet/pull/5300)): List and FixedSizeList compute row lengths via Arrow's `length` kernel plus an `into_parts` null → `-1` rewrite. LargeList and Map subtract adjacent offsets directly into Int32 (LargeList falls back to per-row `try_from` on span overflow). Up to ~12x faster on List / FixedSizeList, ~6.6x on LargeList, and ~6-13x on Map versus the per-row builder loop. Benchmark: `benches/array_size.rs`.
 
 [Spark Expression Support]: ../../user-guide/latest/expressions.md

docs/source/contributor-guide/optimizing_expressions.md

@@ -149,8 +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` 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`                                         |
+| **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; LargeList array path (~6.6x): `offsets().windows(2)` diff into Int32, skips `length`'s Int64→Int32 cast |
+| **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/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                                                                                            |

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::<Vec<i32>>());
 
@@ -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::<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(LargeListArray::new(
         field,
         OffsetBuffer::new(offsets.into()),
         Arc::new(values),
-        Some(nulls),
+        nulls,
     ))
 }
 
@@ -145,11 +145,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),
     );
 
     // Map shapes: no length-kernel support for MapArray, so this path differences

native/spark-expr/src/array_funcs/size.rs

@@ -15,16 +15,19 @@
 // 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};
 use datafusion::common::{exec_err, DataFusionError, Result as DataFusionResult, ScalarValue};
 use datafusion::logical_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
 };
 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";
+
 /// 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<ColumnarValue, DataFusionError> {
@@ -93,10 +96,11 @@ impl ScalarUDFImpl for SparkSizeFunc {
 
 fn spark_size_array(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
     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; compute row lengths from
         // the entry-offset buffer directly.
         DataType::Map(_, _) => spark_size_map(array),
@@ -109,26 +113,21 @@ 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.
+/// 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).
 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:?}");
-        }
-    };
+    // Arrow's `length` returns Int32 for List and builds an Int32Array for
+    // FixedSizeList. LargeList (the only Int64 producer) is routed to the
+    // offset path above. Fail loudly rather than silently paying for a
+    // cast we do not expect.
+    if !matches!(lengths.data_type(), DataType::Int32) {
+        return exec_err!(
+            "unexpected type from length kernel: {:?}",
+            lengths.data_type()
+        );
+    }
 
     // Fast path for the production shape: `CometSize.convert` wraps size() in a
     // `CASE WHEN isnotnull(child)` that filters null rows out before the THEN
@@ -147,6 +146,58 @@ fn spark_size_list_like(array: &ArrayRef) -> Result<ArrayRef, DataFusionError> {
     Ok(rewrite_nulls_to_minus_one(int_lengths))
 }
 
+/// 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<ArrayRef, DataFusionError> {
+    let list = array
+        .as_any()
+        .downcast_ref::<LargeListArray>()
+        .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: Vec<i32> = 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 lengths = Int32Array::new(values.into(), list.nulls().cloned());
+    Ok(rewrite_nulls_to_minus_one(lengths))
+}
+
+/// 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<i64>,
+    nulls: Option<&NullBuffer>,
+) -> Result<Vec<i32>, 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)
+}
+
 /// 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
@@ -204,9 +255,8 @@ fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue, DataFusionErro
             } else {
                 // 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())
-                })?;
+                let len = i32::try_from(array.value_length(0))
+                    .map_err(|_| DataFusionError::Execution(SIZE_OVERFLOW_MSG.to_string()))?;
                 Ok(ScalarValue::Int32(Some(len)))
             }
         }
@@ -240,6 +290,7 @@ fn spark_size_scalar(scalar: &ScalarValue) -> Result<ScalarValue, DataFusionErro
 mod tests {
     use super::*;
     use arrow::array::{Int32Array, ListArray, NullBufferBuilder};
+    use arrow::buffer::NullBuffer;
     use arrow::datatypes::{DataType, Field};
     use std::sync::Arc;
 
@@ -682,6 +733,112 @@ mod tests {
         assert_eq!(result.value(2), 1);
     }
 
+    #[test]
+    fn test_spark_size_large_list_array_no_nulls() {
+        use arrow::array::LargeListArray;
+
+        // Production shape for LargeList: no null buffer, so the offset-diff
+        // fast path in `spark_size_large_list_from_offsets` runs and
+        // `rewrite_nulls_to_minus_one` returns the values buffer untouched.
+        // Mirrors test_spark_size_array_no_nulls for the List path.
+        let value_data = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+        let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, 3, 5, 5, 6].into());
+        let field = Arc::new(Field::new("item", DataType::Int32, true));
+        let large_list_array =
+            LargeListArray::try_new(field, value_offsets, Arc::new(value_data), None).unwrap();
+
+        let array_ref: ArrayRef = Arc::new(large_list_array);
+        let result = spark_size_array(&array_ref).unwrap();
+        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+
+        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_sliced_large_list_array() {
+        // Same sliced-offsets invariant as the List test, for the LargeList helper
+        // that reads `offsets().windows(2)` directly.
+        //
+        // Full array: [[1, 2], [3], [4, 5, 6], null, [7]]
+        use arrow::array::LargeListArray;
+
+        let value_data = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7]);
+        let value_offsets = arrow::buffer::OffsetBuffer::new(vec![0i64, 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 large_list_array = LargeListArray::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(large_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_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_span_overflow_rows_fit() {
+        // Overall span > i32::MAX triggers the checked path, but each per-row
+        // length fits: [i32::MAX, 10]. Pins that the checked helper differences
+        // per row rather than rejecting the whole batch on span overflow. This is
+        // why we loop with try_from instead of erroring outright on the span.
+        let offsets = arrow::buffer::OffsetBuffer::new(
+            vec![0i64, i32::MAX as i64, i32::MAX as i64 + 10].into(),
+        );
+        let values = spark_size_large_list_lengths_checked(&offsets, None).unwrap();
+        assert_eq!(values, vec![i32::MAX, 10]);
+    }
+
+    #[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 lengths = Int32Array::new(values.into(), Some(nulls));
+        let result = rewrite_nulls_to_minus_one(lengths);
+        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(result.value(0), -1);
+        assert_eq!(result.value(1), 0);
+    }
+
     #[test]
     fn test_spark_size_scalar_large_list() {
         use arrow::array::LargeListArray;