kosiew commented on code in PR #23267: URL: https://github.com/apache/datafusion/pull/23267#discussion_r3548696694
########## datafusion/functions-nested/src/array_first.rs: ########## @@ -0,0 +1,427 @@ +// 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. + +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_first function. + +use arrow::{ + array::{ + Array, AsArray, BooleanArray, GenericListArray, OffsetSizeTrait, UInt64Array, + UInt64Builder, new_null_array, + }, + compute::{take, take_arrays}, + datatypes::{DataType, FieldRef}, +}; +use datafusion_common::{ + Result, exec_datafusion_err, exec_err, plan_err, + utils::{adjust_offsets_for_slice, list_values, list_values_row_number}, +}; +use datafusion_expr::{ + ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +use crate::lambda_utils::{ + coerce_single_list_arg, single_list_lambda_parameters, value_lambda_pair, +}; + +make_higher_order_function_expr_and_func!( + ArrayFirst, + array_first, + array lambda, + "returns the first element of an array that satisfies the predicate", + array_first_higher_order_function +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the first element of an array that satisfies the given predicate. Returns null if the array is empty or no element matches. A predicate that returns null for an element is treated as not matching.", + syntax_example = "array_first(array, predicate)", + sql_example = r#"```sql +> select array_first([1, 2, 3, 4], x -> x > 2); ++----------------------------------------+ +| array_first([1,2,3,4],x -> x > 2) | ++----------------------------------------+ +| 3 | ++----------------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument( + name = "predicate", + description = "Lambda predicate that returns a boolean. The first element for which it returns true is returned." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayFirst { + signature: HigherOrderSignature, + aliases: Vec<String>, +} + +impl Default for ArrayFirst { + fn default() -> Self { + Self::new() + } +} + +impl ArrayFirst { + pub fn new() -> Self { + Self { + signature: HigherOrderSignature::exact( + vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())], + Volatility::Immutable, + ), + aliases: vec![String::from("list_first")], + } + } +} + +impl HigherOrderUDFImpl for ArrayFirst { + fn name(&self) -> &str { + "array_first" + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { + coerce_single_list_arg(self.name(), arg_types) + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>], + ) -> Result<LambdaParametersProgress> { + single_list_lambda_parameters(self.name(), fields) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result<FieldRef> { + let (list, _lambda) = value_lambda_pair(self.name(), args.arg_fields)?; + + let element_field = match list.data_type() { + DataType::List(field) | DataType::LargeList(field) => field, + other => { + return plan_err!( + "{} expected a list as first argument, got {other}", + self.name() + ); + } + }; + + // The result is a single element of the array. It is always nullable + // because an empty array (or no matching element) yields null. + Ok(Arc::new( + element_field + .as_ref() + .clone() + .with_name("") + .with_nullable(true), + )) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> { + let (list, lambda) = value_lambda_pair(self.name(), &args.args)?; + + let list_array = list.to_array(args.number_rows)?; + + // Fast path: fully null input. Also required for FixedSizeList which + // can't be handled by clear_null_values when fully null. + if list_array.null_count() == list_array.len() { + return Ok(ColumnarValue::Array(new_null_array( + args.return_type(), + list_array.len(), + ))); + } + + let list_values = list_values(&list_array)?; + + // Evaluate the predicate over every flat element. Captured columns are + // spread to align with the flattened values via list_values_row_number. + let values_param = || Ok(Arc::clone(&list_values)); + + let predicate_results = lambda Review Comment: This evaluates the predicate over every flattened element before picking the first match. That lines up with the existing vectorized higher-order function model, but `array_first` could still read like it short-circuits within each row. Could we add a small regression test or doc note for the same-sublist case where an early element matches, but a later element would error if evaluated? That would make the eager evaluation semantics explicit. Non-blocking, since this PR seems focused on avoiding materializing the filtered array rather than changing predicate evaluation semantics. ########## datafusion/sqllogictest/test_files/array/array_first.slt: ########## @@ -0,0 +1,118 @@ +# 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. + +############# +## array_first Tests +############# + +statement ok +set datafusion.sql_parser.dialect = databricks; + +statement ok +CREATE TABLE t (list array<int>, number int) +AS VALUES +([1, 50], 10), +([4, 50], 40), +([7, 50], 60); + +# basic: returns the first element that matches the predicate +query I +SELECT array_first([1, 2, 3, 4], x -> x > 2); +---- +3 + +# no element matches returns null +query I +SELECT array_first([1, 2, 3], x -> x > 5); +---- +NULL + +# empty array returns null +query I +SELECT array_first(arrow_cast(make_array(), 'List(Int32)'), x -> x > 0); +---- +NULL + +# null array returns null +query I +SELECT array_first(arrow_cast(NULL, 'List(Int32)'), x -> x > 0); +---- +NULL + +# a predicate that returns null for an element is treated as not matching +query I +SELECT array_first([1, 2, NULL, 4], x -> x > 2); +---- +4 + +# the predicate may match a null element, which is returned as null +query I +SELECT array_first(arrow_cast([NULL, 2], 'List(Int32)'), x -> x IS NULL); +---- +NULL + +# predicate always returns null -> no match -> null +query I +SELECT array_first([1, 2, 3], x -> NULL::boolean); +---- +NULL + +# a predicate matching every element returns the first element +query I +SELECT array_first([10, 20, 30], x -> true); +---- +10 + +# string elements +query T +SELECT array_first(['a', 'bb', 'ccc'], x -> length(x) > 1); +---- +bb + +# multiple rows +query I +SELECT array_first(list, x -> x > 5) FROM t; +---- +50 +50 +7 + +# predicate can reference an outer column (last row has no match -> null) +query I +SELECT array_first(list, x -> x > number) FROM t; +---- +50 +50 +NULL + +# large list works Review Comment: Nice coverage for `LargeList`. It might also be worth adding a small `FixedSizeList` or `ListView` case, since `coerce_single_list_arg` advertises that those inputs are normalized for this higher-order function family. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
