1use std::fmt::Display;
21use std::hash::Hash;
22use std::sync::Arc;
23
24use crate::type_coercion::aggregates::NUMERICS;
25use arrow::datatypes::{
26 DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DECIMAL128_MAX_PRECISION, DataType,
27 Decimal128Type, DecimalType, Field, IntervalUnit, TimeUnit,
28};
29use datafusion_common::types::{LogicalType, LogicalTypeRef, NativeType};
30use datafusion_common::utils::ListCoercion;
31use datafusion_common::{Result, internal_err, plan_err};
32use indexmap::IndexSet;
33use itertools::Itertools;
34
35pub const TIMEZONE_WILDCARD: &str = "+TZ";
44
45pub const FIXED_SIZE_LIST_WILDCARD: i32 = i32::MIN;
49
50#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
57pub enum Volatility {
58 Immutable,
66 Stable,
78 Volatile,
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum Arity {
94 Fixed(usize),
96 Variable,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
158pub enum TypeSignature {
159 Variadic(Vec<DataType>),
168 UserDefined,
176 VariadicAny,
178 Uniform(usize, Vec<DataType>),
185 Exact(Vec<DataType>),
189 Coercible(Vec<Coercion>),
197 Comparable(usize),
216 Any(usize),
220 OneOf(Vec<TypeSignature>),
230 ArraySignature(ArrayFunctionSignature),
232 Numeric(usize),
240 String(usize),
251 Nullary,
253}
254
255impl TypeSignature {
256 #[inline]
257 pub fn is_one_of(&self) -> bool {
258 matches!(self, TypeSignature::OneOf(_))
259 }
260
261 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 let has_variable = variants.iter().any(|v| v.arity() == Arity::Variable);
302 if has_variable {
303 return Arity::Variable;
304 }
305 let max_arity = variants
307 .iter()
308 .filter_map(|v| match v.arity() {
309 Arity::Fixed(n) => Some(n),
310 Arity::Variable => None,
311 })
312 .max();
313 match max_arity {
314 Some(n) => Arity::Fixed(n),
315 None => Arity::Variable,
316 }
317 }
318 TypeSignature::Variadic(_)
319 | TypeSignature::VariadicAny
320 | TypeSignature::UserDefined => Arity::Variable,
321 }
322 }
323}
324
325#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Hash)]
334pub enum TypeSignatureClass {
335 Any,
337 Timestamp,
339 Time,
341 Interval,
343 Duration,
345 Native(LogicalTypeRef),
347 Integer,
349 Float,
351 Decimal,
353 Numeric,
355 Binary,
357}
358
359impl Display for TypeSignatureClass {
360 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
361 write!(f, "TypeSignatureClass::{self:?}")
362 }
363}
364
365impl TypeSignatureClass {
366 fn get_example_types(&self) -> Vec<DataType> {
371 match self {
372 TypeSignatureClass::Any => vec![],
375 TypeSignatureClass::Native(l) => get_data_types(l.native()),
376 TypeSignatureClass::Timestamp => {
377 vec![
378 DataType::Timestamp(TimeUnit::Nanosecond, None),
379 DataType::Timestamp(
380 TimeUnit::Nanosecond,
381 Some(TIMEZONE_WILDCARD.into()),
382 ),
383 ]
384 }
385 TypeSignatureClass::Time => {
386 vec![DataType::Time64(TimeUnit::Nanosecond)]
387 }
388 TypeSignatureClass::Interval => {
389 vec![DataType::Interval(IntervalUnit::DayTime)]
390 }
391 TypeSignatureClass::Duration => {
392 vec![DataType::Duration(TimeUnit::Nanosecond)]
393 }
394 TypeSignatureClass::Integer => {
395 vec![DataType::Int64]
396 }
397 TypeSignatureClass::Binary => {
398 vec![DataType::Binary]
399 }
400 TypeSignatureClass::Decimal => vec![Decimal128Type::DEFAULT_TYPE],
401 TypeSignatureClass::Float => vec![DataType::Float64],
402 TypeSignatureClass::Numeric => vec![
403 DataType::Float64,
404 DataType::Int64,
405 Decimal128Type::DEFAULT_TYPE,
406 ],
407 }
408 }
409
410 pub fn matches_native_type(&self, logical_type: &NativeType) -> bool {
412 if logical_type == &NativeType::Null {
413 return true;
414 }
415
416 match self {
417 TypeSignatureClass::Any => true,
418 TypeSignatureClass::Native(t) if t.native() == logical_type => true,
419 TypeSignatureClass::Timestamp if logical_type.is_timestamp() => true,
420 TypeSignatureClass::Time if logical_type.is_time() => true,
421 TypeSignatureClass::Interval if logical_type.is_interval() => true,
422 TypeSignatureClass::Duration if logical_type.is_duration() => true,
423 TypeSignatureClass::Integer if logical_type.is_integer() => true,
424 TypeSignatureClass::Binary if logical_type.is_binary() => true,
425 TypeSignatureClass::Decimal if logical_type.is_decimal() => true,
426 TypeSignatureClass::Float if logical_type.is_float() => true,
427 TypeSignatureClass::Numeric if logical_type.is_numeric() => true,
428 _ => false,
429 }
430 }
431
432 pub fn default_casted_type(
434 &self,
435 native_type: &NativeType,
436 origin_type: &DataType,
437 ) -> Result<DataType> {
438 match self {
439 TypeSignatureClass::Any => Ok(origin_type.to_owned()),
440 TypeSignatureClass::Native(logical_type) => {
441 logical_type.native().default_cast_for(origin_type)
442 }
443 TypeSignatureClass::Timestamp if native_type.is_timestamp() => {
445 Ok(origin_type.to_owned())
446 }
447 TypeSignatureClass::Time if native_type.is_time() => {
448 Ok(origin_type.to_owned())
449 }
450 TypeSignatureClass::Interval if native_type.is_interval() => {
451 Ok(origin_type.to_owned())
452 }
453 TypeSignatureClass::Duration if native_type.is_duration() => {
454 Ok(origin_type.to_owned())
455 }
456 TypeSignatureClass::Integer if native_type.is_integer() => {
457 Ok(origin_type.to_owned())
458 }
459 TypeSignatureClass::Binary if native_type.is_binary() => {
460 Ok(origin_type.to_owned())
461 }
462 TypeSignatureClass::Decimal if native_type.is_decimal() => {
463 Ok(origin_type.to_owned())
464 }
465 TypeSignatureClass::Float if native_type.is_float() => {
466 Ok(origin_type.to_owned())
467 }
468 TypeSignatureClass::Numeric if native_type.is_numeric() => {
469 Ok(origin_type.to_owned())
470 }
471 _ if native_type.is_null() => Ok(origin_type.to_owned()),
472 _ => internal_err!("May miss the matching logic in `matches_native_type`"),
473 }
474 }
475}
476
477#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
478pub enum ArrayFunctionSignature {
479 Array {
481 arguments: Vec<ArrayFunctionArgument>,
483 array_coercion: Option<ListCoercion>,
485 },
486 RecursiveArray,
489 MapArray,
492}
493
494impl Display for ArrayFunctionSignature {
495 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
496 match self {
497 ArrayFunctionSignature::Array { arguments, .. } => {
498 for (idx, argument) in arguments.iter().enumerate() {
499 write!(f, "{argument}")?;
500 if idx != arguments.len() - 1 {
501 write!(f, ", ")?;
502 }
503 }
504 Ok(())
505 }
506 ArrayFunctionSignature::RecursiveArray => {
507 write!(f, "recursive_array")
508 }
509 ArrayFunctionSignature::MapArray => {
510 write!(f, "map_array")
511 }
512 }
513 }
514}
515
516#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
517pub enum ArrayFunctionArgument {
518 Element,
521 Index,
523 Array,
526 String,
528}
529
530impl Display for ArrayFunctionArgument {
531 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
532 match self {
533 ArrayFunctionArgument::Element => {
534 write!(f, "element")
535 }
536 ArrayFunctionArgument::Index => {
537 write!(f, "index")
538 }
539 ArrayFunctionArgument::Array => {
540 write!(f, "array")
541 }
542 ArrayFunctionArgument::String => {
543 write!(f, "string")
544 }
545 }
546 }
547}
548
549impl TypeSignature {
550 pub fn to_string_repr(&self) -> Vec<String> {
551 match self {
552 TypeSignature::Nullary => {
553 vec!["NullAry()".to_string()]
554 }
555 TypeSignature::Variadic(types) => {
556 vec![format!("{}, ..", Self::join_types(types, "/"))]
557 }
558 TypeSignature::Uniform(arg_count, valid_types) => {
559 vec![
560 std::iter::repeat_n(Self::join_types(valid_types, "/"), *arg_count)
561 .collect::<Vec<String>>()
562 .join(", "),
563 ]
564 }
565 TypeSignature::String(num) => {
566 vec![format!("String({num})")]
567 }
568 TypeSignature::Numeric(num) => {
569 vec![format!("Numeric({num})")]
570 }
571 TypeSignature::Comparable(num) => {
572 vec![format!("Comparable({num})")]
573 }
574 TypeSignature::Coercible(coercions) => {
575 vec![Self::join_types(coercions, ", ")]
576 }
577 TypeSignature::Exact(types) => {
578 vec![Self::join_types(types, ", ")]
579 }
580 TypeSignature::Any(arg_count) => {
581 vec![
582 std::iter::repeat_n("Any", *arg_count)
583 .collect::<Vec<&str>>()
584 .join(", "),
585 ]
586 }
587 TypeSignature::UserDefined => {
588 vec!["UserDefined".to_string()]
589 }
590 TypeSignature::VariadicAny => vec!["Any, .., Any".to_string()],
591 TypeSignature::OneOf(sigs) => {
592 sigs.iter().flat_map(|s| s.to_string_repr()).collect()
593 }
594 TypeSignature::ArraySignature(array_signature) => {
595 vec![array_signature.to_string()]
596 }
597 }
598 }
599
600 pub fn to_string_repr_with_names(
626 &self,
627 parameter_names: Option<&[String]>,
628 ) -> Vec<String> {
629 match self {
630 TypeSignature::Exact(types) => {
631 if let Some(names) = parameter_names {
632 vec![
633 names
634 .iter()
635 .zip(types.iter())
636 .map(|(name, typ)| format!("{name}: {typ}"))
637 .collect::<Vec<_>>()
638 .join(", "),
639 ]
640 } else {
641 vec![Self::join_types(types, ", ")]
642 }
643 }
644 TypeSignature::Any(count) => {
645 if let Some(names) = parameter_names {
646 vec![
647 names
648 .iter()
649 .take(*count)
650 .map(|name| format!("{name}: Any"))
651 .collect::<Vec<_>>()
652 .join(", "),
653 ]
654 } else {
655 vec![
656 std::iter::repeat_n("Any", *count)
657 .collect::<Vec<&str>>()
658 .join(", "),
659 ]
660 }
661 }
662 TypeSignature::Uniform(count, types) => {
663 if let Some(names) = parameter_names {
664 let type_str = Self::join_types(types, "/");
665 vec![
666 names
667 .iter()
668 .take(*count)
669 .map(|name| format!("{name}: {type_str}"))
670 .collect::<Vec<_>>()
671 .join(", "),
672 ]
673 } else {
674 self.to_string_repr()
675 }
676 }
677 TypeSignature::Coercible(coercions) => {
678 if let Some(names) = parameter_names {
679 vec![
680 names
681 .iter()
682 .zip(coercions.iter())
683 .map(|(name, coercion)| format!("{name}: {coercion}"))
684 .collect::<Vec<_>>()
685 .join(", "),
686 ]
687 } else {
688 vec![Self::join_types(coercions, ", ")]
689 }
690 }
691 TypeSignature::Comparable(count) => {
692 if let Some(names) = parameter_names {
693 vec![
694 names
695 .iter()
696 .take(*count)
697 .map(|name| format!("{name}: Comparable"))
698 .collect::<Vec<_>>()
699 .join(", "),
700 ]
701 } else {
702 self.to_string_repr()
703 }
704 }
705 TypeSignature::Numeric(count) => {
706 if let Some(names) = parameter_names {
707 vec![
708 names
709 .iter()
710 .take(*count)
711 .map(|name| format!("{name}: Numeric"))
712 .collect::<Vec<_>>()
713 .join(", "),
714 ]
715 } else {
716 self.to_string_repr()
717 }
718 }
719 TypeSignature::String(count) => {
720 if let Some(names) = parameter_names {
721 vec![
722 names
723 .iter()
724 .take(*count)
725 .map(|name| format!("{name}: String"))
726 .collect::<Vec<_>>()
727 .join(", "),
728 ]
729 } else {
730 self.to_string_repr()
731 }
732 }
733 TypeSignature::Nullary => self.to_string_repr(),
734 TypeSignature::ArraySignature(array_sig) => {
735 if let Some(names) = parameter_names {
736 match array_sig {
737 ArrayFunctionSignature::Array { arguments, .. } => {
738 vec![
739 names
740 .iter()
741 .zip(arguments.iter())
742 .map(|(name, arg_type)| format!("{name}: {arg_type}"))
743 .collect::<Vec<_>>()
744 .join(", "),
745 ]
746 }
747 ArrayFunctionSignature::RecursiveArray => {
748 vec![
749 names
750 .iter()
751 .take(1)
752 .map(|name| format!("{name}: recursive_array"))
753 .collect::<Vec<_>>()
754 .join(", "),
755 ]
756 }
757 ArrayFunctionSignature::MapArray => {
758 vec![
759 names
760 .iter()
761 .take(1)
762 .map(|name| format!("{name}: map_array"))
763 .collect::<Vec<_>>()
764 .join(", "),
765 ]
766 }
767 }
768 } else {
769 self.to_string_repr()
770 }
771 }
772 TypeSignature::OneOf(sigs) => sigs
773 .iter()
774 .flat_map(|s| s.to_string_repr_with_names(parameter_names))
775 .collect(),
776 TypeSignature::UserDefined => {
777 if let Some(names) = parameter_names {
778 vec![names.join(", ")]
779 } else {
780 self.to_string_repr()
781 }
782 }
783 TypeSignature::Variadic(_) | TypeSignature::VariadicAny => {
785 self.to_string_repr()
786 }
787 }
788 }
789
790 pub fn join_types<T: Display>(types: &[T], delimiter: &str) -> String {
792 types
793 .iter()
794 .map(|t| t.to_string())
795 .collect::<Vec<String>>()
796 .join(delimiter)
797 }
798
799 pub fn supports_zero_argument(&self) -> bool {
801 match &self {
802 TypeSignature::Exact(vec) => vec.is_empty(),
803 TypeSignature::Nullary => true,
804 TypeSignature::OneOf(types) => types
805 .iter()
806 .any(|type_sig| type_sig.supports_zero_argument()),
807 _ => false,
808 }
809 }
810
811 pub fn used_to_support_zero_arguments(&self) -> bool {
814 match &self {
815 TypeSignature::Any(num) => *num == 0,
816 _ => self.supports_zero_argument(),
817 }
818 }
819
820 #[deprecated(since = "46.0.0", note = "See get_example_types instead")]
821 pub fn get_possible_types(&self) -> Vec<Vec<DataType>> {
822 self.get_example_types()
823 }
824
825 pub fn get_example_types(&self) -> Vec<Vec<DataType>> {
832 match self {
833 TypeSignature::Exact(types) => vec![types.clone()],
834 TypeSignature::OneOf(types) => types
835 .iter()
836 .flat_map(|type_sig| type_sig.get_example_types())
837 .collect(),
838 TypeSignature::Uniform(arg_count, types) => types
839 .iter()
840 .cloned()
841 .map(|data_type| vec![data_type; *arg_count])
842 .collect(),
843 TypeSignature::Coercible(coercions) => coercions
844 .iter()
845 .map(|c| {
846 let mut all_types: IndexSet<DataType> =
847 c.desired_type().get_example_types().into_iter().collect();
848
849 if let Some(implicit_coercion) = c.implicit_coercion() {
850 let allowed_casts: Vec<DataType> = implicit_coercion
851 .allowed_source_types
852 .iter()
853 .flat_map(|t| t.get_example_types())
854 .collect();
855 all_types.extend(allowed_casts);
856 }
857
858 all_types.into_iter().collect::<Vec<_>>()
859 })
860 .multi_cartesian_product()
861 .collect(),
862 TypeSignature::Variadic(types) => types
863 .iter()
864 .cloned()
865 .map(|data_type| vec![data_type])
866 .collect(),
867 TypeSignature::Numeric(arg_count) => NUMERICS
868 .iter()
869 .cloned()
870 .map(|numeric_type| vec![numeric_type; *arg_count])
871 .collect(),
872 TypeSignature::String(arg_count) => get_data_types(&NativeType::String)
873 .into_iter()
874 .map(|dt| vec![dt; *arg_count])
875 .collect::<Vec<_>>(),
876 TypeSignature::Any(_)
878 | TypeSignature::Comparable(_)
879 | TypeSignature::Nullary
880 | TypeSignature::VariadicAny
881 | TypeSignature::ArraySignature(_)
882 | TypeSignature::UserDefined => vec![],
883 }
884 }
885}
886
887fn get_data_types(native_type: &NativeType) -> Vec<DataType> {
888 match native_type {
889 NativeType::Null => vec![DataType::Null],
890 NativeType::Boolean => vec![DataType::Boolean],
891 NativeType::Int8 => vec![DataType::Int8],
892 NativeType::Int16 => vec![DataType::Int16],
893 NativeType::Int32 => vec![DataType::Int32],
894 NativeType::Int64 => vec![DataType::Int64],
895 NativeType::UInt8 => vec![DataType::UInt8],
896 NativeType::UInt16 => vec![DataType::UInt16],
897 NativeType::UInt32 => vec![DataType::UInt32],
898 NativeType::UInt64 => vec![DataType::UInt64],
899 NativeType::Float16 => vec![DataType::Float16],
900 NativeType::Float32 => vec![DataType::Float32],
901 NativeType::Float64 => vec![DataType::Float64],
902 NativeType::Date => vec![DataType::Date32, DataType::Date64],
903 NativeType::Binary => vec![
904 DataType::Binary,
905 DataType::LargeBinary,
906 DataType::BinaryView,
907 ],
908 NativeType::String => {
909 vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]
910 }
911 NativeType::Decimal(precision, scale) => {
912 let mut types = vec![DataType::Decimal256(*precision, *scale)];
914 if *precision <= DECIMAL32_MAX_PRECISION {
915 types.push(DataType::Decimal32(*precision, *scale));
916 }
917 if *precision <= DECIMAL64_MAX_PRECISION {
918 types.push(DataType::Decimal64(*precision, *scale));
919 }
920 if *precision <= DECIMAL128_MAX_PRECISION {
921 types.push(DataType::Decimal128(*precision, *scale));
922 }
923 types
924 }
925 NativeType::Timestamp(time_unit, timezone) => {
926 vec![DataType::Timestamp(*time_unit, timezone.to_owned())]
927 }
928 NativeType::Time(TimeUnit::Second) => vec![DataType::Time32(TimeUnit::Second)],
929 NativeType::Time(TimeUnit::Millisecond) => {
930 vec![DataType::Time32(TimeUnit::Millisecond)]
931 }
932 NativeType::Time(TimeUnit::Microsecond) => {
933 vec![DataType::Time64(TimeUnit::Microsecond)]
934 }
935 NativeType::Time(TimeUnit::Nanosecond) => {
936 vec![DataType::Time64(TimeUnit::Nanosecond)]
937 }
938 NativeType::Duration(time_unit) => vec![DataType::Duration(*time_unit)],
939 NativeType::Interval(interval_unit) => vec![DataType::Interval(*interval_unit)],
940 NativeType::FixedSizeBinary(size) => vec![DataType::FixedSizeBinary(*size)],
941 NativeType::FixedSizeList(logical_field, size) => {
942 get_data_types(logical_field.logical_type.native())
943 .iter()
944 .map(|child_dt| {
945 let field = Field::new(
946 logical_field.name.clone(),
947 child_dt.clone(),
948 logical_field.nullable,
949 );
950 DataType::FixedSizeList(Arc::new(field), *size)
951 })
952 .collect()
953 }
954 NativeType::List(_)
956 | NativeType::Struct(_)
957 | NativeType::Union(_)
958 | NativeType::Map(_) => {
959 vec![]
960 }
961 }
962}
963
964#[derive(Debug, Clone, Eq, PartialOrd)]
989pub enum Coercion {
990 Exact {
992 desired_type: TypeSignatureClass,
994 },
995
996 Implicit {
998 desired_type: TypeSignatureClass,
1000 implicit_coercion: ImplicitCoercion,
1002 },
1003}
1004
1005impl Coercion {
1006 pub fn new_exact(desired_type: TypeSignatureClass) -> Self {
1007 Self::Exact { desired_type }
1008 }
1009
1010 pub fn new_implicit(
1015 desired_type: TypeSignatureClass,
1016 allowed_source_types: Vec<TypeSignatureClass>,
1017 default_casted_type: NativeType,
1018 ) -> Self {
1019 Self::Implicit {
1020 desired_type,
1021 implicit_coercion: ImplicitCoercion {
1022 allowed_source_types,
1023 default_casted_type,
1024 },
1025 }
1026 }
1027
1028 pub fn allowed_source_types(&self) -> &[TypeSignatureClass] {
1029 match self {
1030 Coercion::Exact { .. } => &[],
1031 Coercion::Implicit {
1032 implicit_coercion, ..
1033 } => implicit_coercion.allowed_source_types.as_slice(),
1034 }
1035 }
1036
1037 pub fn default_casted_type(&self) -> Option<&NativeType> {
1038 match self {
1039 Coercion::Exact { .. } => None,
1040 Coercion::Implicit {
1041 implicit_coercion, ..
1042 } => Some(&implicit_coercion.default_casted_type),
1043 }
1044 }
1045
1046 pub fn desired_type(&self) -> &TypeSignatureClass {
1047 match self {
1048 Coercion::Exact { desired_type } => desired_type,
1049 Coercion::Implicit { desired_type, .. } => desired_type,
1050 }
1051 }
1052
1053 pub fn implicit_coercion(&self) -> Option<&ImplicitCoercion> {
1054 match self {
1055 Coercion::Exact { .. } => None,
1056 Coercion::Implicit {
1057 implicit_coercion, ..
1058 } => Some(implicit_coercion),
1059 }
1060 }
1061}
1062
1063impl Display for Coercion {
1064 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1065 write!(f, "Coercion({}", self.desired_type())?;
1066 if let Some(implicit_coercion) = self.implicit_coercion() {
1067 write!(f, ", implicit_coercion={implicit_coercion}",)
1068 } else {
1069 write!(f, ")")
1070 }
1071 }
1072}
1073
1074impl PartialEq for Coercion {
1075 fn eq(&self, other: &Self) -> bool {
1076 self.desired_type() == other.desired_type()
1077 && self.implicit_coercion() == other.implicit_coercion()
1078 }
1079}
1080
1081impl Hash for Coercion {
1082 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1083 self.desired_type().hash(state);
1084 self.implicit_coercion().hash(state);
1085 }
1086}
1087
1088#[derive(Debug, Clone, Eq, PartialOrd)]
1110pub struct ImplicitCoercion {
1111 allowed_source_types: Vec<TypeSignatureClass>,
1113
1114 default_casted_type: NativeType,
1118}
1119
1120impl Display for ImplicitCoercion {
1121 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1122 write!(
1123 f,
1124 "ImplicitCoercion({:?}, default_type={:?})",
1125 self.allowed_source_types, self.default_casted_type
1126 )
1127 }
1128}
1129
1130impl PartialEq for ImplicitCoercion {
1131 fn eq(&self, other: &Self) -> bool {
1132 self.allowed_source_types == other.allowed_source_types
1133 && self.default_casted_type == other.default_casted_type
1134 }
1135}
1136
1137impl Hash for ImplicitCoercion {
1138 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1139 self.allowed_source_types.hash(state);
1140 self.default_casted_type.hash(state);
1141 }
1142}
1143
1144#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
1151pub struct Signature {
1152 pub type_signature: TypeSignature,
1154 pub volatility: Volatility,
1156 pub parameter_names: Option<Vec<String>>,
1163}
1164
1165impl Signature {
1166 pub fn new(type_signature: TypeSignature, volatility: Volatility) -> Self {
1168 Signature {
1169 type_signature,
1170 volatility,
1171 parameter_names: None,
1172 }
1173 }
1174 pub fn variadic(common_types: Vec<DataType>, volatility: Volatility) -> Self {
1176 Self {
1177 type_signature: TypeSignature::Variadic(common_types),
1178 volatility,
1179 parameter_names: None,
1180 }
1181 }
1182 pub fn user_defined(volatility: Volatility) -> Self {
1184 Self {
1185 type_signature: TypeSignature::UserDefined,
1186 volatility,
1187 parameter_names: None,
1188 }
1189 }
1190
1191 pub fn numeric(arg_count: usize, volatility: Volatility) -> Self {
1193 Self {
1194 type_signature: TypeSignature::Numeric(arg_count),
1195 volatility,
1196 parameter_names: None,
1197 }
1198 }
1199
1200 pub fn string(arg_count: usize, volatility: Volatility) -> Self {
1202 Self {
1203 type_signature: TypeSignature::String(arg_count),
1204 volatility,
1205 parameter_names: None,
1206 }
1207 }
1208
1209 pub fn variadic_any(volatility: Volatility) -> Self {
1211 Self {
1212 type_signature: TypeSignature::VariadicAny,
1213 volatility,
1214 parameter_names: None,
1215 }
1216 }
1217 pub fn uniform(
1219 arg_count: usize,
1220 valid_types: Vec<DataType>,
1221 volatility: Volatility,
1222 ) -> Self {
1223 Self {
1224 type_signature: TypeSignature::Uniform(arg_count, valid_types),
1225 volatility,
1226 parameter_names: None,
1227 }
1228 }
1229 pub fn exact(exact_types: Vec<DataType>, volatility: Volatility) -> Self {
1231 Signature {
1232 type_signature: TypeSignature::Exact(exact_types),
1233 volatility,
1234 parameter_names: None,
1235 }
1236 }
1237
1238 pub fn coercible(target_types: Vec<Coercion>, volatility: Volatility) -> Self {
1240 Self {
1241 type_signature: TypeSignature::Coercible(target_types),
1242 volatility,
1243 parameter_names: None,
1244 }
1245 }
1246
1247 pub fn comparable(arg_count: usize, volatility: Volatility) -> Self {
1249 Self {
1250 type_signature: TypeSignature::Comparable(arg_count),
1251 volatility,
1252 parameter_names: None,
1253 }
1254 }
1255
1256 pub fn nullary(volatility: Volatility) -> Self {
1257 Signature {
1258 type_signature: TypeSignature::Nullary,
1259 volatility,
1260 parameter_names: None,
1261 }
1262 }
1263
1264 pub fn any(arg_count: usize, volatility: Volatility) -> Self {
1266 Signature {
1267 type_signature: TypeSignature::Any(arg_count),
1268 volatility,
1269 parameter_names: None,
1270 }
1271 }
1272
1273 pub fn one_of(type_signatures: Vec<TypeSignature>, volatility: Volatility) -> Self {
1275 Signature {
1276 type_signature: TypeSignature::OneOf(type_signatures),
1277 volatility,
1278 parameter_names: None,
1279 }
1280 }
1281
1282 pub fn array_and_element(volatility: Volatility) -> Self {
1284 Signature {
1285 type_signature: TypeSignature::ArraySignature(
1286 ArrayFunctionSignature::Array {
1287 arguments: vec![
1288 ArrayFunctionArgument::Array,
1289 ArrayFunctionArgument::Element,
1290 ],
1291 array_coercion: Some(ListCoercion::FixedSizedListToList),
1292 },
1293 ),
1294 volatility,
1295 parameter_names: None,
1296 }
1297 }
1298
1299 pub fn element_and_array(volatility: Volatility) -> Self {
1301 Signature {
1302 type_signature: TypeSignature::ArraySignature(
1303 ArrayFunctionSignature::Array {
1304 arguments: vec![
1305 ArrayFunctionArgument::Element,
1306 ArrayFunctionArgument::Array,
1307 ],
1308 array_coercion: Some(ListCoercion::FixedSizedListToList),
1309 },
1310 ),
1311 volatility,
1312 parameter_names: None,
1313 }
1314 }
1315
1316 pub fn arrays(
1318 n: usize,
1319 coercion: Option<ListCoercion>,
1320 volatility: Volatility,
1321 ) -> Self {
1322 Signature {
1323 type_signature: TypeSignature::ArraySignature(
1324 ArrayFunctionSignature::Array {
1325 arguments: vec![ArrayFunctionArgument::Array; n],
1326 array_coercion: coercion,
1327 },
1328 ),
1329 volatility,
1330 parameter_names: None,
1331 }
1332 }
1333
1334 pub fn array_and_element_and_optional_index(volatility: Volatility) -> Self {
1336 Signature {
1337 type_signature: TypeSignature::OneOf(vec![
1338 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1339 arguments: vec![
1340 ArrayFunctionArgument::Array,
1341 ArrayFunctionArgument::Element,
1342 ],
1343 array_coercion: Some(ListCoercion::FixedSizedListToList),
1344 }),
1345 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1346 arguments: vec![
1347 ArrayFunctionArgument::Array,
1348 ArrayFunctionArgument::Element,
1349 ArrayFunctionArgument::Index,
1350 ],
1351 array_coercion: Some(ListCoercion::FixedSizedListToList),
1352 }),
1353 ]),
1354 volatility,
1355 parameter_names: None,
1356 }
1357 }
1358
1359 pub fn array_and_index(volatility: Volatility) -> Self {
1361 Signature {
1362 type_signature: TypeSignature::ArraySignature(
1363 ArrayFunctionSignature::Array {
1364 arguments: vec![
1365 ArrayFunctionArgument::Array,
1366 ArrayFunctionArgument::Index,
1367 ],
1368 array_coercion: Some(ListCoercion::FixedSizedListToList),
1369 },
1370 ),
1371 volatility,
1372 parameter_names: None,
1373 }
1374 }
1375
1376 pub fn array(volatility: Volatility) -> Self {
1378 Signature::arrays(1, Some(ListCoercion::FixedSizedListToList), volatility)
1379 }
1380
1381 pub fn with_parameter_names(mut self, names: Vec<impl Into<String>>) -> Result<Self> {
1397 let names = names.into_iter().map(Into::into).collect::<Vec<String>>();
1398 self.validate_parameter_names(&names)?;
1400 self.parameter_names = Some(names);
1401 Ok(self)
1402 }
1403
1404 fn validate_parameter_names(&self, names: &[String]) -> Result<()> {
1406 match self.type_signature.arity() {
1407 Arity::Fixed(expected) => {
1408 if names.len() != expected {
1409 return plan_err!(
1410 "Parameter names count ({}) does not match signature arity ({})",
1411 names.len(),
1412 expected
1413 );
1414 }
1415 }
1416 Arity::Variable => {
1417 if self.type_signature != TypeSignature::UserDefined {
1420 return plan_err!(
1421 "Cannot specify parameter names for variable arity signature: {:?}",
1422 self.type_signature
1423 );
1424 }
1425 }
1426 }
1427
1428 let mut seen = std::collections::HashSet::new();
1429 for name in names {
1430 if !seen.insert(name) {
1431 return plan_err!("Duplicate parameter name: '{}'", name);
1432 }
1433 }
1434
1435 Ok(())
1436 }
1437}
1438
1439#[cfg(test)]
1440mod tests {
1441 use datafusion_common::types::{logical_int32, logical_int64, logical_string};
1442
1443 use super::*;
1444 use crate::signature::{
1445 ArrayFunctionArgument, ArrayFunctionSignature, Coercion, TypeSignatureClass,
1446 };
1447
1448 #[test]
1449 fn supports_zero_argument_tests() {
1450 let positive_cases = vec![
1452 TypeSignature::Exact(vec![]),
1453 TypeSignature::OneOf(vec![
1454 TypeSignature::Exact(vec![DataType::Int8]),
1455 TypeSignature::Nullary,
1456 TypeSignature::Uniform(1, vec![DataType::Int8]),
1457 ]),
1458 TypeSignature::Nullary,
1459 ];
1460
1461 for case in positive_cases {
1462 assert!(
1463 case.supports_zero_argument(),
1464 "Expected {case:?} to support zero arguments"
1465 );
1466 }
1467
1468 let negative_cases = vec![
1470 TypeSignature::Exact(vec![DataType::Utf8]),
1471 TypeSignature::Uniform(1, vec![DataType::Float64]),
1472 TypeSignature::Any(1),
1473 TypeSignature::VariadicAny,
1474 TypeSignature::OneOf(vec![
1475 TypeSignature::Exact(vec![DataType::Int8]),
1476 TypeSignature::Uniform(1, vec![DataType::Int8]),
1477 ]),
1478 ];
1479
1480 for case in negative_cases {
1481 assert!(
1482 !case.supports_zero_argument(),
1483 "Expected {case:?} not to support zero arguments"
1484 );
1485 }
1486 }
1487
1488 #[test]
1489 fn type_signature_partial_ord() {
1490 assert!(TypeSignature::UserDefined < TypeSignature::VariadicAny);
1492 assert!(TypeSignature::UserDefined < TypeSignature::Any(1));
1493
1494 assert!(
1495 TypeSignature::Uniform(1, vec![DataType::Null])
1496 < TypeSignature::Uniform(1, vec![DataType::Boolean])
1497 );
1498 assert!(
1499 TypeSignature::Uniform(1, vec![DataType::Null])
1500 < TypeSignature::Uniform(2, vec![DataType::Null])
1501 );
1502 assert!(
1503 TypeSignature::Uniform(usize::MAX, vec![DataType::Null])
1504 < TypeSignature::Exact(vec![DataType::Null])
1505 );
1506 }
1507
1508 #[test]
1509 fn test_get_possible_types() {
1510 let type_signature = TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]);
1511 let possible_types = type_signature.get_example_types();
1512 assert_eq!(possible_types, vec![vec![DataType::Int32, DataType::Int64]]);
1513
1514 let type_signature = TypeSignature::OneOf(vec![
1515 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1516 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1517 ]);
1518 let possible_types = type_signature.get_example_types();
1519 assert_eq!(
1520 possible_types,
1521 vec![
1522 vec![DataType::Int32, DataType::Int64],
1523 vec![DataType::Float32, DataType::Float64]
1524 ]
1525 );
1526
1527 let type_signature = TypeSignature::OneOf(vec![
1528 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1529 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1530 TypeSignature::Exact(vec![DataType::Utf8]),
1531 ]);
1532 let possible_types = type_signature.get_example_types();
1533 assert_eq!(
1534 possible_types,
1535 vec![
1536 vec![DataType::Int32, DataType::Int64],
1537 vec![DataType::Float32, DataType::Float64],
1538 vec![DataType::Utf8]
1539 ]
1540 );
1541
1542 let type_signature =
1543 TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Int64]);
1544 let possible_types = type_signature.get_example_types();
1545 assert_eq!(
1546 possible_types,
1547 vec![
1548 vec![DataType::Float32, DataType::Float32],
1549 vec![DataType::Int64, DataType::Int64]
1550 ]
1551 );
1552
1553 let type_signature = TypeSignature::Coercible(vec![
1554 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1555 Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
1556 ]);
1557 let possible_types = type_signature.get_example_types();
1558 assert_eq!(
1559 possible_types,
1560 vec![
1561 vec![DataType::Utf8, DataType::Int64],
1562 vec![DataType::LargeUtf8, DataType::Int64],
1563 vec![DataType::Utf8View, DataType::Int64]
1564 ]
1565 );
1566
1567 let type_signature =
1568 TypeSignature::Variadic(vec![DataType::Int32, DataType::Int64]);
1569 let possible_types = type_signature.get_example_types();
1570 assert_eq!(
1571 possible_types,
1572 vec![vec![DataType::Int32], vec![DataType::Int64]]
1573 );
1574
1575 let type_signature = TypeSignature::Numeric(2);
1576 let possible_types = type_signature.get_example_types();
1577 assert_eq!(
1578 possible_types,
1579 vec![
1580 vec![DataType::Int8, DataType::Int8],
1581 vec![DataType::Int16, DataType::Int16],
1582 vec![DataType::Int32, DataType::Int32],
1583 vec![DataType::Int64, DataType::Int64],
1584 vec![DataType::UInt8, DataType::UInt8],
1585 vec![DataType::UInt16, DataType::UInt16],
1586 vec![DataType::UInt32, DataType::UInt32],
1587 vec![DataType::UInt64, DataType::UInt64],
1588 vec![DataType::Float16, DataType::Float16],
1589 vec![DataType::Float32, DataType::Float32],
1590 vec![DataType::Float64, DataType::Float64]
1591 ]
1592 );
1593
1594 let type_signature = TypeSignature::String(2);
1595 let possible_types = type_signature.get_example_types();
1596 assert_eq!(
1597 possible_types,
1598 vec![
1599 vec![DataType::Utf8, DataType::Utf8],
1600 vec![DataType::LargeUtf8, DataType::LargeUtf8],
1601 vec![DataType::Utf8View, DataType::Utf8View]
1602 ]
1603 );
1604 }
1605
1606 #[test]
1607 fn test_signature_with_parameter_names() {
1608 let sig = Signature::exact(
1609 vec![DataType::Int32, DataType::Utf8],
1610 Volatility::Immutable,
1611 )
1612 .with_parameter_names(vec!["count".to_string(), "name".to_string()])
1613 .unwrap();
1614
1615 assert_eq!(
1616 sig.parameter_names,
1617 Some(vec!["count".to_string(), "name".to_string()])
1618 );
1619 assert_eq!(
1620 sig.type_signature,
1621 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8])
1622 );
1623 }
1624
1625 #[test]
1626 fn test_signature_parameter_names_wrong_count() {
1627 let result = Signature::exact(
1628 vec![DataType::Int32, DataType::Utf8],
1629 Volatility::Immutable,
1630 )
1631 .with_parameter_names(vec!["count".to_string()]); assert!(result.is_err());
1634 assert!(
1635 result
1636 .unwrap_err()
1637 .to_string()
1638 .contains("does not match signature arity")
1639 );
1640 }
1641
1642 #[test]
1643 fn test_signature_parameter_names_duplicate() {
1644 let result = Signature::exact(
1645 vec![DataType::Int32, DataType::Int32],
1646 Volatility::Immutable,
1647 )
1648 .with_parameter_names(vec!["count".to_string(), "count".to_string()]);
1649
1650 assert!(result.is_err());
1651 assert!(
1652 result
1653 .unwrap_err()
1654 .to_string()
1655 .contains("Duplicate parameter name")
1656 );
1657 }
1658
1659 #[test]
1660 fn test_signature_parameter_names_variadic() {
1661 let result = Signature::variadic(vec![DataType::Int32], Volatility::Immutable)
1662 .with_parameter_names(vec!["arg".to_string()]);
1663
1664 assert!(result.is_err());
1665 assert!(
1666 result
1667 .unwrap_err()
1668 .to_string()
1669 .contains("variable arity signature")
1670 );
1671 }
1672
1673 #[test]
1674 fn test_signature_without_parameter_names() {
1675 let sig = Signature::exact(
1676 vec![DataType::Int32, DataType::Utf8],
1677 Volatility::Immutable,
1678 );
1679
1680 assert_eq!(sig.parameter_names, None);
1681 }
1682
1683 #[test]
1684 fn test_signature_uniform_with_parameter_names() {
1685 let sig = Signature::uniform(3, vec![DataType::Float64], Volatility::Immutable)
1686 .with_parameter_names(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1687 .unwrap();
1688
1689 assert_eq!(
1690 sig.parameter_names,
1691 Some(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1692 );
1693 }
1694
1695 #[test]
1696 fn test_signature_numeric_with_parameter_names() {
1697 let sig = Signature::numeric(2, Volatility::Immutable)
1698 .with_parameter_names(vec!["a".to_string(), "b".to_string()])
1699 .unwrap();
1700
1701 assert_eq!(
1702 sig.parameter_names,
1703 Some(vec!["a".to_string(), "b".to_string()])
1704 );
1705 }
1706
1707 #[test]
1708 fn test_signature_nullary_with_empty_names() {
1709 let sig = Signature::nullary(Volatility::Immutable)
1710 .with_parameter_names(Vec::<String>::new())
1711 .unwrap();
1712
1713 assert_eq!(sig.parameter_names, Some(vec![]));
1714 }
1715
1716 #[test]
1717 fn test_to_string_repr_with_names_exact() {
1718 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1719
1720 assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
1721
1722 let names = vec!["id".to_string(), "name".to_string()];
1723 assert_eq!(
1724 sig.to_string_repr_with_names(Some(&names)),
1725 vec!["id: Int32, name: Utf8"]
1726 );
1727 }
1728
1729 #[test]
1730 fn test_to_string_repr_with_names_any() {
1731 let sig = TypeSignature::Any(3);
1732
1733 assert_eq!(sig.to_string_repr_with_names(None), vec!["Any, Any, Any"]);
1734
1735 let names = vec!["x".to_string(), "y".to_string(), "z".to_string()];
1736 assert_eq!(
1737 sig.to_string_repr_with_names(Some(&names)),
1738 vec!["x: Any, y: Any, z: Any"]
1739 );
1740 }
1741
1742 #[test]
1743 fn test_to_string_repr_with_names_one_of() {
1744 let sig =
1745 TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]);
1746
1747 assert_eq!(
1748 sig.to_string_repr_with_names(None),
1749 vec!["Any, Any", "Any, Any, Any"]
1750 );
1751
1752 let names = vec![
1753 "str".to_string(),
1754 "start_pos".to_string(),
1755 "length".to_string(),
1756 ];
1757 assert_eq!(
1758 sig.to_string_repr_with_names(Some(&names)),
1759 vec![
1760 "str: Any, start_pos: Any",
1761 "str: Any, start_pos: Any, length: Any"
1762 ]
1763 );
1764 }
1765
1766 #[test]
1767 fn test_to_string_repr_with_names_partial() {
1768 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1770
1771 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1773 assert_eq!(
1774 sig.to_string_repr_with_names(Some(&names)),
1775 vec!["a: Int32, b: Utf8"]
1776 );
1777 }
1778
1779 #[test]
1780 fn test_to_string_repr_with_names_uniform() {
1781 let sig = TypeSignature::Uniform(2, vec![DataType::Float64]);
1782
1783 assert_eq!(
1784 sig.to_string_repr_with_names(None),
1785 vec!["Float64, Float64"]
1786 );
1787
1788 let names = vec!["x".to_string(), "y".to_string()];
1789 assert_eq!(
1790 sig.to_string_repr_with_names(Some(&names)),
1791 vec!["x: Float64, y: Float64"]
1792 );
1793 }
1794
1795 #[test]
1796 fn test_to_string_repr_with_names_coercible() {
1797 let sig = TypeSignature::Coercible(vec![
1798 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1799 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1800 ]);
1801
1802 let names = vec!["a".to_string(), "b".to_string()];
1803 let result = sig.to_string_repr_with_names(Some(&names));
1804 assert_eq!(result.len(), 1);
1806 assert!(result[0].starts_with("a: "));
1807 assert!(result[0].contains(", b: "));
1808 }
1809
1810 #[test]
1811 fn test_to_string_repr_with_names_comparable_numeric_string() {
1812 let comparable = TypeSignature::Comparable(3);
1813 let numeric = TypeSignature::Numeric(2);
1814 let string_sig = TypeSignature::String(2);
1815
1816 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1817
1818 assert_eq!(
1820 comparable.to_string_repr_with_names(Some(&names)),
1821 vec!["a: Comparable, b: Comparable, c: Comparable"]
1822 );
1823 assert_eq!(
1824 numeric.to_string_repr_with_names(Some(&names)),
1825 vec!["a: Numeric, b: Numeric"]
1826 );
1827 assert_eq!(
1828 string_sig.to_string_repr_with_names(Some(&names)),
1829 vec!["a: String, b: String"]
1830 );
1831 }
1832
1833 #[test]
1834 fn test_to_string_repr_with_names_variadic_fallback() {
1835 let variadic = TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]);
1836 let names = vec!["x".to_string()];
1837 assert_eq!(
1838 variadic.to_string_repr_with_names(Some(&names)),
1839 variadic.to_string_repr()
1840 );
1841
1842 let variadic_any = TypeSignature::VariadicAny;
1843 assert_eq!(
1844 variadic_any.to_string_repr_with_names(Some(&names)),
1845 variadic_any.to_string_repr()
1846 );
1847
1848 let user_defined = TypeSignature::UserDefined;
1850 assert_eq!(
1851 user_defined.to_string_repr_with_names(Some(&names)),
1852 vec!["x"]
1853 );
1854 assert_eq!(
1855 user_defined.to_string_repr_with_names(None),
1856 user_defined.to_string_repr()
1857 );
1858 }
1859
1860 #[test]
1861 fn test_to_string_repr_with_names_nullary() {
1862 let sig = TypeSignature::Nullary;
1863 let names = vec!["x".to_string()];
1864
1865 assert_eq!(
1867 sig.to_string_repr_with_names(Some(&names)),
1868 vec!["NullAry()"]
1869 );
1870 assert_eq!(sig.to_string_repr_with_names(None), vec!["NullAry()"]);
1871 }
1872
1873 #[test]
1874 fn test_to_string_repr_with_names_array_signature() {
1875 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1876 arguments: vec![
1877 ArrayFunctionArgument::Array,
1878 ArrayFunctionArgument::Index,
1879 ArrayFunctionArgument::Element,
1880 ],
1881 array_coercion: None,
1882 });
1883
1884 assert_eq!(
1885 sig.to_string_repr_with_names(None),
1886 vec!["array, index, element"]
1887 );
1888
1889 let names = vec!["arr".to_string(), "idx".to_string(), "val".to_string()];
1890 assert_eq!(
1891 sig.to_string_repr_with_names(Some(&names)),
1892 vec!["arr: array, idx: index, val: element"]
1893 );
1894
1895 let recursive =
1896 TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1897 let names = vec!["array".to_string()];
1898 assert_eq!(
1899 recursive.to_string_repr_with_names(Some(&names)),
1900 vec!["array: recursive_array"]
1901 );
1902
1903 let map_array = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1905 let names = vec!["map".to_string()];
1906 assert_eq!(
1907 map_array.to_string_repr_with_names(Some(&names)),
1908 vec!["map: map_array"]
1909 );
1910 }
1911
1912 #[test]
1913 fn test_type_signature_arity_exact() {
1914 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1915 assert_eq!(sig.arity(), Arity::Fixed(2));
1916
1917 let sig = TypeSignature::Exact(vec![]);
1918 assert_eq!(sig.arity(), Arity::Fixed(0));
1919 }
1920
1921 #[test]
1922 fn test_type_signature_arity_uniform() {
1923 let sig = TypeSignature::Uniform(3, vec![DataType::Float64]);
1924 assert_eq!(sig.arity(), Arity::Fixed(3));
1925
1926 let sig = TypeSignature::Uniform(1, vec![DataType::Int32]);
1927 assert_eq!(sig.arity(), Arity::Fixed(1));
1928 }
1929
1930 #[test]
1931 fn test_type_signature_arity_numeric() {
1932 let sig = TypeSignature::Numeric(2);
1933 assert_eq!(sig.arity(), Arity::Fixed(2));
1934 }
1935
1936 #[test]
1937 fn test_type_signature_arity_string() {
1938 let sig = TypeSignature::String(3);
1939 assert_eq!(sig.arity(), Arity::Fixed(3));
1940 }
1941
1942 #[test]
1943 fn test_type_signature_arity_comparable() {
1944 let sig = TypeSignature::Comparable(2);
1945 assert_eq!(sig.arity(), Arity::Fixed(2));
1946 }
1947
1948 #[test]
1949 fn test_type_signature_arity_any() {
1950 let sig = TypeSignature::Any(4);
1951 assert_eq!(sig.arity(), Arity::Fixed(4));
1952 }
1953
1954 #[test]
1955 fn test_type_signature_arity_coercible() {
1956 let sig = TypeSignature::Coercible(vec![
1957 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1958 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1959 ]);
1960 assert_eq!(sig.arity(), Arity::Fixed(2));
1961 }
1962
1963 #[test]
1964 fn test_type_signature_arity_nullary() {
1965 let sig = TypeSignature::Nullary;
1966 assert_eq!(sig.arity(), Arity::Fixed(0));
1967 }
1968
1969 #[test]
1970 fn test_type_signature_arity_array_signature() {
1971 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1973 arguments: vec![ArrayFunctionArgument::Array, ArrayFunctionArgument::Index],
1974 array_coercion: None,
1975 });
1976 assert_eq!(sig.arity(), Arity::Fixed(2));
1977
1978 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1980 arguments: vec![
1981 ArrayFunctionArgument::Array,
1982 ArrayFunctionArgument::Element,
1983 ArrayFunctionArgument::Index,
1984 ],
1985 array_coercion: None,
1986 });
1987 assert_eq!(sig.arity(), Arity::Fixed(3));
1988
1989 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1991 assert_eq!(sig.arity(), Arity::Fixed(1));
1992
1993 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1995 assert_eq!(sig.arity(), Arity::Fixed(1));
1996 }
1997
1998 #[test]
1999 fn test_type_signature_arity_one_of_fixed() {
2000 let sig = TypeSignature::OneOf(vec![
2002 TypeSignature::Exact(vec![DataType::Int32]),
2003 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2004 TypeSignature::Exact(vec![
2005 DataType::Int32,
2006 DataType::Utf8,
2007 DataType::Float64,
2008 ]),
2009 ]);
2010 assert_eq!(sig.arity(), Arity::Fixed(3));
2011 }
2012
2013 #[test]
2014 fn test_type_signature_arity_one_of_variable() {
2015 let sig = TypeSignature::OneOf(vec![
2017 TypeSignature::Exact(vec![DataType::Int32]),
2018 TypeSignature::VariadicAny,
2019 ]);
2020 assert_eq!(sig.arity(), Arity::Variable);
2021 }
2022
2023 #[test]
2024 fn test_type_signature_arity_variadic() {
2025 let sig = TypeSignature::Variadic(vec![DataType::Int32]);
2026 assert_eq!(sig.arity(), Arity::Variable);
2027
2028 let sig = TypeSignature::VariadicAny;
2029 assert_eq!(sig.arity(), Arity::Variable);
2030 }
2031
2032 #[test]
2033 fn test_type_signature_arity_user_defined() {
2034 let sig = TypeSignature::UserDefined;
2035 assert_eq!(sig.arity(), Arity::Variable);
2036 }
2037}