diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6faa9fec4e..33700c4189 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -148,6 +148,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..4438b1a098 --- /dev/null +++ b/native/spark-expr/benches/list_extract.rs @@ -0,0 +1,185 @@ +// 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::{Array, ArrayRef, Int32Array, ListArray, StringArray}; +use arrow::buffer::{NullBuffer, 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, with_nulls: bool) -> ArrayRef { + let offsets = (0..=ROWS) + .map(|i| (i * ELEMENTS_PER_ROW) as i32) + .collect::>(); + let nulls = with_nulls.then(|| (0..ROWS).map(|row| row % 4 != 0).collect::()); + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::new(offsets.into()), + values, + nulls, + )) +} + +fn bench_case( + c: &mut Criterion, + name: &str, + list: ArrayRef, + oob: bool, + default: Option, + with_null_ordinals: bool, +) { + let indices = Arc::new(Int32Array::from_iter((0..ROWS).map(|row| { + if with_null_ordinals && row % 4 == 1 { + None + } else { + Some(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(), list.null_count() > 0), + Field::new("index", DataType::Int32, indices.null_count() > 0), + ])); + 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 int_values = Arc::new(Int32Array::from_iter_values(0..total as i32)); + let string_values = Arc::new(StringArray::from_iter_values( + (0..total).map(|i| format!("value-{i}")), + )); + let ints = list_of(int_values.clone(), false); + let strings = list_of(string_values.clone(), false); + let nullable_ints = list_of(int_values, true); + let nullable_strings = list_of(string_values, true); + + for oob in [false, true] { + let suffix = if oob { "50%-oob" } else { "0%-oob" }; + bench_case( + c, + &format!("list_extract/int32/no-default/{suffix}"), + Arc::clone(&ints), + oob, + None, + false, + ); + bench_case( + c, + &format!("list_extract/int32/null-default/{suffix}"), + Arc::clone(&ints), + oob, + Some(ScalarValue::Int32(None)), + false, + ); + bench_case( + c, + &format!("list_extract/int32/non-null-default/{suffix}"), + Arc::clone(&ints), + oob, + Some(ScalarValue::Int32(Some(0))), + false, + ); + bench_case( + c, + &format!("list_extract/utf8/no-default/{suffix}"), + Arc::clone(&strings), + oob, + None, + false, + ); + bench_case( + c, + &format!("list_extract/utf8/null-default/{suffix}"), + Arc::clone(&strings), + oob, + Some(ScalarValue::Utf8(None)), + false, + ); + bench_case( + c, + &format!("list_extract/utf8/non-null-default/{suffix}"), + Arc::clone(&strings), + oob, + Some(ScalarValue::Utf8(Some(String::new()))), + false, + ); + } + + bench_case( + c, + "list_extract/int32/no-default/25%-null-lists-25%-null-ordinals", + Arc::clone(&nullable_ints), + false, + None, + true, + ); + bench_case( + c, + "list_extract/int32/non-null-default/25%-null-lists-25%-null-ordinals", + nullable_ints, + false, + Some(ScalarValue::Int32(Some(0))), + true, + ); + bench_case( + c, + "list_extract/utf8/no-default/25%-null-lists-25%-null-ordinals", + Arc::clone(&nullable_strings), + false, + None, + true, + ); + bench_case( + c, + "list_extract/utf8/non-null-default/25%-null-lists-25%-null-ordinals", + nullable_strings, + false, + Some(ScalarValue::Utf8(Some(String::new()))), + true, + ); +} + +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..e8e5dd8ef1 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, UInt64Builder, +}; +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,61 +266,128 @@ 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, + } + } +} + +enum RowAction { + Gather(usize), + Null, + Default, +} + +fn resolve_row( + list_is_null: bool, + offset_window: &[O], + index: Option, + fail_on_error: bool, + one_based: bool, + adjust_index: &impl Fn(i32, usize) -> DataFusionResult>, + error_wrapper: &impl Fn(SparkError) -> DataFusionError, +) -> DataFusionResult { + if list_is_null { + return Ok(RowAction::Null); + } + + let Some(index) = index else { + return Ok(RowAction::Null); + }; + + let start = offset_window[0].as_usize(); + let len = offset_window[1].as_usize() - start; + match adjust_index(index, len)? { + Some(index) => Ok(RowAction::Gather(start + index)), + None if fail_on_error => Err(error_wrapper(out_of_bounds_error(one_based, index, len))), + None => Ok(RowAction::Default), + } +} + 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.filter(|default| !default.is_null()) 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 { - // 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)); - } else { - mutable.extend(1, 0, 1); - } - } else { - // index is NULL → result is NULL - mutable.extend_nulls(1); + match resolve_row( + list_array.is_null(row), + offset_window, + index, + fail_on_error, + one_based, + &adjust_index, + error_wrapper, + )? { + RowAction::Gather(index) => mutable.extend(0, index, index + 1), + RowAction::Null => mutable.extend_nulls(1), + RowAction::Default => mutable.extend(1, 0, 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 = UInt64Builder::with_capacity(index_array.len()); + + for (row, (offset_window, index)) in offsets.windows(2).zip(index_array.iter()).enumerate() { + match resolve_row( + list_array.is_null(row), + offset_window, + index, + fail_on_error, + one_based, + &adjust_index, + error_wrapper, + )? { + RowAction::Gather(index) => indices.append_value(index as u64), + RowAction::Null | RowAction::Default => indices.append_null(), + } + } + + let indices = indices.finish(); + Ok(ColumnarValue::Array(take(values.as_ref(), &indices, None)?)) } impl Display for ListExtract { @@ -378,35 +447,35 @@ 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, - false, - false, - |idx, len| zero_based_index(idx, len, &error_wrapper), - &error_wrapper, - )? - else { - unreachable!() - }; + for default_value in [None, Some(ScalarValue::Int32(None))] { + 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(1), None, None]).to_data() - ); + assert_eq!( + &result.to_data(), + &Int32Array::from(vec![Some(1), None, None]).to_data() + ); + } let zero_default = ScalarValue::Int32(Some(0)); 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 +505,80 @@ 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(()) } + + #[test] + fn test_list_extract_without_default_ansi_errors() { + let list = ListArray::from_iter_primitive::(vec![Some(vec![Some(1)])]); + let error_wrapper = |error: SparkError| DataFusionError::from(error); + + for (one_based, index) in [(true, 2), (false, 1)] { + let indices = Int32Array::from(vec![index]); + let error = list_extract( + &list, + &indices, + None, + true, + one_based, + |idx, len| { + if one_based { + one_based_index(idx, len, &error_wrapper) + } else { + zero_based_index(idx, len, &error_wrapper) + } + }, + &error_wrapper, + ) + .unwrap_err(); + let DataFusionError::External(error) = error else { + panic!("expected external Spark error") + }; + let error = error.downcast_ref::().unwrap(); + + if one_based { + assert!(matches!( + error, + SparkError::InvalidElementAtIndex { + index_value: 2, + array_size: 1 + } + )); + } else { + assert!(matches!( + error, + SparkError::InvalidArrayIndex { + index_value: 1, + array_size: 1 + } + )); + } + } + } } diff --git a/spark/src/test/resources/sql-tests/expressions/array/get_array_item.sql b/spark/src/test/resources/sql-tests/expressions/array/get_array_item.sql index c00a21f3df..1fcedd6e85 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/get_array_item.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/get_array_item.sql @@ -30,3 +30,18 @@ SELECT arr[idx] FROM test_get_array_item -- literal arguments query SELECT array(10, 20, 30)[0], array(10, 20, 30)[2], array()[0] + +statement +CREATE TABLE test_get_array_item_complex(nested array>, structs array>, idx int) USING parquet + +statement +INSERT INTO test_get_array_item_complex VALUES + (array(array(1, 2), array(3, 4)), array(named_struct('a', 1, 'b', 'a'), named_struct('a', 2, 'b', 'b')), 1), + (NULL, NULL, 0), + (array(array(5, 6)), array(named_struct('a', 3, 'b', 'c')), NULL) + +query +SELECT nested[idx] FROM test_get_array_item_complex + +query +SELECT structs[idx] FROM test_get_array_item_complex