From 0b0c3a5bf33aa54de3c179e965736555d3f47fae Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 31 Jul 2026 23:33:28 +0800 Subject: [PATCH 1/2] perf: replace list_extract gather with take and zip --- .../src/array_funcs/list_extract.rs | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/native/spark-expr/src/array_funcs/list_extract.rs b/native/spark-expr/src/array_funcs/list_extract.rs index d68784ca70..805eefcf11 100644 --- a/native/spark-expr/src/array_funcs/list_extract.rs +++ b/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, BooleanArray, GenericListArray, Int32Array, OffsetSizeTrait, UInt64Array, +}; +use arrow::compute::{kernels::zip::zip, 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, @@ -275,22 +278,20 @@ fn list_extract( ) -> DataFusionResult { 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()); + let mut indices = Vec::with_capacity(index_array.len()); + let mut use_default = 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) { - mutable.extend_nulls(1); + indices.push(None); + use_default.push(false); } else if let Some(index) = index { if let Some(i) = adjust_index(index, len)? { - mutable.extend(0, start + i, start + i + 1); + indices.push(Some((start + i) as u64)); + use_default.push(false); } else if fail_on_error { // Throw appropriate error based on whether this is element_at (one_based=true) // or GetArrayItem (one_based=false) @@ -309,16 +310,22 @@ fn list_extract( }; return Err(error_wrapper(error)); } else { - mutable.extend(1, 0, 1); + indices.push(None); + use_default.push(true); } } else { // index is NULL → result is NULL - mutable.extend_nulls(1); + indices.push(None); + use_default.push(false); } } - let data = mutable.freeze(); - Ok(ColumnarValue::Array(arrow::array::make_array(data))) + let taken = take(values.as_ref(), &UInt64Array::from(indices), None)?; + Ok(ColumnarValue::Array(zip( + &BooleanArray::from(use_default), + &default_value.to_scalar()?, + &taken, + )?)) } impl Display for ListExtract { @@ -436,13 +443,13 @@ 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 zero_default = ScalarValue::Int32(Some(0)); let error_wrapper = |error: SparkError| DataFusionError::from(error); let ColumnarValue::Array(result) = list_extract( &list, &indices, - &null_default, + &zero_default, false, false, |idx, len| zero_based_index(idx, len, &error_wrapper), From 2a59e089915be8f92291bc7524414fc463ed6562 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 2 Aug 2026 11:34:32 +0800 Subject: [PATCH 2/2] perf: avoid list_extract default regressions --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/list_extract.rs | 122 ++++++++++++++ .../src/array_funcs/list_extract.rs | 155 +++++++++++------- 3 files changed, 226 insertions(+), 55 deletions(-) create mode 100644 native/spark-expr/benches/list_extract.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index c05ae89793..cdbfc32e9e 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -135,6 +135,10 @@ harness = false name = "array_size" harness = false +[[bench]] +name = "list_extract" +harness = false + [[bench]] name = "regexp_extract" harness = false diff --git a/native/spark-expr/benches/list_extract.rs b/native/spark-expr/benches/list_extract.rs new file mode 100644 index 0000000000..150c96bd7b --- /dev/null +++ b/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::>(); + 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, +) { + 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); + 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); diff --git a/native/spark-expr/src/array_funcs/list_extract.rs b/native/spark-expr/src/array_funcs/list_extract.rs index 805eefcf11..8cb0ae253d 100644 --- a/native/spark-expr/src/array_funcs/list_extract.rs +++ b/native/spark-expr/src/array_funcs/list_extract.rs @@ -16,9 +16,9 @@ // under the License. use arrow::array::{ - Array, BooleanArray, GenericListArray, Int32Array, OffsetSizeTrait, UInt64Array, + Array, GenericListArray, Int32Array, MutableArrayData, OffsetSizeTrait, UInt64Array, }; -use arrow::compute::{kernels::zip::zip, take}; +use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow::{datatypes::ArrowNativeType, record_batch::RecordBatch}; use datafusion::common::{ @@ -146,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) => { @@ -158,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); @@ -179,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, @@ -193,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, @@ -267,10 +266,73 @@ 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( list_array: &GenericListArray, index_array: &Int32Array, - default_value: &ScalarValue, + default_value: Option<&ScalarValue>, + fail_on_error: bool, + one_based: bool, + adjust_index: impl Fn(i32, usize) -> DataFusionResult>, + error_wrapper: &impl Fn(SparkError) -> DataFusionError, +) -> DataFusionResult { + 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() { + let start = offset_window[0].as_usize(); + let len = offset_window[1].as_usize() - start; + + if list_array.is_null(row) { + mutable.extend_nulls(1); + } else if let Some(index) = index { + if let Some(i) = adjust_index(index, len)? { + mutable.extend(0, start + i, start + i + 1); + } else if fail_on_error { + return Err(error_wrapper(out_of_bounds_error(one_based, index, len))); + } else { + mutable.extend(1, 0, 1); + } + } else { + mutable.extend_nulls(1); + } + } + + Ok(ColumnarValue::Array(arrow::array::make_array( + mutable.freeze(), + ))) +} + +fn list_extract_without_default( + list_array: &GenericListArray, + index_array: &Int32Array, fail_on_error: bool, one_based: bool, adjust_index: impl Fn(i32, usize) -> DataFusionResult>, @@ -279,7 +341,6 @@ fn list_extract( let values = list_array.values(); let offsets = list_array.offsets(); let mut indices = Vec::with_capacity(index_array.len()); - let mut use_default = 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(); @@ -287,44 +348,24 @@ fn list_extract( if list_array.is_null(row) { indices.push(None); - use_default.push(false); } else if let Some(index) = index { if let Some(i) = adjust_index(index, len)? { indices.push(Some((start + i) as u64)); - use_default.push(false); } 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 { indices.push(None); - use_default.push(true); } } else { // index is NULL → result is NULL indices.push(None); - use_default.push(false); } } - let taken = take(values.as_ref(), &UInt64Array::from(indices), None)?; - Ok(ColumnarValue::Array(zip( - &BooleanArray::from(use_default), - &default_value.to_scalar()?, - &taken, + Ok(ColumnarValue::Array(take( + values.as_ref(), + &UInt64Array::from(indices), + None, )?)) } @@ -385,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), @@ -413,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), @@ -443,26 +482,32 @@ mod test { ]); let indices = Int32Array::from(vec![Some(0), Some(1), Some(2), Some(0), Some(0), None]); - let zero_default = ScalarValue::Int32(Some(0)); let error_wrapper = |error: SparkError| DataFusionError::from(error); - let ColumnarValue::Array(result) = list_extract( - &list, - &indices, - &zero_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(()) } }