gstvg commented on code in PR #18921: URL: https://github.com/apache/datafusion/pull/18921#discussion_r2976030740
########## datafusion/functions-nested/src/array_transform.rs: ########## @@ -0,0 +1,253 @@ +// 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. + +//! [`LambdaUDF`] definitions for array_transform function. + +use arrow::{ + array::{Array, ArrayRef, AsArray, FixedSizeListArray, LargeListArray, ListArray}, + datatypes::{DataType, Field, FieldRef}, +}; +use datafusion_common::{ + Result, exec_err, plan_err, + utils::{list_values, take_function_args}, +}; +use datafusion_expr::{ + ColumnarValue, Documentation, LambdaFunctionArgs, LambdaReturnFieldArgs, + LambdaSignature, LambdaUDF, ValueOrLambda, Volatility, +}; +use datafusion_macros::user_doc; +use std::{any::Any, fmt::Debug, sync::Arc}; + +make_udlf_expr_and_func!( + ArrayTransform, + array_transform, + array lambda, + "transforms the values of a array", + array_transform_udlf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "transforms the values of a array", + syntax_example = "array_transform(array, x -> x*2)", + sql_example = r#"```sql +> select array_transform([1, 2, 3, 4, 5], x -> x*2); ++-------------------------------------------+ +| array_transform([1, 2, 3, 4, 5], x -> x*2) | ++-------------------------------------------+ +| [2, 4, 6, 8, 10] | ++-------------------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument(name = "lambda", description = "Lambda") +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayTransform { + signature: LambdaSignature, + aliases: Vec<String>, +} + +impl Default for ArrayTransform { + fn default() -> Self { + Self::new() + } +} + +impl ArrayTransform { + pub fn new() -> Self { + Self { + signature: LambdaSignature::user_defined(Volatility::Immutable), + aliases: vec![String::from("list_transform")], + } + } +} + +impl LambdaUDF for ArrayTransform { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "array_transform" + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &LambdaSignature { + &self.signature + } + + fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { + let list = if arg_types.len() == 1 { + &arg_types[0] + } else { + return plan_err!( + "{} function requires 1 value arguments, got {}", + self.name(), + arg_types.len() + ); + }; + + let coerced = match list { + DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) => list.clone(), + DataType::ListView(field) => DataType::List(Arc::clone(field)), + DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)), + _ => { + return plan_err!( + "{} expected a list as first argument, got {}", + self.name(), + list + ); + } + }; + + Ok(vec![coerced]) + } + + fn lambdas_parameters( + &self, + args: &[ValueOrLambda<FieldRef, ()>], + ) -> Result<Vec<Option<Vec<Field>>>> { + let (list, _lambda) = value_lambda_pair(self.name(), args)?; + + let field = match list.data_type() { + DataType::List(field) => field, + DataType::LargeList(field) => field, + DataType::FixedSizeList(field, _) => field, Review Comment: ListView/LargeListView are currently [handled via type coercion](https://github.com/apache/datafusion/pull/18921/changes#diff-e464d894d5a0fe37392c40f9e54f67590220578259a8dc41fa57b329f6f0d343R114-R115). This makes things simpler and if the list view is casted to normal list multiple times on a projection, CSE will dedup it. I can only imagine few cases where it's more efficiently to handle views directly, but I'm not sure if they are worth handling: 1. Every value is referenced by at least one view, regardless of ordering, so we can transform them directly without packing them first, but that I guess that's uncommon and that check is a bit costly for the false cases. 2. All referenced values are packed on order. Is cheaper to check, but if that would make sense here, it likely would make sense on arrow_cast ListView->List instead, as it allow zero-copy casting, and we would get it for free here as well as CSE. 3. The sublists avg size is very high to the point is efficient to call the lamba body physical expr for every sublist instead of packing them together first, and the output type of the transformation is cheaper to move around than the inputs, because we would need to concat the transformed sublists into a single array, e.g: `array_transform(list_view_of_strings, str -> length(str))`. Instead of packing `list_view_of_strings`, which is expensive for strings, we could call the lambda body for each sublist, and then concat the resulting int32 arrays, which is cheaper than packing strings. And for maps, I thought it would be a different function, like `map_transform(map_col, (k, v, i) -> ...)`. Shoud we support it here too? -- 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]
