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 !matches!(self.type_signature, TypeSignature::UserDefined) {
1420 return plan_err!(
1421 "Cannot specify parameter names for variable arity signature: {:?}",
1422 self.type_signature
1423 );
1424 }
1425 }
1426 }
1427
1428 let mut seen = std::collections::HashSet::new();
1429 for name in names {
1430 if !seen.insert(name) {
1431 return plan_err!("Duplicate parameter name: '{}'", name);
1432 }
1433 }
1434
1435 Ok(())
1436 }
1437}
1438
1439#[cfg(test)]
1440mod tests {
1441 use datafusion_common::types::{logical_int32, logical_int64, logical_string};
1442
1443 use super::*;
1444 use crate::signature::{
1445 ArrayFunctionArgument, ArrayFunctionSignature, Coercion, TypeSignatureClass,
1446 };
1447
1448 #[test]
1449 fn supports_zero_argument_tests() {
1450 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::Float32, DataType::Float32],
1589 vec![DataType::Float64, DataType::Float64]
1590 ]
1591 );
1592
1593 let type_signature = TypeSignature::String(2);
1594 let possible_types = type_signature.get_example_types();
1595 assert_eq!(
1596 possible_types,
1597 vec![
1598 vec![DataType::Utf8, DataType::Utf8],
1599 vec![DataType::LargeUtf8, DataType::LargeUtf8],
1600 vec![DataType::Utf8View, DataType::Utf8View]
1601 ]
1602 );
1603 }
1604
1605 #[test]
1606 fn test_signature_with_parameter_names() {
1607 let sig = Signature::exact(
1608 vec![DataType::Int32, DataType::Utf8],
1609 Volatility::Immutable,
1610 )
1611 .with_parameter_names(vec!["count".to_string(), "name".to_string()])
1612 .unwrap();
1613
1614 assert_eq!(
1615 sig.parameter_names,
1616 Some(vec!["count".to_string(), "name".to_string()])
1617 );
1618 assert_eq!(
1619 sig.type_signature,
1620 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8])
1621 );
1622 }
1623
1624 #[test]
1625 fn test_signature_parameter_names_wrong_count() {
1626 let result = Signature::exact(
1627 vec![DataType::Int32, DataType::Utf8],
1628 Volatility::Immutable,
1629 )
1630 .with_parameter_names(vec!["count".to_string()]); assert!(result.is_err());
1633 assert!(
1634 result
1635 .unwrap_err()
1636 .to_string()
1637 .contains("does not match signature arity")
1638 );
1639 }
1640
1641 #[test]
1642 fn test_signature_parameter_names_duplicate() {
1643 let result = Signature::exact(
1644 vec![DataType::Int32, DataType::Int32],
1645 Volatility::Immutable,
1646 )
1647 .with_parameter_names(vec!["count".to_string(), "count".to_string()]);
1648
1649 assert!(result.is_err());
1650 assert!(
1651 result
1652 .unwrap_err()
1653 .to_string()
1654 .contains("Duplicate parameter name")
1655 );
1656 }
1657
1658 #[test]
1659 fn test_signature_parameter_names_variadic() {
1660 let result = Signature::variadic(vec![DataType::Int32], Volatility::Immutable)
1661 .with_parameter_names(vec!["arg".to_string()]);
1662
1663 assert!(result.is_err());
1664 assert!(
1665 result
1666 .unwrap_err()
1667 .to_string()
1668 .contains("variable arity signature")
1669 );
1670 }
1671
1672 #[test]
1673 fn test_signature_without_parameter_names() {
1674 let sig = Signature::exact(
1675 vec![DataType::Int32, DataType::Utf8],
1676 Volatility::Immutable,
1677 );
1678
1679 assert_eq!(sig.parameter_names, None);
1680 }
1681
1682 #[test]
1683 fn test_signature_uniform_with_parameter_names() {
1684 let sig = Signature::uniform(3, vec![DataType::Float64], Volatility::Immutable)
1685 .with_parameter_names(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1686 .unwrap();
1687
1688 assert_eq!(
1689 sig.parameter_names,
1690 Some(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1691 );
1692 }
1693
1694 #[test]
1695 fn test_signature_numeric_with_parameter_names() {
1696 let sig = Signature::numeric(2, Volatility::Immutable)
1697 .with_parameter_names(vec!["a".to_string(), "b".to_string()])
1698 .unwrap();
1699
1700 assert_eq!(
1701 sig.parameter_names,
1702 Some(vec!["a".to_string(), "b".to_string()])
1703 );
1704 }
1705
1706 #[test]
1707 fn test_signature_nullary_with_empty_names() {
1708 let sig = Signature::nullary(Volatility::Immutable)
1709 .with_parameter_names(Vec::<String>::new())
1710 .unwrap();
1711
1712 assert_eq!(sig.parameter_names, Some(vec![]));
1713 }
1714
1715 #[test]
1716 fn test_to_string_repr_with_names_exact() {
1717 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1718
1719 assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
1720
1721 let names = vec!["id".to_string(), "name".to_string()];
1722 assert_eq!(
1723 sig.to_string_repr_with_names(Some(&names)),
1724 vec!["id: Int32, name: Utf8"]
1725 );
1726 }
1727
1728 #[test]
1729 fn test_to_string_repr_with_names_any() {
1730 let sig = TypeSignature::Any(3);
1731
1732 assert_eq!(sig.to_string_repr_with_names(None), vec!["Any, Any, Any"]);
1733
1734 let names = vec!["x".to_string(), "y".to_string(), "z".to_string()];
1735 assert_eq!(
1736 sig.to_string_repr_with_names(Some(&names)),
1737 vec!["x: Any, y: Any, z: Any"]
1738 );
1739 }
1740
1741 #[test]
1742 fn test_to_string_repr_with_names_one_of() {
1743 let sig =
1744 TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]);
1745
1746 assert_eq!(
1747 sig.to_string_repr_with_names(None),
1748 vec!["Any, Any", "Any, Any, Any"]
1749 );
1750
1751 let names = vec![
1752 "str".to_string(),
1753 "start_pos".to_string(),
1754 "length".to_string(),
1755 ];
1756 assert_eq!(
1757 sig.to_string_repr_with_names(Some(&names)),
1758 vec![
1759 "str: Any, start_pos: Any",
1760 "str: Any, start_pos: Any, length: Any"
1761 ]
1762 );
1763 }
1764
1765 #[test]
1766 fn test_to_string_repr_with_names_partial() {
1767 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1769
1770 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1772 assert_eq!(
1773 sig.to_string_repr_with_names(Some(&names)),
1774 vec!["a: Int32, b: Utf8"]
1775 );
1776 }
1777
1778 #[test]
1779 fn test_to_string_repr_with_names_uniform() {
1780 let sig = TypeSignature::Uniform(2, vec![DataType::Float64]);
1781
1782 assert_eq!(
1783 sig.to_string_repr_with_names(None),
1784 vec!["Float64, Float64"]
1785 );
1786
1787 let names = vec!["x".to_string(), "y".to_string()];
1788 assert_eq!(
1789 sig.to_string_repr_with_names(Some(&names)),
1790 vec!["x: Float64, y: Float64"]
1791 );
1792 }
1793
1794 #[test]
1795 fn test_to_string_repr_with_names_coercible() {
1796 let sig = TypeSignature::Coercible(vec![
1797 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1798 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1799 ]);
1800
1801 let names = vec!["a".to_string(), "b".to_string()];
1802 let result = sig.to_string_repr_with_names(Some(&names));
1803 assert_eq!(result.len(), 1);
1805 assert!(result[0].starts_with("a: "));
1806 assert!(result[0].contains(", b: "));
1807 }
1808
1809 #[test]
1810 fn test_to_string_repr_with_names_comparable_numeric_string() {
1811 let comparable = TypeSignature::Comparable(3);
1812 let numeric = TypeSignature::Numeric(2);
1813 let string_sig = TypeSignature::String(2);
1814
1815 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1816
1817 assert_eq!(
1819 comparable.to_string_repr_with_names(Some(&names)),
1820 vec!["a: Comparable, b: Comparable, c: Comparable"]
1821 );
1822 assert_eq!(
1823 numeric.to_string_repr_with_names(Some(&names)),
1824 vec!["a: Numeric, b: Numeric"]
1825 );
1826 assert_eq!(
1827 string_sig.to_string_repr_with_names(Some(&names)),
1828 vec!["a: String, b: String"]
1829 );
1830 }
1831
1832 #[test]
1833 fn test_to_string_repr_with_names_variadic_fallback() {
1834 let variadic = TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]);
1835 let names = vec!["x".to_string()];
1836 assert_eq!(
1837 variadic.to_string_repr_with_names(Some(&names)),
1838 variadic.to_string_repr()
1839 );
1840
1841 let variadic_any = TypeSignature::VariadicAny;
1842 assert_eq!(
1843 variadic_any.to_string_repr_with_names(Some(&names)),
1844 variadic_any.to_string_repr()
1845 );
1846
1847 let user_defined = TypeSignature::UserDefined;
1849 assert_eq!(
1850 user_defined.to_string_repr_with_names(Some(&names)),
1851 vec!["x"]
1852 );
1853 assert_eq!(
1854 user_defined.to_string_repr_with_names(None),
1855 user_defined.to_string_repr()
1856 );
1857 }
1858
1859 #[test]
1860 fn test_to_string_repr_with_names_nullary() {
1861 let sig = TypeSignature::Nullary;
1862 let names = vec!["x".to_string()];
1863
1864 assert_eq!(
1866 sig.to_string_repr_with_names(Some(&names)),
1867 vec!["NullAry()"]
1868 );
1869 assert_eq!(sig.to_string_repr_with_names(None), vec!["NullAry()"]);
1870 }
1871
1872 #[test]
1873 fn test_to_string_repr_with_names_array_signature() {
1874 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1875 arguments: vec![
1876 ArrayFunctionArgument::Array,
1877 ArrayFunctionArgument::Index,
1878 ArrayFunctionArgument::Element,
1879 ],
1880 array_coercion: None,
1881 });
1882
1883 assert_eq!(
1884 sig.to_string_repr_with_names(None),
1885 vec!["array, index, element"]
1886 );
1887
1888 let names = vec!["arr".to_string(), "idx".to_string(), "val".to_string()];
1889 assert_eq!(
1890 sig.to_string_repr_with_names(Some(&names)),
1891 vec!["arr: array, idx: index, val: element"]
1892 );
1893
1894 let recursive =
1895 TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1896 let names = vec!["array".to_string()];
1897 assert_eq!(
1898 recursive.to_string_repr_with_names(Some(&names)),
1899 vec!["array: recursive_array"]
1900 );
1901
1902 let map_array = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1904 let names = vec!["map".to_string()];
1905 assert_eq!(
1906 map_array.to_string_repr_with_names(Some(&names)),
1907 vec!["map: map_array"]
1908 );
1909 }
1910
1911 #[test]
1912 fn test_type_signature_arity_exact() {
1913 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1914 assert_eq!(sig.arity(), Arity::Fixed(2));
1915
1916 let sig = TypeSignature::Exact(vec![]);
1917 assert_eq!(sig.arity(), Arity::Fixed(0));
1918 }
1919
1920 #[test]
1921 fn test_type_signature_arity_uniform() {
1922 let sig = TypeSignature::Uniform(3, vec![DataType::Float64]);
1923 assert_eq!(sig.arity(), Arity::Fixed(3));
1924
1925 let sig = TypeSignature::Uniform(1, vec![DataType::Int32]);
1926 assert_eq!(sig.arity(), Arity::Fixed(1));
1927 }
1928
1929 #[test]
1930 fn test_type_signature_arity_numeric() {
1931 let sig = TypeSignature::Numeric(2);
1932 assert_eq!(sig.arity(), Arity::Fixed(2));
1933 }
1934
1935 #[test]
1936 fn test_type_signature_arity_string() {
1937 let sig = TypeSignature::String(3);
1938 assert_eq!(sig.arity(), Arity::Fixed(3));
1939 }
1940
1941 #[test]
1942 fn test_type_signature_arity_comparable() {
1943 let sig = TypeSignature::Comparable(2);
1944 assert_eq!(sig.arity(), Arity::Fixed(2));
1945 }
1946
1947 #[test]
1948 fn test_type_signature_arity_any() {
1949 let sig = TypeSignature::Any(4);
1950 assert_eq!(sig.arity(), Arity::Fixed(4));
1951 }
1952
1953 #[test]
1954 fn test_type_signature_arity_coercible() {
1955 let sig = TypeSignature::Coercible(vec![
1956 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1957 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1958 ]);
1959 assert_eq!(sig.arity(), Arity::Fixed(2));
1960 }
1961
1962 #[test]
1963 fn test_type_signature_arity_nullary() {
1964 let sig = TypeSignature::Nullary;
1965 assert_eq!(sig.arity(), Arity::Fixed(0));
1966 }
1967
1968 #[test]
1969 fn test_type_signature_arity_array_signature() {
1970 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1972 arguments: vec![ArrayFunctionArgument::Array, ArrayFunctionArgument::Index],
1973 array_coercion: None,
1974 });
1975 assert_eq!(sig.arity(), Arity::Fixed(2));
1976
1977 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1979 arguments: vec![
1980 ArrayFunctionArgument::Array,
1981 ArrayFunctionArgument::Element,
1982 ArrayFunctionArgument::Index,
1983 ],
1984 array_coercion: None,
1985 });
1986 assert_eq!(sig.arity(), Arity::Fixed(3));
1987
1988 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1990 assert_eq!(sig.arity(), Arity::Fixed(1));
1991
1992 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1994 assert_eq!(sig.arity(), Arity::Fixed(1));
1995 }
1996
1997 #[test]
1998 fn test_type_signature_arity_one_of_fixed() {
1999 let sig = TypeSignature::OneOf(vec![
2001 TypeSignature::Exact(vec![DataType::Int32]),
2002 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2003 TypeSignature::Exact(vec![
2004 DataType::Int32,
2005 DataType::Utf8,
2006 DataType::Float64,
2007 ]),
2008 ]);
2009 assert_eq!(sig.arity(), Arity::Fixed(3));
2010 }
2011
2012 #[test]
2013 fn test_type_signature_arity_one_of_variable() {
2014 let sig = TypeSignature::OneOf(vec![
2016 TypeSignature::Exact(vec![DataType::Int32]),
2017 TypeSignature::VariadicAny,
2018 ]);
2019 assert_eq!(sig.arity(), Arity::Variable);
2020 }
2021
2022 #[test]
2023 fn test_type_signature_arity_variadic() {
2024 let sig = TypeSignature::Variadic(vec![DataType::Int32]);
2025 assert_eq!(sig.arity(), Arity::Variable);
2026
2027 let sig = TypeSignature::VariadicAny;
2028 assert_eq!(sig.arity(), Arity::Variable);
2029 }
2030
2031 #[test]
2032 fn test_type_signature_arity_user_defined() {
2033 let sig = TypeSignature::UserDefined;
2034 assert_eq!(sig.arity(), Arity::Variable);
2035 }
2036}