perf: compute spark_size list lengths with Arrow length kernel

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                                                                                            |

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);