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

datafusion_expr_common/
signature.rs
