perf: optimize list_extract without defaults using Arrow take

native/spark-expr/Cargo.toml

@@ -143,6 +143,10 @@ harness = false
 name = "array_size"
 harness = false
 
+[[bench]]
+name = "list_extract"
+harness = false
+
 [[bench]]
 name = "regexp_extract"
 harness = false

native/spark-expr/benches/list_extract.rs

@@ -0,0 +1,122 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
+use datafusion::physical_expr::expressions::{Column, Literal};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::{create_query_context_map, ListExtract};
+use std::hint::black_box;
+use std::sync::Arc;
+
+const ROWS: usize = 8192;
+const ELEMENTS_PER_ROW: usize = 5;
+
+fn list_of(values: ArrayRef) -> ArrayRef {
+    let offsets = (0..=ROWS)
+        .map(|i| (i * ELEMENTS_PER_ROW) as i32)
+        .collect::<Vec<_>>();
+    let field = Arc::new(Field::new("item", values.data_type().clone(), true));
+    Arc::new(ListArray::new(
+        field,
+        OffsetBuffer::new(offsets.into()),
+        values,
+        None,
+    ))
+}
+
+fn bench_case(
+    c: &mut Criterion,
+    name: &str,
+    list: ArrayRef,
+    oob: bool,
+    default: Option<ScalarValue>,
+) {
+    let indices = Arc::new(Int32Array::from_iter_values((0..ROWS).map(|row| {
+        if oob && row % 2 == 0 {
+            ELEMENTS_PER_ROW as i32 + 1
+        } else {
+            3
+        }
+    })));
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("list", list.data_type().clone(), false),
+        Field::new("index", DataType::Int32, false),
+    ]));
+    let batch = RecordBatch::try_new(schema, vec![list, indices]).unwrap();
+    let default = default.map(|value| Arc::new(Literal::new(value)) as Arc<dyn PhysicalExpr>);
+    let expr = ListExtract::new(
+        Arc::new(Column::new("list", 0)),
+        Arc::new(Column::new("index", 1)),
+        default,
+        true,
+        false,
+        None,
+        create_query_context_map(),
+    );
+
+    c.bench_function(name, |b| {
+        b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap()))
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let total = ROWS * ELEMENTS_PER_ROW;
+    let ints = list_of(Arc::new(Int32Array::from_iter_values(0..total as i32)));
+    let strings = list_of(Arc::new(StringArray::from_iter_values(
+        (0..total).map(|i| format!("value-{i}")),
+    )));
+
+    for oob in [false, true] {
+        let suffix = if oob { "50%-oob" } else { "0%-oob" };
+        bench_case(
+            c,
+            &format!("list_extract/int32/null-default/{suffix}"),
+            Arc::clone(&ints),
+            oob,
+            None,
+        );
+        bench_case(
+            c,
+            &format!("list_extract/int32/non-null-default/{suffix}"),
+            Arc::clone(&ints),
+            oob,
+            Some(ScalarValue::Int32(Some(0))),
+        );
+        bench_case(
+            c,
+            &format!("list_extract/utf8/null-default/{suffix}"),
+            Arc::clone(&strings),
+            oob,
+            None,
+        );
+        bench_case(
+            c,
+            &format!("list_extract/utf8/non-null-default/{suffix}"),
+            Arc::clone(&strings),
+            oob,
+            Some(ScalarValue::Utf8(Some(String::new()))),
+        );
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

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

@@ -15,9 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, GenericListArray, Int32Array, OffsetSizeTrait};
+use arrow::array::{
+    Array, GenericListArray, Int32Array, MutableArrayData, OffsetSizeTrait, UInt64Array,
+};
+use arrow::compute::take;
 use arrow::datatypes::{DataType, FieldRef, Schema};
-use arrow::{array::MutableArrayData, datatypes::ArrowNativeType, record_batch::RecordBatch};
+use arrow::{datatypes::ArrowNativeType, record_batch::RecordBatch};
 use datafusion::common::{
     cast::{as_int32_array, as_large_list_array, as_list_array},
     internal_err, DataFusionError, Result as DataFusionResult, ScalarValue,
@@ -143,7 +146,7 @@ impl PhysicalExpr for ListExtract {
             .default_value
             .as_ref()
             .map(|d| {
-                d.evaluate(batch).map(|value| match value {
+                d.evaluate(batch).and_then(|value| match value {
                     ColumnarValue::Scalar(scalar)
                         if !scalar.data_type().equals_datatype(&element_type) =>
                     {
@@ -155,8 +158,7 @@ impl PhysicalExpr for ListExtract {
                     ))),
                 })
             })
-            .transpose()?
-            .unwrap_or(element_type.try_into())?;
+            .transpose()?;
 
         // Create error wrapper closure that has access to self
         let error_wrapper = |error: SparkError| self.wrap_error_with_context(error);
@@ -176,7 +178,7 @@ impl PhysicalExpr for ListExtract {
                 list_extract(
                     list_array,
                     index_array,
-                    &default_value,
+                    default_value.as_ref(),
                     self.fail_on_error,
                     self.one_based,
                     adjust_index,
@@ -190,7 +192,7 @@ impl PhysicalExpr for ListExtract {
                 list_extract(
                     list_array,
                     index_array,
-                    &default_value,
+                    default_value.as_ref(),
                     self.fail_on_error,
                     self.one_based,
                     adjust_index,
@@ -264,22 +266,44 @@ fn zero_based_index(
     }
 }
 
+fn out_of_bounds_error(one_based: bool, index: i32, len: usize) -> SparkError {
+    if one_based {
+        SparkError::InvalidElementAtIndex {
+            index_value: index,
+            array_size: len as i32,
+        }
+    } else {
+        SparkError::InvalidArrayIndex {
+            index_value: index,
+            array_size: len as i32,
+        }
+    }
+}
+
 fn list_extract<O: OffsetSizeTrait>(
     list_array: &GenericListArray<O>,
     index_array: &Int32Array,
-    default_value: &ScalarValue,
+    default_value: Option<&ScalarValue>,
     fail_on_error: bool,
     one_based: bool,
     adjust_index: impl Fn(i32, usize) -> DataFusionResult<Option<usize>>,
     error_wrapper: &impl Fn(SparkError) -> DataFusionError,
 ) -> DataFusionResult<ColumnarValue> {
+    let Some(default_value) = default_value else {
+        return list_extract_without_default(
+            list_array,
+            index_array,
+            fail_on_error,
+            one_based,
+            adjust_index,
+            error_wrapper,
+        );
+    };
+
     let values = list_array.values();
     let offsets = list_array.offsets();
-
     let data = values.to_data();
-
     let default_data = default_value.to_array()?.to_data();
-
     let mut mutable = MutableArrayData::new(vec![&data, &default_data], true, index_array.len());
 
     for (row, (offset_window, index)) in offsets.windows(2).zip(index_array.iter()).enumerate() {
@@ -292,33 +316,57 @@ fn list_extract<O: OffsetSizeTrait>(
             if let Some(i) = adjust_index(index, len)? {
                 mutable.extend(0, start + i, start + i + 1);
             } else if fail_on_error {
-                // Throw appropriate error based on whether this is element_at (one_based=true)
-                // or GetArrayItem (one_based=false)
-                let error = if one_based {
-                    // element_at function
-                    SparkError::InvalidElementAtIndex {
-                        index_value: index,
-                        array_size: len as i32,
-                    }
-                } else {
-                    // GetArrayItem (arr[index])
-                    SparkError::InvalidArrayIndex {
-                        index_value: index,
-                        array_size: len as i32,
-                    }
-                };
-                return Err(error_wrapper(error));
+                return Err(error_wrapper(out_of_bounds_error(one_based, index, len)));
             } else {
                 mutable.extend(1, 0, 1);
             }
         } else {
-            // index is NULL → result is NULL
             mutable.extend_nulls(1);
         }
     }
 
-    let data = mutable.freeze();
-    Ok(ColumnarValue::Array(arrow::array::make_array(data)))
+    Ok(ColumnarValue::Array(arrow::array::make_array(
+        mutable.freeze(),
+    )))
+}
+
+fn list_extract_without_default<O: OffsetSizeTrait>(
+    list_array: &GenericListArray<O>,
+    index_array: &Int32Array,
+    fail_on_error: bool,
+    one_based: bool,
+    adjust_index: impl Fn(i32, usize) -> DataFusionResult<Option<usize>>,
+    error_wrapper: &impl Fn(SparkError) -> DataFusionError,
+) -> DataFusionResult<ColumnarValue> {
+    let values = list_array.values();
+    let offsets = list_array.offsets();
+    let mut indices = Vec::with_capacity(index_array.len());
+
+    for (row, (offset_window, index)) in offsets.windows(2).zip(index_array.iter()).enumerate() {
+        let start = offset_window[0].as_usize();
+        let len = offset_window[1].as_usize() - start;
+
+        if list_array.is_null(row) {
+            indices.push(None);
+        } else if let Some(index) = index {
+            if let Some(i) = adjust_index(index, len)? {
+                indices.push(Some((start + i) as u64));
+            } else if fail_on_error {
+                return Err(error_wrapper(out_of_bounds_error(one_based, index, len)));
+            } else {
+                indices.push(None);
+            }
+        } else {
+            // index is NULL → result is NULL
+            indices.push(None);
+        }
+    }
+
+    Ok(ColumnarValue::Array(take(
+        values.as_ref(),
+        &UInt64Array::from(indices),
+        None,
+    )?))
 }
 
 impl Display for ListExtract {
@@ -378,15 +426,13 @@ mod test {
         ]);
         let indices = Int32Array::from(vec![0, 0, 0]);
 
-        let null_default = ScalarValue::Int32(None);
-
         // Simple error wrapper for tests - just converts SparkError to DataFusionError
         let error_wrapper = |error: SparkError| DataFusionError::from(error);
 
         let ColumnarValue::Array(result) = list_extract(
             &list,
             &indices,
-            &null_default,
+            None,
             false,
             false,
             |idx, len| zero_based_index(idx, len, &error_wrapper),
@@ -406,7 +452,7 @@ mod test {
         let ColumnarValue::Array(result) = list_extract(
             &list,
             &indices,
-            &zero_default,
+            Some(&zero_default),
             false,
             false,
             |idx, len| zero_based_index(idx, len, &error_wrapper),
@@ -436,26 +482,32 @@ mod test {
         ]);
         let indices = Int32Array::from(vec![Some(0), Some(1), Some(2), Some(0), Some(0), None]);
 
-        let null_default = ScalarValue::Int32(None);
         let error_wrapper = |error: SparkError| DataFusionError::from(error);
 
-        let ColumnarValue::Array(result) = list_extract(
-            &list,
-            &indices,
-            &null_default,
-            false,
-            false,
-            |idx, len| zero_based_index(idx, len, &error_wrapper),
-            &error_wrapper,
-        )?
-        else {
-            unreachable!()
-        };
+        for default_value in [
+            None,
+            Some(ScalarValue::Int32(None)),
+            Some(ScalarValue::Int32(Some(0))),
+        ] {
+            let ColumnarValue::Array(result) = list_extract(
+                &list,
+                &indices,
+                default_value.as_ref(),
+                false,
+                false,
+                |idx, len| zero_based_index(idx, len, &error_wrapper),
+                &error_wrapper,
+            )?
+            else {
+                unreachable!()
+            };
 
-        assert_eq!(
-            &result.to_data(),
-            &Int32Array::from(vec![Some(10), Some(20), Some(30), Some(1), None, None]).to_data()
-        );
+            assert_eq!(
+                &result.to_data(),
+                &Int32Array::from(vec![Some(10), Some(20), Some(30), Some(1), None, None])
+                    .to_data()
+            );
+        }
         Ok(())
     }
 }