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 pub fn get_example_types(&self) -> Vec<Vec<DataType>> {
890 match self {
891 TypeSignature::Exact(types) => vec![types.clone()],
892 TypeSignature::OneOf(types) => types
893 .iter()
894 .flat_map(|type_sig| type_sig.get_example_types())
895 .collect(),
896 TypeSignature::Uniform(arg_count, types) => types
897 .iter()
898 .cloned()
899 .map(|data_type| vec![data_type; *arg_count])
900 .collect(),
901 TypeSignature::Coercible(coercions) => coercions
902 .iter()
903 .map(|c| {
904 let mut all_types: IndexSet<DataType> =
905 c.desired_type().get_example_types().into_iter().collect();
906
907 if let Some(implicit_coercion) = c.implicit_coercion() {
908 let allowed_casts: Vec<DataType> = implicit_coercion
909 .allowed_source_types
910 .iter()
911 .flat_map(|t| t.get_example_types())
912 .collect();
913 all_types.extend(allowed_casts);
914 }
915
916 all_types.into_iter().collect::<Vec<_>>()
917 })
918 .multi_cartesian_product()
919 .collect(),
920 TypeSignature::Variadic(types) => types
921 .iter()
922 .cloned()
923 .map(|data_type| vec![data_type])
924 .collect(),
925 TypeSignature::Numeric(arg_count) => NUMERICS
926 .iter()
927 .cloned()
928 .map(|numeric_type| vec![numeric_type; *arg_count])
929 .collect(),
930 TypeSignature::String(arg_count) => get_data_types(&NativeType::String)
931 .into_iter()
932 .map(|dt| vec![dt; *arg_count])
933 .collect::<Vec<_>>(),
934 TypeSignature::Any(_)
936 | TypeSignature::Comparable(_)
937 | TypeSignature::Nullary
938 | TypeSignature::VariadicAny
939 | TypeSignature::ArraySignature(_)
940 | TypeSignature::UserDefined => vec![],
941 }
942 }
943}
944
945fn get_data_types(native_type: &NativeType) -> Vec<DataType> {
946 match native_type {
947 NativeType::Null => vec![DataType::Null],
948 NativeType::Boolean => vec![DataType::Boolean],
949 NativeType::Int8 => vec![DataType::Int8],
950 NativeType::Int16 => vec![DataType::Int16],
951 NativeType::Int32 => vec![DataType::Int32],
952 NativeType::Int64 => vec![DataType::Int64],
953 NativeType::UInt8 => vec![DataType::UInt8],
954 NativeType::UInt16 => vec![DataType::UInt16],
955 NativeType::UInt32 => vec![DataType::UInt32],
956 NativeType::UInt64 => vec![DataType::UInt64],
957 NativeType::Float16 => vec![DataType::Float16],
958 NativeType::Float32 => vec![DataType::Float32],
959 NativeType::Float64 => vec![DataType::Float64],
960 NativeType::Date => vec![DataType::Date32, DataType::Date64],
961 NativeType::Binary => vec![
962 DataType::Binary,
963 DataType::LargeBinary,
964 DataType::BinaryView,
965 ],
966 NativeType::String => {
967 vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]
968 }
969 NativeType::Decimal(precision, scale) => {
970 let mut types = vec![DataType::Decimal256(*precision, *scale)];
972 if *precision <= DECIMAL32_MAX_PRECISION {
973 types.push(DataType::Decimal32(*precision, *scale));
974 }
975 if *precision <= DECIMAL64_MAX_PRECISION {
976 types.push(DataType::Decimal64(*precision, *scale));
977 }
978 if *precision <= DECIMAL128_MAX_PRECISION {
979 types.push(DataType::Decimal128(*precision, *scale));
980 }
981 types
982 }
983 NativeType::Timestamp(time_unit, timezone) => {
984 vec![DataType::Timestamp(*time_unit, timezone.to_owned())]
985 }
986 NativeType::Time(TimeUnit::Second) => vec![DataType::Time32(TimeUnit::Second)],
987 NativeType::Time(TimeUnit::Millisecond) => {
988 vec![DataType::Time32(TimeUnit::Millisecond)]
989 }
990 NativeType::Time(TimeUnit::Microsecond) => {
991 vec![DataType::Time64(TimeUnit::Microsecond)]
992 }
993 NativeType::Time(TimeUnit::Nanosecond) => {
994 vec![DataType::Time64(TimeUnit::Nanosecond)]
995 }
996 NativeType::Duration(time_unit) => vec![DataType::Duration(*time_unit)],
997 NativeType::Interval(interval_unit) => vec![DataType::Interval(*interval_unit)],
998 NativeType::FixedSizeBinary(size) => vec![DataType::FixedSizeBinary(*size)],
999 NativeType::FixedSizeList(logical_field, size) => {
1000 get_data_types(logical_field.logical_type.native())
1001 .iter()
1002 .map(|child_dt| {
1003 let field = Field::new(
1004 logical_field.name.clone(),
1005 child_dt.clone(),
1006 logical_field.nullable,
1007 );
1008 DataType::FixedSizeList(Arc::new(field), *size)
1009 })
1010 .collect()
1011 }
1012 NativeType::List(_)
1014 | NativeType::Struct(_)
1015 | NativeType::Union(_)
1016 | NativeType::Map(_) => {
1017 vec![]
1018 }
1019 }
1020}
1021
1022#[derive(Debug, Clone, Eq, PartialOrd)]
1047pub enum Coercion {
1048 Exact {
1050 desired_type: TypeSignatureClass,
1052 encoding_preservation: EncodingPreservation,
1054 },
1055
1056 Implicit {
1058 desired_type: TypeSignatureClass,
1060 implicit_coercion: ImplicitCoercion,
1062 encoding_preservation: EncodingPreservation,
1064 },
1065}
1066
1067#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)]
1070pub struct EncodingPreservation {
1071 preserve_dictionary: bool,
1072}
1073
1074impl EncodingPreservation {
1075 pub const fn dictionary() -> Self {
1077 Self {
1078 preserve_dictionary: true,
1079 }
1080 }
1081
1082 pub const fn with_dictionary(mut self) -> Self {
1084 self.preserve_dictionary = true;
1085 self
1086 }
1087
1088 pub const fn preserve_dictionary(self) -> bool {
1090 self.preserve_dictionary
1091 }
1092}
1093
1094impl Coercion {
1095 pub fn new_exact(desired_type: TypeSignatureClass) -> Self {
1096 Self::Exact {
1097 desired_type,
1098 encoding_preservation: EncodingPreservation::default(),
1099 }
1100 }
1101
1102 pub fn new_implicit(
1107 desired_type: TypeSignatureClass,
1108 allowed_source_types: Vec<TypeSignatureClass>,
1109 default_casted_type: NativeType,
1110 ) -> Self {
1111 Self::Implicit {
1112 desired_type,
1113 implicit_coercion: ImplicitCoercion {
1114 allowed_source_types,
1115 default_casted_type,
1116 },
1117 encoding_preservation: EncodingPreservation::default(),
1118 }
1119 }
1120
1121 pub fn with_encoding_preservation(
1122 mut self,
1123 encoding_preservation: EncodingPreservation,
1124 ) -> Self {
1125 match &mut self {
1126 Coercion::Exact {
1127 encoding_preservation: current,
1128 ..
1129 }
1130 | Coercion::Implicit {
1131 encoding_preservation: current,
1132 ..
1133 } => *current = encoding_preservation,
1134 }
1135 self
1136 }
1137
1138 pub fn encoding_preservation(&self) -> EncodingPreservation {
1139 match self {
1140 Coercion::Exact {
1141 encoding_preservation,
1142 ..
1143 }
1144 | Coercion::Implicit {
1145 encoding_preservation,
1146 ..
1147 } => *encoding_preservation,
1148 }
1149 }
1150
1151 pub fn allowed_source_types(&self) -> &[TypeSignatureClass] {
1152 match self {
1153 Coercion::Exact { .. } => &[],
1154 Coercion::Implicit {
1155 implicit_coercion, ..
1156 } => implicit_coercion.allowed_source_types.as_slice(),
1157 }
1158 }
1159
1160 pub fn default_casted_type(&self) -> Option<&NativeType> {
1161 match self {
1162 Coercion::Exact { .. } => None,
1163 Coercion::Implicit {
1164 implicit_coercion, ..
1165 } => Some(&implicit_coercion.default_casted_type),
1166 }
1167 }
1168
1169 pub fn desired_type(&self) -> &TypeSignatureClass {
1170 match self {
1171 Coercion::Exact { desired_type, .. } => desired_type,
1172 Coercion::Implicit { desired_type, .. } => desired_type,
1173 }
1174 }
1175
1176 pub fn implicit_coercion(&self) -> Option<&ImplicitCoercion> {
1177 match self {
1178 Coercion::Exact { .. } => None,
1179 Coercion::Implicit {
1180 implicit_coercion, ..
1181 } => Some(implicit_coercion),
1182 }
1183 }
1184}
1185
1186impl Display for Coercion {
1187 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1188 write!(f, "{}", self.desired_type())
1189 }
1190}
1191
1192impl PartialEq for Coercion {
1193 fn eq(&self, other: &Self) -> bool {
1194 self.desired_type() == other.desired_type()
1195 && self.implicit_coercion() == other.implicit_coercion()
1196 && self.encoding_preservation() == other.encoding_preservation()
1197 }
1198}
1199
1200impl Hash for Coercion {
1201 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1202 self.desired_type().hash(state);
1203 self.implicit_coercion().hash(state);
1204 self.encoding_preservation().hash(state);
1205 }
1206}
1207
1208#[derive(Debug, Clone, Eq, PartialOrd)]
1230pub struct ImplicitCoercion {
1231 allowed_source_types: Vec<TypeSignatureClass>,
1233
1234 default_casted_type: NativeType,
1238}
1239
1240impl Display for ImplicitCoercion {
1241 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1242 write!(f, "ImplicitCoercion(")?;
1243 for (i, source_type) in self.allowed_source_types.iter().enumerate() {
1244 if i > 0 {
1245 write!(f, ", ")?;
1246 }
1247 write!(f, "{source_type}")?;
1248 }
1249 write!(f, "; default={}", self.default_casted_type)
1250 }
1251}
1252
1253impl PartialEq for ImplicitCoercion {
1254 fn eq(&self, other: &Self) -> bool {
1255 self.allowed_source_types == other.allowed_source_types
1256 && self.default_casted_type == other.default_casted_type
1257 }
1258}
1259
1260impl Hash for ImplicitCoercion {
1261 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
1262 self.allowed_source_types.hash(state);
1263 self.default_casted_type.hash(state);
1264 }
1265}
1266
1267#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
1274pub struct Signature {
1275 pub type_signature: TypeSignature,
1277 pub volatility: Volatility,
1279 pub parameter_names: Option<Vec<String>>,
1286}
1287
1288impl Signature {
1289 pub fn new(type_signature: TypeSignature, volatility: Volatility) -> Self {
1291 Signature {
1292 type_signature,
1293 volatility,
1294 parameter_names: None,
1295 }
1296 }
1297 pub fn variadic(common_types: Vec<DataType>, volatility: Volatility) -> Self {
1299 Self {
1300 type_signature: TypeSignature::Variadic(common_types),
1301 volatility,
1302 parameter_names: None,
1303 }
1304 }
1305 pub fn user_defined(volatility: Volatility) -> Self {
1307 Self {
1308 type_signature: TypeSignature::UserDefined,
1309 volatility,
1310 parameter_names: None,
1311 }
1312 }
1313
1314 pub fn numeric(arg_count: usize, volatility: Volatility) -> Self {
1316 Self {
1317 type_signature: TypeSignature::Numeric(arg_count),
1318 volatility,
1319 parameter_names: None,
1320 }
1321 }
1322
1323 pub fn string(arg_count: usize, volatility: Volatility) -> Self {
1325 Self {
1326 type_signature: TypeSignature::String(arg_count),
1327 volatility,
1328 parameter_names: None,
1329 }
1330 }
1331
1332 pub fn variadic_any(volatility: Volatility) -> Self {
1334 Self {
1335 type_signature: TypeSignature::VariadicAny,
1336 volatility,
1337 parameter_names: None,
1338 }
1339 }
1340 pub fn uniform(
1342 arg_count: usize,
1343 valid_types: Vec<DataType>,
1344 volatility: Volatility,
1345 ) -> Self {
1346 Self {
1347 type_signature: TypeSignature::Uniform(arg_count, valid_types),
1348 volatility,
1349 parameter_names: None,
1350 }
1351 }
1352 pub fn exact(exact_types: Vec<DataType>, volatility: Volatility) -> Self {
1354 Signature {
1355 type_signature: TypeSignature::Exact(exact_types),
1356 volatility,
1357 parameter_names: None,
1358 }
1359 }
1360
1361 pub fn coercible(target_types: Vec<Coercion>, volatility: Volatility) -> Self {
1363 Self {
1364 type_signature: TypeSignature::Coercible(target_types),
1365 volatility,
1366 parameter_names: None,
1367 }
1368 }
1369
1370 pub fn comparable(arg_count: usize, volatility: Volatility) -> Self {
1372 Self {
1373 type_signature: TypeSignature::Comparable(arg_count),
1374 volatility,
1375 parameter_names: None,
1376 }
1377 }
1378
1379 pub fn nullary(volatility: Volatility) -> Self {
1380 Signature {
1381 type_signature: TypeSignature::Nullary,
1382 volatility,
1383 parameter_names: None,
1384 }
1385 }
1386
1387 pub fn any(arg_count: usize, volatility: Volatility) -> Self {
1389 Signature {
1390 type_signature: TypeSignature::Any(arg_count),
1391 volatility,
1392 parameter_names: None,
1393 }
1394 }
1395
1396 pub fn one_of(type_signatures: Vec<TypeSignature>, volatility: Volatility) -> Self {
1398 Signature {
1399 type_signature: TypeSignature::OneOf(type_signatures),
1400 volatility,
1401 parameter_names: None,
1402 }
1403 }
1404
1405 pub fn array_and_element(volatility: Volatility) -> Self {
1407 Signature {
1408 type_signature: TypeSignature::ArraySignature(
1409 ArrayFunctionSignature::Array {
1410 arguments: vec![
1411 ArrayFunctionArgument::Array,
1412 ArrayFunctionArgument::Element,
1413 ],
1414 array_coercion: Some(ListCoercion::FixedSizedListToList),
1415 },
1416 ),
1417 volatility,
1418 parameter_names: None,
1419 }
1420 }
1421
1422 pub fn element_and_array(volatility: Volatility) -> Self {
1424 Signature {
1425 type_signature: TypeSignature::ArraySignature(
1426 ArrayFunctionSignature::Array {
1427 arguments: vec![
1428 ArrayFunctionArgument::Element,
1429 ArrayFunctionArgument::Array,
1430 ],
1431 array_coercion: Some(ListCoercion::FixedSizedListToList),
1432 },
1433 ),
1434 volatility,
1435 parameter_names: None,
1436 }
1437 }
1438
1439 pub fn arrays(
1441 n: usize,
1442 coercion: Option<ListCoercion>,
1443 volatility: Volatility,
1444 ) -> Self {
1445 Signature {
1446 type_signature: TypeSignature::ArraySignature(
1447 ArrayFunctionSignature::Array {
1448 arguments: vec![ArrayFunctionArgument::Array; n],
1449 array_coercion: coercion,
1450 },
1451 ),
1452 volatility,
1453 parameter_names: None,
1454 }
1455 }
1456
1457 pub fn array_and_element_and_optional_index(volatility: Volatility) -> Self {
1459 Signature {
1460 type_signature: TypeSignature::OneOf(vec![
1461 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1462 arguments: vec![
1463 ArrayFunctionArgument::Array,
1464 ArrayFunctionArgument::Element,
1465 ],
1466 array_coercion: Some(ListCoercion::FixedSizedListToList),
1467 }),
1468 TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
1469 arguments: vec![
1470 ArrayFunctionArgument::Array,
1471 ArrayFunctionArgument::Element,
1472 ArrayFunctionArgument::Index,
1473 ],
1474 array_coercion: Some(ListCoercion::FixedSizedListToList),
1475 }),
1476 ]),
1477 volatility,
1478 parameter_names: None,
1479 }
1480 }
1481
1482 pub fn array_and_index(volatility: Volatility) -> Self {
1484 Signature {
1485 type_signature: TypeSignature::ArraySignature(
1486 ArrayFunctionSignature::Array {
1487 arguments: vec![
1488 ArrayFunctionArgument::Array,
1489 ArrayFunctionArgument::Index,
1490 ],
1491 array_coercion: Some(ListCoercion::FixedSizedListToList),
1492 },
1493 ),
1494 volatility,
1495 parameter_names: None,
1496 }
1497 }
1498
1499 pub fn array(volatility: Volatility) -> Self {
1501 Signature::arrays(1, Some(ListCoercion::FixedSizedListToList), volatility)
1502 }
1503
1504 pub fn with_parameter_names(mut self, names: Vec<impl Into<String>>) -> Result<Self> {
1520 let names = names.into_iter().map(Into::into).collect::<Vec<String>>();
1521 self.validate_parameter_names(&names)?;
1523 self.parameter_names = Some(names);
1524 Ok(self)
1525 }
1526
1527 fn validate_parameter_names(&self, names: &[String]) -> Result<()> {
1529 match self.type_signature.arity() {
1530 Arity::Fixed(expected) => {
1531 if names.len() != expected {
1532 return plan_err!(
1533 "Parameter names count ({}) does not match signature arity ({})",
1534 names.len(),
1535 expected
1536 );
1537 }
1538 }
1539 Arity::Variable => {
1540 if self.type_signature != TypeSignature::UserDefined {
1543 return plan_err!(
1544 "Cannot specify parameter names for variable arity signature: {:?}",
1545 self.type_signature
1546 );
1547 }
1548 }
1549 }
1550
1551 let mut seen = std::collections::HashSet::new();
1552 for name in names {
1553 if !seen.insert(name) {
1554 return plan_err!("Duplicate parameter name: '{}'", name);
1555 }
1556 }
1557
1558 Ok(())
1559 }
1560}
1561
1562#[cfg(test)]
1563mod tests {
1564 use datafusion_common::types::{
1565 NativeType, logical_float64, logical_int32, logical_int64, logical_string,
1566 };
1567
1568 use super::*;
1569 use crate::signature::{
1570 ArrayFunctionArgument, ArrayFunctionSignature, Coercion, TypeSignatureClass,
1571 };
1572
1573 #[test]
1574 fn supports_zero_argument_tests() {
1575 let positive_cases = vec![
1577 TypeSignature::Exact(vec![]),
1578 TypeSignature::OneOf(vec![
1579 TypeSignature::Exact(vec![DataType::Int8]),
1580 TypeSignature::Nullary,
1581 TypeSignature::Uniform(1, vec![DataType::Int8]),
1582 ]),
1583 TypeSignature::Nullary,
1584 ];
1585
1586 for case in positive_cases {
1587 assert!(
1588 case.supports_zero_argument(),
1589 "Expected {case:?} to support zero arguments"
1590 );
1591 }
1592
1593 let negative_cases = vec![
1595 TypeSignature::Exact(vec![DataType::Utf8]),
1596 TypeSignature::Uniform(1, vec![DataType::Float64]),
1597 TypeSignature::Any(1),
1598 TypeSignature::VariadicAny,
1599 TypeSignature::OneOf(vec![
1600 TypeSignature::Exact(vec![DataType::Int8]),
1601 TypeSignature::Uniform(1, vec![DataType::Int8]),
1602 ]),
1603 ];
1604
1605 for case in negative_cases {
1606 assert!(
1607 !case.supports_zero_argument(),
1608 "Expected {case:?} not to support zero arguments"
1609 );
1610 }
1611 }
1612
1613 #[test]
1614 fn type_signature_partial_ord() {
1615 assert!(TypeSignature::UserDefined < TypeSignature::VariadicAny);
1617 assert!(TypeSignature::UserDefined < TypeSignature::Any(1));
1618
1619 assert!(
1620 TypeSignature::Uniform(1, vec![DataType::Null])
1621 < TypeSignature::Uniform(1, vec![DataType::Boolean])
1622 );
1623 assert!(
1624 TypeSignature::Uniform(1, vec![DataType::Null])
1625 < TypeSignature::Uniform(2, vec![DataType::Null])
1626 );
1627 assert!(
1628 TypeSignature::Uniform(usize::MAX, vec![DataType::Null])
1629 < TypeSignature::Exact(vec![DataType::Null])
1630 );
1631 }
1632
1633 #[test]
1634 fn test_get_possible_types() {
1635 let type_signature = TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]);
1636 let possible_types = type_signature.get_example_types();
1637 assert_eq!(possible_types, vec![vec![DataType::Int32, DataType::Int64]]);
1638
1639 let type_signature = TypeSignature::OneOf(vec![
1640 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1641 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1642 ]);
1643 let possible_types = type_signature.get_example_types();
1644 assert_eq!(
1645 possible_types,
1646 vec![
1647 vec![DataType::Int32, DataType::Int64],
1648 vec![DataType::Float32, DataType::Float64]
1649 ]
1650 );
1651
1652 let type_signature = TypeSignature::OneOf(vec![
1653 TypeSignature::Exact(vec![DataType::Int32, DataType::Int64]),
1654 TypeSignature::Exact(vec![DataType::Float32, DataType::Float64]),
1655 TypeSignature::Exact(vec![DataType::Utf8]),
1656 ]);
1657 let possible_types = type_signature.get_example_types();
1658 assert_eq!(
1659 possible_types,
1660 vec![
1661 vec![DataType::Int32, DataType::Int64],
1662 vec![DataType::Float32, DataType::Float64],
1663 vec![DataType::Utf8]
1664 ]
1665 );
1666
1667 let type_signature =
1668 TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Int64]);
1669 let possible_types = type_signature.get_example_types();
1670 assert_eq!(
1671 possible_types,
1672 vec![
1673 vec![DataType::Float32, DataType::Float32],
1674 vec![DataType::Int64, DataType::Int64]
1675 ]
1676 );
1677
1678 let type_signature = TypeSignature::Coercible(vec![
1679 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
1680 Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
1681 ]);
1682 let possible_types = type_signature.get_example_types();
1683 assert_eq!(
1684 possible_types,
1685 vec![
1686 vec![DataType::Utf8, DataType::Int64],
1687 vec![DataType::LargeUtf8, DataType::Int64],
1688 vec![DataType::Utf8View, DataType::Int64]
1689 ]
1690 );
1691
1692 let type_signature =
1693 TypeSignature::Variadic(vec![DataType::Int32, DataType::Int64]);
1694 let possible_types = type_signature.get_example_types();
1695 assert_eq!(
1696 possible_types,
1697 vec![vec![DataType::Int32], vec![DataType::Int64]]
1698 );
1699
1700 let type_signature = TypeSignature::Numeric(2);
1701 let possible_types = type_signature.get_example_types();
1702 assert_eq!(
1703 possible_types,
1704 vec![
1705 vec![DataType::Int8, DataType::Int8],
1706 vec![DataType::Int16, DataType::Int16],
1707 vec![DataType::Int32, DataType::Int32],
1708 vec![DataType::Int64, DataType::Int64],
1709 vec![DataType::UInt8, DataType::UInt8],
1710 vec![DataType::UInt16, DataType::UInt16],
1711 vec![DataType::UInt32, DataType::UInt32],
1712 vec![DataType::UInt64, DataType::UInt64],
1713 vec![DataType::Float16, DataType::Float16],
1714 vec![DataType::Float32, DataType::Float32],
1715 vec![DataType::Float64, DataType::Float64]
1716 ]
1717 );
1718
1719 let type_signature = TypeSignature::String(2);
1720 let possible_types = type_signature.get_example_types();
1721 assert_eq!(
1722 possible_types,
1723 vec![
1724 vec![DataType::Utf8, DataType::Utf8],
1725 vec![DataType::LargeUtf8, DataType::LargeUtf8],
1726 vec![DataType::Utf8View, DataType::Utf8View]
1727 ]
1728 );
1729 }
1730
1731 #[test]
1732 fn test_signature_with_parameter_names() {
1733 let sig = Signature::exact(
1734 vec![DataType::Int32, DataType::Utf8],
1735 Volatility::Immutable,
1736 )
1737 .with_parameter_names(vec!["count".to_string(), "name".to_string()])
1738 .unwrap();
1739
1740 assert_eq!(
1741 sig.parameter_names,
1742 Some(vec!["count".to_string(), "name".to_string()])
1743 );
1744 assert_eq!(
1745 sig.type_signature,
1746 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8])
1747 );
1748 }
1749
1750 #[test]
1751 fn test_signature_parameter_names_wrong_count() {
1752 let result = Signature::exact(
1753 vec![DataType::Int32, DataType::Utf8],
1754 Volatility::Immutable,
1755 )
1756 .with_parameter_names(vec!["count".to_string()]); assert!(result.is_err());
1759 assert!(
1760 result
1761 .unwrap_err()
1762 .to_string()
1763 .contains("does not match signature arity")
1764 );
1765 }
1766
1767 #[test]
1768 fn test_signature_parameter_names_duplicate() {
1769 let result = Signature::exact(
1770 vec![DataType::Int32, DataType::Int32],
1771 Volatility::Immutable,
1772 )
1773 .with_parameter_names(vec!["count".to_string(), "count".to_string()]);
1774
1775 assert!(result.is_err());
1776 assert!(
1777 result
1778 .unwrap_err()
1779 .to_string()
1780 .contains("Duplicate parameter name")
1781 );
1782 }
1783
1784 #[test]
1785 fn test_signature_parameter_names_variadic() {
1786 let result = Signature::variadic(vec![DataType::Int32], Volatility::Immutable)
1787 .with_parameter_names(vec!["arg".to_string()]);
1788
1789 assert!(result.is_err());
1790 assert!(
1791 result
1792 .unwrap_err()
1793 .to_string()
1794 .contains("variable arity signature")
1795 );
1796 }
1797
1798 #[test]
1799 fn test_signature_without_parameter_names() {
1800 let sig = Signature::exact(
1801 vec![DataType::Int32, DataType::Utf8],
1802 Volatility::Immutable,
1803 );
1804
1805 assert_eq!(sig.parameter_names, None);
1806 }
1807
1808 #[test]
1809 fn test_signature_uniform_with_parameter_names() {
1810 let sig = Signature::uniform(3, vec![DataType::Float64], Volatility::Immutable)
1811 .with_parameter_names(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1812 .unwrap();
1813
1814 assert_eq!(
1815 sig.parameter_names,
1816 Some(vec!["x".to_string(), "y".to_string(), "z".to_string()])
1817 );
1818 }
1819
1820 #[test]
1821 fn test_signature_numeric_with_parameter_names() {
1822 let sig = Signature::numeric(2, Volatility::Immutable)
1823 .with_parameter_names(vec!["a".to_string(), "b".to_string()])
1824 .unwrap();
1825
1826 assert_eq!(
1827 sig.parameter_names,
1828 Some(vec!["a".to_string(), "b".to_string()])
1829 );
1830 }
1831
1832 #[test]
1833 fn test_signature_nullary_with_empty_names() {
1834 let sig = Signature::nullary(Volatility::Immutable)
1835 .with_parameter_names(Vec::<String>::new())
1836 .unwrap();
1837
1838 assert_eq!(sig.parameter_names, Some(vec![]));
1839 }
1840
1841 #[test]
1842 fn test_to_string_repr_with_names_exact() {
1843 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1844
1845 assert_eq!(sig.to_string_repr_with_names(None), vec!["Int32, Utf8"]);
1846
1847 let names = vec!["id".to_string(), "name".to_string()];
1848 assert_eq!(
1849 sig.to_string_repr_with_names(Some(&names)),
1850 vec!["id: Int32, name: Utf8"]
1851 );
1852 }
1853
1854 #[test]
1855 fn test_to_string_repr_with_names_any() {
1856 let sig = TypeSignature::Any(3);
1857
1858 assert_eq!(sig.to_string_repr_with_names(None), vec!["Any, Any, Any"]);
1859
1860 let names = vec!["x".to_string(), "y".to_string(), "z".to_string()];
1861 assert_eq!(
1862 sig.to_string_repr_with_names(Some(&names)),
1863 vec!["x: Any, y: Any, z: Any"]
1864 );
1865 }
1866
1867 #[test]
1868 fn test_to_string_repr_with_names_one_of() {
1869 let sig =
1870 TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]);
1871
1872 assert_eq!(
1873 sig.to_string_repr_with_names(None),
1874 vec!["Any, Any", "Any, Any, Any"]
1875 );
1876
1877 let names = vec![
1878 "str".to_string(),
1879 "start_pos".to_string(),
1880 "length".to_string(),
1881 ];
1882 assert_eq!(
1883 sig.to_string_repr_with_names(Some(&names)),
1884 vec![
1885 "str: Any, start_pos: Any",
1886 "str: Any, start_pos: Any, length: Any"
1887 ]
1888 );
1889 }
1890
1891 #[test]
1892 fn test_to_string_repr_with_names_partial() {
1893 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
1895
1896 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1898 assert_eq!(
1899 sig.to_string_repr_with_names(Some(&names)),
1900 vec!["a: Int32, b: Utf8"]
1901 );
1902 }
1903
1904 #[test]
1905 fn test_to_string_repr_with_names_uniform() {
1906 let sig = TypeSignature::Uniform(2, vec![DataType::Float64]);
1907
1908 assert_eq!(
1909 sig.to_string_repr_with_names(None),
1910 vec!["Float64, Float64"]
1911 );
1912
1913 let names = vec!["x".to_string(), "y".to_string()];
1914 assert_eq!(
1915 sig.to_string_repr_with_names(Some(&names)),
1916 vec!["x: Float64, y: Float64"]
1917 );
1918 }
1919
1920 #[test]
1921 fn test_to_string_repr_with_names_coercible() {
1922 let sig = TypeSignature::Coercible(vec![
1923 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1924 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
1925 ]);
1926
1927 let names = vec!["a".to_string(), "b".to_string()];
1928 let result = sig.to_string_repr_with_names(Some(&names));
1929 assert_eq!(result.len(), 1);
1931 assert!(result[0].starts_with("a: "));
1932 assert!(result[0].contains(", b: "));
1933 }
1934
1935 #[test]
1936 fn test_to_string_repr_with_names_comparable_numeric_string() {
1937 let comparable = TypeSignature::Comparable(3);
1938 let numeric = TypeSignature::Numeric(2);
1939 let string_sig = TypeSignature::String(2);
1940
1941 let names = vec!["a".to_string(), "b".to_string(), "c".to_string()];
1942
1943 assert_eq!(
1945 comparable.to_string_repr_with_names(Some(&names)),
1946 vec!["a: Comparable, b: Comparable, c: Comparable"]
1947 );
1948 assert_eq!(
1949 numeric.to_string_repr_with_names(Some(&names)),
1950 vec!["a: Numeric, b: Numeric"]
1951 );
1952 assert_eq!(
1953 string_sig.to_string_repr_with_names(Some(&names)),
1954 vec!["a: String, b: String"]
1955 );
1956 }
1957
1958 #[test]
1959 fn test_to_string_repr_with_names_variadic_fallback() {
1960 let variadic = TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]);
1961 let names = vec!["x".to_string()];
1962 assert_eq!(
1963 variadic.to_string_repr_with_names(Some(&names)),
1964 variadic.to_string_repr()
1965 );
1966
1967 let variadic_any = TypeSignature::VariadicAny;
1968 assert_eq!(
1969 variadic_any.to_string_repr_with_names(Some(&names)),
1970 variadic_any.to_string_repr()
1971 );
1972
1973 let user_defined = TypeSignature::UserDefined;
1975 assert_eq!(
1976 user_defined.to_string_repr_with_names(Some(&names)),
1977 vec!["x"]
1978 );
1979 assert_eq!(
1980 user_defined.to_string_repr_with_names(None),
1981 user_defined.to_string_repr()
1982 );
1983 }
1984
1985 #[test]
1986 fn test_to_string_repr_with_names_nullary() {
1987 let sig = TypeSignature::Nullary;
1988 let names = vec!["x".to_string()];
1989
1990 assert_eq!(
1992 sig.to_string_repr_with_names(Some(&names)),
1993 vec!["NullAry()"]
1994 );
1995 assert_eq!(sig.to_string_repr_with_names(None), vec!["NullAry()"]);
1996 }
1997
1998 #[test]
1999 fn test_to_string_repr_with_names_array_signature() {
2000 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
2001 arguments: vec![
2002 ArrayFunctionArgument::Array,
2003 ArrayFunctionArgument::Index,
2004 ArrayFunctionArgument::Element,
2005 ],
2006 array_coercion: None,
2007 });
2008
2009 assert_eq!(
2010 sig.to_string_repr_with_names(None),
2011 vec!["array, index, element"]
2012 );
2013
2014 let names = vec!["arr".to_string(), "idx".to_string(), "val".to_string()];
2015 assert_eq!(
2016 sig.to_string_repr_with_names(Some(&names)),
2017 vec!["arr: array, idx: index, val: element"]
2018 );
2019
2020 let recursive =
2021 TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
2022 let names = vec!["array".to_string()];
2023 assert_eq!(
2024 recursive.to_string_repr_with_names(Some(&names)),
2025 vec!["array: recursive_array"]
2026 );
2027
2028 let map_array = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
2030 let names = vec!["map".to_string()];
2031 assert_eq!(
2032 map_array.to_string_repr_with_names(Some(&names)),
2033 vec!["map: map_array"]
2034 );
2035 }
2036
2037 #[test]
2038 fn test_type_signature_arity_exact() {
2039 let sig = TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]);
2040 assert_eq!(sig.arity(), Arity::Fixed(2));
2041
2042 let sig = TypeSignature::Exact(vec![]);
2043 assert_eq!(sig.arity(), Arity::Fixed(0));
2044 }
2045
2046 #[test]
2047 fn test_type_signature_arity_uniform() {
2048 let sig = TypeSignature::Uniform(3, vec![DataType::Float64]);
2049 assert_eq!(sig.arity(), Arity::Fixed(3));
2050
2051 let sig = TypeSignature::Uniform(1, vec![DataType::Int32]);
2052 assert_eq!(sig.arity(), Arity::Fixed(1));
2053 }
2054
2055 #[test]
2056 fn test_type_signature_arity_numeric() {
2057 let sig = TypeSignature::Numeric(2);
2058 assert_eq!(sig.arity(), Arity::Fixed(2));
2059 }
2060
2061 #[test]
2062 fn test_type_signature_arity_string() {
2063 let sig = TypeSignature::String(3);
2064 assert_eq!(sig.arity(), Arity::Fixed(3));
2065 }
2066
2067 #[test]
2068 fn test_type_signature_arity_comparable() {
2069 let sig = TypeSignature::Comparable(2);
2070 assert_eq!(sig.arity(), Arity::Fixed(2));
2071 }
2072
2073 #[test]
2074 fn test_type_signature_arity_any() {
2075 let sig = TypeSignature::Any(4);
2076 assert_eq!(sig.arity(), Arity::Fixed(4));
2077 }
2078
2079 #[test]
2080 fn test_type_signature_arity_coercible() {
2081 let sig = TypeSignature::Coercible(vec![
2082 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
2083 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
2084 ]);
2085 assert_eq!(sig.arity(), Arity::Fixed(2));
2086 }
2087
2088 #[test]
2089 fn test_type_signature_arity_nullary() {
2090 let sig = TypeSignature::Nullary;
2091 assert_eq!(sig.arity(), Arity::Fixed(0));
2092 }
2093
2094 #[test]
2095 fn test_type_signature_arity_array_signature() {
2096 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
2098 arguments: vec![ArrayFunctionArgument::Array, ArrayFunctionArgument::Index],
2099 array_coercion: None,
2100 });
2101 assert_eq!(sig.arity(), Arity::Fixed(2));
2102
2103 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
2105 arguments: vec![
2106 ArrayFunctionArgument::Array,
2107 ArrayFunctionArgument::Element,
2108 ArrayFunctionArgument::Index,
2109 ],
2110 array_coercion: None,
2111 });
2112 assert_eq!(sig.arity(), Arity::Fixed(3));
2113
2114 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::RecursiveArray);
2116 assert_eq!(sig.arity(), Arity::Fixed(1));
2117
2118 let sig = TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray);
2120 assert_eq!(sig.arity(), Arity::Fixed(1));
2121 }
2122
2123 #[test]
2124 fn test_type_signature_arity_one_of_fixed() {
2125 let sig = TypeSignature::OneOf(vec![
2127 TypeSignature::Exact(vec![DataType::Int32]),
2128 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2129 TypeSignature::Exact(vec![
2130 DataType::Int32,
2131 DataType::Utf8,
2132 DataType::Float64,
2133 ]),
2134 ]);
2135 assert_eq!(sig.arity(), Arity::Fixed(3));
2136 }
2137
2138 #[test]
2139 fn test_type_signature_arity_one_of_variable() {
2140 let sig = TypeSignature::OneOf(vec![
2142 TypeSignature::Exact(vec![DataType::Int32]),
2143 TypeSignature::VariadicAny,
2144 ]);
2145 assert_eq!(sig.arity(), Arity::Variable);
2146 }
2147
2148 #[test]
2149 fn test_type_signature_arity_variadic() {
2150 let sig = TypeSignature::Variadic(vec![DataType::Int32]);
2151 assert_eq!(sig.arity(), Arity::Variable);
2152
2153 let sig = TypeSignature::VariadicAny;
2154 assert_eq!(sig.arity(), Arity::Variable);
2155 }
2156
2157 #[test]
2158 fn test_type_signature_arity_user_defined() {
2159 let sig = TypeSignature::UserDefined;
2160 assert_eq!(sig.arity(), Arity::Variable);
2161 }
2162
2163 #[test]
2164 fn test_type_signature_display() {
2165 use insta::assert_snapshot;
2166
2167 assert_snapshot!(TypeSignature::Nullary, @"Nullary");
2168 assert_snapshot!(TypeSignature::Any(2), @"Any(2)");
2169 assert_snapshot!(TypeSignature::Numeric(3), @"Numeric(3)");
2170 assert_snapshot!(TypeSignature::String(1), @"String(1)");
2171 assert_snapshot!(TypeSignature::Comparable(2), @"Comparable(2)");
2172 assert_snapshot!(TypeSignature::VariadicAny, @"VariadicAny");
2173 assert_snapshot!(TypeSignature::UserDefined, @"UserDefined");
2174
2175 assert_snapshot!(
2176 TypeSignature::Exact(vec![DataType::Int32, DataType::Utf8]),
2177 @"Exact(Int32, Utf8)"
2178 );
2179 assert_snapshot!(
2180 TypeSignature::Variadic(vec![DataType::Utf8, DataType::LargeUtf8]),
2181 @"Variadic(Utf8, LargeUtf8)"
2182 );
2183 assert_snapshot!(
2184 TypeSignature::Uniform(2, vec![DataType::Float32, DataType::Float64]),
2185 @"Uniform(2, [Float32, Float64])"
2186 );
2187
2188 assert_snapshot!(
2189 TypeSignature::Coercible(vec![
2190 Coercion::new_exact(TypeSignatureClass::Native(logical_float64())),
2191 Coercion::new_exact(TypeSignatureClass::Native(logical_int32())),
2192 ]),
2193 @"Coercible(Float64, Int32)"
2194 );
2195
2196 assert_snapshot!(
2197 TypeSignature::OneOf(vec![
2198 TypeSignature::Nullary,
2199 TypeSignature::VariadicAny,
2200 ]),
2201 @"OneOf(Nullary, VariadicAny)"
2202 );
2203 }
2204
2205 #[test]
2206 fn test_type_signature_class_display() {
2207 use insta::assert_snapshot;
2208
2209 assert_snapshot!(TypeSignatureClass::Any, @"Any");
2210 assert_snapshot!(TypeSignatureClass::Numeric, @"Numeric");
2211 assert_snapshot!(TypeSignatureClass::Integer, @"Integer");
2212 assert_snapshot!(TypeSignatureClass::Float, @"Float");
2213 assert_snapshot!(TypeSignatureClass::Decimal, @"Decimal");
2214 assert_snapshot!(TypeSignatureClass::Timestamp, @"Timestamp");
2215 assert_snapshot!(TypeSignatureClass::Time, @"Time");
2216 assert_snapshot!(TypeSignatureClass::Interval, @"Interval");
2217 assert_snapshot!(TypeSignatureClass::Duration, @"Duration");
2218 assert_snapshot!(TypeSignatureClass::Binary, @"Binary");
2219 assert_snapshot!(TypeSignatureClass::Native(logical_int32()), @"Int32");
2220 assert_snapshot!(TypeSignatureClass::Native(logical_string()), @"String");
2221 }
2222
2223 #[test]
2224 fn test_coercion_display() {
2225 use insta::assert_snapshot;
2226
2227 let exact_int = Coercion::new_exact(TypeSignatureClass::Native(logical_int32()));
2228 assert_snapshot!(exact_int, @"Int32");
2229
2230 let exact_numeric = Coercion::new_exact(TypeSignatureClass::Numeric);
2231 assert_snapshot!(exact_numeric, @"Numeric");
2232
2233 let implicit = Coercion::new_implicit(
2234 TypeSignatureClass::Native(logical_float64()),
2235 vec![TypeSignatureClass::Numeric],
2236 NativeType::Float64,
2237 );
2238 assert_snapshot!(implicit, @"Float64");
2239
2240 let implicit_with_multiple_sources = Coercion::new_implicit(
2241 TypeSignatureClass::Native(logical_int64()),
2242 vec![TypeSignatureClass::Integer, TypeSignatureClass::Numeric],
2243 NativeType::Int64,
2244 );
2245 assert_snapshot!(implicit_with_multiple_sources, @"Int64");
2246 }
2247
2248 #[test]
2249 fn test_coercion_encoding_preservation_affects_equality() {
2250 assert!(!EncodingPreservation::default().preserve_dictionary());
2251 let preserve_dictionary = EncodingPreservation::dictionary();
2252 assert!(preserve_dictionary.preserve_dictionary());
2253
2254 let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string()));
2255 let preserving = default
2256 .clone()
2257 .with_encoding_preservation(preserve_dictionary);
2258
2259 assert_ne!(default, preserving);
2260 }
2261
2262 #[test]
2263 fn test_to_string_repr_coercible() {
2264 use insta::assert_snapshot;
2265
2266 let sig = TypeSignature::Coercible(vec![
2268 Coercion::new_implicit(
2269 TypeSignatureClass::Native(logical_float64()),
2270 vec![TypeSignatureClass::Numeric],
2271 NativeType::Float64,
2272 ),
2273 Coercion::new_implicit(
2274 TypeSignatureClass::Native(logical_int64()),
2275 vec![TypeSignatureClass::Integer],
2276 NativeType::Int64,
2277 ),
2278 ]);
2279 let repr = sig.to_string_repr();
2280 assert_eq!(repr.len(), 1);
2281 assert_snapshot!(repr[0], @"Float64, Int64");
2282 }
2283
2284 #[test]
2285 fn test_to_string_repr_coercible_exact() {
2286 use insta::assert_snapshot;
2287
2288 let sig = TypeSignature::Coercible(vec![
2289 Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
2290 Coercion::new_exact(TypeSignatureClass::Native(logical_int64())),
2291 ]);
2292 let repr = sig.to_string_repr();
2293 assert_eq!(repr.len(), 1);
2294 assert_snapshot!(repr[0], @"String, Int64");
2295 }
2296}