diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 775311935c..f91aa58ca0 100644 --- a/native/spark-expr/Cargo.toml +++ b/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 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 d68784ca70..8cb0ae253d 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, 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( 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() { @@ -292,33 +316,57 @@ fn list_extract( 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( + list_array: &GenericListArray, + index_array: &Int32Array, + fail_on_error: bool, + one_based: bool, + adjust_index: impl Fn(i32, usize) -> DataFusionResult>, + error_wrapper: &impl Fn(SparkError) -> DataFusionError, +) -> DataFusionResult { + 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(()) } }