datafusion_expr_common/
signature.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
18//! Function signatures: [`Volatility`], [`Signature`] and [`TypeSignature`]
19
20use std::fmt::Display;
21use std::hash::Hash;
22use std::sync::Arc;
23
24use crate::type_coercion::aggregates::NUMERICS;
25use arrow::datatypes::{
26    DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DECIMAL128_MAX_PRECISION, DataType,
27    Decimal128Type, DecimalType, Field, IntervalUnit, TimeUnit,
28};
29use datafusion_common::types::{LogicalType, LogicalTypeRef, NativeType};
30use datafusion_common::utils::ListCoercion;
31use datafusion_common::{Result, internal_err, plan_err};
32use indexmap::IndexSet;
33use itertools::Itertools;
34
35/// Constant that is used as a placeholder for any valid timezone.
36/// This is used where a function can accept a timestamp type with any
37/// valid timezone, it exists to avoid the need to enumerate all possible
38/// timezones. See [`TypeSignature`] for more details.
39///
40/// Type coercion always ensures that functions will be executed using
41/// timestamp arrays that have a valid time zone. Functions must never
42/// return results with this timezone.
43pub const TIMEZONE_WILDCARD: &str = "+TZ";
44
45/// Constant that is used as a placeholder for any valid fixed size list.
46/// This is used where a function can accept a fixed size list type with any
47/// valid length. It exists to avoid the need to enumerate all possible fixed size list lengths.
48pub const FIXED_SIZE_LIST_WILDCARD: i32 = i32::MIN;
49
50/// How a function's output changes with respect to a fixed input
51///
52/// The volatility of a function determines eligibility for certain
53/// optimizations. You should always define your function to have the strictest
54/// possible volatility to maximize performance and avoid unexpected
55/// results.
56#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
57pub enum Volatility {
58    /// Always returns the same output when given the same input.
59    ///
60    /// DataFusion will inline immutable functions during planning.
61    ///
62    /// For example, the `abs` function is immutable, so `abs(-1)` will be
63    /// evaluated and replaced  with `1` during planning rather than invoking
64    /// the function at runtime.
65    Immutable,
66    /// May return different values given the same input across different
67    /// queries but must return the same value for a given input within a query.
68    ///
69    /// For example, the `now()` function is stable, because the query `select
70    /// col1, now() from t1`, will return different results each time it is run,
71    /// but within the same query, the output of the `now()` function has the
72    /// same value for each output row.
73    ///
74    /// DataFusion will inline `Stable` functions when possible. For example,
75    /// `Stable` functions are inlined when planning a query for execution, but
76    /// not in View definitions or prepared statements.
77    Stable,
78    /// May change the return value from evaluation to evaluation.
79    ///
80    /// Multiple invocations of a volatile function may return different results
81    /// when used in the same query on different rows. An example of this is the
82    /// `random()` function.
83    ///
84    /// DataFusion can not evaluate such functions during planning or push these
85    /// predicates into scans. In the query `select col1, random() from t1`,
86    /// `random()` function will be evaluated for each output row, resulting in
87    /// a unique random value for each row.
88    Volatile,
89}
90
91/// Represents the arity (number of arguments) of a function signature
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum Arity {
94    /// Fixed number of arguments
95    Fixed(usize),
96    /// Variable number of arguments (e.g., Variadic, VariadicAny, UserDefined)
97    Variable,
98}
99
100/// The types of arguments for which a function has implementations.
101///
102/// [`TypeSignature`] **DOES NOT** define the types that a user query could call the
103/// function with. DataFusion will automatically coerce (cast) argument types to
104/// one of the supported function signatures, if possible.
105///
106/// # Overview
107/// Functions typically provide implementations for a small number of different
108/// argument [`DataType`]s, rather than all possible combinations. If a user
109/// calls a function with arguments that do not match any of the declared types,
110/// DataFusion will attempt to automatically coerce (add casts to) function
111/// arguments so they match the [`TypeSignature`]. See the [`type_coercion`] module
112/// for more details
113///
114/// # Example: Numeric Functions
115/// For example, a function like `cos` may only provide an implementation for
116/// [`DataType::Float64`]. When users call `cos` with a different argument type,
117/// such as `cos(int_column)`, and type coercion automatically adds a cast such
118/// as `cos(CAST int_column AS DOUBLE)` during planning.
119///
120/// [`type_coercion`]: crate::type_coercion
121///
122/// ## Example: Strings
123///
124/// There are several different string types in Arrow, such as
125/// [`DataType::Utf8`], [`DataType::LargeUtf8`], and [`DataType::Utf8View`].
126///
127/// Some functions may have specialized implementations for these types, while others
128/// may be able to handle only one of them. For example, a function that
129/// only works with [`DataType::Utf8View`] would have the following signature:
130///
131/// ```
132/// # use arrow::datatypes::DataType;
133/// # use datafusion_expr_common::signature::{TypeSignature};
134/// // Declares the function must be invoked with a single argument of type `Utf8View`.
135/// // if a user calls the function with `Utf8` or `LargeUtf8`, DataFusion will
136/// // automatically add a cast to `Utf8View` during planning.
137/// let type_signature = TypeSignature::Exact(vec![DataType::Utf8View]);
138/// ```
139///
140/// # Example: Timestamps
141///
142/// Types to match are represented using Arrow's [`DataType`].  [`DataType::Timestamp`] has an optional variable
143/// timezone specification. To specify a function can handle a timestamp with *ANY* timezone, use
144/// the [`TIMEZONE_WILDCARD`]. For example:
145///
146/// ```
147/// # use arrow::datatypes::{DataType, TimeUnit};
148/// # use datafusion_expr_common::signature::{TIMEZONE_WILDCARD, TypeSignature};
149/// let type_signature = TypeSignature::Exact(vec![
150///     // A nanosecond precision timestamp with ANY timezone
151///     // matches  Timestamp(Nanosecond, Some("+0:00"))
152///     // matches  Timestamp(Nanosecond, Some("+5:00"))
153///     // does not match  Timestamp(Nanosecond, None)
154///     DataType::Timestamp(TimeUnit::Nanosecond, Some(TIMEZONE_WILDCARD.into())),
155/// ]);
156/// ```
157#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
158pub enum TypeSignature {
159    /// One or more arguments of a common type out of a list of valid types.
160    ///
161    /// For functions that take no arguments (e.g. `random()` see [`TypeSignature::Nullary`]).
162    ///
163    /// # Examples
164    ///
165    /// A function such as `concat` is `Variadic(vec![DataType::Utf8,
166    /// DataType::LargeUtf8])`
167    Variadic(Vec<DataType>),
168    /// The acceptable signature and coercions rules are special for this
169    /// function.
170    ///
171    /// If this signature is specified,
172    /// DataFusion will call [`ScalarUDFImpl::coerce_types`] to prepare argument types.
173    ///
174    /// [`ScalarUDFImpl::coerce_types`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.coerce_types
175    UserDefined,
176    /// One or more arguments with arbitrary types
177    VariadicAny,
178    /// One or more arguments of an arbitrary but equal type out of a list of valid types.
179    ///
180    /// # Examples
181    ///
182    /// 1. A function of one argument of f64 is `Uniform(1, vec![DataType::Float64])`
183    /// 2. A function of one argument of f64 or f32 is `Uniform(1, vec![DataType::Float32, DataType::Float64])`
184    Uniform(usize, Vec<DataType>),
185    /// One or more arguments with exactly the specified types in order.
186    ///
187    /// For functions that take no arguments (e.g. `random()`) use [`TypeSignature::Nullary`].
188    Exact(Vec<DataType>),
189    /// One or more arguments belonging to the [`TypeSignatureClass`], in order.
190    ///
191    /// [`Coercion`] contains not only the desired type but also the allowed
192    /// casts. For example, if you expect a function has string type, but you
193    /// also allow it to be casted from binary type.
194    ///
195    /// For functions that take no arguments (e.g. `random()`) see [`TypeSignature::Nullary`].
196    Coercible(Vec<Coercion>),
197    /// One or more arguments coercible to a single, comparable type.
198    ///
199    /// Each argument will be coerced to a single type using the
200    /// coercion rules described in [`comparison_coercion_numeric`].
201    ///
202    /// # Examples
203    ///
204    /// If the `nullif(1, 2)` function is called with `i32` and `i64` arguments
205    /// the types will both be coerced to `i64` before the function is invoked.
206    ///
207    /// If the `nullif('1', 2)` function is called with `Utf8` and `i64` arguments
208    /// the types will both be coerced to `Utf8` before the function is invoked.
209    ///
210    /// Note:
211    /// - For functions that take no arguments (e.g. `random()` see [`TypeSignature::Nullary`]).
212    /// - If all arguments have type [`DataType::Null`], they are coerced to `Utf8`
213    ///
214    /// [`comparison_coercion_numeric`]: crate::type_coercion::binary::comparison_coercion_numeric
215    Comparable(usize),
216    /// One or more arguments of arbitrary types.
217    ///
218    /// For functions that take no arguments (e.g. `random()`) use [`TypeSignature::Nullary`].
219    Any(usize),
220    /// Matches exactly one of a list of [`TypeSignature`]s.
221    ///
222    /// Coercion is attempted to match the signatures in order, and stops after
223    /// the first success, if any.
224    ///
225    /// # Examples
226    ///
227    /// Since `make_array` takes 0 or more arguments with arbitrary types, its `TypeSignature`
228    /// is `OneOf(vec![Any(0), VariadicAny])`.
229    OneOf(Vec<TypeSignature>),
230    /// A function that has an [`ArrayFunctionSignature`]
231    ArraySignature(ArrayFunctionSignature),
232    /// One or more arguments of numeric types, coerced to a common numeric type.
233    ///
234    /// See [`NativeType::is_numeric`] to know which type is considered numeric
235    ///
236    /// For functions that take no arguments (e.g. `random()`) use [`TypeSignature::Nullary`].
237    ///
238    /// [`NativeType::is_numeric`]: datafusion_common::types::NativeType::is_numeric
239    Numeric(usize),
240    /// One or arguments of all the same string types.
241    ///
242    /// The precedence of type from high to low is Utf8View, LargeUtf8 and Utf8.
243    /// Null is considered as `Utf8` by default
244    /// Dictionary with string value type is also handled.
245    ///
246    /// For example, if a function is called with (utf8, large_utf8), all
247    /// arguments will be coerced to  `LargeUtf8`
248    ///
249    /// For functions that take no arguments (e.g. `random()` use [`TypeSignature::Nullary`]).
250    String(usize),
251    /// No arguments
252    Nullary,
253}
254
255impl TypeSignature {
256    #[inline]
257    pub fn is_one_of(&self) -> bool {
258        matches!(self, TypeSignature::OneOf(_))
259    }
260
261    /// Returns the arity (expected number of arguments) for this type signature.
262    ///
263    /// Returns `Arity::Fixed(n)` for signatures with a specific argument count,
264    /// or `Arity::Variable` for variable-arity signatures like `Variadic`, `VariadicAny`, `UserDefined`.
265    ///
266    /// # Examples
267    ///
268    /// ```
269    /// # use datafusion_expr_common::signature::{TypeSignature, Arity};
270    /// # use arrow::datatypes::DataType;
271    /// // Exact signature has fixed arity
272    /// let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
273    /// assert_eq!(sig.arity(), Arity::Fixed(2));
274    ///
275    /// // Variadic signature has variable arity
276    /// let sig = TypeSignature::VariadicAny;
277    /// assert_eq!(sig.arity(), Arity::Variable);
278    /// ```
279    pub fn arity(&self) -> Arity {
280        match self {
281            TypeSignature::Exact(types) => Arity::Fixed(types.len()),
282            TypeSignature::Uniform(count, _) => Arity::Fixed(*count),
283            TypeSignature::Numeric(count) => Arity::Fixed(*count),
284            TypeSignature::String(count) => Arity::Fixed(*count),
285            TypeSignature::Comparable(count) => Arity::Fixed(*count),
286            TypeSignature::Any(count) => Arity::Fixed(*count),
287            TypeSignature::Coercible(types) => Arity::Fixed(types.len()),
288            TypeSignature::Nullary => Arity::Fixed(0),
289            TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
290                arguments,
291                ..
292            }) => Arity::Fixed(arguments.len()),
293            TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray) => {
294                Arity::Fixed(1)
295            }
296            TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray) => {
297                Arity::Fixed(1)
298            }
299            TypeSignature::OneOf(variants) => {
300                // If any variant is Variable, the whole OneOf is Variable
301                let has_variable = variants.iter().any(|v| v.arity() == Arity::Variable);
302                if has_variable {
303                    return Arity::Variable;
304                }
305                // Otherwise, get max arity from all fixed arity variants
306                let max_arity = variants
307                    .iter()
308                    .filter_map(|v| match v.arity() {
309                        Arity::Fixed(n) => Some(n),
310                        Arity::Variable => None,
311                    })
312                    .max();
313                match max_arity {
314                    Some(n) => Arity::Fixed(n),
315                    None => Arity::Variable,
316                }
317            }
318            TypeSignature::Variadic(_)
319            | TypeSignature::VariadicAny
320            | TypeSignature::UserDefined => Arity::Variable,
321        }
322    }
323}
324
325/// Represents the class of types that can be used in a function signature.
326///
327/// This is used to specify what types are valid for function arguments in a more flexible way than
328/// just listing specific DataTypes. For example, TypeSignatureClass::Timestamp matches any timestamp
329/// type regardless of timezone or precision.
330///
331/// Used primarily with [`TypeSignature::Coercible`] to define function signatures that can accept
332/// arguments that can be coerced to a particular class of types.
333#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Hash)]
334pub enum TypeSignatureClass {
335    /// Allows an arbitrary type argument without coercing the argument.
336    Any,
337    /// Timestamps, allowing arbitrary (or no) timezones
338    Timestamp,
339    /// All time types
340    Time,
341    /// All interval types
342    Interval,
343    /// All duration types
344    Duration,
345    /// A specific native type
346    Native(LogicalTypeRef),
347    /// Signed and unsigned integers
348    Integer,
349    /// All float types
350    Float,
351    /// All decimal types, allowing arbitrary precision & scale
352    Decimal,
353    /// Integers, floats and decimals
354    Numeric,
355    /// Encompasses both the native Binary/LargeBinary types as well as arbitrarily sized FixedSizeBinary types
356    Binary,
357}
358
359impl Display for TypeSignatureClass {
360    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
361        write!(f, "TypeSignatureClass::{self:?}")
362    }
363}
364
365impl TypeSignatureClass {
366    /// Get example acceptable types for this `TypeSignatureClass`
367    ///
368    /// This is used for `information_schema` and can be used to generate
369    /// documentation or error messages.
370    fn get_example_types(&self) -> Vec<DataType> {
371        match self {
372            // TODO: might be too much info to return every single type here
373            //       maybe https://github.com/apache/datafusion/issues/14761 will help here?
374            TypeSignatureClass::Any => vec![],
375            TypeSignatureClass::Native(l) => get_data_types(l.native()),
376            TypeSignatureClass::Timestamp => {
377                vec![
378                    DataType::Timestamp(TimeUnit::Nanosecond, None),
379                    DataType::Timestamp(
380                        TimeUnit::Nanosecond,
381                        Some(TIMEZONE_WILDCARD.into()),
382                    ),
383                ]
384            }
385            TypeSignatureClass::Time => {
386                vec![DataType::Time64(TimeUnit::Nanosecond)]
387            }
388            TypeSignatureClass::Interval => {
389                vec![DataType::Interval(IntervalUnit::DayTime)]
390            }
391            TypeSignatureClass::Duration => {
392                vec![DataType::Duration(TimeUnit::Nanosecond)]
393            }
394            TypeSignatureClass::Integer => {
395                vec![DataType::Int64]
396            }
397            TypeSignatureClass::Binary => {
398                vec![DataType::Binary]
399            }
400            TypeSignatureClass::Decimal => vec![Decimal128Type::DEFAULT_TYPE],
401            TypeSignatureClass::Float => vec![DataType::Float64],
402            TypeSignatureClass::Numeric => vec![
403                DataType::Float64,
404                DataType::Int64,
405                Decimal128Type::DEFAULT_TYPE,
406            ],
407        }
408    }
409
410    /// Does the specified `NativeType` match this type signature class?
411    pub fn matches_native_type(&self, logical_type: &NativeType) -> bool {
412        if logical_type == &NativeType::Null {
413            return true;
414        }
415
416        match self {
417            TypeSignatureClass::Any => true,
418            TypeSignatureClass::Native(t) if t.native() == logical_type => true,
419            TypeSignatureClass::Timestamp if logical_type.is_timestamp() => true,
420            TypeSignatureClass::Time if logical_type.is_time() => true,
421            TypeSignatureClass::Interval if logical_type.is_interval() => true,
422            TypeSignatureClass::Duration if logical_type.is_duration() => true,
423            TypeSignatureClass::Integer if logical_type.is_integer() => true,
424            TypeSignatureClass::Binary if logical_type.is_binary() => true,
425            TypeSignatureClass::Decimal if logical_type.is_decimal() => true,
426            TypeSignatureClass::Float if logical_type.is_float() => true,
427            TypeSignatureClass::Numeric if logical_type.is_numeric() => true,
428            _ => false,
429        }
430    }
431
432    /// What type would `origin_type` be casted to when casting to the specified native type?
433    pub fn default_casted_type(
434        &self,
435        native_type: &NativeType,
436        origin_type: &DataType,
437    ) -> Result<DataType> {
438        match self {
439            TypeSignatureClass::Any => Ok(origin_type.to_owned()),
440            TypeSignatureClass::Native(logical_type) => {
441                logical_type.native().default_cast_for(origin_type)
442            }
443            // If the given type is already a timestamp, we don't change the unit and timezone
444            TypeSignatureClass::Timestamp if native_type.is_timestamp() => {
445                Ok(origin_type.to_owned())
446            }
447            TypeSignatureClass::Time if native_type.is_time() => {
448                Ok(origin_type.to_owned())
449            }
450            TypeSignatureClass::Interval if native_type.is_interval() => {
451                Ok(origin_type.to_owned())
452            }
453            TypeSignatureClass::Duration if native_type.is_duration() => {
454                Ok(origin_type.to_owned())
455            }
456            TypeSignatureClass::Integer if native_type.is_integer() => {
457                Ok(origin_type.to_owned())
458            }
459            TypeSignatureClass::Binary if native_type.is_binary() => {
460                Ok(origin_type.to_owned())
461            }
462            TypeSignatureClass::Decimal if native_type.is_decimal() => {
463                Ok(origin_type.to_owned())
464            }
465            TypeSignatureClass::Float if native_type.is_float() => {
466                Ok(origin_type.to_owned())
467            }
468            TypeSignatureClass::Numeric if native_type.is_numeric() => {
469                Ok(origin_type.to_owned())
470            }
471            _ if native_type.is_null() => Ok(origin_type.to_owned()),
472            _ => internal_err!("May miss the matching logic in `matches_native_type`"),
473        }
474    }
475}
476
477#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
478pub enum ArrayFunctionSignature {
479    /// A function takes at least one List/LargeList/FixedSizeList argument.
480    Array {
481        /// A full list of the arguments accepted by this function.
482        arguments: Vec<ArrayFunctionArgument>,
483        /// Additional information about how array arguments should be coerced.
484        array_coercion: Option<ListCoercion>,
485    },
486    /// A function takes a single argument that must be a List/LargeList/FixedSizeList
487    /// which gets coerced to List, with element type recursively coerced to List too if it is list-like.
488    RecursiveArray,
489    /// Specialized Signature for MapArray
490    /// The function takes a single argument that must be a MapArray
491    MapArray,
492}
493
494impl Display for ArrayFunctionSignature {
495    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
496        match self {
497            ArrayFunctionSignature::Array { arguments, .. } => {
498                for (idx, argument) in arguments.iter().enumerate() {
499                    write!(f, "{argument}")?;
500                    if idx != arguments.len() - 1 {
501                        write!(f, ", ")?;
502                    }
503                }
504                Ok(())
505            }
506            ArrayFunctionSignature::RecursiveArray => {
507                write!(f, "recursive_array")
508            }
509            ArrayFunctionSignature::MapArray => {
510                write!(f, "map_array")
511            }
512        }
513    }
514}
515
516#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
517pub enum ArrayFunctionArgument {
518    /// A non-list or list argument. The list dimensions should be one less than the Array's list
519    /// dimensions.
520    Element,
521    /// An Int64 index argument.
522    Index,
523    /// An argument of type List/LargeList/FixedSizeList. All Array arguments must be coercible
524    /// to the same type.
525    Array,
526    // A Utf8 argument.
527    String,
528}
529
530impl Display for ArrayFunctionArgument {
531    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
532        match self {
533            ArrayFunctionArgument::Element => {
534                write!(f, "element")
535            }
536            ArrayFunctionArgument::Index => {
537                write!(f, "index")
538            }
539            ArrayFunctionArgument::Array => {
540                write!(f, "array")
541            }
542            ArrayFunctionArgument::String => {
543                write!(f, "string")
544            }
545        }
546    }
547}
548
549impl TypeSignature {
550    pub fn to_string_repr(&self) -> Vec<String> {
551        match self {
552            TypeSignature::Nullary => {
553                vec!["NullAry()".to_string()]
554            }
555            TypeSignature::Variadic(types) => {
556                vec![format!("{}, ..", Self::join_types(types, "/"))]
557            }
558            TypeSignature::Uniform(arg_count, valid_types) => {
559                vec![
560                    std::iter::repeat_n(Self::join_types(valid_types, "/"), *arg_count)
561                        .collect::<Vec<String>>()
562                        .join(", "),
563                ]
564            }
565            TypeSignature::String(num) => {
566                vec![format!("String({num})")]
567            }
568            TypeSignature::Numeric(num) => {
569                vec![format!("Numeric({num})")]
570            }
571            TypeSignature::Comparable(num) => {
572                vec![format!("Comparable({num})")]
573            }
574            TypeSignature::Coercible(coercions) => {
575                vec![Self::join_types(coercions, ", ")]
576            }
577            TypeSignature::Exact(types) => {
578                vec![Self::join_types(types, ", ")]
579            }
580            TypeSignature::Any(arg_count) => {
581                vec![
582                    std::iter::repeat_n("Any", *arg_count)
583                        .collect::<Vec<&str>>()
584                        .join(", "),
585                ]
586            }
587            TypeSignature::UserDefined => {
588                vec!["UserDefined".to_string()]
589            }
590            TypeSignature::VariadicAny => vec!["Any, .., Any".to_string()],
591            TypeSignature::OneOf(sigs) => {
592                sigs.iter().flat_map(|s| s.to_string_repr()).collect()
593            }
594            TypeSignature::ArraySignature(array_signature) => {
595                vec![array_signature.to_string()]
596            }
597        }
598    }
599
600    /// Return string representation of the function signature with parameter names.
601    ///
602    /// This method is similar to [`Self::to_string_repr`] but uses parameter names
603    /// instead of types when available. This is useful for generating more helpful
604    /// error messages.
605    ///
606    /// # Arguments
607    /// * `parameter_names` - Optional slice of parameter names. When provided, these
608    ///   names will be used instead of type names in the output.
609    ///
610    /// # Examples
611    /// ```
612    /// # use datafusion_expr_common::signature::TypeSignature;
613    /// # use arrow::datatypes::DataType;
614    /// let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
615    ///
616    /// // Without names: shows types only
617    /// assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
618    ///
619    /// // With names: shows parameter names with types
620    /// assert_eq!(
621    ///     sig.to_string_repr_with_names(Some(&["id".to_string(), "name".to_string()])),
622    ///     vec!["id: Int32, name: Utf8"]
623    /// );
624    /// ```
625    pub fn to_string_repr_with_names(
626        &self,
627        parameter_names: Option<&[String]>,
628    ) -> Vec<String> {
629        match self {
630            TypeSignature::Exact(types) => {
631                if let Some(names) = parameter_names {
632                    vec![
633                        names
634                            .iter()
635                            .zip(types.iter())
636                            .map(|(name, typ)| format!("{name}: {typ}"))
637                            .collect::<Vec<_>>()
638                            .join(", "),
639                    ]
640                } else {
641                    vec![Self::join_types(types, ", ")]
642                }
643            }
644            TypeSignature::Any(count) => {
645                if let Some(names) = parameter_names {
646                    vec![
647                        names
648                            .iter()
649                            .take(*count)
650                            .map(|name| format!("{name}: Any"))
651                            .collect::<Vec<_>>()
652                            .join(", "),
653                    ]
654                } else {
655                    vec![
656                        std::iter::repeat_n("Any", *count)
657                            .collect::<Vec<&str>>()
658                            .join(", "),
659                    ]
660                }
661            }
662            TypeSignature::Uniform(count, types) => {
663                if let Some(names) = parameter_names {
664                    let type_str = Self::join_types(types, "/");
665                    vec![
666                        names
667                            .iter()
668                            .take(*count)
669                            .map(|name| format!("{name}: {type_str}"))
670                            .collect::<Vec<_>>()
671                            .join(", "),
672                    ]
673                } else {
674                    self.to_string_repr()
675                }
676            }
677            TypeSignature::Coercible(coercions) => {
678                if let Some(names) = parameter_names {
679                    vec![
680                        names
681                            .iter()
682                            .zip(coercions.iter())
683                            .map(|(name, coercion)| format!("{name}: {coercion}"))
684                            .collect::<Vec<_>>()
685                            .join(", "),
686                    ]
687                } else {
688                    vec![Self::join_types(coercions, ", ")]
689                }
690            }
691            TypeSignature::Comparable(count) => {
692                if let Some(names) = parameter_names {
693                    vec![
694                        names
695                            .iter()
696                            .take(*count)
697                            .map(|name| format!("{name}: Comparable"))
698                            .collect::<Vec<_>>()
699                            .join(", "),
700                    ]
701                } else {
702                    self.to_string_repr()
703                }
704            }
705            TypeSignature::Numeric(count) => {
706                if let Some(names) = parameter_names {
707                    vec![
708                        names
709                            .iter()
710                            .take(*count)
711                            .map(|name| format!("{name}: Numeric"))
712                            .collect::<Vec<_>>()
713                            .join(", "),
714                    ]
715                } else {
716                    self.to_string_repr()
717                }
718            }
719            TypeSignature::String(count) => {
720                if let Some(names) = parameter_names {
721                    vec![
722                        names
723                            .iter()
724                            .take(*count)
725                            .map(|name| format!("{name}: String"))
726                            .collect::<Vec<_>>()
727                            .join(", "),
728                    ]
729                } else {
730                    self.to_string_repr()
731                }
732            }
733            TypeSignature::Nullary => self.to_string_repr(),
734            TypeSignature::ArraySignature(array_sig) => {
735                if let Some(names) = parameter_names {
736                    match array_sig {
737                        ArrayFunctionSignature::Array { arguments, .. } => {
738                            vec![
739                                names
740                                    .iter()
741                                    .zip(arguments.iter())
742                                    .map(|(name, arg_type)| format!("{name}: {arg_type}"))
743                                    .collect::<Vec<_>>()
744                                    .join(", "),
745                            ]
746                        }
747                        ArrayFunctionSignature::RecursiveArray => {
748                            vec![
749                                names
750                                    .iter()
751                                    .take(1)
752                                    .map(|name| format!("{name}: recursive_array"))
753                                    .collect::<Vec<_>>()
754                                    .join(", "),
755                            ]
756                        }
757                        ArrayFunctionSignature::MapArray => {
758                            vec![
759                                names
760                                    .iter()
761                                    .take(1)
762                                    .map(|name| format!("{name}: map_array"))
763                                    .collect::<Vec<_>>()
764                                    .join(", "),
765                            ]
766                        }
767                    }
768                } else {
769                    self.to_string_repr()
770                }
771            }
772            TypeSignature::OneOf(sigs) => sigs
773                .iter()
774                .flat_map(|s| s.to_string_repr_with_names(parameter_names))
775                .collect(),
776            TypeSignature::UserDefined => {
777                if let Some(names) = parameter_names {
778                    vec![names.join(", ")]
779                } else {
780                    self.to_string_repr()
781                }
782            }
783            // Variable arity signatures cannot use parameter names
784            TypeSignature::Variadic(_) | TypeSignature::VariadicAny => {
785                self.to_string_repr()
786            }
787        }
788    }
789
790    /// Helper function to join types with specified delimiter.
791    pub fn join_types<T: Display>(types: &[T], delimiter: &str) -> String {
792        types
793            .iter()
794            .map(|t| t.to_string())
795            .collect::<Vec<String>>()
796            .join(delimiter)
797    }
798
799    /// Check whether 0 input argument is valid for given `TypeSignature`
800    pub fn supports_zero_argument(&self) -> bool {
801        match &self {
802            TypeSignature::Exact(vec) => vec.is_empty(),
803            TypeSignature::Nullary => true,
804            TypeSignature::OneOf(types) => types
805                .iter()
806                .any(|type_sig| type_sig.supports_zero_argument()),
807            _ => false,
808        }
809    }
810
811    /// Returns true if the signature currently supports or used to supported 0
812    /// input arguments in a previous version of DataFusion.
813    pub fn used_to_support_zero_arguments(&self) -> bool {
814        match &self {
815            TypeSignature::Any(num) => *num == 0,
816            _ => self.supports_zero_argument(),
817        }
818    }
819
820    #[deprecated(since = "46.0.0", note = "See get_example_types instead")]
821    pub fn get_possible_types(&self) -> Vec<Vec<DataType>> {
822        self.get_example_types()
823    }
824
825    /// Return example acceptable types for this `TypeSignature`'
826    ///
827    /// Returns a `Vec<DataType>` for each argument to the function
828    ///
829    /// This is used for `information_schema` and can be used to generate
830    /// documentation or error messages.
831    pub fn get_example_types(&self) -> Vec<Vec<DataType>> {
832        match self {
833            TypeSignature::Exact(types) => vec![types.clone()],
834            TypeSignature::OneOf(types) => types
835                .iter()
836                .flat_map(|type_sig| type_sig.get_example_types())
837                .collect(),
838            TypeSignature::Uniform(arg_count, types) => types
839                .iter()
840                .cloned()
841                .map(|data_type| vec![data_type; *arg_count])
842                .collect(),
843            TypeSignature::Coercible(coercions) => coercions
844                .iter()
845                .map(|c| {
846                    let mut all_types: IndexSet<DataType> =
847                        c.desired_type().get_example_types().into_iter().collect();
848
849                    if let Some(implicit_coercion) = c.implicit_coercion() {
850                        let allowed_casts: Vec<DataType> = implicit_coercion
851                            .allowed_source_types
852                            .iter()
853                            .flat_map(|t| t.get_example_types())
854                            .collect();
855                        all_types.extend(allowed_casts);
856                    }
857
858                    all_types.into_iter().collect::<Vec<_>>()
859                })
860                .multi_cartesian_product()
861                .collect(),
862            TypeSignature::Variadic(types) => types
863                .iter()
864                .cloned()
865                .map(|data_type| vec![data_type])
866                .collect(),
867            TypeSignature::Numeric(arg_count) => NUMERICS
868                .iter()
869                .cloned()
870                .map(|numeric_type| vec![numeric_type; *arg_count])
871                .collect(),
872            TypeSignature::String(arg_count) => get_data_types(&NativeType::String)
873                .into_iter()
874                .map(|dt| vec![dt; *arg_count])
875                .collect::<Vec<_>>(),
876            // TODO: Implement for other types
877            TypeSignature::Any(_)
878            | TypeSignature::Comparable(_)
879            | TypeSignature::Nullary
880            | TypeSignature::VariadicAny
881            | TypeSignature::ArraySignature(_)
882            | TypeSignature::UserDefined => vec![],
883        }
884    }
885}
886
887fn get_data_types(native_type: &NativeType) -> Vec<DataType> {
888    match native_type {
889        NativeType::Null => vec![DataType::Null],
890        NativeType::Boolean => vec![DataType::Boolean],
891        NativeType::Int8 => vec![DataType::Int8],
892        NativeType::Int16 => vec![DataType::Int16],
893        NativeType::Int32 => vec![DataType::Int32],
894        NativeType::Int64 => vec![DataType::Int64],
895        NativeType::UInt8 => vec![DataType::UInt8],
896        NativeType::UInt16 => vec![DataType::UInt16],
897        NativeType::UInt32 => vec![DataType::UInt32],
898        NativeType::UInt64 => vec![DataType::UInt64],
899        NativeType::Float16 => vec![DataType::Float16],
900        NativeType::Float32 => vec![DataType::Float32],
901        NativeType::Float64 => vec![DataType::Float64],
902        NativeType::Date => vec![DataType::Date32, DataType::Date64],
903        NativeType::Binary => vec![
904            DataType::Binary,
905            DataType::LargeBinary,
906            DataType::BinaryView,
907        ],
908        NativeType::String => {
909            vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]
910        }
911        NativeType::Decimal(precision, scale) => {
912            // We assume incoming NativeType is valid already, in terms of precision & scale
913            let mut types = vec![DataType::Decimal256(*precision, *scale)];
914            if *precision <= DECIMAL32_MAX_PRECISION {
915                types.push(DataType::Decimal32(*precision, *scale));
916            }
917            if *precision <= DECIMAL64_MAX_PRECISION {
918                types.push(DataType::Decimal64(*precision, *scale));
919            }
920            if *precision <= DECIMAL128_MAX_PRECISION {
921                types.push(DataType::Decimal128(*precision, *scale));
922            }
923            types
924        }
925        NativeType::Timestamp(time_unit, timezone) => {
926            vec![DataType::Timestamp(*time_unit, timezone.to_owned())]
927        }
928        NativeType::Time(TimeUnit::Second) => vec![DataType::Time32(TimeUnit::Second)],
929        NativeType::Time(TimeUnit::Millisecond) => {
930            vec![DataType::Time32(TimeUnit::Millisecond)]
931        }
932        NativeType::Time(TimeUnit::Microsecond) => {
933            vec![DataType::Time64(TimeUnit::Microsecond)]
934        }
935        NativeType::Time(TimeUnit::Nanosecond) => {
936            vec![DataType::Time64(TimeUnit::Nanosecond)]
937        }
938        NativeType::Duration(time_unit) => vec![DataType::Duration(*time_unit)],
939        NativeType::Interval(interval_unit) => vec![DataType::Interval(*interval_unit)],
940        NativeType::FixedSizeBinary(size) => vec![DataType::FixedSizeBinary(*size)],
941        NativeType::FixedSizeList(logical_field, size) => {
942            get_data_types(logical_field.logical_type.native())
943                .iter()
944                .map(|child_dt| {
945                    let field = Field::new(
946                        logical_field.name.clone(),
947                        child_dt.clone(),
948                        logical_field.nullable,
949                    );
950                    DataType::FixedSizeList(Arc::new(field), *size)
951                })
952                .collect()
953        }
954        // TODO: implement for nested types
955        NativeType::List(_)
956        | NativeType::Struct(_)
957        | NativeType::Union(_)
958        | NativeType::Map(_) => {
959            vec![]
960        }
961    }
962}
963
964/// Represents type coercion rules for function arguments, specifying both the desired type
965/// and optional implicit coercion rules for source types.
966///
967/// # Examples
968///
969/// ```
970/// use datafusion_common::types::{logical_binary, logical_string, NativeType};
971/// use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
972///
973/// // Exact coercion that only accepts timestamp types
974/// let exact = Coercion::new_exact(TypeSignatureClass::Timestamp);
975///
976/// // Implicit coercion that accepts string types but can coerce from binary types
977/// let implicit = Coercion::new_implicit(
978///     TypeSignatureClass::Native(logical_string()),
979///     vec![TypeSignatureClass::Native(logical_binary())],
980///     NativeType::String,
981/// );
982/// ```
983///
984/// There are two variants:
985///
986/// * `Exact` - Only accepts arguments that exactly match the desired type
987/// * `Implicit` - Accepts the desired type and can coerce from specified source types
988#[derive(Debug, Clone, Eq, PartialOrd)]
989pub enum Coercion {
990    /// Coercion that only accepts arguments exactly matching the desired type.
991    Exact {
992        /// The required type for the argument
993        desired_type: TypeSignatureClass,
994    },
995
996    /// Coercion that accepts the desired type and can implicitly coerce from other types.
997    Implicit {
998        /// The primary desired type for the argument
999        desired_type: TypeSignatureClass,
1000        /// Rules for implicit coercion from other types
1001        implicit_coercion: ImplicitCoercion,
1002    },
1003}
1004
1005impl Coercion {
1006    pub fn new_exact(desired_type: TypeSignatureClass) -> Self {
1007        Self::Exact { desired_type }
1008    }
1009
1010    /// Create a new coercion with implicit coercion rules.
1011    ///
1012    /// `allowed_source_types` defines the possible types that can be coerced to `desired_type`.
1013    /// `default_casted_type` is the default type to be used for coercion if we cast from other types via `allowed_source_types`.
1014    pub fn new_implicit(
1015        desired_type: TypeSignatureClass,
1016        allowed_source_types: Vec<TypeSignatureClass>,
1017        default_casted_type: NativeType,
1018    ) -> Self {
1019        Self::Implicit {
1020            desired_type,
1021            implicit_coercion: ImplicitCoercion {
1022                allowed_source_types,
1023                default_casted_type,
1024            },
1025        }
1026    }
1027
1028    pub fn allowed_source_types(&self) -> &[TypeSignatureClass] {
1029        match self {
1030            Coercion::Exact { .. } => &[],
1031            Coercion::Implicit {
1032                implicit_coercion, ..
1033            } => implicit_coercion.allowed_source_types.as_slice(),
1034        }
1035    }
1036
1037    pub fn default_casted_type(&self) -> Option<&NativeType> {
1038        match self {
1039            Coercion::Exact { .. } => None,
1040            Coercion::Implicit {
1041                implicit_coercion, ..
1042            } => Some(&implicit_coercion.default_casted_type),
1043        }
1044    }
1045
1046    pub fn desired_type(&self) -> &TypeSignatureClass {
1047        match self {
1048            Coercion::Exact { desired_type } => desired_type,
1049            Coercion::Implicit { desired_type, .. } => desired_type,
1050        }
1051    }
1052
1053    pub fn implicit_coercion(&self) -> Option<&ImplicitCoercion> {
1054        match self {
1055            Coercion::Exact { .. } => None,
1056            Coercion::Implicit {
1057                implicit_coercion, ..
1058            } => Some(implicit_coercion),
1059        }
1060    }
1061}
1062
1063impl Display for Coercion {
1064    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1065        write!(f, "Coercion({}", self.desired_type())?;
1066        if let Some(implicit_coercion) = self.implicit_coercion() {
1067            write!(f, ", implicit_coercion={implicit_coercion}",)
1068        } else {
1069            write!(f, ")")
1070        }
1071    }
1072}
1073
1074impl PartialEq for Coercion {
1075    fn eq(&self, other: &Self) -> bool {
1076        self.desired_type() == other.desired_type()
1077            && self.implicit_coercion() == other.implicit_coercion()
1078    }
1079}
1080
1081impl Hash for Coercion {
1082    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1083        self.desired_type().hash(state);
1084        self.implicit_coercion().hash(state);
1085    }
1086}
1087
1088/// Defines rules for implicit type coercion, specifying which source types can be
1089/// coerced and the default type to use when coercing.
1090///
1091/// This is used by functions to specify which types they can accept via implicit
1092/// coercion in addition to their primary desired type.
1093///
1094/// # Examples
1095///
1096/// ```
1097/// use arrow::datatypes::TimeUnit;
1098///
1099/// use datafusion_expr_common::signature::{Coercion, ImplicitCoercion, TypeSignatureClass};
1100/// use datafusion_common::types::{NativeType, logical_binary};
1101///
1102/// // Allow coercing from binary types to timestamp, coerce to specific timestamp unit and timezone
1103/// let implicit = Coercion::new_implicit(
1104///     TypeSignatureClass::Timestamp,
1105///     vec![TypeSignatureClass::Native(logical_binary())],
1106///     NativeType::Timestamp(TimeUnit::Second, None),
1107/// );
1108/// ```
1109#[derive(Debug, Clone, Eq, PartialOrd)]
1110pub struct ImplicitCoercion {
1111    /// The types that can be coerced from via implicit casting
1112    allowed_source_types: Vec<TypeSignatureClass>,
1113
1114    /// The default type to use when coercing from allowed source types.
1115    /// This is particularly important for types like Timestamp that have multiple
1116    /// possible configurations (different time units and timezones).
1117    default_casted_type: NativeType,
1118}
1119
1120impl Display for ImplicitCoercion {
1121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1122        write!(
1123            f,
1124            "ImplicitCoercion({:?}, default_type={:?})",
1125            self.allowed_source_types, self.default_casted_type
1126        )
1127    }
1128}
1129
1130impl PartialEq for ImplicitCoercion {
1131    fn eq(&self, other: &Self) -> bool {
1132        self.allowed_source_types == other.allowed_source_types
1133            && self.default_casted_type == other.default_casted_type
1134    }
1135}
1136
1137impl Hash for ImplicitCoercion {
1138    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1139        self.allowed_source_types.hash(state);
1140        self.default_casted_type.hash(state);
1141    }
1142}
1143
1144/// Provides  information necessary for calling a function.
1145///
1146/// - [`TypeSignature`] defines the argument types that a function has implementations
1147///   for.
1148///
1149/// - [`Volatility`] defines how the output of the function changes with the input.
1150#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
1151pub struct Signature {
1152    /// The data types that the function accepts. See [TypeSignature] for more information.
1153    pub type_signature: TypeSignature,
1154    /// The volatility of the function. See [Volatility] for more information.
1155    pub volatility: Volatility,
1156    /// Optional parameter names for the function arguments.
1157    ///
1158    /// If provided, enables named argument notation for function calls (e.g., `func(a => 1, b => 2)`).
1159    /// The length must match the number of arguments defined by `type_signature`.
1160    ///
1161    /// Defaults to `None`, meaning only positional arguments are supported.
1162    pub parameter_names: Option<Vec<String>>,
1163}
1164
1165impl Signature {
1166    /// Creates a new Signature from a given type signature and volatility.
1167    pub fn new(type_signature: TypeSignature, volatility: Volatility) -> Self {
1168        Signature {
1169            type_signature,
1170            volatility,
1171            parameter_names: None,
1172        }
1173    }
1174    /// An arbitrary number of arguments with the same type, from those listed in `common_types`.
1175    pub fn variadic(common_types: Vec<DataType>, volatility: Volatility) -> Self {
1176        Self {
1177            type_signature: TypeSignature::Variadic(common_types),
1178            volatility,
1179            parameter_names: None,
1180        }
1181    }
1182    /// User-defined coercion rules for the function.
1183    pub fn user_defined(volatility: Volatility) -> Self {
1184        Self {
1185            type_signature: TypeSignature::UserDefined,
1186            volatility,
1187            parameter_names: None,
1188        }
1189    }
1190
1191    /// A specified number of numeric arguments
1192    pub fn numeric(arg_count: usize, volatility: Volatility) -> Self {
1193        Self {
1194            type_signature: TypeSignature::Numeric(arg_count),
1195            volatility,
1196            parameter_names: None,
1197        }
1198    }
1199
1200    /// A specified number of string arguments
1201    pub fn string(arg_count: usize, volatility: Volatility) -> Self {
1202        Self {
1203            type_signature: TypeSignature::String(arg_count),
1204            volatility,
1205            parameter_names: None,
1206        }
1207    }
1208
1209    /// An arbitrary number of arguments of any type.
1210    pub fn variadic_any(volatility: Volatility) -> Self {
1211        Self {
1212            type_signature: TypeSignature::VariadicAny,
1213            volatility,
1214            parameter_names: None,
1215        }
1216    }
1217    /// A fixed number of arguments of the same type, from those listed in `valid_types`.
1218    pub fn uniform(
1219        arg_count: usize,
1220        valid_types: Vec<DataType>,
1221        volatility: Volatility,
1222    ) -> Self {
1223        Self {
1224            type_signature: TypeSignature::Uniform(arg_count, valid_types),
1225            volatility,
1226            parameter_names: None,
1227        }
1228    }
1229    /// Exactly matches the types in `exact_types`, in order.
1230    pub fn exact(exact_types: Vec<DataType>, volatility: Volatility) -> Self {
1231        Signature {
1232            type_signature: TypeSignature::Exact(exact_types),
1233            volatility,
1234            parameter_names: None,
1235        }
1236    }
1237
1238    /// Target coerce types in order
1239    pub fn coercible(target_types: Vec<Coercion>, volatility: Volatility) -> Self {
1240        Self {
1241            type_signature: TypeSignature::Coercible(target_types),
1242            volatility,
1243            parameter_names: None,
1244        }
1245    }
1246
1247    /// Used for function that expects comparable data types, it will try to coerced all the types into single final one.
1248    pub fn comparable(arg_count: usize, volatility: Volatility) -> Self {
1249        Self {
1250            type_signature: TypeSignature::Comparable(arg_count),
1251            volatility,
1252            parameter_names: None,
1253        }
1254    }
1255
1256    pub fn nullary(volatility: Volatility) -> Self {
1257        Signature {
1258            type_signature: TypeSignature::Nullary,
1259            volatility,
1260            parameter_names: None,
1261        }
1262    }
1263
1264    /// A specified number of arguments of any type
1265    pub fn any(arg_count: usize, volatility: Volatility) -> Self {
1266        Signature {
1267            type_signature: TypeSignature::Any(arg_count),
1268            volatility,
1269            parameter_names: None,
1270        }
1271    }
1272
1273    /// Any one of a list of [TypeSignature]s.
1274    pub fn one_of(type_signatures: Vec<TypeSignature>, volatility: Volatility) -> Self {
1275        Signature {
1276            type_signature: TypeSignature::OneOf(type_signatures),
1277            volatility,
1278            parameter_names: None,
1279        }
1280    }
1281
1282    /// Specialized [Signature] for ArrayAppend and similar functions.
1283    pub fn array_and_element(volatility: Volatility) -> Self {
1284        Signature {
1285            type_signature: TypeSignature::ArraySignature(
1286                ArrayFunctionSignature::Array {
1287                    arguments: vec![
1288                        ArrayFunctionArgument::Array,
1289                        ArrayFunctionArgument::Element,
1290                    ],
1291                    array_coercion: Some(ListCoercion::FixedSizedListToList),
1292                },
1293            ),
1294            volatility,
1295            parameter_names: None,
1296        }
1297    }
1298
1299    /// Specialized [Signature] for ArrayPrepend and similar functions.
1300    pub fn element_and_array(volatility: Volatility) -> Self {
1301        Signature {
1302            type_signature: TypeSignature::ArraySignature(
1303                ArrayFunctionSignature::Array {
1304                    arguments: vec![
1305                        ArrayFunctionArgument::Element,
1306                        ArrayFunctionArgument::Array,
1307                    ],
1308                    array_coercion: Some(ListCoercion::FixedSizedListToList),
1309                },
1310            ),
1311            volatility,
1312            parameter_names: None,
1313        }
1314    }
1315
1316    /// Specialized [Signature] for functions that take a fixed number of arrays.
1317    pub fn arrays(
1318        n: usize,
1319        coercion: Option<ListCoercion>,
1320        volatility: Volatility,
1321    ) -> Self {
1322        Signature {
1323            type_signature: TypeSignature::ArraySignature(
1324                ArrayFunctionSignature::Array {
1325                    arguments: vec![ArrayFunctionArgument::Array; n],
1326                    array_coercion: coercion,
1327                },
1328            ),
1329            volatility,
1330            parameter_names: None,
1331        }
1332    }
1333
1334    /// Specialized [Signature] for Array functions with an optional index.
1335    pub fn array_and_element_and_optional_index(volatility: Volatility) -> Self {
1336        Signature {
1337            type_signature: TypeSignature::OneOf(vec![
1338                TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1339                    arguments: vec![
1340                        ArrayFunctionArgument::Array,
1341                        ArrayFunctionArgument::Element,
1342                    ],
1343                    array_coercion: Some(ListCoercion::FixedSizedListToList),
1344                }),
1345                TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1346                    arguments: vec![
1347                        ArrayFunctionArgument::Array,
1348                        ArrayFunctionArgument::Element,
1349                        ArrayFunctionArgument::Index,
1350                    ],
1351                    array_coercion: Some(ListCoercion::FixedSizedListToList),
1352                }),
1353            ]),
1354            volatility,
1355            parameter_names: None,
1356        }
1357    }
1358
1359    /// Specialized [Signature] for ArrayElement and similar functions.
1360    pub fn array_and_index(volatility: Volatility) -> Self {
1361        Signature {
1362            type_signature: TypeSignature::ArraySignature(
1363                ArrayFunctionSignature::Array {
1364                    arguments: vec![
1365                        ArrayFunctionArgument::Array,
1366                        ArrayFunctionArgument::Index,
1367                    ],
1368                    array_coercion: Some(ListCoercion::FixedSizedListToList),
1369                },
1370            ),
1371            volatility,
1372            parameter_names: None,
1373        }
1374    }
1375
1376    /// Specialized [Signature] for ArrayEmpty and similar functions.
1377    pub fn array(volatility: Volatility) -> Self {
1378        Signature::arrays(1, Some(ListCoercion::FixedSizedListToList), volatility)
1379    }
1380
1381    /// Add parameter names to this signature, enabling named argument notation.
1382    ///
1383    /// # Example
1384    /// ```
1385    /// # use datafusion_expr_common::signature::{Signature, Volatility};
1386    /// # use arrow::datatypes::DataType;
1387    /// let sig =
1388    ///     Signature::exact(vec![DataType::Int32, DataType::Utf8], Volatility::Immutable)
1389    ///         .with_parameter_names(vec!["count".to_string(), "name".to_string()]);
1390    /// ```
1391    ///
1392    /// # Errors
1393    /// Returns an error if the number of parameter names doesn't match the signature's arity.
1394    /// For signatures with variable arity (e.g., `Variadic`, `VariadicAny`), parameter names
1395    /// cannot be specified.
1396    pub fn with_parameter_names(mut self, names: Vec<impl Into<String>>) -> Result<Self> {
1397        let names = names.into_iter().map(Into::into).collect::<Vec<String>>();
1398        // Validate that the number of names matches the signature
1399        self.validate_parameter_names(&names)?;
1400        self.parameter_names = Some(names);
1401        Ok(self)
1402    }
1403
1404    /// Validate that parameter names are compatible with this signature
1405    fn validate_parameter_names(&self, names: &[String]) -> Result<()> {
1406        match self.type_signature.arity() {
1407            Arity::Fixed(expected) => {
1408                if names.len() != expected {
1409                    return plan_err!(
1410                        "Parameter names count ({}) does not match signature arity ({})",
1411                        names.len(),
1412                        expected
1413                    );
1414                }
1415            }
1416            Arity::Variable => {
1417                // For UserDefined signatures, allow parameter names
1418                // The function implementer is responsible for validating the names match the actual arguments
1419                if !matches!(self.type_signature, TypeSignature::UserDefined) {
1420                    return plan_err!(
1421                        "Cannot specify parameter names for variable arity signature: {:?}",
1422                        self.type_signature
1423                    );
1424                }
1425            }
1426        }
1427
1428        let mut seen = std::collections::HashSet::new();
1429        for name in names {
1430            if !seen.insert(name) {
1431                return plan_err!("Duplicate parameter name: '{}'", name);
1432            }
1433        }
1434
1435        Ok(())
1436    }
1437}
1438
1439#[cfg(test)]
1440mod tests {
1441    use datafusion_common::types::{logical_int32, logical_int64, logical_string};
1442
1443    use super::*;
1444    use crate::signature::{
1445        ArrayFunctionArgument, ArrayFunctionSignature, Coercion, TypeSignatureClass,
1446    };
1447
1448    #[test]
1449    fn supports_zero_argument_tests() {
1450        // Testing `TypeSignature`s which supports 0 arg
1451        let positive_cases = vec![
1452            TypeSignature::Exact(vec![]),
1453            TypeSignature::OneOf(vec![
1454                TypeSignature::Exact(vec![DataType::Int8]),
1455                TypeSignature::Nullary,
1456                TypeSignature::Uniform(1, vec![DataType::Int8]),
1457            ]),
1458            TypeSignature::Nullary,
1459        ];
1460
1461        for case in positive_cases {
1462            assert!(
1463                case.supports_zero_argument(),
1464                "Expected {case:?} to support zero arguments"
1465            );
1466        }
1467
1468        // Testing `TypeSignature`s which doesn't support 0 arg
1469        let negative_cases = vec![
1470            TypeSignature::Exact(vec![DataType::Utf8]),
1471            TypeSignature::Uniform(1, vec![DataType::Float64]),
1472            TypeSignature::Any(1),
1473            TypeSignature::VariadicAny,
1474            TypeSignature::OneOf(vec![
1475                TypeSignature::Exact(vec![DataType::Int8]),
1476                TypeSignature::Uniform(1, vec![DataType::Int8]),
1477            ]),
1478        ];
1479
1480        for case in negative_cases {
1481            assert!(
1482                !case.supports_zero_argument(),
1483                "Expected {case:?} not to support zero arguments"
1484            );
1485        }
1486    }
1487
1488    #[test]
1489    fn type_signature_partial_ord() {
1490        // Test validates that partial ord is defined for TypeSignature and Signature.
1491        assert!(TypeSignature::UserDefined < TypeSignature::VariadicAny);
1492        assert!(TypeSignature::UserDefined < TypeSignature::Any(1));
1493
1494        assert!(
1495            TypeSignature::Uniform(1, vec![DataType::Null])
1496                < TypeSignature::Uniform(1, vec![DataType::Boolean])
1497        );
1498        assert!(
1499            TypeSignature::Uniform(1, vec![DataType::Null])
1500                < TypeSignature::Uniform(2, vec![DataType::Null])
1501        );
1502        assert!(
1503            TypeSignature::Uniform(usize::MAX, vec![DataType::Null])
1504                < TypeSignature::Exact(vec![DataType::Null])
1505        );
1506    }
1507
1508    #[test]
1509    fn test_get_possible_types() {
1510        let type_signature = TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]);
1511        let possible_types = type_signature.get_example_types();
1512        assert_eq!(possible_types, vec![vec![DataType::Int32, DataType::Int64]]);
1513
1514        let type_signature = TypeSignature::OneOf(vec![
1515            TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1516            TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1517        ]);
1518        let possible_types = type_signature.get_example_types();
1519        assert_eq!(
1520            possible_types,
1521            vec![
1522                vec![DataType::Int32, DataType::Int64],
1523                vec![DataType::Float32, DataType::Float64]
1524            ]
1525        );
1526
1527        let type_signature = TypeSignature::OneOf(vec![
1528            TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1529            TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1530            TypeSignature::Exact(vec![DataType::Utf8]),
1531        ]);
1532        let possible_types = type_signature.get_example_types();
1533        assert_eq!(
1534            possible_types,
1535            vec![
1536                vec![DataType::Int32, DataType::Int64],
1537                vec![DataType::Float32, DataType::Float64],
1538                vec![DataType::Utf8]
1539            ]
1540        );
1541
1542        let type_signature =
1543            TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Int64]);
1544        let possible_types = type_signature.get_example_types();
1545        assert_eq!(
1546            possible_types,
1547            vec![
1548                vec![DataType::Float32, DataType::Float32],
1549                vec![DataType::Int64, DataType::Int64]
1550            ]
1551        );
1552
1553        let type_signature = TypeSignature::Coercible(vec![
1554            Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1555            Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
1556        ]);
1557        let possible_types = type_signature.get_example_types();
1558        assert_eq!(
1559            possible_types,
1560            vec![
1561                vec![DataType::Utf8, DataType::Int64],
1562                vec![DataType::LargeUtf8, DataType::Int64],
1563                vec![DataType::Utf8View, DataType::Int64]
1564            ]
1565        );
1566
1567        let type_signature =
1568            TypeSignature::Variadic(vec![DataType::Int32, DataType::Int64]);
1569        let possible_types = type_signature.get_example_types();
1570        assert_eq!(
1571            possible_types,
1572            vec![vec![DataType::Int32], vec![DataType::Int64]]
1573        );
1574
1575        let type_signature = TypeSignature::Numeric(2);
1576        let possible_types = type_signature.get_example_types();
1577        assert_eq!(
1578            possible_types,
1579            vec![
1580                vec![DataType::Int8, DataType::Int8],
1581                vec![DataType::Int16, DataType::Int16],
1582                vec![DataType::Int32, DataType::Int32],
1583                vec![DataType::Int64, DataType::Int64],
1584                vec![DataType::UInt8, DataType::UInt8],
1585                vec![DataType::UInt16, DataType::UInt16],
1586                vec![DataType::UInt32, DataType::UInt32],
1587                vec![DataType::UInt64, DataType::UInt64],
1588                vec![DataType::Float32, DataType::Float32],
1589                vec![DataType::Float64, DataType::Float64]
1590            ]
1591        );
1592
1593        let type_signature = TypeSignature::String(2);
1594        let possible_types = type_signature.get_example_types();
1595        assert_eq!(
1596            possible_types,
1597            vec![
1598                vec![DataType::Utf8, DataType::Utf8],
1599                vec![DataType::LargeUtf8, DataType::LargeUtf8],
1600                vec![DataType::Utf8View, DataType::Utf8View]
1601            ]
1602        );
1603    }
1604
1605    #[test]
1606    fn test_signature_with_parameter_names() {
1607        let sig = Signature::exact(
1608            vec![DataType::Int32, DataType::Utf8],
1609            Volatility::Immutable,
1610        )
1611        .with_parameter_names(vec!["count".to_string(), "name".to_string()])
1612        .unwrap();
1613
1614        assert_eq!(
1615            sig.parameter_names,
1616            Some(vec!["count".to_string(), "name".to_string()])
1617        );
1618        assert_eq!(
1619            sig.type_signature,
1620            TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8])
1621        );
1622    }
1623
1624    #[test]
1625    fn test_signature_parameter_names_wrong_count() {
1626        let result = Signature::exact(
1627            vec![DataType::Int32, DataType::Utf8],
1628            Volatility::Immutable,
1629        )
1630        .with_parameter_names(vec!["count".to_string()]); // Only 1 name for 2 args
1631
1632        assert!(result.is_err());
1633        assert!(
1634            result
1635                .unwrap_err()
1636                .to_string()
1637                .contains("does not match signature arity")
1638        );
1639    }
1640
1641    #[test]
1642    fn test_signature_parameter_names_duplicate() {
1643        let result = Signature::exact(
1644            vec![DataType::Int32, DataType::Int32],
1645            Volatility::Immutable,
1646        )
1647        .with_parameter_names(vec!["count".to_string(), "count".to_string()]);
1648
1649        assert!(result.is_err());
1650        assert!(
1651            result
1652                .unwrap_err()
1653                .to_string()
1654                .contains("Duplicate parameter name")
1655        );
1656    }
1657
1658    #[test]
1659    fn test_signature_parameter_names_variadic() {
1660        let result = Signature::variadic(vec![DataType::Int32], Volatility::Immutable)
1661            .with_parameter_names(vec!["arg".to_string()]);
1662
1663        assert!(result.is_err());
1664        assert!(
1665            result
1666                .unwrap_err()
1667                .to_string()
1668                .contains("variable arity signature")
1669        );
1670    }
1671
1672    #[test]
1673    fn test_signature_without_parameter_names() {
1674        let sig = Signature::exact(
1675            vec![DataType::Int32, DataType::Utf8],
1676            Volatility::Immutable,
1677        );
1678
1679        assert_eq!(sig.parameter_names, None);
1680    }
1681
1682    #[test]
1683    fn test_signature_uniform_with_parameter_names() {
1684        let sig = Signature::uniform(3, vec![DataType::Float64], Volatility::Immutable)
1685            .with_parameter_names(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1686            .unwrap();
1687
1688        assert_eq!(
1689            sig.parameter_names,
1690            Some(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1691        );
1692    }
1693
1694    #[test]
1695    fn test_signature_numeric_with_parameter_names() {
1696        let sig = Signature::numeric(2, Volatility::Immutable)
1697            .with_parameter_names(vec!["a".to_string(), "b".to_string()])
1698            .unwrap();
1699
1700        assert_eq!(
1701            sig.parameter_names,
1702            Some(vec!["a".to_string(), "b".to_string()])
1703        );
1704    }
1705
1706    #[test]
1707    fn test_signature_nullary_with_empty_names() {
1708        let sig = Signature::nullary(Volatility::Immutable)
1709            .with_parameter_names(Vec::<String>::new())
1710            .unwrap();
1711
1712        assert_eq!(sig.parameter_names, Some(vec![]));
1713    }
1714
1715    #[test]
1716    fn test_to_string_repr_with_names_exact() {
1717        let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1718
1719        assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
1720
1721        let names = vec!["id".to_string(), "name".to_string()];
1722        assert_eq!(
1723            sig.to_string_repr_with_names(Some(&names)),
1724            vec!["id: Int32, name: Utf8"]
1725        );
1726    }
1727
1728    #[test]
1729    fn test_to_string_repr_with_names_any() {
1730        let sig = TypeSignature::Any(3);
1731
1732        assert_eq!(sig.to_string_repr_with_names(None), vec!["Any, Any, Any"]);
1733
1734        let names = vec!["x".to_string(), "y".to_string(), "z".to_string()];
1735        assert_eq!(
1736            sig.to_string_repr_with_names(Some(&names)),
1737            vec!["x: Any, y: Any, z: Any"]
1738        );
1739    }
1740
1741    #[test]
1742    fn test_to_string_repr_with_names_one_of() {
1743        let sig =
1744            TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]);
1745
1746        assert_eq!(
1747            sig.to_string_repr_with_names(None),
1748            vec!["Any, Any", "Any, Any, Any"]
1749        );
1750
1751        let names = vec![
1752            "str".to_string(),
1753            "start_pos".to_string(),
1754            "length".to_string(),
1755        ];
1756        assert_eq!(
1757            sig.to_string_repr_with_names(Some(&names)),
1758            vec![
1759                "str: Any, start_pos: Any",
1760                "str: Any, start_pos: Any, length: Any"
1761            ]
1762        );
1763    }
1764
1765    #[test]
1766    fn test_to_string_repr_with_names_partial() {
1767        // This simulates providing max arity names for a OneOf signature
1768        let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1769
1770        // Provide 3 names for 2-parameter signature (extra name is ignored via zip)
1771        let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1772        assert_eq!(
1773            sig.to_string_repr_with_names(Some(&names)),
1774            vec!["a: Int32, b: Utf8"]
1775        );
1776    }
1777
1778    #[test]
1779    fn test_to_string_repr_with_names_uniform() {
1780        let sig = TypeSignature::Uniform(2, vec![DataType::Float64]);
1781
1782        assert_eq!(
1783            sig.to_string_repr_with_names(None),
1784            vec!["Float64, Float64"]
1785        );
1786
1787        let names = vec!["x".to_string(), "y".to_string()];
1788        assert_eq!(
1789            sig.to_string_repr_with_names(Some(&names)),
1790            vec!["x: Float64, y: Float64"]
1791        );
1792    }
1793
1794    #[test]
1795    fn test_to_string_repr_with_names_coercible() {
1796        let sig = TypeSignature::Coercible(vec![
1797            Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1798            Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1799        ]);
1800
1801        let names = vec!["a".to_string(), "b".to_string()];
1802        let result = sig.to_string_repr_with_names(Some(&names));
1803        // Check that it contains the parameter names with type annotations
1804        assert_eq!(result.len(), 1);
1805        assert!(result[0].starts_with("a: "));
1806        assert!(result[0].contains(", b: "));
1807    }
1808
1809    #[test]
1810    fn test_to_string_repr_with_names_comparable_numeric_string() {
1811        let comparable = TypeSignature::Comparable(3);
1812        let numeric = TypeSignature::Numeric(2);
1813        let string_sig = TypeSignature::String(2);
1814
1815        let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1816
1817        // All should show parameter names with type annotations
1818        assert_eq!(
1819            comparable.to_string_repr_with_names(Some(&names)),
1820            vec!["a: Comparable, b: Comparable, c: Comparable"]
1821        );
1822        assert_eq!(
1823            numeric.to_string_repr_with_names(Some(&names)),
1824            vec!["a: Numeric, b: Numeric"]
1825        );
1826        assert_eq!(
1827            string_sig.to_string_repr_with_names(Some(&names)),
1828            vec!["a: String, b: String"]
1829        );
1830    }
1831
1832    #[test]
1833    fn test_to_string_repr_with_names_variadic_fallback() {
1834        let variadic = TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]);
1835        let names = vec!["x".to_string()];
1836        assert_eq!(
1837            variadic.to_string_repr_with_names(Some(&names)),
1838            variadic.to_string_repr()
1839        );
1840
1841        let variadic_any = TypeSignature::VariadicAny;
1842        assert_eq!(
1843            variadic_any.to_string_repr_with_names(Some(&names)),
1844            variadic_any.to_string_repr()
1845        );
1846
1847        // UserDefined now shows parameter names when available
1848        let user_defined = TypeSignature::UserDefined;
1849        assert_eq!(
1850            user_defined.to_string_repr_with_names(Some(&names)),
1851            vec!["x"]
1852        );
1853        assert_eq!(
1854            user_defined.to_string_repr_with_names(None),
1855            user_defined.to_string_repr()
1856        );
1857    }
1858
1859    #[test]
1860    fn test_to_string_repr_with_names_nullary() {
1861        let sig = TypeSignature::Nullary;
1862        let names = vec!["x".to_string()];
1863
1864        // Should return empty representation, names don't apply
1865        assert_eq!(
1866            sig.to_string_repr_with_names(Some(&names)),
1867            vec!["NullAry()"]
1868        );
1869        assert_eq!(sig.to_string_repr_with_names(None), vec!["NullAry()"]);
1870    }
1871
1872    #[test]
1873    fn test_to_string_repr_with_names_array_signature() {
1874        let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1875            arguments: vec![
1876                ArrayFunctionArgument::Array,
1877                ArrayFunctionArgument::Index,
1878                ArrayFunctionArgument::Element,
1879            ],
1880            array_coercion: None,
1881        });
1882
1883        assert_eq!(
1884            sig.to_string_repr_with_names(None),
1885            vec!["array, index, element"]
1886        );
1887
1888        let names = vec!["arr".to_string(), "idx".to_string(), "val".to_string()];
1889        assert_eq!(
1890            sig.to_string_repr_with_names(Some(&names)),
1891            vec!["arr: array, idx: index, val: element"]
1892        );
1893
1894        let recursive =
1895            TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1896        let names = vec!["array".to_string()];
1897        assert_eq!(
1898            recursive.to_string_repr_with_names(Some(&names)),
1899            vec!["array: recursive_array"]
1900        );
1901
1902        // Test MapArray (1 argument)
1903        let map_array = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1904        let names = vec!["map".to_string()];
1905        assert_eq!(
1906            map_array.to_string_repr_with_names(Some(&names)),
1907            vec!["map: map_array"]
1908        );
1909    }
1910
1911    #[test]
1912    fn test_type_signature_arity_exact() {
1913        let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1914        assert_eq!(sig.arity(), Arity::Fixed(2));
1915
1916        let sig = TypeSignature::Exact(vec![]);
1917        assert_eq!(sig.arity(), Arity::Fixed(0));
1918    }
1919
1920    #[test]
1921    fn test_type_signature_arity_uniform() {
1922        let sig = TypeSignature::Uniform(3, vec![DataType::Float64]);
1923        assert_eq!(sig.arity(), Arity::Fixed(3));
1924
1925        let sig = TypeSignature::Uniform(1, vec![DataType::Int32]);
1926        assert_eq!(sig.arity(), Arity::Fixed(1));
1927    }
1928
1929    #[test]
1930    fn test_type_signature_arity_numeric() {
1931        let sig = TypeSignature::Numeric(2);
1932        assert_eq!(sig.arity(), Arity::Fixed(2));
1933    }
1934
1935    #[test]
1936    fn test_type_signature_arity_string() {
1937        let sig = TypeSignature::String(3);
1938        assert_eq!(sig.arity(), Arity::Fixed(3));
1939    }
1940
1941    #[test]
1942    fn test_type_signature_arity_comparable() {
1943        let sig = TypeSignature::Comparable(2);
1944        assert_eq!(sig.arity(), Arity::Fixed(2));
1945    }
1946
1947    #[test]
1948    fn test_type_signature_arity_any() {
1949        let sig = TypeSignature::Any(4);
1950        assert_eq!(sig.arity(), Arity::Fixed(4));
1951    }
1952
1953    #[test]
1954    fn test_type_signature_arity_coercible() {
1955        let sig = TypeSignature::Coercible(vec![
1956            Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1957            Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1958        ]);
1959        assert_eq!(sig.arity(), Arity::Fixed(2));
1960    }
1961
1962    #[test]
1963    fn test_type_signature_arity_nullary() {
1964        let sig = TypeSignature::Nullary;
1965        assert_eq!(sig.arity(), Arity::Fixed(0));
1966    }
1967
1968    #[test]
1969    fn test_type_signature_arity_array_signature() {
1970        // Test Array variant with 2 arguments
1971        let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1972            arguments: vec![ArrayFunctionArgument::Array, ArrayFunctionArgument::Index],
1973            array_coercion: None,
1974        });
1975        assert_eq!(sig.arity(), Arity::Fixed(2));
1976
1977        // Test Array variant with 3 arguments
1978        let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1979            arguments: vec![
1980                ArrayFunctionArgument::Array,
1981                ArrayFunctionArgument::Element,
1982                ArrayFunctionArgument::Index,
1983            ],
1984            array_coercion: None,
1985        });
1986        assert_eq!(sig.arity(), Arity::Fixed(3));
1987
1988        // Test RecursiveArray variant
1989        let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1990        assert_eq!(sig.arity(), Arity::Fixed(1));
1991
1992        // Test MapArray variant
1993        let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1994        assert_eq!(sig.arity(), Arity::Fixed(1));
1995    }
1996
1997    #[test]
1998    fn test_type_signature_arity_one_of_fixed() {
1999        // OneOf with all fixed arity variants should return max arity
2000        let sig = TypeSignature::OneOf(vec![
2001            TypeSignature::Exact(vec![DataType::Int32]),
2002            TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2003            TypeSignature::Exact(vec![
2004                DataType::Int32,
2005                DataType::Utf8,
2006                DataType::Float64,
2007            ]),
2008        ]);
2009        assert_eq!(sig.arity(), Arity::Fixed(3));
2010    }
2011
2012    #[test]
2013    fn test_type_signature_arity_one_of_variable() {
2014        // OneOf with variable arity variant should return Variable
2015        let sig = TypeSignature::OneOf(vec![
2016            TypeSignature::Exact(vec![DataType::Int32]),
2017            TypeSignature::VariadicAny,
2018        ]);
2019        assert_eq!(sig.arity(), Arity::Variable);
2020    }
2021
2022    #[test]
2023    fn test_type_signature_arity_variadic() {
2024        let sig = TypeSignature::Variadic(vec![DataType::Int32]);
2025        assert_eq!(sig.arity(), Arity::Variable);
2026
2027        let sig = TypeSignature::VariadicAny;
2028        assert_eq!(sig.arity(), Arity::Variable);
2029    }
2030
2031    #[test]
2032    fn test_type_signature_arity_user_defined() {
2033        let sig = TypeSignature::UserDefined;
2034        assert_eq!(sig.arity(), Arity::Variable);
2035    }
2036}