Skip to main content

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