1use std::fmt::Display;
21use std::hash::Hash;
22use std::sync::Arc;
23
24use arrow::datatypes::{
25 DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DECIMAL128_MAX_PRECISION, DataType,
26 Decimal128Type, DecimalType, Field, IntervalUnit, TimeUnit,
27};
28use datafusion_common::types::{LogicalType, LogicalTypeRef, NativeType};
29use datafusion_common::utils::ListCoercion;
30use datafusion_common::{Result, internal_err, plan_err};
31use indexmap::IndexSet;
32use itertools::Itertools;
33
34pub const TIMEZONE_WILDCARD: &str = "+TZ";
43
44pub const FIXED_SIZE_LIST_WILDCARD: i32 = i32::MIN;
48
49#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
56pub enum Volatility {
57 Immutable,
65 Stable,
77 Volatile,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum Arity {
93 Fixed(usize),
95 Variable,
97}
98
99#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
157pub enum TypeSignature {
158 Variadic(Vec<DataType>),
167 UserDefined,
175 VariadicAny,
177 Uniform(usize, Vec<DataType>),
184 Exact(Vec<DataType>),
188 Coercible(Vec<Coercion>),
196 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
325impl Display for TypeSignature {
326 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
327 match self {
328 TypeSignature::Variadic(types) => {
329 write!(f, "Variadic({})", types.iter().join(", "))
330 }
331 TypeSignature::UserDefined => write!(f, "UserDefined"),
332 TypeSignature::VariadicAny => write!(f, "VariadicAny"),
333 TypeSignature::Uniform(count, types) => {
334 write!(f, "Uniform({count}, [{}])", types.iter().join(", "))
335 }
336 TypeSignature::Exact(types) => {
337 write!(f, "Exact({})", types.iter().join(", "))
338 }
339 TypeSignature::Coercible(coercions) => {
340 write!(f, "Coercible({})", coercions.iter().join(", "))
341 }
342 TypeSignature::Comparable(count) => write!(f, "Comparable({count})"),
343 TypeSignature::Any(count) => write!(f, "Any({count})"),
344 TypeSignature::OneOf(sigs) => {
345 write!(f, "OneOf(")?;
346 for (i, sig) in sigs.iter().enumerate() {
347 if i > 0 {
348 write!(f, ", ")?;
349 }
350 write!(f, "{sig}")?;
351 }
352 write!(f, ")")
353 }
354 TypeSignature::ArraySignature(sig) => write!(f, "ArraySignature({sig})"),
355 TypeSignature::Numeric(count) => write!(f, "Numeric({count})"),
356 TypeSignature::String(count) => write!(f, "String({count})"),
357 TypeSignature::Nullary => write!(f, "Nullary"),
358 }
359 }
360}
361
362#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Hash)]
371pub enum TypeSignatureClass {
372 Any,
374 Timestamp,
376 Time,
378 Interval,
380 Duration,
382 Native(LogicalTypeRef),
384 Integer,
386 Float,
388 Decimal,
390 Numeric,
392 Binary,
394}
395
396impl Display for TypeSignatureClass {
397 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
398 match self {
399 Self::Any => write!(f, "Any"),
400 Self::Timestamp => write!(f, "Timestamp"),
401 Self::Time => write!(f, "Time"),
402 Self::Interval => write!(f, "Interval"),
403 Self::Duration => write!(f, "Duration"),
404 Self::Native(logical_type) => write!(f, "{logical_type}"),
405 Self::Integer => write!(f, "Integer"),
406 Self::Float => write!(f, "Float"),
407 Self::Decimal => write!(f, "Decimal"),
408 Self::Numeric => write!(f, "Numeric"),
409 Self::Binary => write!(f, "Binary"),
410 }
411 }
412}
413
414impl TypeSignatureClass {
415 fn get_example_types(&self) -> Vec<DataType> {
420 match self {
421 TypeSignatureClass::Any => vec![],
424 TypeSignatureClass::Native(l) => get_data_types(l.native()),
425 TypeSignatureClass::Timestamp => {
426 vec![
427 DataType::Timestamp(TimeUnit::Nanosecond, None),
428 DataType::Timestamp(
429 TimeUnit::Nanosecond,
430 Some(TIMEZONE_WILDCARD.into()),
431 ),
432 ]
433 }
434 TypeSignatureClass::Time => {
435 vec![DataType::Time64(TimeUnit::Nanosecond)]
436 }
437 TypeSignatureClass::Interval => {
438 vec![DataType::Interval(IntervalUnit::DayTime)]
439 }
440 TypeSignatureClass::Duration => {
441 vec![DataType::Duration(TimeUnit::Nanosecond)]
442 }
443 TypeSignatureClass::Integer => {
444 vec![DataType::Int64]
445 }
446 TypeSignatureClass::Binary => {
447 vec![DataType::Binary]
448 }
449 TypeSignatureClass::Decimal => vec![Decimal128Type::DEFAULT_TYPE],
450 TypeSignatureClass::Float => vec![DataType::Float64],
451 TypeSignatureClass::Numeric => vec![
452 DataType::Float64,
453 DataType::Int64,
454 Decimal128Type::DEFAULT_TYPE,
455 ],
456 }
457 }
458
459 pub fn matches_native_type(&self, logical_type: &NativeType) -> bool {
461 if logical_type == &NativeType::Null {
462 return true;
463 }
464
465 match self {
466 TypeSignatureClass::Any => true,
467 TypeSignatureClass::Native(t) if t.native() == logical_type => true,
468 TypeSignatureClass::Timestamp if logical_type.is_timestamp() => true,
469 TypeSignatureClass::Time if logical_type.is_time() => true,
470 TypeSignatureClass::Interval if logical_type.is_interval() => true,
471 TypeSignatureClass::Duration if logical_type.is_duration() => true,
472 TypeSignatureClass::Integer if logical_type.is_integer() => true,
473 TypeSignatureClass::Binary if logical_type.is_binary() => true,
474 TypeSignatureClass::Decimal if logical_type.is_decimal() => true,
475 TypeSignatureClass::Float if logical_type.is_float() => true,
476 TypeSignatureClass::Numeric if logical_type.is_numeric() => true,
477 _ => false,
478 }
479 }
480
481 pub fn default_casted_type(
483 &self,
484 native_type: &NativeType,
485 origin_type: &DataType,
486 ) -> Result<DataType> {
487 match self {
488 TypeSignatureClass::Any => Ok(origin_type.to_owned()),
489 TypeSignatureClass::Native(logical_type) => {
490 logical_type.native().default_cast_for(origin_type)
491 }
492 TypeSignatureClass::Timestamp if native_type.is_timestamp() => {
494 Ok(origin_type.to_owned())
495 }
496 TypeSignatureClass::Time if native_type.is_time() => {
497 Ok(origin_type.to_owned())
498 }
499 TypeSignatureClass::Interval if native_type.is_interval() => {
500 Ok(origin_type.to_owned())
501 }
502 TypeSignatureClass::Duration if native_type.is_duration() => {
503 Ok(origin_type.to_owned())
504 }
505 TypeSignatureClass::Integer if native_type.is_integer() => {
506 Ok(origin_type.to_owned())
507 }
508 TypeSignatureClass::Binary if native_type.is_binary() => {
509 Ok(origin_type.to_owned())
510 }
511 TypeSignatureClass::Decimal if native_type.is_decimal() => {
512 Ok(origin_type.to_owned())
513 }
514 TypeSignatureClass::Float if native_type.is_float() => {
515 Ok(origin_type.to_owned())
516 }
517 TypeSignatureClass::Numeric if native_type.is_numeric() => {
518 Ok(origin_type.to_owned())
519 }
520 _ if native_type.is_null() => Ok(origin_type.to_owned()),
521 _ => internal_err!("May miss the matching logic in `matches_native_type`"),
522 }
523 }
524}
525
526#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
527pub enum ArrayFunctionSignature {
528 Array {
530 arguments: Vec<ArrayFunctionArgument>,
532 array_coercion: Option<ListCoercion>,
534 },
535 RecursiveArray,
538 MapArray,
541}
542
543impl Display for ArrayFunctionSignature {
544 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
545 match self {
546 ArrayFunctionSignature::Array { arguments, .. } => {
547 for (idx, argument) in arguments.iter().enumerate() {
548 write!(f, "{argument}")?;
549 if idx != arguments.len() - 1 {
550 write!(f, ", ")?;
551 }
552 }
553 Ok(())
554 }
555 ArrayFunctionSignature::RecursiveArray => {
556 write!(f, "recursive_array")
557 }
558 ArrayFunctionSignature::MapArray => {
559 write!(f, "map_array")
560 }
561 }
562 }
563}
564
565#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
566pub enum ArrayFunctionArgument {
567 Element,
570 Index,
572 Array,
575 String,
577}
578
579impl Display for ArrayFunctionArgument {
580 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
581 match self {
582 ArrayFunctionArgument::Element => {
583 write!(f, "element")
584 }
585 ArrayFunctionArgument::Index => {
586 write!(f, "index")
587 }
588 ArrayFunctionArgument::Array => {
589 write!(f, "array")
590 }
591 ArrayFunctionArgument::String => {
592 write!(f, "string")
593 }
594 }
595 }
596}
597
598static NUMERICS: &[DataType] = &[
599 DataType::Int8,
600 DataType::Int16,
601 DataType::Int32,
602 DataType::Int64,
603 DataType::UInt8,
604 DataType::UInt16,
605 DataType::UInt32,
606 DataType::UInt64,
607 DataType::Float16,
608 DataType::Float32,
609 DataType::Float64,
610];
611
612impl TypeSignature {
613 pub fn to_string_repr(&self) -> Vec<String> {
614 match self {
615 TypeSignature::Nullary => {
616 vec!["NullAry()".to_string()]
617 }
618 TypeSignature::Variadic(types) => {
619 vec![format!("{}, ..", Self::join_types(types, "/"))]
620 }
621 TypeSignature::Uniform(arg_count, valid_types) => {
622 vec![
623 std::iter::repeat_n(Self::join_types(valid_types, "/"), *arg_count)
624 .collect::<Vec<String>>()
625 .join(", "),
626 ]
627 }
628 TypeSignature::String(num) => {
629 vec![format!("String({num})")]
630 }
631 TypeSignature::Numeric(num) => {
632 vec![format!("Numeric({num})")]
633 }
634 TypeSignature::Comparable(num) => {
635 vec![format!("Comparable({num})")]
636 }
637 TypeSignature::Coercible(coercions) => {
638 vec![Self::join_types(coercions, ", ")]
639 }
640 TypeSignature::Exact(types) => {
641 vec![Self::join_types(types, ", ")]
642 }
643 TypeSignature::Any(arg_count) => {
644 vec![
645 std::iter::repeat_n("Any", *arg_count)
646 .collect::<Vec<&str>>()
647 .join(", "),
648 ]
649 }
650 TypeSignature::UserDefined => {
651 vec!["UserDefined".to_string()]
652 }
653 TypeSignature::VariadicAny => vec!["Any, .., Any".to_string()],
654 TypeSignature::OneOf(sigs) => {
655 sigs.iter().flat_map(|s| s.to_string_repr()).collect()
656 }
657 TypeSignature::ArraySignature(array_signature) => {
658 vec![array_signature.to_string()]
659 }
660 }
661 }
662
663 pub fn to_string_repr_with_names(
689 &self,
690 parameter_names: Option<&[String]>,
691 ) -> Vec<String> {
692 match self {
693 TypeSignature::Exact(types) => {
694 if let Some(names) = parameter_names {
695 vec![
696 names
697 .iter()
698 .zip(types.iter())
699 .map(|(name, typ)| format!("{name}: {typ}"))
700 .collect::<Vec<_>>()
701 .join(", "),
702 ]
703 } else {
704 vec![Self::join_types(types, ", ")]
705 }
706 }
707 TypeSignature::Any(count) => {
708 if let Some(names) = parameter_names {
709 vec![
710 names
711 .iter()
712 .take(*count)
713 .map(|name| format!("{name}: Any"))
714 .collect::<Vec<_>>()
715 .join(", "),
716 ]
717 } else {
718 vec![
719 std::iter::repeat_n("Any", *count)
720 .collect::<Vec<&str>>()
721 .join(", "),
722 ]
723 }
724 }
725 TypeSignature::Uniform(count, types) => {
726 if let Some(names) = parameter_names {
727 let type_str = Self::join_types(types, "/");
728 vec![
729 names
730 .iter()
731 .take(*count)
732 .map(|name| format!("{name}: {type_str}"))
733 .collect::<Vec<_>>()
734 .join(", "),
735 ]
736 } else {
737 self.to_string_repr()
738 }
739 }
740 TypeSignature::Coercible(coercions) => {
741 if let Some(names) = parameter_names {
742 vec![
743 names
744 .iter()
745 .zip(coercions.iter())
746 .map(|(name, coercion)| format!("{name}: {coercion}"))
747 .collect::<Vec<_>>()
748 .join(", "),
749 ]
750 } else {
751 vec![Self::join_types(coercions, ", ")]
752 }
753 }
754 TypeSignature::Comparable(count) => {
755 if let Some(names) = parameter_names {
756 vec![
757 names
758 .iter()
759 .take(*count)
760 .map(|name| format!("{name}: Comparable"))
761 .collect::<Vec<_>>()
762 .join(", "),
763 ]
764 } else {
765 self.to_string_repr()
766 }
767 }
768 TypeSignature::Numeric(count) => {
769 if let Some(names) = parameter_names {
770 vec![
771 names
772 .iter()
773 .take(*count)
774 .map(|name| format!("{name}: Numeric"))
775 .collect::<Vec<_>>()
776 .join(", "),
777 ]
778 } else {
779 self.to_string_repr()
780 }
781 }
782 TypeSignature::String(count) => {
783 if let Some(names) = parameter_names {
784 vec![
785 names
786 .iter()
787 .take(*count)
788 .map(|name| format!("{name}: String"))
789 .collect::<Vec<_>>()
790 .join(", "),
791 ]
792 } else {
793 self.to_string_repr()
794 }
795 }
796 TypeSignature::Nullary => self.to_string_repr(),
797 TypeSignature::ArraySignature(array_sig) => {
798 if let Some(names) = parameter_names {
799 match array_sig {
800 ArrayFunctionSignature::Array { arguments, .. } => {
801 vec![
802 names
803 .iter()
804 .zip(arguments.iter())
805 .map(|(name, arg_type)| format!("{name}: {arg_type}"))
806 .collect::<Vec<_>>()
807 .join(", "),
808 ]
809 }
810 ArrayFunctionSignature::RecursiveArray => {
811 vec![
812 names
813 .iter()
814 .take(1)
815 .map(|name| format!("{name}: recursive_array"))
816 .collect::<Vec<_>>()
817 .join(", "),
818 ]
819 }
820 ArrayFunctionSignature::MapArray => {
821 vec![
822 names
823 .iter()
824 .take(1)
825 .map(|name| format!("{name}: map_array"))
826 .collect::<Vec<_>>()
827 .join(", "),
828 ]
829 }
830 }
831 } else {
832 self.to_string_repr()
833 }
834 }
835 TypeSignature::OneOf(sigs) => sigs
836 .iter()
837 .flat_map(|s| s.to_string_repr_with_names(parameter_names))
838 .collect(),
839 TypeSignature::UserDefined => {
840 if let Some(names) = parameter_names {
841 vec![names.join(", ")]
842 } else {
843 self.to_string_repr()
844 }
845 }
846 TypeSignature::Variadic(_) | TypeSignature::VariadicAny => {
848 self.to_string_repr()
849 }
850 }
851 }
852
853 pub fn join_types<T: Display>(types: &[T], delimiter: &str) -> String {
855 types
856 .iter()
857 .map(|t| t.to_string())
858 .collect::<Vec<String>>()
859 .join(delimiter)
860 }
861
862 pub fn supports_zero_argument(&self) -> bool {
864 match &self {
865 TypeSignature::Exact(vec) => vec.is_empty(),
866 TypeSignature::Nullary => true,
867 TypeSignature::OneOf(types) => types
868 .iter()
869 .any(|type_sig| type_sig.supports_zero_argument()),
870 _ => false,
871 }
872 }
873
874 pub fn used_to_support_zero_arguments(&self) -> bool {
877 match &self {
878 TypeSignature::Any(num) => *num == 0,
879 _ => self.supports_zero_argument(),
880 }
881 }
882
883 #[deprecated(since = "46.0.0", note = "See get_example_types instead")]
884 pub fn get_possible_types(&self) -> Vec<Vec<DataType>> {
885 self.get_example_types()
886 }
887
888 pub fn get_example_types(&self) -> Vec<Vec<DataType>> {
895 match self {
896 TypeSignature::Exact(types) => vec![types.clone()],
897 TypeSignature::OneOf(types) => types
898 .iter()
899 .flat_map(|type_sig| type_sig.get_example_types())
900 .collect(),
901 TypeSignature::Uniform(arg_count, types) => types
902 .iter()
903 .cloned()
904 .map(|data_type| vec![data_type; *arg_count])
905 .collect(),
906 TypeSignature::Coercible(coercions) => coercions
907 .iter()
908 .map(|c| {
909 let mut all_types: IndexSet<DataType> =
910 c.desired_type().get_example_types().into_iter().collect();
911
912 if let Some(implicit_coercion) = c.implicit_coercion() {
913 let allowed_casts: Vec<DataType> = implicit_coercion
914 .allowed_source_types
915 .iter()
916 .flat_map(|t| t.get_example_types())
917 .collect();
918 all_types.extend(allowed_casts);
919 }
920
921 all_types.into_iter().collect::<Vec<_>>()
922 })
923 .multi_cartesian_product()
924 .collect(),
925 TypeSignature::Variadic(types) => types
926 .iter()
927 .cloned()
928 .map(|data_type| vec![data_type])
929 .collect(),
930 TypeSignature::Numeric(arg_count) => NUMERICS
931 .iter()
932 .cloned()
933 .map(|numeric_type| vec![numeric_type; *arg_count])
934 .collect(),
935 TypeSignature::String(arg_count) => get_data_types(&NativeType::String)
936 .into_iter()
937 .map(|dt| vec![dt; *arg_count])
938 .collect::<Vec<_>>(),
939 TypeSignature::Any(_)
941 | TypeSignature::Comparable(_)
942 | TypeSignature::Nullary
943 | TypeSignature::VariadicAny
944 | TypeSignature::ArraySignature(_)
945 | TypeSignature::UserDefined => vec![],
946 }
947 }
948}
949
950fn get_data_types(native_type: &NativeType) -> Vec<DataType> {
951 match native_type {
952 NativeType::Null => vec![DataType::Null],
953 NativeType::Boolean => vec![DataType::Boolean],
954 NativeType::Int8 => vec![DataType::Int8],
955 NativeType::Int16 => vec![DataType::Int16],
956 NativeType::Int32 => vec![DataType::Int32],
957 NativeType::Int64 => vec![DataType::Int64],
958 NativeType::UInt8 => vec![DataType::UInt8],
959 NativeType::UInt16 => vec![DataType::UInt16],
960 NativeType::UInt32 => vec![DataType::UInt32],
961 NativeType::UInt64 => vec![DataType::UInt64],
962 NativeType::Float16 => vec![DataType::Float16],
963 NativeType::Float32 => vec![DataType::Float32],
964 NativeType::Float64 => vec![DataType::Float64],
965 NativeType::Date => vec![DataType::Date32, DataType::Date64],
966 NativeType::Binary => vec![
967 DataType::Binary,
968 DataType::LargeBinary,
969 DataType::BinaryView,
970 ],
971 NativeType::String => {
972 vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]
973 }
974 NativeType::Decimal(precision, scale) => {
975 let mut types = vec![DataType::Decimal256(*precision, *scale)];
977 if *precision <= DECIMAL32_MAX_PRECISION {
978 types.push(DataType::Decimal32(*precision, *scale));
979 }
980 if *precision <= DECIMAL64_MAX_PRECISION {
981 types.push(DataType::Decimal64(*precision, *scale));
982 }
983 if *precision <= DECIMAL128_MAX_PRECISION {
984 types.push(DataType::Decimal128(*precision, *scale));
985 }
986 types
987 }
988 NativeType::Timestamp(time_unit, timezone) => {
989 vec![DataType::Timestamp(*time_unit, timezone.to_owned())]
990 }
991 NativeType::Time(TimeUnit::Second) => vec![DataType::Time32(TimeUnit::Second)],
992 NativeType::Time(TimeUnit::Millisecond) => {
993 vec![DataType::Time32(TimeUnit::Millisecond)]
994 }
995 NativeType::Time(TimeUnit::Microsecond) => {
996 vec![DataType::Time64(TimeUnit::Microsecond)]
997 }
998 NativeType::Time(TimeUnit::Nanosecond) => {
999 vec![DataType::Time64(TimeUnit::Nanosecond)]
1000 }
1001 NativeType::Duration(time_unit) => vec![DataType::Duration(*time_unit)],
1002 NativeType::Interval(interval_unit) => vec![DataType::Interval(*interval_unit)],
1003 NativeType::FixedSizeBinary(size) => vec![DataType::FixedSizeBinary(*size)],
1004 NativeType::FixedSizeList(logical_field, size) => {
1005 get_data_types(logical_field.logical_type.native())
1006 .iter()
1007 .map(|child_dt| {
1008 let field = Field::new(
1009 logical_field.name.clone(),
1010 child_dt.clone(),
1011 logical_field.nullable,
1012 );
1013 DataType::FixedSizeList(Arc::new(field), *size)
1014 })
1015 .collect()
1016 }
1017 NativeType::List(_)
1019 | NativeType::Struct(_)
1020 | NativeType::Union(_)
1021 | NativeType::Map(_) => {
1022 vec![]
1023 }
1024 }
1025}
1026
1027#[derive(Debug, Clone, Eq, PartialOrd)]
1052pub enum Coercion {
1053 Exact {
1055 desired_type: TypeSignatureClass,
1057 },
1058
1059 Implicit {
1061 desired_type: TypeSignatureClass,
1063 implicit_coercion: ImplicitCoercion,
1065 },
1066}
1067
1068impl Coercion {
1069 pub fn new_exact(desired_type: TypeSignatureClass) -> Self {
1070 Self::Exact { desired_type }
1071 }
1072
1073 pub fn new_implicit(
1078 desired_type: TypeSignatureClass,
1079 allowed_source_types: Vec<TypeSignatureClass>,
1080 default_casted_type: NativeType,
1081 ) -> Self {
1082 Self::Implicit {
1083 desired_type,
1084 implicit_coercion: ImplicitCoercion {
1085 allowed_source_types,
1086 default_casted_type,
1087 },
1088 }
1089 }
1090
1091 pub fn allowed_source_types(&self) -> &[TypeSignatureClass] {
1092 match self {
1093 Coercion::Exact { .. } => &[],
1094 Coercion::Implicit {
1095 implicit_coercion, ..
1096 } => implicit_coercion.allowed_source_types.as_slice(),
1097 }
1098 }
1099
1100 pub fn default_casted_type(&self) -> Option<&NativeType> {
1101 match self {
1102 Coercion::Exact { .. } => None,
1103 Coercion::Implicit {
1104 implicit_coercion, ..
1105 } => Some(&implicit_coercion.default_casted_type),
1106 }
1107 }
1108
1109 pub fn desired_type(&self) -> &TypeSignatureClass {
1110 match self {
1111 Coercion::Exact { desired_type } => desired_type,
1112 Coercion::Implicit { desired_type, .. } => desired_type,
1113 }
1114 }
1115
1116 pub fn implicit_coercion(&self) -> Option<&ImplicitCoercion> {
1117 match self {
1118 Coercion::Exact { .. } => None,
1119 Coercion::Implicit {
1120 implicit_coercion, ..
1121 } => Some(implicit_coercion),
1122 }
1123 }
1124}
1125
1126impl Display for Coercion {
1127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1128 write!(f, "{}", self.desired_type())
1129 }
1130}
1131
1132impl PartialEq for Coercion {
1133 fn eq(&self, other: &Self) -> bool {
1134 self.desired_type() == other.desired_type()
1135 && self.implicit_coercion() == other.implicit_coercion()
1136 }
1137}
1138
1139impl Hash for Coercion {
1140 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1141 self.desired_type().hash(state);
1142 self.implicit_coercion().hash(state);
1143 }
1144}
1145
1146#[derive(Debug, Clone, Eq, PartialOrd)]
1168pub struct ImplicitCoercion {
1169 allowed_source_types: Vec<TypeSignatureClass>,
1171
1172 default_casted_type: NativeType,
1176}
1177
1178impl Display for ImplicitCoercion {
1179 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1180 write!(f, "ImplicitCoercion(")?;
1181 for (i, source_type) in self.allowed_source_types.iter().enumerate() {
1182 if i > 0 {
1183 write!(f, ", ")?;
1184 }
1185 write!(f, "{source_type}")?;
1186 }
1187 write!(f, "; default={}", self.default_casted_type)
1188 }
1189}
1190
1191impl PartialEq for ImplicitCoercion {
1192 fn eq(&self, other: &Self) -> bool {
1193 self.allowed_source_types == other.allowed_source_types
1194 && self.default_casted_type == other.default_casted_type
1195 }
1196}
1197
1198impl Hash for ImplicitCoercion {
1199 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1200 self.allowed_source_types.hash(state);
1201 self.default_casted_type.hash(state);
1202 }
1203}
1204
1205#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
1212pub struct Signature {
1213 pub type_signature: TypeSignature,
1215 pub volatility: Volatility,
1217 pub parameter_names: Option<Vec<String>>,
1224}
1225
1226impl Signature {
1227 pub fn new(type_signature: TypeSignature, volatility: Volatility) -> Self {
1229 Signature {
1230 type_signature,
1231 volatility,
1232 parameter_names: None,
1233 }
1234 }
1235 pub fn variadic(common_types: Vec<DataType>, volatility: Volatility) -> Self {
1237 Self {
1238 type_signature: TypeSignature::Variadic(common_types),
1239 volatility,
1240 parameter_names: None,
1241 }
1242 }
1243 pub fn user_defined(volatility: Volatility) -> Self {
1245 Self {
1246 type_signature: TypeSignature::UserDefined,
1247 volatility,
1248 parameter_names: None,
1249 }
1250 }
1251
1252 pub fn numeric(arg_count: usize, volatility: Volatility) -> Self {
1254 Self {
1255 type_signature: TypeSignature::Numeric(arg_count),
1256 volatility,
1257 parameter_names: None,
1258 }
1259 }
1260
1261 pub fn string(arg_count: usize, volatility: Volatility) -> Self {
1263 Self {
1264 type_signature: TypeSignature::String(arg_count),
1265 volatility,
1266 parameter_names: None,
1267 }
1268 }
1269
1270 pub fn variadic_any(volatility: Volatility) -> Self {
1272 Self {
1273 type_signature: TypeSignature::VariadicAny,
1274 volatility,
1275 parameter_names: None,
1276 }
1277 }
1278 pub fn uniform(
1280 arg_count: usize,
1281 valid_types: Vec<DataType>,
1282 volatility: Volatility,
1283 ) -> Self {
1284 Self {
1285 type_signature: TypeSignature::Uniform(arg_count, valid_types),
1286 volatility,
1287 parameter_names: None,
1288 }
1289 }
1290 pub fn exact(exact_types: Vec<DataType>, volatility: Volatility) -> Self {
1292 Signature {
1293 type_signature: TypeSignature::Exact(exact_types),
1294 volatility,
1295 parameter_names: None,
1296 }
1297 }
1298
1299 pub fn coercible(target_types: Vec<Coercion>, volatility: Volatility) -> Self {
1301 Self {
1302 type_signature: TypeSignature::Coercible(target_types),
1303 volatility,
1304 parameter_names: None,
1305 }
1306 }
1307
1308 pub fn comparable(arg_count: usize, volatility: Volatility) -> Self {
1310 Self {
1311 type_signature: TypeSignature::Comparable(arg_count),
1312 volatility,
1313 parameter_names: None,
1314 }
1315 }
1316
1317 pub fn nullary(volatility: Volatility) -> Self {
1318 Signature {
1319 type_signature: TypeSignature::Nullary,
1320 volatility,
1321 parameter_names: None,
1322 }
1323 }
1324
1325 pub fn any(arg_count: usize, volatility: Volatility) -> Self {
1327 Signature {
1328 type_signature: TypeSignature::Any(arg_count),
1329 volatility,
1330 parameter_names: None,
1331 }
1332 }
1333
1334 pub fn one_of(type_signatures: Vec<TypeSignature>, volatility: Volatility) -> Self {
1336 Signature {
1337 type_signature: TypeSignature::OneOf(type_signatures),
1338 volatility,
1339 parameter_names: None,
1340 }
1341 }
1342
1343 pub fn array_and_element(volatility: Volatility) -> Self {
1345 Signature {
1346 type_signature: TypeSignature::ArraySignature(
1347 ArrayFunctionSignature::Array {
1348 arguments: vec![
1349 ArrayFunctionArgument::Array,
1350 ArrayFunctionArgument::Element,
1351 ],
1352 array_coercion: Some(ListCoercion::FixedSizedListToList),
1353 },
1354 ),
1355 volatility,
1356 parameter_names: None,
1357 }
1358 }
1359
1360 pub fn element_and_array(volatility: Volatility) -> Self {
1362 Signature {
1363 type_signature: TypeSignature::ArraySignature(
1364 ArrayFunctionSignature::Array {
1365 arguments: vec![
1366 ArrayFunctionArgument::Element,
1367 ArrayFunctionArgument::Array,
1368 ],
1369 array_coercion: Some(ListCoercion::FixedSizedListToList),
1370 },
1371 ),
1372 volatility,
1373 parameter_names: None,
1374 }
1375 }
1376
1377 pub fn arrays(
1379 n: usize,
1380 coercion: Option<ListCoercion>,
1381 volatility: Volatility,
1382 ) -> Self {
1383 Signature {
1384 type_signature: TypeSignature::ArraySignature(
1385 ArrayFunctionSignature::Array {
1386 arguments: vec![ArrayFunctionArgument::Array; n],
1387 array_coercion: coercion,
1388 },
1389 ),
1390 volatility,
1391 parameter_names: None,
1392 }
1393 }
1394
1395 pub fn array_and_element_and_optional_index(volatility: Volatility) -> Self {
1397 Signature {
1398 type_signature: TypeSignature::OneOf(vec![
1399 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1400 arguments: vec![
1401 ArrayFunctionArgument::Array,
1402 ArrayFunctionArgument::Element,
1403 ],
1404 array_coercion: Some(ListCoercion::FixedSizedListToList),
1405 }),
1406 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1407 arguments: vec![
1408 ArrayFunctionArgument::Array,
1409 ArrayFunctionArgument::Element,
1410 ArrayFunctionArgument::Index,
1411 ],
1412 array_coercion: Some(ListCoercion::FixedSizedListToList),
1413 }),
1414 ]),
1415 volatility,
1416 parameter_names: None,
1417 }
1418 }
1419
1420 pub fn array_and_index(volatility: Volatility) -> Self {
1422 Signature {
1423 type_signature: TypeSignature::ArraySignature(
1424 ArrayFunctionSignature::Array {
1425 arguments: vec![
1426 ArrayFunctionArgument::Array,
1427 ArrayFunctionArgument::Index,
1428 ],
1429 array_coercion: Some(ListCoercion::FixedSizedListToList),
1430 },
1431 ),
1432 volatility,
1433 parameter_names: None,
1434 }
1435 }
1436
1437 pub fn array(volatility: Volatility) -> Self {
1439 Signature::arrays(1, Some(ListCoercion::FixedSizedListToList), volatility)
1440 }
1441
1442 pub fn with_parameter_names(mut self, names: Vec<impl Into<String>>) -> Result<Self> {
1458 let names = names.into_iter().map(Into::into).collect::<Vec<String>>();
1459 self.validate_parameter_names(&names)?;
1461 self.parameter_names = Some(names);
1462 Ok(self)
1463 }
1464
1465 fn validate_parameter_names(&self, names: &[String]) -> Result<()> {
1467 match self.type_signature.arity() {
1468 Arity::Fixed(expected) => {
1469 if names.len() != expected {
1470 return plan_err!(
1471 "Parameter names count ({}) does not match signature arity ({})",
1472 names.len(),
1473 expected
1474 );
1475 }
1476 }
1477 Arity::Variable => {
1478 if self.type_signature != TypeSignature::UserDefined {
1481 return plan_err!(
1482 "Cannot specify parameter names for variable arity signature: {:?}",
1483 self.type_signature
1484 );
1485 }
1486 }
1487 }
1488
1489 let mut seen = std::collections::HashSet::new();
1490 for name in names {
1491 if !seen.insert(name) {
1492 return plan_err!("Duplicate parameter name: '{}'", name);
1493 }
1494 }
1495
1496 Ok(())
1497 }
1498}
1499
1500#[cfg(test)]
1501mod tests {
1502 use datafusion_common::types::{
1503 NativeType, logical_float64, logical_int32, logical_int64, logical_string,
1504 };
1505
1506 use super::*;
1507 use crate::signature::{
1508 ArrayFunctionArgument, ArrayFunctionSignature, Coercion, TypeSignatureClass,
1509 };
1510
1511 #[test]
1512 fn supports_zero_argument_tests() {
1513 let positive_cases = vec![
1515 TypeSignature::Exact(vec![]),
1516 TypeSignature::OneOf(vec![
1517 TypeSignature::Exact(vec![DataType::Int8]),
1518 TypeSignature::Nullary,
1519 TypeSignature::Uniform(1, vec![DataType::Int8]),
1520 ]),
1521 TypeSignature::Nullary,
1522 ];
1523
1524 for case in positive_cases {
1525 assert!(
1526 case.supports_zero_argument(),
1527 "Expected {case:?} to support zero arguments"
1528 );
1529 }
1530
1531 let negative_cases = vec![
1533 TypeSignature::Exact(vec![DataType::Utf8]),
1534 TypeSignature::Uniform(1, vec![DataType::Float64]),
1535 TypeSignature::Any(1),
1536 TypeSignature::VariadicAny,
1537 TypeSignature::OneOf(vec![
1538 TypeSignature::Exact(vec![DataType::Int8]),
1539 TypeSignature::Uniform(1, vec![DataType::Int8]),
1540 ]),
1541 ];
1542
1543 for case in negative_cases {
1544 assert!(
1545 !case.supports_zero_argument(),
1546 "Expected {case:?} not to support zero arguments"
1547 );
1548 }
1549 }
1550
1551 #[test]
1552 fn type_signature_partial_ord() {
1553 assert!(TypeSignature::UserDefined < TypeSignature::VariadicAny);
1555 assert!(TypeSignature::UserDefined < TypeSignature::Any(1));
1556
1557 assert!(
1558 TypeSignature::Uniform(1, vec![DataType::Null])
1559 < TypeSignature::Uniform(1, vec![DataType::Boolean])
1560 );
1561 assert!(
1562 TypeSignature::Uniform(1, vec![DataType::Null])
1563 < TypeSignature::Uniform(2, vec![DataType::Null])
1564 );
1565 assert!(
1566 TypeSignature::Uniform(usize::MAX, vec![DataType::Null])
1567 < TypeSignature::Exact(vec![DataType::Null])
1568 );
1569 }
1570
1571 #[test]
1572 fn test_get_possible_types() {
1573 let type_signature = TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]);
1574 let possible_types = type_signature.get_example_types();
1575 assert_eq!(possible_types, vec![vec![DataType::Int32, DataType::Int64]]);
1576
1577 let type_signature = TypeSignature::OneOf(vec![
1578 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1579 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1580 ]);
1581 let possible_types = type_signature.get_example_types();
1582 assert_eq!(
1583 possible_types,
1584 vec![
1585 vec![DataType::Int32, DataType::Int64],
1586 vec![DataType::Float32, DataType::Float64]
1587 ]
1588 );
1589
1590 let type_signature = TypeSignature::OneOf(vec![
1591 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1592 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1593 TypeSignature::Exact(vec![DataType::Utf8]),
1594 ]);
1595 let possible_types = type_signature.get_example_types();
1596 assert_eq!(
1597 possible_types,
1598 vec![
1599 vec![DataType::Int32, DataType::Int64],
1600 vec![DataType::Float32, DataType::Float64],
1601 vec![DataType::Utf8]
1602 ]
1603 );
1604
1605 let type_signature =
1606 TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Int64]);
1607 let possible_types = type_signature.get_example_types();
1608 assert_eq!(
1609 possible_types,
1610 vec![
1611 vec![DataType::Float32, DataType::Float32],
1612 vec![DataType::Int64, DataType::Int64]
1613 ]
1614 );
1615
1616 let type_signature = TypeSignature::Coercible(vec![
1617 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1618 Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
1619 ]);
1620 let possible_types = type_signature.get_example_types();
1621 assert_eq!(
1622 possible_types,
1623 vec![
1624 vec![DataType::Utf8, DataType::Int64],
1625 vec![DataType::LargeUtf8, DataType::Int64],
1626 vec![DataType::Utf8View, DataType::Int64]
1627 ]
1628 );
1629
1630 let type_signature =
1631 TypeSignature::Variadic(vec![DataType::Int32, DataType::Int64]);
1632 let possible_types = type_signature.get_example_types();
1633 assert_eq!(
1634 possible_types,
1635 vec![vec![DataType::Int32], vec![DataType::Int64]]
1636 );
1637
1638 let type_signature = TypeSignature::Numeric(2);
1639 let possible_types = type_signature.get_example_types();
1640 assert_eq!(
1641 possible_types,
1642 vec![
1643 vec![DataType::Int8, DataType::Int8],
1644 vec![DataType::Int16, DataType::Int16],
1645 vec![DataType::Int32, DataType::Int32],
1646 vec![DataType::Int64, DataType::Int64],
1647 vec![DataType::UInt8, DataType::UInt8],
1648 vec![DataType::UInt16, DataType::UInt16],
1649 vec![DataType::UInt32, DataType::UInt32],
1650 vec![DataType::UInt64, DataType::UInt64],
1651 vec![DataType::Float16, DataType::Float16],
1652 vec![DataType::Float32, DataType::Float32],
1653 vec![DataType::Float64, DataType::Float64]
1654 ]
1655 );
1656
1657 let type_signature = TypeSignature::String(2);
1658 let possible_types = type_signature.get_example_types();
1659 assert_eq!(
1660 possible_types,
1661 vec![
1662 vec![DataType::Utf8, DataType::Utf8],
1663 vec![DataType::LargeUtf8, DataType::LargeUtf8],
1664 vec![DataType::Utf8View, DataType::Utf8View]
1665 ]
1666 );
1667 }
1668
1669 #[test]
1670 fn test_signature_with_parameter_names() {
1671 let sig = Signature::exact(
1672 vec![DataType::Int32, DataType::Utf8],
1673 Volatility::Immutable,
1674 )
1675 .with_parameter_names(vec!["count".to_string(), "name".to_string()])
1676 .unwrap();
1677
1678 assert_eq!(
1679 sig.parameter_names,
1680 Some(vec!["count".to_string(), "name".to_string()])
1681 );
1682 assert_eq!(
1683 sig.type_signature,
1684 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8])
1685 );
1686 }
1687
1688 #[test]
1689 fn test_signature_parameter_names_wrong_count() {
1690 let result = Signature::exact(
1691 vec![DataType::Int32, DataType::Utf8],
1692 Volatility::Immutable,
1693 )
1694 .with_parameter_names(vec!["count".to_string()]); assert!(result.is_err());
1697 assert!(
1698 result
1699 .unwrap_err()
1700 .to_string()
1701 .contains("does not match signature arity")
1702 );
1703 }
1704
1705 #[test]
1706 fn test_signature_parameter_names_duplicate() {
1707 let result = Signature::exact(
1708 vec![DataType::Int32, DataType::Int32],
1709 Volatility::Immutable,
1710 )
1711 .with_parameter_names(vec!["count".to_string(), "count".to_string()]);
1712
1713 assert!(result.is_err());
1714 assert!(
1715 result
1716 .unwrap_err()
1717 .to_string()
1718 .contains("Duplicate parameter name")
1719 );
1720 }
1721
1722 #[test]
1723 fn test_signature_parameter_names_variadic() {
1724 let result = Signature::variadic(vec![DataType::Int32], Volatility::Immutable)
1725 .with_parameter_names(vec!["arg".to_string()]);
1726
1727 assert!(result.is_err());
1728 assert!(
1729 result
1730 .unwrap_err()
1731 .to_string()
1732 .contains("variable arity signature")
1733 );
1734 }
1735
1736 #[test]
1737 fn test_signature_without_parameter_names() {
1738 let sig = Signature::exact(
1739 vec![DataType::Int32, DataType::Utf8],
1740 Volatility::Immutable,
1741 );
1742
1743 assert_eq!(sig.parameter_names, None);
1744 }
1745
1746 #[test]
1747 fn test_signature_uniform_with_parameter_names() {
1748 let sig = Signature::uniform(3, vec![DataType::Float64], Volatility::Immutable)
1749 .with_parameter_names(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1750 .unwrap();
1751
1752 assert_eq!(
1753 sig.parameter_names,
1754 Some(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1755 );
1756 }
1757
1758 #[test]
1759 fn test_signature_numeric_with_parameter_names() {
1760 let sig = Signature::numeric(2, Volatility::Immutable)
1761 .with_parameter_names(vec!["a".to_string(), "b".to_string()])
1762 .unwrap();
1763
1764 assert_eq!(
1765 sig.parameter_names,
1766 Some(vec!["a".to_string(), "b".to_string()])
1767 );
1768 }
1769
1770 #[test]
1771 fn test_signature_nullary_with_empty_names() {
1772 let sig = Signature::nullary(Volatility::Immutable)
1773 .with_parameter_names(Vec::<String>::new())
1774 .unwrap();
1775
1776 assert_eq!(sig.parameter_names, Some(vec![]));
1777 }
1778
1779 #[test]
1780 fn test_to_string_repr_with_names_exact() {
1781 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1782
1783 assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
1784
1785 let names = vec!["id".to_string(), "name".to_string()];
1786 assert_eq!(
1787 sig.to_string_repr_with_names(Some(&names)),
1788 vec!["id: Int32, name: Utf8"]
1789 );
1790 }
1791
1792 #[test]
1793 fn test_to_string_repr_with_names_any() {
1794 let sig = TypeSignature::Any(3);
1795
1796 assert_eq!(sig.to_string_repr_with_names(None), vec!["Any, Any, Any"]);
1797
1798 let names = vec!["x".to_string(), "y".to_string(), "z".to_string()];
1799 assert_eq!(
1800 sig.to_string_repr_with_names(Some(&names)),
1801 vec!["x: Any, y: Any, z: Any"]
1802 );
1803 }
1804
1805 #[test]
1806 fn test_to_string_repr_with_names_one_of() {
1807 let sig =
1808 TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]);
1809
1810 assert_eq!(
1811 sig.to_string_repr_with_names(None),
1812 vec!["Any, Any", "Any, Any, Any"]
1813 );
1814
1815 let names = vec![
1816 "str".to_string(),
1817 "start_pos".to_string(),
1818 "length".to_string(),
1819 ];
1820 assert_eq!(
1821 sig.to_string_repr_with_names(Some(&names)),
1822 vec![
1823 "str: Any, start_pos: Any",
1824 "str: Any, start_pos: Any, length: Any"
1825 ]
1826 );
1827 }
1828
1829 #[test]
1830 fn test_to_string_repr_with_names_partial() {
1831 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1833
1834 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1836 assert_eq!(
1837 sig.to_string_repr_with_names(Some(&names)),
1838 vec!["a: Int32, b: Utf8"]
1839 );
1840 }
1841
1842 #[test]
1843 fn test_to_string_repr_with_names_uniform() {
1844 let sig = TypeSignature::Uniform(2, vec![DataType::Float64]);
1845
1846 assert_eq!(
1847 sig.to_string_repr_with_names(None),
1848 vec!["Float64, Float64"]
1849 );
1850
1851 let names = vec!["x".to_string(), "y".to_string()];
1852 assert_eq!(
1853 sig.to_string_repr_with_names(Some(&names)),
1854 vec!["x: Float64, y: Float64"]
1855 );
1856 }
1857
1858 #[test]
1859 fn test_to_string_repr_with_names_coercible() {
1860 let sig = TypeSignature::Coercible(vec![
1861 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1862 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1863 ]);
1864
1865 let names = vec!["a".to_string(), "b".to_string()];
1866 let result = sig.to_string_repr_with_names(Some(&names));
1867 assert_eq!(result.len(), 1);
1869 assert!(result[0].starts_with("a: "));
1870 assert!(result[0].contains(", b: "));
1871 }
1872
1873 #[test]
1874 fn test_to_string_repr_with_names_comparable_numeric_string() {
1875 let comparable = TypeSignature::Comparable(3);
1876 let numeric = TypeSignature::Numeric(2);
1877 let string_sig = TypeSignature::String(2);
1878
1879 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1880
1881 assert_eq!(
1883 comparable.to_string_repr_with_names(Some(&names)),
1884 vec!["a: Comparable, b: Comparable, c: Comparable"]
1885 );
1886 assert_eq!(
1887 numeric.to_string_repr_with_names(Some(&names)),
1888 vec!["a: Numeric, b: Numeric"]
1889 );
1890 assert_eq!(
1891 string_sig.to_string_repr_with_names(Some(&names)),
1892 vec!["a: String, b: String"]
1893 );
1894 }
1895
1896 #[test]
1897 fn test_to_string_repr_with_names_variadic_fallback() {
1898 let variadic = TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]);
1899 let names = vec!["x".to_string()];
1900 assert_eq!(
1901 variadic.to_string_repr_with_names(Some(&names)),
1902 variadic.to_string_repr()
1903 );
1904
1905 let variadic_any = TypeSignature::VariadicAny;
1906 assert_eq!(
1907 variadic_any.to_string_repr_with_names(Some(&names)),
1908 variadic_any.to_string_repr()
1909 );
1910
1911 let user_defined = TypeSignature::UserDefined;
1913 assert_eq!(
1914 user_defined.to_string_repr_with_names(Some(&names)),
1915 vec!["x"]
1916 );
1917 assert_eq!(
1918 user_defined.to_string_repr_with_names(None),
1919 user_defined.to_string_repr()
1920 );
1921 }
1922
1923 #[test]
1924 fn test_to_string_repr_with_names_nullary() {
1925 let sig = TypeSignature::Nullary;
1926 let names = vec!["x".to_string()];
1927
1928 assert_eq!(
1930 sig.to_string_repr_with_names(Some(&names)),
1931 vec!["NullAry()"]
1932 );
1933 assert_eq!(sig.to_string_repr_with_names(None), vec!["NullAry()"]);
1934 }
1935
1936 #[test]
1937 fn test_to_string_repr_with_names_array_signature() {
1938 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1939 arguments: vec![
1940 ArrayFunctionArgument::Array,
1941 ArrayFunctionArgument::Index,
1942 ArrayFunctionArgument::Element,
1943 ],
1944 array_coercion: None,
1945 });
1946
1947 assert_eq!(
1948 sig.to_string_repr_with_names(None),
1949 vec!["array, index, element"]
1950 );
1951
1952 let names = vec!["arr".to_string(), "idx".to_string(), "val".to_string()];
1953 assert_eq!(
1954 sig.to_string_repr_with_names(Some(&names)),
1955 vec!["arr: array, idx: index, val: element"]
1956 );
1957
1958 let recursive =
1959 TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
1960 let names = vec!["array".to_string()];
1961 assert_eq!(
1962 recursive.to_string_repr_with_names(Some(&names)),
1963 vec!["array: recursive_array"]
1964 );
1965
1966 let map_array = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
1968 let names = vec!["map".to_string()];
1969 assert_eq!(
1970 map_array.to_string_repr_with_names(Some(&names)),
1971 vec!["map: map_array"]
1972 );
1973 }
1974
1975 #[test]
1976 fn test_type_signature_arity_exact() {
1977 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1978 assert_eq!(sig.arity(), Arity::Fixed(2));
1979
1980 let sig = TypeSignature::Exact(vec![]);
1981 assert_eq!(sig.arity(), Arity::Fixed(0));
1982 }
1983
1984 #[test]
1985 fn test_type_signature_arity_uniform() {
1986 let sig = TypeSignature::Uniform(3, vec![DataType::Float64]);
1987 assert_eq!(sig.arity(), Arity::Fixed(3));
1988
1989 let sig = TypeSignature::Uniform(1, vec![DataType::Int32]);
1990 assert_eq!(sig.arity(), Arity::Fixed(1));
1991 }
1992
1993 #[test]
1994 fn test_type_signature_arity_numeric() {
1995 let sig = TypeSignature::Numeric(2);
1996 assert_eq!(sig.arity(), Arity::Fixed(2));
1997 }
1998
1999 #[test]
2000 fn test_type_signature_arity_string() {
2001 let sig = TypeSignature::String(3);
2002 assert_eq!(sig.arity(), Arity::Fixed(3));
2003 }
2004
2005 #[test]
2006 fn test_type_signature_arity_comparable() {
2007 let sig = TypeSignature::Comparable(2);
2008 assert_eq!(sig.arity(), Arity::Fixed(2));
2009 }
2010
2011 #[test]
2012 fn test_type_signature_arity_any() {
2013 let sig = TypeSignature::Any(4);
2014 assert_eq!(sig.arity(), Arity::Fixed(4));
2015 }
2016
2017 #[test]
2018 fn test_type_signature_arity_coercible() {
2019 let sig = TypeSignature::Coercible(vec![
2020 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
2021 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
2022 ]);
2023 assert_eq!(sig.arity(), Arity::Fixed(2));
2024 }
2025
2026 #[test]
2027 fn test_type_signature_arity_nullary() {
2028 let sig = TypeSignature::Nullary;
2029 assert_eq!(sig.arity(), Arity::Fixed(0));
2030 }
2031
2032 #[test]
2033 fn test_type_signature_arity_array_signature() {
2034 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
2036 arguments: vec![ArrayFunctionArgument::Array, ArrayFunctionArgument::Index],
2037 array_coercion: None,
2038 });
2039 assert_eq!(sig.arity(), Arity::Fixed(2));
2040
2041 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
2043 arguments: vec![
2044 ArrayFunctionArgument::Array,
2045 ArrayFunctionArgument::Element,
2046 ArrayFunctionArgument::Index,
2047 ],
2048 array_coercion: None,
2049 });
2050 assert_eq!(sig.arity(), Arity::Fixed(3));
2051
2052 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
2054 assert_eq!(sig.arity(), Arity::Fixed(1));
2055
2056 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
2058 assert_eq!(sig.arity(), Arity::Fixed(1));
2059 }
2060
2061 #[test]
2062 fn test_type_signature_arity_one_of_fixed() {
2063 let sig = TypeSignature::OneOf(vec![
2065 TypeSignature::Exact(vec![DataType::Int32]),
2066 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2067 TypeSignature::Exact(vec![
2068 DataType::Int32,
2069 DataType::Utf8,
2070 DataType::Float64,
2071 ]),
2072 ]);
2073 assert_eq!(sig.arity(), Arity::Fixed(3));
2074 }
2075
2076 #[test]
2077 fn test_type_signature_arity_one_of_variable() {
2078 let sig = TypeSignature::OneOf(vec![
2080 TypeSignature::Exact(vec![DataType::Int32]),
2081 TypeSignature::VariadicAny,
2082 ]);
2083 assert_eq!(sig.arity(), Arity::Variable);
2084 }
2085
2086 #[test]
2087 fn test_type_signature_arity_variadic() {
2088 let sig = TypeSignature::Variadic(vec![DataType::Int32]);
2089 assert_eq!(sig.arity(), Arity::Variable);
2090
2091 let sig = TypeSignature::VariadicAny;
2092 assert_eq!(sig.arity(), Arity::Variable);
2093 }
2094
2095 #[test]
2096 fn test_type_signature_arity_user_defined() {
2097 let sig = TypeSignature::UserDefined;
2098 assert_eq!(sig.arity(), Arity::Variable);
2099 }
2100
2101 #[test]
2102 fn test_type_signature_display() {
2103 use insta::assert_snapshot;
2104
2105 assert_snapshot!(TypeSignature::Nullary, @"Nullary");
2106 assert_snapshot!(TypeSignature::Any(2), @"Any(2)");
2107 assert_snapshot!(TypeSignature::Numeric(3), @"Numeric(3)");
2108 assert_snapshot!(TypeSignature::String(1), @"String(1)");
2109 assert_snapshot!(TypeSignature::Comparable(2), @"Comparable(2)");
2110 assert_snapshot!(TypeSignature::VariadicAny, @"VariadicAny");
2111 assert_snapshot!(TypeSignature::UserDefined, @"UserDefined");
2112
2113 assert_snapshot!(
2114 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2115 @"Exact(Int32, Utf8)"
2116 );
2117 assert_snapshot!(
2118 TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]),
2119 @"Variadic(Utf8, LargeUtf8)"
2120 );
2121 assert_snapshot!(
2122 TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Float64]),
2123 @"Uniform(2, [Float32, Float64])"
2124 );
2125
2126 assert_snapshot!(
2127 TypeSignature::Coercible(vec![
2128 Coercion::new_exact(TypeSignatureClass::Native(logical_float64())),
2129 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
2130 ]),
2131 @"Coercible(Float64, Int32)"
2132 );
2133
2134 assert_snapshot!(
2135 TypeSignature::OneOf(vec![
2136 TypeSignature::Nullary,
2137 TypeSignature::VariadicAny,
2138 ]),
2139 @"OneOf(Nullary, VariadicAny)"
2140 );
2141 }
2142
2143 #[test]
2144 fn test_type_signature_class_display() {
2145 use insta::assert_snapshot;
2146
2147 assert_snapshot!(TypeSignatureClass::Any, @"Any");
2148 assert_snapshot!(TypeSignatureClass::Numeric, @"Numeric");
2149 assert_snapshot!(TypeSignatureClass::Integer, @"Integer");
2150 assert_snapshot!(TypeSignatureClass::Float, @"Float");
2151 assert_snapshot!(TypeSignatureClass::Decimal, @"Decimal");
2152 assert_snapshot!(TypeSignatureClass::Timestamp, @"Timestamp");
2153 assert_snapshot!(TypeSignatureClass::Time, @"Time");
2154 assert_snapshot!(TypeSignatureClass::Interval, @"Interval");
2155 assert_snapshot!(TypeSignatureClass::Duration, @"Duration");
2156 assert_snapshot!(TypeSignatureClass::Binary, @"Binary");
2157 assert_snapshot!(TypeSignatureClass::Native(logical_int32()), @"Int32");
2158 assert_snapshot!(TypeSignatureClass::Native(logical_string()), @"String");
2159 }
2160
2161 #[test]
2162 fn test_coercion_display() {
2163 use insta::assert_snapshot;
2164
2165 let exact_int = Coercion::new_exact(TypeSignatureClass::Native(logical_int32()));
2166 assert_snapshot!(exact_int, @"Int32");
2167
2168 let exact_numeric = Coercion::new_exact(TypeSignatureClass::Numeric);
2169 assert_snapshot!(exact_numeric, @"Numeric");
2170
2171 let implicit = Coercion::new_implicit(
2172 TypeSignatureClass::Native(logical_float64()),
2173 vec![TypeSignatureClass::Numeric],
2174 NativeType::Float64,
2175 );
2176 assert_snapshot!(implicit, @"Float64");
2177
2178 let implicit_with_multiple_sources = Coercion::new_implicit(
2179 TypeSignatureClass::Native(logical_int64()),
2180 vec![TypeSignatureClass::Integer, TypeSignatureClass::Numeric],
2181 NativeType::Int64,
2182 );
2183 assert_snapshot!(implicit_with_multiple_sources, @"Int64");
2184 }
2185
2186 #[test]
2187 fn test_to_string_repr_coercible() {
2188 use insta::assert_snapshot;
2189
2190 let sig = TypeSignature::Coercible(vec![
2192 Coercion::new_implicit(
2193 TypeSignatureClass::Native(logical_float64()),
2194 vec![TypeSignatureClass::Numeric],
2195 NativeType::Float64,
2196 ),
2197 Coercion::new_implicit(
2198 TypeSignatureClass::Native(logical_int64()),
2199 vec![TypeSignatureClass::Integer],
2200 NativeType::Int64,
2201 ),
2202 ]);
2203 let repr = sig.to_string_repr();
2204 assert_eq!(repr.len(), 1);
2205 assert_snapshot!(repr[0], @"Float64, Int64");
2206 }
2207
2208 #[test]
2209 fn test_to_string_repr_coercible_exact() {
2210 use insta::assert_snapshot;
2211
2212 let sig = TypeSignature::Coercible(vec![
2213 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
2214 Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
2215 ]);
2216 let repr = sig.to_string_repr();
2217 assert_eq!(repr.len(), 1);
2218 assert_snapshot!(repr[0], @"String, Int64");
2219 }
2220}