Skip to main content

datafusion_expr/type_coercion/
functions.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use super::binary::binary_numeric_coercion;
19use crate::{
20    AggregateUDF, HigherOrderTypeSignature, HigherOrderUDF, ScalarUDF, Signature,
21    TypeSignature, ValueOrLambda, WindowUDF,
22};
23use arrow::datatypes::{Field, FieldRef};
24use arrow::{
25    compute::can_cast_types,
26    datatypes::{DataType, TimeUnit},
27};
28use datafusion_common::internal_datafusion_err;
29use datafusion_common::types::LogicalType;
30use datafusion_common::utils::{
31    ListCoercion, base_type, coerced_fixed_size_list_to_list,
32};
33use datafusion_common::{
34    Result, exec_err, internal_err, plan_err, types::NativeType, utils::list_ndims,
35};
36use datafusion_expr_common::signature::{
37    ArrayFunctionArgument, EncodingPreservation, TypeSignatureClass,
38};
39use datafusion_expr_common::type_coercion::binary::type_union_resolution;
40use datafusion_expr_common::{
41    signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD},
42    type_coercion::binary::comparison_coercion,
43    type_coercion::binary::string_coercion,
44};
45use itertools::Itertools as _;
46use std::sync::Arc;
47
48/// Extension trait to unify common functionality between [`ScalarUDF`], [`AggregateUDF`]
49/// and [`WindowUDF`] for use by signature coercion functions.
50pub trait UDFCoercionExt {
51    /// Should delegate to [`ScalarUDF::name`], [`AggregateUDF::name`] or [`WindowUDF::name`].
52    fn name(&self) -> &str;
53    /// Should delegate to [`ScalarUDF::signature`], [`AggregateUDF::signature`]
54    /// or [`WindowUDF::signature`].
55    fn signature(&self) -> &Signature;
56    /// Should delegate to [`ScalarUDF::coerce_types`], [`AggregateUDF::coerce_types`]
57    /// or [`WindowUDF::coerce_types`].
58    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>>;
59}
60
61impl UDFCoercionExt for ScalarUDF {
62    fn name(&self) -> &str {
63        self.name()
64    }
65
66    fn signature(&self) -> &Signature {
67        self.signature()
68    }
69
70    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
71        self.coerce_types(arg_types)
72    }
73}
74
75impl UDFCoercionExt for AggregateUDF {
76    fn name(&self) -> &str {
77        self.name()
78    }
79
80    fn signature(&self) -> &Signature {
81        self.signature()
82    }
83
84    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
85        self.coerce_types(arg_types)
86    }
87}
88
89impl UDFCoercionExt for WindowUDF {
90    fn name(&self) -> &str {
91        self.name()
92    }
93
94    fn signature(&self) -> &Signature {
95        self.signature()
96    }
97
98    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
99        self.coerce_types(arg_types)
100    }
101}
102
103/// Performs type coercion for UDF arguments.
104///
105/// Returns the data types to which each argument must be coerced to
106/// match `signature`.
107///
108/// For more details on coercion in general, please see the
109/// [`type_coercion`](crate::type_coercion) module.
110pub fn fields_with_udf<F: UDFCoercionExt>(
111    current_fields: &[FieldRef],
112    func: &F,
113) -> Result<Vec<FieldRef>> {
114    let signature = func.signature();
115    let type_signature = &signature.type_signature;
116
117    if current_fields.is_empty() && type_signature != &TypeSignature::UserDefined {
118        if type_signature.supports_zero_argument() {
119            return Ok(vec![]);
120        } else if type_signature.used_to_support_zero_arguments() {
121            // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
122            return plan_err!(
123                "'{}' does not support zero arguments. Use TypeSignature::Nullary for zero arguments",
124                func.name()
125            );
126        } else {
127            return plan_err!("'{}' does not support zero arguments", func.name());
128        }
129    }
130    let current_types = current_fields
131        .iter()
132        .map(|f| f.data_type())
133        .cloned()
134        .collect::<Vec<_>>();
135
136    let valid_types = get_valid_types_with_udf(type_signature, &current_types, func)?;
137    if valid_types
138        .iter()
139        .any(|data_type| data_type == &current_types)
140    {
141        return Ok(current_fields.to_vec());
142    }
143
144    let updated_types =
145        try_coerce_types(func.name(), valid_types, &current_types, type_signature)?;
146
147    Ok(current_fields
148        .iter()
149        .zip(updated_types)
150        .map(|(current_field, new_type)| {
151            current_field.as_ref().clone().with_data_type(new_type)
152        })
153        .map(Arc::new)
154        .collect())
155}
156
157/// Performs type coercion for higher order function arguments.
158///
159/// For value arguments, returns the field to which each
160/// argument must be coerced to match `signature`.
161/// For lambda arguments, returns a clone of the associated data
162///
163/// Note this does not invokes [crate::HigherOrderUDFImpl::coerce_values_for_lambdas].
164/// If that's required, use [value_fields_with_higher_order_udf_and_lambdas]
165/// instead
166///
167/// For more details on coercion in general, please see the
168/// [`type_coercion`](crate::type_coercion) module.
169pub fn value_fields_with_higher_order_udf<L: Clone>(
170    current_fields: &[ValueOrLambda<FieldRef, L>],
171    func: &HigherOrderUDF,
172) -> Result<Vec<ValueOrLambda<FieldRef, L>>> {
173    match func.signature().type_signature {
174        HigherOrderTypeSignature::UserDefined => {
175            let arg_types = current_fields
176                .iter()
177                .filter_map(|p| match p {
178                    ValueOrLambda::Value(field) => Some(field.data_type().clone()),
179                    ValueOrLambda::Lambda(_) => None,
180                })
181                .collect::<Vec<_>>();
182
183            let coerced_types = func.coerce_value_types(&arg_types)?;
184
185            if coerced_types.len() != arg_types.len() {
186                return plan_err!(
187                    "{} coerce_value_types should have returned {} items but returned {}",
188                    func.name(),
189                    arg_types.len(),
190                    coerced_types.len()
191                );
192            }
193
194            // coerced_types has been partitioned from current_fields
195            // and refers only to values and not to lambdas, so instead
196            // of zipping them, we iterate over current_fields and only
197            // consume from coerced_types when a given argument is a value
198            // to reconstruct the arguments list with the correct order
199            // this supports any value and lambda positioning including
200            // multiple lambdas interleaved with values
201            let mut coerced_types = coerced_types.into_iter();
202
203            current_fields
204                .iter()
205                .map(|current_field| match current_field {
206                    ValueOrLambda::Value(field) => {
207                        let data_type = coerced_types.next().ok_or_else(|| {
208                            internal_datafusion_err!(
209                                "coerced_types len should have been checked above"
210                            )
211                        })?;
212
213                        Ok(ValueOrLambda::Value(Arc::new(
214                            field.as_ref().clone().with_data_type(data_type),
215                        )))
216                    }
217                    ValueOrLambda::Lambda(lambda) => {
218                        Ok(ValueOrLambda::Lambda(lambda.clone()))
219                    }
220                })
221                .collect()
222        }
223        HigherOrderTypeSignature::VariadicAny => Ok(current_fields.to_vec()),
224        HigherOrderTypeSignature::Any(number) => {
225            if current_fields.len() != number {
226                return plan_err!(
227                    "The function '{}' expected {number} arguments but received {}",
228                    func.name(),
229                    current_fields.len()
230                );
231            }
232
233            Ok(current_fields.to_vec())
234        }
235        HigherOrderTypeSignature::Exact(ref expected) => {
236            if current_fields.len() != expected.len() {
237                let name = func.name();
238                let expected_len = expected.len();
239                let actual_len = current_fields.len();
240                return plan_err!(
241                    "The function '{name}' expected {expected_len} argument(s) but received {actual_len}"
242                );
243            }
244
245            for (i, (actual, expected)) in
246                current_fields.iter().zip(expected.iter()).enumerate()
247            {
248                match (actual, expected) {
249                    (ValueOrLambda::Value(_), ValueOrLambda::Value(_)) => {}
250                    (ValueOrLambda::Lambda(_), ValueOrLambda::Lambda(_)) => {}
251                    (ValueOrLambda::Value(_), ValueOrLambda::Lambda(_)) => {
252                        let name = func.name();
253                        return plan_err!(
254                            "The function '{name}' expected a lambda at position {i} but received a value"
255                        );
256                    }
257                    (ValueOrLambda::Lambda(_), ValueOrLambda::Value(_)) => {
258                        let name = func.name();
259                        return plan_err!(
260                            "The function '{name}' expected a value at position {i} but received a lambda"
261                        );
262                    }
263                }
264            }
265
266            let arg_types = current_fields
267                .iter()
268                .filter_map(|p| match p {
269                    ValueOrLambda::Value(field) => Some(field.data_type().clone()),
270                    ValueOrLambda::Lambda(_) => None,
271                })
272                .collect::<Vec<_>>();
273
274            let coerced_types = func.coerce_value_types(&arg_types)?;
275
276            if coerced_types.len() != arg_types.len() {
277                return plan_err!(
278                    "{} coerce_value_types should have returned {} items but returned {}",
279                    func.name(),
280                    arg_types.len(),
281                    coerced_types.len()
282                );
283            }
284
285            let mut coerced_types = coerced_types.into_iter();
286
287            current_fields
288                .iter()
289                .map(|current_field| match current_field {
290                    ValueOrLambda::Value(field) => {
291                        let data_type = coerced_types.next().ok_or_else(|| {
292                            internal_datafusion_err!(
293                                "coerced_types len should have been checked above"
294                            )
295                        })?;
296
297                        Ok(ValueOrLambda::Value(Arc::new(
298                            field.as_ref().clone().with_data_type(data_type),
299                        )))
300                    }
301                    ValueOrLambda::Lambda(lambda) => {
302                        Ok(ValueOrLambda::Lambda(lambda.clone()))
303                    }
304                })
305                .collect()
306        }
307    }
308}
309
310/// Performs type coercion for higher order function arguments,
311/// including those defined by [crate::HigherOrderUDFImpl::coerce_values_for_lambdas],
312/// if it returns `Some(...)` instead of the default `None`. Note that
313/// compared to [value_fields_with_higher_order_udf], this function requires
314/// the [ValueOrLambda::Lambda] variant to contain the output field of the lambda.
315///
316/// For value arguments, returns the field to which each
317/// argument must be coerced to match `signature`.
318/// For lambda arguments, returns a clone of the output field
319///
320/// For more details on coercion in general, please see the
321/// [`type_coercion`](crate::type_coercion) module.
322pub fn value_fields_with_higher_order_udf_and_lambdas(
323    current_fields: &[ValueOrLambda<FieldRef, FieldRef>],
324    func: &HigherOrderUDF,
325) -> Result<Vec<ValueOrLambda<FieldRef, FieldRef>>> {
326    let mut new_fields = value_fields_with_higher_order_udf(current_fields, func)?;
327
328    let new_types = new_fields
329        .iter()
330        .map(|f| match f {
331            ValueOrLambda::Value(f) => ValueOrLambda::Value(f.data_type().clone()),
332            ValueOrLambda::Lambda(f) => ValueOrLambda::Lambda(f.data_type().clone()),
333        })
334        .collect::<Vec<_>>();
335
336    if let Some(new_value_types) = func.coerce_values_for_lambdas(&new_types)? {
337        let mut new_value_types = new_value_types.into_iter();
338
339        let value_types_count = new_types
340            .iter()
341            .filter(|e| matches!(e, ValueOrLambda::Value(_)))
342            .count();
343
344        if new_value_types.len() != value_types_count {
345            return plan_err!(
346                "{} coerce_values_for_lambdas returned {} values but {value_types_count} expected",
347                func.name(),
348                new_value_types.len()
349            );
350        }
351
352        for new_field in &mut new_fields {
353            match new_field {
354                ValueOrLambda::Value(value) => {
355                    let coerce_to = new_value_types.next().ok_or_else(|| {
356                        internal_datafusion_err!(
357                            "new_value_types len should have been checked above"
358                        )
359                    })?;
360
361                    if value.data_type() != &coerce_to {
362                        Arc::make_mut(value).set_data_type(coerce_to);
363                    }
364                }
365                ValueOrLambda::Lambda(_) => {}
366            }
367        }
368    };
369
370    Ok(new_fields)
371}
372
373/// Performs type coercion for scalar function arguments.
374///
375/// Returns the data types to which each argument must be coerced to
376/// match `signature`.
377///
378/// For more details on coercion in general, please see the
379/// [`type_coercion`](crate::type_coercion) module.
380#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
381pub fn data_types_with_scalar_udf(
382    current_types: &[DataType],
383    func: &ScalarUDF,
384) -> Result<Vec<DataType>> {
385    let current_fields = current_types
386        .iter()
387        .map(|dt| Arc::new(Field::new("f", dt.clone(), true)))
388        .collect::<Vec<_>>();
389    Ok(fields_with_udf(&current_fields, func)?
390        .iter()
391        .map(|f| f.data_type().clone())
392        .collect())
393}
394
395/// Performs type coercion for aggregate function arguments.
396///
397/// Returns the fields to which each argument must be coerced to
398/// match `signature`.
399///
400/// For more details on coercion in general, please see the
401/// [`type_coercion`](crate::type_coercion) module.
402#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
403pub fn fields_with_aggregate_udf(
404    current_fields: &[FieldRef],
405    func: &AggregateUDF,
406) -> Result<Vec<FieldRef>> {
407    fields_with_udf(current_fields, func)
408}
409
410/// Performs type coercion for window function arguments.
411///
412/// Returns the data types to which each argument must be coerced to
413/// match `signature`.
414///
415/// For more details on coercion in general, please see the
416/// [`type_coercion`](crate::type_coercion) module.
417#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
418pub fn fields_with_window_udf(
419    current_fields: &[FieldRef],
420    func: &WindowUDF,
421) -> Result<Vec<FieldRef>> {
422    fields_with_udf(current_fields, func)
423}
424
425/// Performs type coercion for function arguments.
426///
427/// Returns the data types to which each argument must be coerced to
428/// match `signature`.
429///
430/// For more details on coercion in general, please see the
431/// [`type_coercion`](crate::type_coercion) module.
432#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
433pub fn data_types(
434    function_name: impl AsRef<str>,
435    current_types: &[DataType],
436    signature: &Signature,
437) -> Result<Vec<DataType>> {
438    let type_signature = &signature.type_signature;
439
440    if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
441        if type_signature.supports_zero_argument() {
442            return Ok(vec![]);
443        } else if type_signature.used_to_support_zero_arguments() {
444            // Special error to help during upgrade: https://github.com/apache/datafusion/issues/13763
445            return plan_err!(
446                "function '{}' has signature {type_signature} which does not support zero arguments. Use TypeSignature::Nullary for zero arguments",
447                function_name.as_ref()
448            );
449        } else {
450            return plan_err!(
451                "Function '{}' has signature {type_signature} which does not support zero arguments",
452                function_name.as_ref()
453            );
454        }
455    }
456
457    let valid_types =
458        get_valid_types(function_name.as_ref(), type_signature, current_types)?;
459    if valid_types
460        .iter()
461        .any(|data_type| data_type == current_types)
462    {
463        return Ok(current_types.to_vec());
464    }
465
466    try_coerce_types(
467        function_name.as_ref(),
468        valid_types,
469        current_types,
470        type_signature,
471    )
472}
473
474fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
475    match type_signature {
476        TypeSignature::OneOf(type_signatures) => {
477            type_signatures.iter().all(is_well_supported_signature)
478        }
479        TypeSignature::UserDefined
480        | TypeSignature::Numeric(_)
481        | TypeSignature::String(_)
482        | TypeSignature::Coercible(_)
483        | TypeSignature::Any(_)
484        | TypeSignature::Nullary
485        | TypeSignature::Comparable(_) => true,
486        TypeSignature::Variadic(_)
487        | TypeSignature::VariadicAny
488        | TypeSignature::Uniform(_, _)
489        | TypeSignature::Exact(_)
490        | TypeSignature::ArraySignature(_) => false,
491    }
492}
493
494fn try_coerce_types(
495    function_name: &str,
496    valid_types: Vec<Vec<DataType>>,
497    current_types: &[DataType],
498    type_signature: &TypeSignature,
499) -> Result<Vec<DataType>> {
500    let mut valid_types = valid_types;
501
502    // Well-supported signature that returns exact valid types.
503    if !valid_types.is_empty() && is_well_supported_signature(type_signature) {
504        // There may be many valid types if valid signature is OneOf
505        // Otherwise, there should be only one valid type
506        if !type_signature.is_one_of() {
507            assert_eq!(valid_types.len(), 1);
508        }
509
510        let valid_types = valid_types.swap_remove(0);
511        if let Some(t) = maybe_data_types_without_coercion(&valid_types, current_types) {
512            return Ok(t);
513        }
514    } else {
515        // TODO: Deprecate this branch after all signatures are well-supported (aka coercion has happened already)
516        // Try and coerce the argument types to match the signature, returning the
517        // coerced types from the first matching signature.
518        for valid_types in valid_types {
519            if let Some(types) = maybe_data_types(&valid_types, current_types) {
520                return Ok(types);
521            }
522        }
523    }
524
525    // none possible -> Error
526    plan_err!(
527        "Failed to coerce arguments to satisfy a call to '{function_name}' function: coercion from {} to the signature {type_signature} failed",
528        current_types.iter().join(", ")
529    )
530}
531
532fn get_valid_types_with_udf<F: UDFCoercionExt>(
533    signature: &TypeSignature,
534    current_types: &[DataType],
535    func: &F,
536) -> Result<Vec<Vec<DataType>>> {
537    let valid_types = match signature {
538        TypeSignature::UserDefined => match func.coerce_types(current_types) {
539            Ok(coerced_types) => vec![coerced_types],
540            Err(e) => {
541                return exec_err!(
542                    "Function '{}' user-defined coercion failed with: {}",
543                    func.name(),
544                    e.strip_backtrace()
545                );
546            }
547        },
548        TypeSignature::OneOf(signatures) => {
549            let mut res = vec![];
550            let mut errors = vec![];
551            for sig in signatures {
552                match get_valid_types_with_udf(sig, current_types, func) {
553                    Ok(valid_types) => {
554                        res.extend(valid_types);
555                    }
556                    Err(e) => {
557                        errors.push(e.to_string());
558                    }
559                }
560            }
561
562            // Every signature failed, return the joined error
563            if res.is_empty() {
564                return internal_err!(
565                    "Function '{}' failed to match any signature, errors: {}",
566                    func.name(),
567                    errors.join(",")
568                );
569            } else {
570                res
571            }
572        }
573        _ => get_valid_types(func.name(), signature, current_types)?,
574    };
575
576    Ok(valid_types)
577}
578
579/// Returns a Vec of all possible valid argument types for the given signature.
580fn get_valid_types(
581    function_name: &str,
582    signature: &TypeSignature,
583    current_types: &[DataType],
584) -> Result<Vec<Vec<DataType>>> {
585    fn array_valid_types(
586        function_name: &str,
587        current_types: &[DataType],
588        arguments: &[ArrayFunctionArgument],
589        array_coercion: Option<&ListCoercion>,
590    ) -> Result<Vec<Vec<DataType>>> {
591        fn rebuild_array_type(
592            current_type: &DataType,
593            element_type: &DataType,
594            nullable: bool,
595            large_list: bool,
596            fixed_size: Option<i32>,
597        ) -> DataType {
598            // Preserve the original list field when possible so field name or
599            // metadata differences do not introduce otherwise unnecessary casts.
600            let field = match current_type {
601                DataType::List(field)
602                | DataType::LargeList(field)
603                | DataType::FixedSizeList(field, _) => Some(Arc::new(
604                    field
605                        .as_ref()
606                        .clone()
607                        .with_data_type(element_type.clone())
608                        .with_nullable(nullable),
609                )),
610                _ => None,
611            };
612
613            if large_list {
614                field.map_or_else(
615                    || DataType::new_large_list(element_type.clone(), nullable),
616                    DataType::LargeList,
617                )
618            } else if let Some(size) = fixed_size {
619                field.map_or_else(
620                    || {
621                        DataType::new_fixed_size_list(
622                            element_type.clone(),
623                            size,
624                            nullable,
625                        )
626                    },
627                    |field| DataType::FixedSizeList(field, size),
628                )
629            } else {
630                field.map_or_else(
631                    || DataType::new_list(element_type.clone(), nullable),
632                    DataType::List,
633                )
634            }
635        }
636
637        if current_types.len() != arguments.len() {
638            return Ok(vec![vec![]]);
639        }
640
641        let mut large_list = false;
642        let mut fixed_size = array_coercion != Some(&ListCoercion::FixedSizedListToList);
643        let mut list_sizes = Vec::with_capacity(arguments.len());
644        let mut element_types = Vec::with_capacity(arguments.len());
645        let mut nested_item_nullability = Vec::with_capacity(arguments.len());
646        for (argument, current_type) in arguments.iter().zip(current_types.iter()) {
647            match argument {
648                ArrayFunctionArgument::Index | ArrayFunctionArgument::String => {
649                    nested_item_nullability.push(None);
650                }
651                ArrayFunctionArgument::Element => {
652                    element_types.push(current_type.clone());
653                    nested_item_nullability.push(None);
654                }
655                ArrayFunctionArgument::Array => match current_type {
656                    DataType::Null => {
657                        element_types.push(DataType::Null);
658                        nested_item_nullability.push(None);
659                    }
660                    DataType::List(field) | DataType::ListView(field) => {
661                        element_types.push(field.data_type().clone());
662                        nested_item_nullability.push(Some(field.is_nullable()));
663                        fixed_size = false;
664                    }
665                    DataType::LargeList(field) | DataType::LargeListView(field) => {
666                        element_types.push(field.data_type().clone());
667                        nested_item_nullability.push(Some(field.is_nullable()));
668                        large_list = true;
669                        fixed_size = false;
670                    }
671                    DataType::FixedSizeList(field, size) => {
672                        element_types.push(field.data_type().clone());
673                        nested_item_nullability.push(Some(field.is_nullable()));
674                        list_sizes.push(*size)
675                    }
676                    arg_type => {
677                        plan_err!("{function_name} does not support type {arg_type}")?
678                    }
679                },
680            }
681        }
682
683        debug_assert_eq!(nested_item_nullability.len(), arguments.len());
684
685        let Some(element_type) = type_union_resolution(&element_types) else {
686            return Ok(vec![vec![]]);
687        };
688
689        if !fixed_size {
690            list_sizes.clear()
691        };
692
693        let mut list_sizes = list_sizes.into_iter();
694        let valid_types = arguments
695            .iter()
696            .zip(current_types.iter())
697            .zip(nested_item_nullability)
698            .map(|((argument_type, current_type), is_nested_item_nullable)| {
699                match argument_type {
700                    ArrayFunctionArgument::Index => DataType::Int64,
701                    ArrayFunctionArgument::String => DataType::Utf8,
702                    ArrayFunctionArgument::Element => element_type.clone(),
703                    // TODO: support maintaining ListView types here
704                    // https://github.com/apache/datafusion/issues/21777
705                    ArrayFunctionArgument::Array => {
706                        if current_type.is_null() {
707                            DataType::Null
708                        } else {
709                            rebuild_array_type(
710                                current_type,
711                                &element_type,
712                                is_nested_item_nullable.unwrap_or(true),
713                                large_list,
714                                list_sizes.next(),
715                            )
716                        }
717                    }
718                }
719            });
720
721        Ok(vec![valid_types.collect()])
722    }
723
724    fn recursive_array(array_type: &DataType) -> Option<DataType> {
725        match array_type {
726            DataType::List(_)
727            | DataType::LargeList(_)
728            | DataType::ListView(_)
729            | DataType::LargeListView(_)
730            | DataType::FixedSizeList(_, _) => {
731                let array_type = coerced_fixed_size_list_to_list(array_type);
732                Some(array_type)
733            }
734            _ => None,
735        }
736    }
737
738    fn function_length_check(
739        function_name: &str,
740        length: usize,
741        expected_length: usize,
742    ) -> Result<()> {
743        if length != expected_length {
744            return plan_err!(
745                "Function '{function_name}' expects {expected_length} arguments but received {length}"
746            );
747        }
748        Ok(())
749    }
750
751    let valid_types = match signature {
752        TypeSignature::Variadic(valid_types) => valid_types
753            .iter()
754            .map(|valid_type| vec![valid_type.clone(); current_types.len()])
755            .collect(),
756        TypeSignature::String(number) => {
757            function_length_check(function_name, current_types.len(), *number)?;
758
759            let mut new_types = Vec::with_capacity(current_types.len());
760            for data_type in current_types.iter() {
761                let logical_data_type: NativeType = data_type.into();
762                if logical_data_type == NativeType::String {
763                    new_types.push(data_type.to_owned());
764                } else if logical_data_type == NativeType::Null {
765                    // TODO: Switch to Utf8View if all the string functions supports Utf8View
766                    new_types.push(DataType::Utf8);
767                } else {
768                    return plan_err!(
769                        "Function '{function_name}' expects String but received {logical_data_type}"
770                    );
771                }
772            }
773
774            // Find the common string type for the given types
775            fn find_common_type(
776                function_name: &str,
777                lhs_type: &DataType,
778                rhs_type: &DataType,
779            ) -> Result<DataType> {
780                match (lhs_type, rhs_type) {
781                    (DataType::Dictionary(_, lhs), DataType::Dictionary(_, rhs)) => {
782                        find_common_type(function_name, lhs, rhs)
783                    }
784                    (DataType::Dictionary(_, v), other)
785                    | (other, DataType::Dictionary(_, v)) => {
786                        find_common_type(function_name, v, other)
787                    }
788                    _ => {
789                        if let Some(coerced_type) = string_coercion(lhs_type, rhs_type) {
790                            Ok(coerced_type)
791                        } else {
792                            plan_err!(
793                                "Function '{function_name}' could not coerce {lhs_type} and {rhs_type} to a common string type"
794                            )
795                        }
796                    }
797                }
798            }
799
800            // Length checked above, safe to unwrap
801            let mut coerced_type = new_types.first().unwrap().to_owned();
802            for t in new_types.iter().skip(1) {
803                coerced_type = find_common_type(function_name, &coerced_type, t)?;
804            }
805
806            fn base_type_or_default_type(data_type: &DataType) -> DataType {
807                if let DataType::Dictionary(_, v) = data_type {
808                    base_type_or_default_type(v)
809                } else {
810                    data_type.to_owned()
811                }
812            }
813
814            vec![vec![base_type_or_default_type(&coerced_type); *number]]
815        }
816        TypeSignature::Numeric(number) => {
817            function_length_check(function_name, current_types.len(), *number)?;
818
819            // Find common numeric type among given types except string
820            let mut valid_type = current_types.first().unwrap().to_owned();
821            for t in current_types.iter().skip(1) {
822                let logical_data_type: NativeType = t.into();
823                if logical_data_type == NativeType::Null {
824                    continue;
825                }
826
827                if !logical_data_type.is_numeric() {
828                    return plan_err!(
829                        "Function '{function_name}' expects Numeric but received {logical_data_type}"
830                    );
831                }
832
833                if let Some(coerced_type) = binary_numeric_coercion(&valid_type, t) {
834                    valid_type = coerced_type;
835                } else {
836                    return plan_err!(
837                        "For function '{function_name}' {valid_type} and {t} are not coercible to a common numeric type"
838                    );
839                }
840            }
841
842            let logical_data_type: NativeType = valid_type.clone().into();
843            // Fallback to default type if we don't know which type to coerced to
844            // f64 is chosen since most of the math functions utilize Signature::numeric,
845            // and their default type is double precision
846            if logical_data_type == NativeType::Null {
847                valid_type = DataType::Float64;
848            } else if !logical_data_type.is_numeric() {
849                return plan_err!(
850                    "Function '{function_name}' expects Numeric but received {logical_data_type}"
851                );
852            }
853
854            vec![vec![valid_type; *number]]
855        }
856        TypeSignature::Comparable(num) => {
857            function_length_check(function_name, current_types.len(), *num)?;
858            let mut target_type = current_types[0].to_owned();
859            for data_type in current_types.iter().skip(1) {
860                if let Some(dt) = comparison_coercion(&target_type, data_type) {
861                    target_type = dt;
862                } else {
863                    return plan_err!(
864                        "For function '{function_name}' {target_type} and {data_type} is not comparable"
865                    );
866                }
867            }
868            // Convert null to String type.
869            if target_type.is_null() {
870                vec![vec![DataType::Utf8View; *num]]
871            } else {
872                vec![vec![target_type; *num]]
873            }
874        }
875        TypeSignature::Coercible(param_types) => {
876            function_length_check(function_name, current_types.len(), param_types.len())?;
877
878            fn coercion_value_type<'a>(
879                current_type: &'a DataType,
880                desired_type: &TypeSignatureClass,
881            ) -> &'a DataType {
882                if matches!(desired_type, TypeSignatureClass::Any) {
883                    return current_type;
884                }
885
886                match current_type {
887                    DataType::Dictionary(_, value_type) => {
888                        coercion_value_type(value_type, desired_type)
889                    }
890                    _ => current_type,
891                }
892            }
893
894            fn preserve_encoding(
895                current_type: &DataType,
896                casted_type: DataType,
897                desired_type: &TypeSignatureClass,
898                encoding_preservation: EncodingPreservation,
899            ) -> DataType {
900                if matches!(desired_type, TypeSignatureClass::Any) {
901                    return casted_type;
902                }
903
904                match current_type {
905                    DataType::Dictionary(key_type, value_type) => {
906                        let casted_type = preserve_encoding(
907                            value_type,
908                            casted_type,
909                            desired_type,
910                            encoding_preservation,
911                        );
912                        if encoding_preservation.preserve_dictionary() {
913                            DataType::Dictionary(key_type.clone(), Box::new(casted_type))
914                        } else {
915                            casted_type
916                        }
917                    }
918                    _ => casted_type,
919                }
920            }
921
922            let mut new_types = Vec::with_capacity(current_types.len());
923            for (current_type, param) in current_types.iter().zip(param_types.iter()) {
924                let current_native_type: NativeType = current_type.into();
925                let encoding_preservation = param.encoding_preservation();
926                let coercion_value_type =
927                    coercion_value_type(current_type, param.desired_type());
928
929                if param
930                    .desired_type()
931                    .matches_native_type(&current_native_type)
932                {
933                    let casted_type = param
934                        .desired_type()
935                        .default_casted_type(&current_native_type, coercion_value_type)?;
936
937                    new_types.push(preserve_encoding(
938                        current_type,
939                        casted_type,
940                        param.desired_type(),
941                        encoding_preservation,
942                    ));
943                } else if param
944                    .allowed_source_types()
945                    .iter()
946                    .any(|t| t.matches_native_type(&current_native_type))
947                {
948                    // If the condition is met which means `implicit coercion`` is provided so we can safely unwrap
949                    let default_casted_type = param.default_casted_type().unwrap();
950                    let casted_type =
951                        default_casted_type.default_cast_for(coercion_value_type)?;
952                    new_types.push(preserve_encoding(
953                        current_type,
954                        casted_type,
955                        param.desired_type(),
956                        encoding_preservation,
957                    ));
958                } else {
959                    let hint = if matches!(current_native_type, NativeType::Binary) {
960                        "\n\nHint: Binary types are not automatically coerced to String. Use CAST(column AS VARCHAR) to convert Binary data to String."
961                    } else {
962                        ""
963                    };
964                    return plan_err!(
965                        "Function '{function_name}' requires {}, but received {} (DataType: {}).{hint}",
966                        param.desired_type(),
967                        current_native_type,
968                        current_type
969                    );
970                }
971            }
972
973            vec![new_types]
974        }
975        TypeSignature::Uniform(number, valid_types) => {
976            if *number == 0 {
977                return plan_err!(
978                    "The function '{function_name}' expected at least one argument"
979                );
980            }
981
982            valid_types
983                .iter()
984                .map(|valid_type| vec![valid_type.clone(); *number])
985                .collect()
986        }
987        TypeSignature::UserDefined => {
988            return internal_err!(
989                "Function '{function_name}' user-defined signature should be handled by function-specific coerce_types"
990            );
991        }
992        TypeSignature::VariadicAny => {
993            if current_types.is_empty() {
994                return plan_err!(
995                    "Function '{function_name}' expected at least one argument but received 0"
996                );
997            }
998            vec![current_types.to_vec()]
999        }
1000        TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
1001        TypeSignature::ArraySignature(function_signature) => match function_signature {
1002            ArrayFunctionSignature::Array {
1003                arguments,
1004                array_coercion,
1005            } => array_valid_types(
1006                function_name,
1007                current_types,
1008                arguments,
1009                array_coercion.as_ref(),
1010            )?,
1011            ArrayFunctionSignature::RecursiveArray => {
1012                if current_types.len() != 1 {
1013                    return Ok(vec![vec![]]);
1014                }
1015                recursive_array(&current_types[0])
1016                    .map_or_else(|| vec![vec![]], |array_type| vec![vec![array_type]])
1017            }
1018            ArrayFunctionSignature::MapArray => {
1019                if current_types.len() != 1 {
1020                    return Ok(vec![vec![]]);
1021                }
1022
1023                match &current_types[0] {
1024                    DataType::Map(_, _) => vec![vec![current_types[0].clone()]],
1025                    _ => vec![vec![]],
1026                }
1027            }
1028        },
1029        TypeSignature::Nullary => {
1030            if !current_types.is_empty() {
1031                return plan_err!(
1032                    "The function '{function_name}' expected zero argument but received {}",
1033                    current_types.len()
1034                );
1035            }
1036            vec![vec![]]
1037        }
1038        TypeSignature::Any(number) => {
1039            if current_types.is_empty() {
1040                return plan_err!(
1041                    "The function '{function_name}' expected at least one argument but received 0"
1042                );
1043            }
1044
1045            if current_types.len() != *number {
1046                return plan_err!(
1047                    "The function '{function_name}' expected {number} arguments but received {}",
1048                    current_types.len()
1049                );
1050            }
1051            vec![current_types.to_vec()]
1052        }
1053        TypeSignature::OneOf(types) => types
1054            .iter()
1055            .filter_map(|t| get_valid_types(function_name, t, current_types).ok())
1056            .flatten()
1057            .collect::<Vec<_>>(),
1058    };
1059
1060    Ok(valid_types)
1061}
1062
1063/// Try to coerce the current argument types to match the given `valid_types`.
1064///
1065/// For example, if a function `func` accepts arguments of  `(int64, int64)`,
1066/// but was called with `(int32, int64)`, this function could match the
1067/// valid_types by coercing the first argument to `int64`, and would return
1068/// `Some([int64, int64])`.
1069fn maybe_data_types(
1070    valid_types: &[DataType],
1071    current_types: &[DataType],
1072) -> Option<Vec<DataType>> {
1073    if valid_types.len() != current_types.len() {
1074        return None;
1075    }
1076
1077    let mut new_type = Vec::with_capacity(valid_types.len());
1078    for (i, valid_type) in valid_types.iter().enumerate() {
1079        let current_type = &current_types[i];
1080
1081        if current_type == valid_type {
1082            new_type.push(current_type.clone())
1083        } else {
1084            // attempt to coerce.
1085            // TODO: Replace with `can_cast_types` after failing cases are resolved
1086            // (they need new signature that returns exactly valid types instead of list of possible valid types).
1087            let coerced_type = coerced_from(valid_type, current_type)?;
1088            new_type.push(coerced_type)
1089        }
1090    }
1091    Some(new_type)
1092}
1093
1094/// Check if the current argument types can be coerced to match the given `valid_types`
1095/// unlike `maybe_data_types`, this function does not coerce the types.
1096/// TODO: I think this function should replace `maybe_data_types` after signature are well-supported.
1097fn maybe_data_types_without_coercion(
1098    valid_types: &[DataType],
1099    current_types: &[DataType],
1100) -> Option<Vec<DataType>> {
1101    if valid_types.len() != current_types.len() {
1102        return None;
1103    }
1104
1105    let mut new_type = Vec::with_capacity(valid_types.len());
1106    for (i, valid_type) in valid_types.iter().enumerate() {
1107        let current_type = &current_types[i];
1108
1109        if current_type == valid_type {
1110            new_type.push(current_type.clone())
1111        } else if can_cast_types(current_type, valid_type) {
1112            // validate the valid type is castable from the current type
1113            new_type.push(valid_type.clone())
1114        } else {
1115            return None;
1116        }
1117    }
1118    Some(new_type)
1119}
1120
1121/// Return true if a value of type `type_from` can be coerced
1122/// (losslessly converted) into a value of `type_to`
1123///
1124/// See the module level documentation for more detail on coercion.
1125#[deprecated(since = "53.0.0", note = "Unused internal function")]
1126pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
1127    if type_into == type_from {
1128        return true;
1129    }
1130    if let Some(coerced) = coerced_from(type_into, type_from) {
1131        return coerced == *type_into;
1132    }
1133    false
1134}
1135
1136/// Find the coerced type for the given `type_into` and `type_from`.
1137/// Returns `None` if coercion is not possible.
1138///
1139/// Expect uni-directional coercion, for example, i32 is coerced to i64, but i64 is not coerced to i32.
1140///
1141/// Unlike [crate::binary::comparison_coercion], the coerced type is usually `wider` for lossless conversion.
1142fn coerced_from<'a>(
1143    type_into: &'a DataType,
1144    type_from: &'a DataType,
1145) -> Option<DataType> {
1146    use self::DataType::*;
1147
1148    // match Dictionary first
1149    match (type_into, type_from) {
1150        // coerced dictionary first
1151        (_, Dictionary(_, value_type))
1152            if coerced_from(type_into, value_type).is_some() =>
1153        {
1154            Some(type_into.clone())
1155        }
1156        (Dictionary(_, value_type), _)
1157            if coerced_from(value_type, type_from).is_some() =>
1158        {
1159            Some(type_into.clone())
1160        }
1161        // coerced into type_into
1162        (Int8, Null | Int8) => Some(type_into.clone()),
1163        (Int16, Null | Int8 | Int16 | UInt8) => Some(type_into.clone()),
1164        (Int32, Null | Int8 | Int16 | Int32 | UInt8 | UInt16) => Some(type_into.clone()),
1165        (Int64, Null | Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32) => {
1166            Some(type_into.clone())
1167        }
1168        (UInt8, Null | UInt8) => Some(type_into.clone()),
1169        (UInt16, Null | UInt8 | UInt16) => Some(type_into.clone()),
1170        (UInt32, Null | UInt8 | UInt16 | UInt32) => Some(type_into.clone()),
1171        (UInt64, Null | UInt8 | UInt16 | UInt32 | UInt64) => Some(type_into.clone()),
1172        (Float16, Null | Int8 | Int16 | UInt8 | UInt16 | Float16) => {
1173            Some(type_into.clone())
1174        }
1175        (
1176            Float32,
1177            Null | Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64
1178            | Float16 | Float32,
1179        ) => Some(type_into.clone()),
1180        (
1181            Float64,
1182            Null
1183            | Int8
1184            | Int16
1185            | Int32
1186            | Int64
1187            | UInt8
1188            | UInt16
1189            | UInt32
1190            | UInt64
1191            | Float16
1192            | Float32
1193            | Float64
1194            | Decimal32(_, _)
1195            | Decimal64(_, _)
1196            | Decimal128(_, _)
1197            | Decimal256(_, _),
1198        ) => Some(type_into.clone()),
1199        (
1200            Timestamp(TimeUnit::Nanosecond, None),
1201            Null | Timestamp(_, None) | Date32 | Date64 | Utf8 | LargeUtf8 | Utf8View,
1202        ) => Some(type_into.clone()),
1203        (Interval(_), Null | Utf8 | LargeUtf8 | Utf8View) => Some(type_into.clone()),
1204        // Any type can be coerced into strings
1205        (Utf8 | LargeUtf8 | Utf8View, _) => Some(type_into.clone()),
1206        // We can go into a BinaryView from a Binary or LargeBinary
1207        (BinaryView, Binary | LargeBinary | Null) => Some(type_into.clone()),
1208        (Null, _) if can_cast_types(type_from, type_into) => Some(type_into.clone()),
1209
1210        (List(_), FixedSizeList(_, _)) => Some(type_into.clone()),
1211
1212        // Only accept list and largelist with the same number of dimensions unless the type is Null.
1213        // List or LargeList with different dimensions should be handled in TypeSignature or other places before this
1214        (List(_) | LargeList(_) | ListView(_) | LargeListView(_), _)
1215            if base_type(type_from).is_null()
1216                || list_ndims(type_from) == list_ndims(type_into) =>
1217        {
1218            Some(type_into.clone())
1219        }
1220        // should be able to coerce wildcard fixed size list to non wildcard fixed size list
1221        (
1222            FixedSizeList(f_into, FIXED_SIZE_LIST_WILDCARD),
1223            FixedSizeList(f_from, size_from),
1224        ) => match coerced_from(f_into.data_type(), f_from.data_type()) {
1225            Some(data_type) if &data_type != f_into.data_type() => {
1226                let new_field =
1227                    Arc::new(f_into.as_ref().clone().with_data_type(data_type));
1228                Some(FixedSizeList(new_field, *size_from))
1229            }
1230            Some(_) => Some(FixedSizeList(Arc::clone(f_into), *size_from)),
1231            _ => None,
1232        },
1233        (Timestamp(unit, Some(tz)), _) if tz.as_ref() == TIMEZONE_WILDCARD => {
1234            match type_from {
1235                Timestamp(_, Some(from_tz)) => {
1236                    Some(Timestamp(*unit, Some(Arc::clone(from_tz))))
1237                }
1238                Null | Date32 | Utf8 | LargeUtf8 | Timestamp(_, None) => {
1239                    // In the absence of any other information assume the time zone is "+00" (UTC).
1240                    Some(Timestamp(*unit, Some("+00".into())))
1241                }
1242                _ => None,
1243            }
1244        }
1245        (Timestamp(_, Some(_)), Null | Timestamp(_, _) | Date32 | Utf8 | LargeUtf8) => {
1246            Some(type_into.clone())
1247        }
1248        // Null can be coerced to any target type, provided the cast is valid.
1249        // This mirrors null_coercion() in binary comparison coercion
1250        // (expr-common/src/type_coercion/binary.rs) and is the symmetric
1251        // counterpart of the (Null, _) arm above. Without this, untyped
1252        // placeholders ($1, $foo) inside function calls fail signature matching
1253        // because their Null type doesn't match any Exact(...) variant.
1254        (_, Null) if can_cast_types(type_from, type_into) => Some(type_into.clone()),
1255        _ => None,
1256    }
1257}
1258
1259#[cfg(test)]
1260mod tests {
1261    use crate::{
1262        HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature,
1263        HigherOrderUDFImpl, Volatility,
1264    };
1265
1266    use super::*;
1267    use arrow::datatypes::IntervalUnit;
1268    use datafusion_common::{
1269        assert_contains,
1270        types::{logical_binary, logical_int64, logical_string},
1271    };
1272    use datafusion_expr_common::{
1273        columnar_value::ColumnarValue,
1274        signature::{Coercion, EncodingPreservation, TypeSignatureClass},
1275    };
1276
1277    #[test]
1278    fn test_string_conversion() {
1279        let cases = vec![
1280            (DataType::Utf8View, DataType::Utf8),
1281            (DataType::Utf8View, DataType::LargeUtf8),
1282            (DataType::Utf8View, DataType::Null),
1283        ];
1284
1285        for case in cases {
1286            assert_eq!(coerced_from(&case.0, &case.1), Some(case.0));
1287        }
1288    }
1289
1290    #[test]
1291    fn test_binary_conversion() {
1292        let cases = vec![
1293            (DataType::BinaryView, DataType::Binary),
1294            (DataType::BinaryView, DataType::LargeBinary),
1295            (DataType::BinaryView, DataType::Null),
1296        ];
1297
1298        for case in cases {
1299            assert_eq!(coerced_from(&case.0, &case.1), Some(case.0));
1300        }
1301    }
1302
1303    #[test]
1304    fn test_coerced_from_null() {
1305        // Null should coerce to Interval (the motivating case)
1306        assert_eq!(
1307            coerced_from(
1308                &DataType::Interval(IntervalUnit::MonthDayNano),
1309                &DataType::Null
1310            ),
1311            Some(DataType::Interval(IntervalUnit::MonthDayNano))
1312        );
1313
1314        // Null should coerce to Date32
1315        assert_eq!(
1316            coerced_from(&DataType::Date32, &DataType::Null),
1317            Some(DataType::Date32)
1318        );
1319
1320        // Null should coerce to Timestamp with timezone
1321        assert_eq!(
1322            coerced_from(
1323                &DataType::Timestamp(TimeUnit::Microsecond, Some("+00".into())),
1324                &DataType::Null
1325            ),
1326            Some(DataType::Timestamp(
1327                TimeUnit::Microsecond,
1328                Some("+00".into())
1329            ))
1330        );
1331    }
1332
1333    #[test]
1334    fn test_maybe_data_types() {
1335        // this vec contains: arg1, arg2, expected result
1336        let cases = vec![
1337            // 2 entries, same values
1338            (
1339                vec![DataType::UInt8, DataType::UInt16],
1340                vec![DataType::UInt8, DataType::UInt16],
1341                Some(vec![DataType::UInt8, DataType::UInt16]),
1342            ),
1343            // 2 entries, can coerce values
1344            (
1345                vec![DataType::UInt16, DataType::UInt16],
1346                vec![DataType::UInt8, DataType::UInt16],
1347                Some(vec![DataType::UInt16, DataType::UInt16]),
1348            ),
1349            // 0 entries, all good
1350            (vec![], vec![], Some(vec![])),
1351            // 2 entries, can't coerce
1352            (
1353                vec![DataType::Boolean, DataType::UInt16],
1354                vec![DataType::UInt8, DataType::UInt16],
1355                None,
1356            ),
1357            // u32 -> u16 is possible
1358            (
1359                vec![DataType::Boolean, DataType::UInt32],
1360                vec![DataType::Boolean, DataType::UInt16],
1361                Some(vec![DataType::Boolean, DataType::UInt32]),
1362            ),
1363            // UTF8 -> Timestamp
1364            (
1365                vec![
1366                    DataType::Timestamp(TimeUnit::Nanosecond, None),
1367                    DataType::Timestamp(TimeUnit::Nanosecond, Some("+TZ".into())),
1368                    DataType::Timestamp(TimeUnit::Nanosecond, Some("+01".into())),
1369                ],
1370                vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
1371                Some(vec![
1372                    DataType::Timestamp(TimeUnit::Nanosecond, None),
1373                    DataType::Timestamp(TimeUnit::Nanosecond, Some("+00".into())),
1374                    DataType::Timestamp(TimeUnit::Nanosecond, Some("+01".into())),
1375                ]),
1376            ),
1377        ];
1378
1379        for case in cases {
1380            assert_eq!(maybe_data_types(&case.0, &case.1), case.2)
1381        }
1382    }
1383
1384    #[test]
1385    fn test_get_valid_types_numeric() -> Result<()> {
1386        let get_valid_types_flatten =
1387            |function_name: &str,
1388             signature: &TypeSignature,
1389             current_types: &[DataType]| {
1390                get_valid_types(function_name, signature, current_types)
1391                    .unwrap()
1392                    .into_iter()
1393                    .flatten()
1394                    .collect::<Vec<_>>()
1395            };
1396
1397        // Trivial case.
1398        let got = get_valid_types_flatten(
1399            "test",
1400            &TypeSignature::Numeric(1),
1401            &[DataType::Int32],
1402        );
1403        assert_eq!(got, [DataType::Int32]);
1404
1405        // Args are coerced into a common numeric type.
1406        let got = get_valid_types_flatten(
1407            "test",
1408            &TypeSignature::Numeric(2),
1409            &[DataType::Int32, DataType::Int64],
1410        );
1411        assert_eq!(got, [DataType::Int64, DataType::Int64]);
1412
1413        // Args are coerced into a common numeric type, specifically, int would be coerced to float.
1414        let got = get_valid_types_flatten(
1415            "test",
1416            &TypeSignature::Numeric(3),
1417            &[DataType::Int32, DataType::Int64, DataType::Float64],
1418        );
1419        assert_eq!(
1420            got,
1421            [DataType::Float64, DataType::Float64, DataType::Float64]
1422        );
1423
1424        // Cannot coerce args to a common numeric type.
1425        let got = get_valid_types(
1426            "test",
1427            &TypeSignature::Numeric(2),
1428            &[DataType::Int32, DataType::Utf8],
1429        )
1430        .unwrap_err();
1431        assert_contains!(
1432            got.to_string(),
1433            "Function 'test' expects Numeric but received String"
1434        );
1435
1436        // Fallbacks to float64 if the arg is of type null.
1437        let got = get_valid_types_flatten(
1438            "test",
1439            &TypeSignature::Numeric(1),
1440            &[DataType::Null],
1441        );
1442        assert_eq!(got, [DataType::Float64]);
1443
1444        // Rejects non-numeric arg.
1445        let got = get_valid_types(
1446            "test",
1447            &TypeSignature::Numeric(1),
1448            &[DataType::Timestamp(TimeUnit::Second, None)],
1449        )
1450        .unwrap_err();
1451        assert_contains!(
1452            got.to_string(),
1453            "Function 'test' expects Numeric but received Timestamp(s)"
1454        );
1455
1456        Ok(())
1457    }
1458
1459    #[test]
1460    fn test_get_valid_types_one_of() -> Result<()> {
1461        let signature =
1462            TypeSignature::OneOf(vec![TypeSignature::Any(1), TypeSignature::Any(2)]);
1463
1464        let invalid_types = get_valid_types(
1465            "test",
1466            &signature,
1467            &[DataType::Int32, DataType::Int32, DataType::Int32],
1468        )?;
1469        assert_eq!(invalid_types.len(), 0);
1470
1471        let args = vec![DataType::Int32, DataType::Int32];
1472        let valid_types = get_valid_types("test", &signature, &args)?;
1473        assert_eq!(valid_types.len(), 1);
1474        assert_eq!(valid_types[0], args);
1475
1476        let args = vec![DataType::Int32];
1477        let valid_types = get_valid_types("test", &signature, &args)?;
1478        assert_eq!(valid_types.len(), 1);
1479        assert_eq!(valid_types[0], args);
1480
1481        Ok(())
1482    }
1483
1484    #[test]
1485    fn test_get_valid_types_length_check() -> Result<()> {
1486        let signature = TypeSignature::Numeric(1);
1487
1488        let err = get_valid_types("test", &signature, &[]).unwrap_err();
1489        assert_contains!(
1490            err.to_string(),
1491            "Function 'test' expects 1 arguments but received 0"
1492        );
1493
1494        let err = get_valid_types(
1495            "test",
1496            &signature,
1497            &[DataType::Int32, DataType::Int32, DataType::Int32],
1498        )
1499        .unwrap_err();
1500        assert_contains!(
1501            err.to_string(),
1502            "Function 'test' expects 1 arguments but received 3"
1503        );
1504
1505        Ok(())
1506    }
1507
1508    struct MockUdf(Signature);
1509
1510    impl UDFCoercionExt for MockUdf {
1511        fn name(&self) -> &str {
1512            "test"
1513        }
1514        fn signature(&self) -> &Signature {
1515            &self.0
1516        }
1517        fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
1518            unimplemented!()
1519        }
1520    }
1521
1522    #[test]
1523    fn test_fixed_list_wildcard_coerce() -> Result<()> {
1524        let inner = Arc::new(Field::new_list_field(DataType::Int32, false));
1525        // able to coerce for any size
1526        let current_fields = vec![Arc::new(Field::new(
1527            "t",
1528            DataType::FixedSizeList(Arc::clone(&inner), 2),
1529            true,
1530        ))];
1531
1532        let signature = Signature::exact(
1533            vec![DataType::FixedSizeList(
1534                Arc::clone(&inner),
1535                FIXED_SIZE_LIST_WILDCARD,
1536            )],
1537            Volatility::Stable,
1538        );
1539
1540        let coerced_fields = fields_with_udf(&current_fields, &MockUdf(signature))?;
1541        assert_eq!(coerced_fields, current_fields);
1542
1543        // make sure it can't coerce to a different size
1544        let signature = Signature::exact(
1545            vec![DataType::FixedSizeList(Arc::clone(&inner), 3)],
1546            Volatility::Stable,
1547        );
1548        let coerced_fields = fields_with_udf(&current_fields, &MockUdf(signature));
1549        assert!(coerced_fields.is_err());
1550
1551        // make sure it works with the same type.
1552        let signature = Signature::exact(
1553            vec![DataType::FixedSizeList(Arc::clone(&inner), 2)],
1554            Volatility::Stable,
1555        );
1556        let coerced_fields =
1557            fields_with_udf(&current_fields, &MockUdf(signature)).unwrap();
1558        assert_eq!(coerced_fields, current_fields);
1559
1560        Ok(())
1561    }
1562
1563    #[test]
1564    fn test_nested_wildcard_fixed_size_lists() -> Result<()> {
1565        let type_into = DataType::FixedSizeList(
1566            Arc::new(Field::new_list_field(
1567                DataType::FixedSizeList(
1568                    Arc::new(Field::new_list_field(DataType::Int32, false)),
1569                    FIXED_SIZE_LIST_WILDCARD,
1570                ),
1571                false,
1572            )),
1573            FIXED_SIZE_LIST_WILDCARD,
1574        );
1575
1576        let type_from = DataType::FixedSizeList(
1577            Arc::new(Field::new_list_field(
1578                DataType::FixedSizeList(
1579                    Arc::new(Field::new_list_field(DataType::Int8, false)),
1580                    4,
1581                ),
1582                false,
1583            )),
1584            3,
1585        );
1586
1587        assert_eq!(
1588            coerced_from(&type_into, &type_from),
1589            Some(DataType::FixedSizeList(
1590                Arc::new(Field::new_list_field(
1591                    DataType::FixedSizeList(
1592                        Arc::new(Field::new_list_field(DataType::Int32, false)),
1593                        4,
1594                    ),
1595                    false,
1596                )),
1597                3,
1598            ))
1599        );
1600
1601        Ok(())
1602    }
1603
1604    #[test]
1605    fn test_coerced_from_dictionary() {
1606        let type_into =
1607            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::UInt32));
1608        let type_from = DataType::Int64;
1609        assert_eq!(coerced_from(&type_into, &type_from), None);
1610
1611        let type_from =
1612            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::UInt32));
1613        let type_into = DataType::Int64;
1614        assert_eq!(
1615            coerced_from(&type_into, &type_from),
1616            Some(type_into.clone())
1617        );
1618    }
1619
1620    #[test]
1621    fn test_get_valid_types_array_and_array() -> Result<()> {
1622        let function = "array_and_array";
1623        let signature = Signature::arrays(
1624            2,
1625            Some(ListCoercion::FixedSizedListToList),
1626            Volatility::Immutable,
1627        );
1628
1629        let data_types = vec![
1630            DataType::new_list(DataType::Int32, true),
1631            DataType::new_large_list(DataType::Float64, true),
1632        ];
1633        assert_eq!(
1634            get_valid_types(function, &signature.type_signature, &data_types)?,
1635            vec![vec![
1636                DataType::new_large_list(DataType::Float64, true),
1637                DataType::new_large_list(DataType::Float64, true),
1638            ]]
1639        );
1640
1641        let data_types = vec![
1642            DataType::new_fixed_size_list(DataType::Int64, 3, true),
1643            DataType::new_fixed_size_list(DataType::Int32, 5, true),
1644        ];
1645        assert_eq!(
1646            get_valid_types(function, &signature.type_signature, &data_types)?,
1647            vec![vec![
1648                DataType::new_list(DataType::Int64, true),
1649                DataType::new_list(DataType::Int64, true),
1650            ]]
1651        );
1652
1653        let data_types = vec![
1654            DataType::new_fixed_size_list(DataType::Null, 3, true),
1655            DataType::new_large_list(DataType::Utf8, true),
1656        ];
1657        assert_eq!(
1658            get_valid_types(function, &signature.type_signature, &data_types)?,
1659            vec![vec![
1660                DataType::new_large_list(DataType::Utf8, true),
1661                DataType::new_large_list(DataType::Utf8, true),
1662            ]]
1663        );
1664
1665        let data_types = vec![
1666            DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1667            DataType::new_list(DataType::Int32, true),
1668        ];
1669        assert_eq!(
1670            get_valid_types(function, &signature.type_signature, &data_types)?,
1671            vec![vec![
1672                DataType::new_list(DataType::Int32, true),
1673                DataType::new_list(DataType::Int32, true),
1674            ]]
1675        );
1676
1677        let data_types = vec![
1678            DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1679            DataType::new_list(DataType::Int32, true),
1680        ];
1681        assert_eq!(
1682            get_valid_types(function, &signature.type_signature, &data_types)?,
1683            vec![vec![
1684                DataType::new_large_list(DataType::Int32, true),
1685                DataType::new_large_list(DataType::Int32, true),
1686            ]]
1687        );
1688
1689        let data_types = vec![
1690            DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1691            DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1692        ];
1693        assert_eq!(
1694            get_valid_types(function, &signature.type_signature, &data_types)?,
1695            vec![vec![
1696                DataType::new_list(DataType::Int32, true),
1697                DataType::new_list(DataType::Int32, true),
1698            ]]
1699        );
1700
1701        let data_types = vec![
1702            DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1703            DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1704        ];
1705        assert_eq!(
1706            get_valid_types(function, &signature.type_signature, &data_types)?,
1707            vec![vec![
1708                DataType::new_large_list(DataType::Int32, true),
1709                DataType::new_large_list(DataType::Int32, true),
1710            ]]
1711        );
1712
1713        Ok(())
1714    }
1715
1716    #[test]
1717    fn test_get_valid_types_array_and_element() -> Result<()> {
1718        let function = "array_and_element";
1719        let signature = Signature::array_and_element(Volatility::Immutable);
1720
1721        let data_types =
1722            vec![DataType::new_list(DataType::Int32, true), DataType::Float64];
1723        assert_eq!(
1724            get_valid_types(function, &signature.type_signature, &data_types)?,
1725            vec![vec![
1726                DataType::new_list(DataType::Float64, true),
1727                DataType::Float64,
1728            ]]
1729        );
1730
1731        let data_types = vec![
1732            DataType::new_large_list(DataType::Int32, true),
1733            DataType::Null,
1734        ];
1735        assert_eq!(
1736            get_valid_types(function, &signature.type_signature, &data_types)?,
1737            vec![vec![
1738                DataType::new_large_list(DataType::Int32, true),
1739                DataType::Int32,
1740            ]]
1741        );
1742
1743        let data_types = vec![
1744            DataType::new_fixed_size_list(DataType::Null, 3, true),
1745            DataType::Utf8,
1746        ];
1747        assert_eq!(
1748            get_valid_types(function, &signature.type_signature, &data_types)?,
1749            vec![vec![
1750                DataType::new_list(DataType::Utf8, true),
1751                DataType::Utf8,
1752            ]]
1753        );
1754
1755        Ok(())
1756    }
1757
1758    #[test]
1759    fn test_get_valid_types_array_and_index_preserves_list_field_name() -> Result<()> {
1760        let struct_fields = vec![
1761            Field::new("id", DataType::Utf8, true),
1762            Field::new("prim", DataType::Boolean, true),
1763        ];
1764        let current_type = DataType::List(Arc::new(Field::new(
1765            "element",
1766            DataType::Struct(struct_fields.into()),
1767            true,
1768        )));
1769        let signature = Signature::array_and_index(Volatility::Immutable);
1770
1771        assert_eq!(
1772            get_valid_types(
1773                "array_element",
1774                &signature.type_signature,
1775                &[current_type.clone(), DataType::Int64],
1776            )?,
1777            vec![vec![current_type, DataType::Int64]]
1778        );
1779
1780        Ok(())
1781    }
1782
1783    #[test]
1784    fn test_get_valid_types_element_and_array() -> Result<()> {
1785        let function = "element_and_array";
1786        let signature = Signature::element_and_array(Volatility::Immutable);
1787
1788        let data_types = vec![
1789            DataType::new_large_list(DataType::Null, false),
1790            DataType::new_list(DataType::new_list(DataType::Int64, true), true),
1791        ];
1792        assert_eq!(
1793            get_valid_types(function, &signature.type_signature, &data_types)?,
1794            vec![vec![
1795                DataType::new_large_list(DataType::Int64, true),
1796                DataType::new_list(DataType::new_large_list(DataType::Int64, true), true),
1797            ]]
1798        );
1799
1800        Ok(())
1801    }
1802
1803    #[test]
1804    fn test_coercible_nulls() -> Result<()> {
1805        fn null_input(coercion: Coercion) -> Result<Vec<DataType>> {
1806            fields_with_udf(
1807                &[Field::new("field", DataType::Null, true).into()],
1808                &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1809            )
1810            .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1811        }
1812
1813        // Casts Null to Int64 if we use TypeSignatureClass::Native
1814        let output = null_input(Coercion::new_exact(TypeSignatureClass::Native(
1815            logical_int64(),
1816        )))?;
1817        assert_eq!(vec![DataType::Int64], output);
1818
1819        let output = null_input(Coercion::new_implicit(
1820            TypeSignatureClass::Native(logical_int64()),
1821            vec![],
1822            NativeType::Int64,
1823        ))?;
1824        assert_eq!(vec![DataType::Int64], output);
1825
1826        // Null gets passed through if we use TypeSignatureClass apart from Native
1827        let output = null_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
1828        assert_eq!(vec![DataType::Null], output);
1829
1830        let output = null_input(Coercion::new_implicit(
1831            TypeSignatureClass::Integer,
1832            vec![],
1833            NativeType::Int64,
1834        ))?;
1835        assert_eq!(vec![DataType::Null], output);
1836
1837        Ok(())
1838    }
1839
1840    #[test]
1841    fn test_coercible_dictionary() -> Result<()> {
1842        let dictionary =
1843            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64));
1844        fn dictionary_input(coercion: Coercion) -> Result<Vec<DataType>> {
1845            fields_with_udf(
1846                &[Field::new(
1847                    "field",
1848                    DataType::Dictionary(
1849                        Box::new(DataType::Int8),
1850                        Box::new(DataType::Int64),
1851                    ),
1852                    true,
1853                )
1854                .into()],
1855                &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1856            )
1857            .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1858        }
1859
1860        // Casts Dictionary to Int64 if we use TypeSignatureClass::Native
1861        let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Native(
1862            logical_int64(),
1863        )))?;
1864        assert_eq!(vec![DataType::Int64], output);
1865
1866        let output = dictionary_input(Coercion::new_implicit(
1867            TypeSignatureClass::Native(logical_int64()),
1868            vec![],
1869            NativeType::Int64,
1870        ))?;
1871        assert_eq!(vec![DataType::Int64], output);
1872
1873        // Any always preserves the original physical type
1874        let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Any))?;
1875        assert_eq!(vec![dictionary.clone()], output);
1876
1877        let output = dictionary_input(
1878            Coercion::new_exact(TypeSignatureClass::Any)
1879                .with_encoding_preservation(EncodingPreservation::dictionary()),
1880        )?;
1881        assert_eq!(vec![dictionary.clone()], output);
1882
1883        // Typed non-Native classes materialize dictionaries by default
1884        let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
1885        assert_eq!(vec![DataType::Int64], output);
1886
1887        let output = dictionary_input(Coercion::new_implicit(
1888            TypeSignatureClass::Integer,
1889            vec![],
1890            NativeType::Int64,
1891        ))?;
1892        assert_eq!(vec![DataType::Int64], output);
1893
1894        // Typed non-Native classes preserve dictionaries only when requested
1895        let output = dictionary_input(
1896            Coercion::new_exact(TypeSignatureClass::Integer)
1897                .with_encoding_preservation(EncodingPreservation::dictionary()),
1898        )?;
1899        assert_eq!(vec![dictionary], output);
1900
1901        Ok(())
1902    }
1903
1904    #[test]
1905    fn test_coercible_dictionary_preserves_encoding() -> Result<()> {
1906        fn dictionary_input(
1907            value_type: DataType,
1908            coercion: Coercion,
1909        ) -> Result<Vec<DataType>> {
1910            fields_with_udf(
1911                &[Field::new(
1912                    "field",
1913                    DataType::Dictionary(Box::new(DataType::Int8), Box::new(value_type)),
1914                    true,
1915                )
1916                .into()],
1917                &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1918            )
1919            .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1920        }
1921
1922        let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
1923            .with_encoding_preservation(EncodingPreservation::dictionary());
1924
1925        assert_eq!(
1926            dictionary_input(DataType::LargeUtf8, coercion.clone())?,
1927            vec![DataType::Dictionary(
1928                Box::new(DataType::Int8),
1929                Box::new(DataType::LargeUtf8),
1930            )]
1931        );
1932        assert_eq!(
1933            dictionary_input(
1934                DataType::BinaryView,
1935                Coercion::new_implicit(
1936                    TypeSignatureClass::Native(logical_string()),
1937                    vec![TypeSignatureClass::Native(logical_binary())],
1938                    NativeType::String,
1939                )
1940                .with_encoding_preservation(EncodingPreservation::dictionary()),
1941            )?,
1942            vec![DataType::Dictionary(
1943                Box::new(DataType::Int8),
1944                Box::new(DataType::Utf8View),
1945            )]
1946        );
1947        // Contrast: without encoding_preservation, Native strips dictionary entirely
1948        assert_eq!(
1949            dictionary_input(
1950                DataType::Int32,
1951                Coercion::new_implicit(
1952                    TypeSignatureClass::Native(logical_int64()),
1953                    vec![TypeSignatureClass::Integer],
1954                    NativeType::Int64,
1955                ),
1956            )?,
1957            vec![DataType::Int64]
1958        );
1959        // With encoding_preservation, dictionary wrapper is preserved, value coerced
1960        assert_eq!(
1961            dictionary_input(
1962                DataType::Int32,
1963                Coercion::new_implicit(
1964                    TypeSignatureClass::Native(logical_int64()),
1965                    vec![TypeSignatureClass::Integer],
1966                    NativeType::Int64,
1967                )
1968                .with_encoding_preservation(EncodingPreservation::dictionary()),
1969            )?,
1970            vec![DataType::Dictionary(
1971                Box::new(DataType::Int8),
1972                Box::new(DataType::Int64),
1973            )]
1974        );
1975        // Without encoding_preservation, non-Native classes materialize dictionaries
1976        assert_eq!(
1977            dictionary_input(
1978                DataType::Int32,
1979                Coercion::new_implicit(
1980                    TypeSignatureClass::Integer,
1981                    vec![],
1982                    NativeType::Int64,
1983                ),
1984            )?,
1985            vec![DataType::Int32]
1986        );
1987        // With encoding_preservation, non-Native classes preserve dictionaries
1988        assert_eq!(
1989            dictionary_input(
1990                DataType::Int32,
1991                Coercion::new_implicit(
1992                    TypeSignatureClass::Integer,
1993                    vec![],
1994                    NativeType::Int64,
1995                )
1996                .with_encoding_preservation(EncodingPreservation::dictionary()),
1997            )?,
1998            vec![DataType::Dictionary(
1999                Box::new(DataType::Int8),
2000                Box::new(DataType::Int32),
2001            )]
2002        );
2003
2004        Ok(())
2005    }
2006
2007    #[test]
2008    fn test_coercible_nested_dictionary() -> Result<()> {
2009        let nested_dictionary = DataType::Dictionary(
2010            Box::new(DataType::Int8),
2011            Box::new(DataType::Dictionary(
2012                Box::new(DataType::Int16),
2013                Box::new(DataType::Int32),
2014            )),
2015        );
2016        let nested_dictionary_input = |coercion| -> Result<Vec<DataType>> {
2017            fields_with_udf(
2018                &[Field::new("field", nested_dictionary.clone(), true).into()],
2019                &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
2020            )
2021            .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
2022        };
2023
2024        // Without preservation, recursively unwrap dictionaries to the unchanged leaf.
2025        let output =
2026            nested_dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
2027        assert_eq!(vec![DataType::Int32], output);
2028
2029        // With preservation, restore the complete dictionary stack around the leaf.
2030        let output = nested_dictionary_input(
2031            Coercion::new_exact(TypeSignatureClass::Integer)
2032                .with_encoding_preservation(EncodingPreservation::dictionary()),
2033        )?;
2034        assert_eq!(vec![nested_dictionary.clone()], output);
2035
2036        let int64_coercion = || {
2037            Coercion::new_implicit(
2038                TypeSignatureClass::Native(logical_int64()),
2039                vec![TypeSignatureClass::Integer],
2040                NativeType::Int64,
2041            )
2042        };
2043
2044        // Without preservation, materialize the coerced leaf type.
2045        let output = nested_dictionary_input(int64_coercion())?;
2046        assert_eq!(vec![DataType::Int64], output);
2047
2048        // With preservation, restore the complete dictionary stack around the coerced leaf.
2049        let output = nested_dictionary_input(
2050            int64_coercion()
2051                .with_encoding_preservation(EncodingPreservation::dictionary()),
2052        )?;
2053        assert_eq!(
2054            vec![DataType::Dictionary(
2055                Box::new(DataType::Int8),
2056                Box::new(DataType::Dictionary(
2057                    Box::new(DataType::Int16),
2058                    Box::new(DataType::Int64),
2059                )),
2060            )],
2061            output
2062        );
2063
2064        Ok(())
2065    }
2066
2067    #[test]
2068    fn test_coercible_run_end_encoded() -> Result<()> {
2069        let run_end_encoded = DataType::RunEndEncoded(
2070            Field::new("run_ends", DataType::Int16, false).into(),
2071            Field::new("values", DataType::Int64, true).into(),
2072        );
2073        fn run_end_encoded_input(coercion: Coercion) -> Result<Vec<DataType>> {
2074            fields_with_udf(
2075                &[Field::new(
2076                    "field",
2077                    DataType::RunEndEncoded(
2078                        Field::new("run_ends", DataType::Int16, false).into(),
2079                        Field::new("values", DataType::Int64, true).into(),
2080                    ),
2081                    true,
2082                )
2083                .into()],
2084                &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
2085            )
2086            .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
2087        }
2088
2089        // Casts REE to Int64 if we use TypeSignatureClass::Native
2090        let output = run_end_encoded_input(Coercion::new_exact(
2091            TypeSignatureClass::Native(logical_int64()),
2092        ))?;
2093        assert_eq!(vec![DataType::Int64], output);
2094
2095        let output = run_end_encoded_input(Coercion::new_implicit(
2096            TypeSignatureClass::Native(logical_int64()),
2097            vec![],
2098            NativeType::Int64,
2099        ))?;
2100        assert_eq!(vec![DataType::Int64], output);
2101
2102        // REE gets passed through if we use TypeSignatureClass apart from Native
2103        let output =
2104            run_end_encoded_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
2105        assert_eq!(vec![run_end_encoded.clone()], output);
2106
2107        let output = run_end_encoded_input(Coercion::new_implicit(
2108            TypeSignatureClass::Integer,
2109            vec![],
2110            NativeType::Int64,
2111        ))?;
2112        assert_eq!(vec![run_end_encoded.clone()], output);
2113
2114        Ok(())
2115    }
2116
2117    #[test]
2118    fn test_get_valid_types_coercible_binary() -> Result<()> {
2119        let signature = Signature::coercible(
2120            vec![Coercion::new_exact(TypeSignatureClass::Native(
2121                logical_binary(),
2122            ))],
2123            Volatility::Immutable,
2124        );
2125
2126        // Binary types should stay their original selves
2127        for t in [
2128            DataType::Binary,
2129            DataType::BinaryView,
2130            DataType::LargeBinary,
2131        ] {
2132            assert_eq!(
2133                get_valid_types("", &signature.type_signature, std::slice::from_ref(&t))?,
2134                vec![vec![t]]
2135            );
2136        }
2137
2138        Ok(())
2139    }
2140
2141    #[test]
2142    fn test_get_valid_types_fixed_size_arrays() -> Result<()> {
2143        let function = "fixed_size_arrays";
2144        let signature = Signature::arrays(2, None, Volatility::Immutable);
2145
2146        let data_types = vec![
2147            DataType::new_fixed_size_list(DataType::Int64, 3, true),
2148            DataType::new_fixed_size_list(DataType::Int32, 5, true),
2149        ];
2150        assert_eq!(
2151            get_valid_types(function, &signature.type_signature, &data_types)?,
2152            vec![vec![
2153                DataType::new_fixed_size_list(DataType::Int64, 3, true),
2154                DataType::new_fixed_size_list(DataType::Int64, 5, true),
2155            ]]
2156        );
2157
2158        let data_types = vec![
2159            DataType::new_fixed_size_list(DataType::Int64, 3, true),
2160            DataType::new_list(DataType::Int32, true),
2161        ];
2162        assert_eq!(
2163            get_valid_types(function, &signature.type_signature, &data_types)?,
2164            vec![vec![
2165                DataType::new_list(DataType::Int64, true),
2166                DataType::new_list(DataType::Int64, true),
2167            ]]
2168        );
2169
2170        let data_types = vec![
2171            DataType::new_fixed_size_list(DataType::Utf8, 3, true),
2172            DataType::new_list(DataType::new_list(DataType::Int32, true), true),
2173        ];
2174        assert_eq!(
2175            get_valid_types(function, &signature.type_signature, &data_types)?,
2176            vec![vec![]]
2177        );
2178
2179        let data_types = vec![
2180            DataType::new_fixed_size_list(DataType::Int64, 3, false),
2181            DataType::new_list(DataType::Int32, false),
2182        ];
2183        assert_eq!(
2184            get_valid_types(function, &signature.type_signature, &data_types)?,
2185            vec![vec![
2186                DataType::new_list(DataType::Int64, false),
2187                DataType::new_list(DataType::Int64, false),
2188            ]]
2189        );
2190
2191        Ok(())
2192    }
2193
2194    #[derive(Debug, PartialEq, Eq, Hash)]
2195    struct MockHigherOrderUDF {
2196        signature: HigherOrderSignature,
2197        coerced_value_types: Vec<DataType>,
2198    }
2199
2200    impl HigherOrderUDFImpl for MockHigherOrderUDF {
2201        fn name(&self) -> &str {
2202            "mock_higher_order_function"
2203        }
2204
2205        fn signature(&self) -> &HigherOrderSignature {
2206            &self.signature
2207        }
2208
2209        fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
2210            if arg_types.len() != 1 {
2211                return plan_err!(
2212                    "mock_higher_order_function expects 1 value arguments, got {}",
2213                    arg_types.len()
2214                );
2215            }
2216            Ok(self.coerced_value_types.clone())
2217        }
2218
2219        fn coerce_values_for_lambdas(
2220            &self,
2221            fields: &[ValueOrLambda<DataType, DataType>],
2222        ) -> Result<Option<Vec<DataType>>> {
2223            // thoerical impl of array_reduce without finish
2224            let [
2225                ValueOrLambda::Value(list),
2226                ValueOrLambda::Value(_initial),
2227                ValueOrLambda::Lambda(merge),
2228            ] = fields
2229            else {
2230                unreachable!()
2231            };
2232
2233            Ok(Some(vec![list.clone(), merge.clone()]))
2234        }
2235
2236        fn lambda_parameters(
2237            &self,
2238            _step: usize,
2239            _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
2240        ) -> Result<crate::LambdaParametersProgress> {
2241            unimplemented!("mock_higher_order_function")
2242        }
2243
2244        fn return_field_from_args(
2245            &self,
2246            _args: HigherOrderReturnFieldArgs,
2247        ) -> Result<FieldRef> {
2248            unimplemented!("mock_higher_order_function")
2249        }
2250
2251        fn invoke_with_args(
2252            &self,
2253            _args: HigherOrderFunctionArgs,
2254        ) -> Result<ColumnarValue> {
2255            unimplemented!("mock_higher_order_function")
2256        }
2257    }
2258
2259    #[test]
2260    fn test_higher_order_function_user_defined_type_coercion() {
2261        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2262            signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2263            coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)],
2264        });
2265
2266        let new_fields = value_fields_with_higher_order_udf(
2267            &[
2268                ValueOrLambda::Value(Arc::new(Field::new_list(
2269                    "",
2270                    Field::new_list_field(DataType::Int32, false),
2271                    false,
2272                ))),
2273                ValueOrLambda::Lambda(()),
2274            ],
2275            &fun,
2276        )
2277        .unwrap();
2278
2279        // from List(Int32) to LargeList(Int32)
2280        assert_eq!(
2281            new_fields,
2282            vec![
2283                ValueOrLambda::Value(Arc::new(Field::new_large_list(
2284                    "",
2285                    Field::new_list_field(DataType::Int32, false),
2286                    false
2287                ))),
2288                ValueOrLambda::Lambda(()),
2289            ]
2290        )
2291    }
2292
2293    #[test]
2294    fn test_higher_order_function_coerce_values_for_lambdas() {
2295        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2296            signature: HigherOrderSignature::variadic_any(Volatility::Immutable),
2297            coerced_value_types: vec![],
2298        });
2299
2300        let new_fields = value_fields_with_higher_order_udf_and_lambdas(
2301            &[
2302                ValueOrLambda::Value(Arc::new(Field::new_list(
2303                    "",
2304                    Field::new_list_field(DataType::Float32, true),
2305                    true,
2306                ))),
2307                ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, true))),
2308                ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
2309            ],
2310            &fun,
2311        )
2312        .unwrap();
2313
2314        // second parameter from Int32 to Float32
2315        assert_eq!(
2316            new_fields,
2317            vec![
2318                ValueOrLambda::Value(Arc::new(Field::new_list(
2319                    "",
2320                    Field::new_list_field(DataType::Float32, true),
2321                    true,
2322                ))),
2323                ValueOrLambda::Value(Arc::new(Field::new("", DataType::Float32, true))),
2324                ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
2325            ]
2326        )
2327    }
2328
2329    #[test]
2330    fn test_higher_order_function_user_defined_type_coercion_bad_args() {
2331        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2332            signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2333            coerced_value_types: vec![DataType::Int32],
2334        });
2335
2336        let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
2337
2338        assert_contains!(
2339            err.to_string(),
2340            "mock_higher_order_function expects 1 value arguments, got 0"
2341        );
2342    }
2343
2344    #[test]
2345    fn test_higher_order_function_faulty_user_defined_type_coercion() {
2346        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2347            signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2348            coerced_value_types: vec![DataType::Int32, DataType::Int32],
2349        });
2350
2351        let err = value_fields_with_higher_order_udf::<()>(
2352            &[ValueOrLambda::Value(Arc::new(Field::new(
2353                "",
2354                DataType::Int32,
2355                false,
2356            )))],
2357            &fun,
2358        )
2359        .unwrap_err();
2360
2361        assert_contains!(
2362            err.to_string(),
2363            "mock_higher_order_function coerce_value_types should have returned 1 items but returned 2"
2364        );
2365    }
2366
2367    #[test]
2368    fn test_higher_order_function_any_signature() {
2369        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2370            signature: HigherOrderSignature::any(1, Volatility::Immutable),
2371            coerced_value_types: vec![],
2372        });
2373
2374        let new_fields =
2375            value_fields_with_higher_order_udf(&[ValueOrLambda::Lambda(())], &fun)
2376                .unwrap();
2377
2378        // no coercion, just number of args checked
2379        assert_eq!(new_fields, vec![ValueOrLambda::Lambda(())])
2380    }
2381
2382    #[test]
2383    fn test_higher_order_function_any_signature_bad_args() {
2384        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2385            signature: HigherOrderSignature::any(1, Volatility::Immutable),
2386            coerced_value_types: vec![],
2387        });
2388
2389        let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
2390
2391        assert_contains!(
2392            err.to_string(),
2393            "The function 'mock_higher_order_function' expected 1 arguments but received 0"
2394        );
2395    }
2396
2397    #[test]
2398    fn test_higher_order_function_exact_signature() {
2399        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2400            signature: HigherOrderSignature::exact(
2401                vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2402                Volatility::Immutable,
2403            ),
2404            coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)],
2405        });
2406
2407        let new_fields = value_fields_with_higher_order_udf(
2408            &[
2409                ValueOrLambda::Value(Arc::new(Field::new_list(
2410                    "",
2411                    Field::new_list_field(DataType::Int32, false),
2412                    false,
2413                ))),
2414                ValueOrLambda::Lambda(()),
2415            ],
2416            &fun,
2417        )
2418        .unwrap();
2419
2420        // type coercion applied: List(Int32) -> LargeList(Int32)
2421        assert_eq!(
2422            new_fields,
2423            vec![
2424                ValueOrLambda::Value(Arc::new(Field::new_large_list(
2425                    "",
2426                    Field::new_list_field(DataType::Int32, false),
2427                    false
2428                ))),
2429                ValueOrLambda::Lambda(()),
2430            ]
2431        )
2432    }
2433
2434    #[test]
2435    fn test_higher_order_function_exact_signature_wrong_value_count() {
2436        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2437            signature: HigherOrderSignature::exact(
2438                vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2439                Volatility::Immutable,
2440            ),
2441            coerced_value_types: vec![],
2442        });
2443
2444        let err = value_fields_with_higher_order_udf::<()>(
2445            &[ValueOrLambda::Lambda(()), ValueOrLambda::Lambda(())],
2446            &fun,
2447        )
2448        .unwrap_err();
2449
2450        assert_contains!(
2451            err.to_string(),
2452            "expected a value at position 0 but received a lambda"
2453        );
2454    }
2455
2456    #[test]
2457    fn test_higher_order_function_exact_signature_wrong_lambda_count() {
2458        let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2459            signature: HigherOrderSignature::exact(
2460                vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2461                Volatility::Immutable,
2462            ),
2463            coerced_value_types: vec![],
2464        });
2465
2466        let err = value_fields_with_higher_order_udf::<()>(
2467            &[
2468                ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, false))),
2469                ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, false))),
2470            ],
2471            &fun,
2472        )
2473        .unwrap_err();
2474
2475        assert_contains!(
2476            err.to_string(),
2477            "expected a lambda at position 1 but received a value"
2478        );
2479    }
2480}