1use super::binary::binary_numeric_coercion;
19use crate::{
20 AggregateUDF, HigherOrderTypeSignature, HigherOrderUDF, ScalarUDF, Signature,
21 TypeSignature, ValueOrLambda, WindowUDF,
22};
23use arrow::datatypes::{Field, FieldRef};
24use arrow::{
25 compute::can_cast_types,
26 datatypes::{DataType, TimeUnit},
27};
28use datafusion_common::internal_datafusion_err;
29use datafusion_common::types::LogicalType;
30use datafusion_common::utils::{
31 ListCoercion, base_type, coerced_fixed_size_list_to_list,
32};
33use datafusion_common::{
34 Result, exec_err, internal_err, plan_err, types::NativeType, utils::list_ndims,
35};
36use datafusion_expr_common::signature::{
37 ArrayFunctionArgument, EncodingPreservation, TypeSignatureClass,
38};
39use datafusion_expr_common::type_coercion::binary::type_union_resolution;
40use datafusion_expr_common::{
41 signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD},
42 type_coercion::binary::comparison_coercion,
43 type_coercion::binary::string_coercion,
44};
45use itertools::Itertools as _;
46use std::sync::Arc;
47
48pub trait UDFCoercionExt {
51 fn name(&self) -> &str;
53 fn signature(&self) -> &Signature;
56 fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>>;
59}
60
61impl UDFCoercionExt for ScalarUDF {
62 fn name(&self) -> &str {
63 self.name()
64 }
65
66 fn signature(&self) -> &Signature {
67 self.signature()
68 }
69
70 fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
71 self.coerce_types(arg_types)
72 }
73}
74
75impl UDFCoercionExt for AggregateUDF {
76 fn name(&self) -> &str {
77 self.name()
78 }
79
80 fn signature(&self) -> &Signature {
81 self.signature()
82 }
83
84 fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
85 self.coerce_types(arg_types)
86 }
87}
88
89impl UDFCoercionExt for WindowUDF {
90 fn name(&self) -> &str {
91 self.name()
92 }
93
94 fn signature(&self) -> &Signature {
95 self.signature()
96 }
97
98 fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
99 self.coerce_types(arg_types)
100 }
101}
102
103pub fn fields_with_udf<F: UDFCoercionExt>(
111 current_fields: &[FieldRef],
112 func: &F,
113) -> Result<Vec<FieldRef>> {
114 let signature = func.signature();
115 let type_signature = &signature.type_signature;
116
117 if current_fields.is_empty() && type_signature != &TypeSignature::UserDefined {
118 if type_signature.supports_zero_argument() {
119 return Ok(vec![]);
120 } else if type_signature.used_to_support_zero_arguments() {
121 return plan_err!(
123 "'{}' does not support zero arguments. Use TypeSignature::Nullary for zero arguments",
124 func.name()
125 );
126 } else {
127 return plan_err!("'{}' does not support zero arguments", func.name());
128 }
129 }
130 let current_types = current_fields
131 .iter()
132 .map(|f| f.data_type())
133 .cloned()
134 .collect::<Vec<_>>();
135
136 let valid_types = get_valid_types_with_udf(type_signature, ¤t_types, func)?;
137 if valid_types
138 .iter()
139 .any(|data_type| data_type == ¤t_types)
140 {
141 return Ok(current_fields.to_vec());
142 }
143
144 let updated_types =
145 try_coerce_types(func.name(), valid_types, ¤t_types, type_signature)?;
146
147 Ok(current_fields
148 .iter()
149 .zip(updated_types)
150 .map(|(current_field, new_type)| {
151 current_field.as_ref().clone().with_data_type(new_type)
152 })
153 .map(Arc::new)
154 .collect())
155}
156
157pub fn value_fields_with_higher_order_udf<L: Clone>(
170 current_fields: &[ValueOrLambda<FieldRef, L>],
171 func: &HigherOrderUDF,
172) -> Result<Vec<ValueOrLambda<FieldRef, L>>> {
173 match func.signature().type_signature {
174 HigherOrderTypeSignature::UserDefined => {
175 let arg_types = current_fields
176 .iter()
177 .filter_map(|p| match p {
178 ValueOrLambda::Value(field) => Some(field.data_type().clone()),
179 ValueOrLambda::Lambda(_) => None,
180 })
181 .collect::<Vec<_>>();
182
183 let coerced_types = func.coerce_value_types(&arg_types)?;
184
185 if coerced_types.len() != arg_types.len() {
186 return plan_err!(
187 "{} coerce_value_types should have returned {} items but returned {}",
188 func.name(),
189 arg_types.len(),
190 coerced_types.len()
191 );
192 }
193
194 let mut coerced_types = coerced_types.into_iter();
202
203 current_fields
204 .iter()
205 .map(|current_field| match current_field {
206 ValueOrLambda::Value(field) => {
207 let data_type = coerced_types.next().ok_or_else(|| {
208 internal_datafusion_err!(
209 "coerced_types len should have been checked above"
210 )
211 })?;
212
213 Ok(ValueOrLambda::Value(Arc::new(
214 field.as_ref().clone().with_data_type(data_type),
215 )))
216 }
217 ValueOrLambda::Lambda(lambda) => {
218 Ok(ValueOrLambda::Lambda(lambda.clone()))
219 }
220 })
221 .collect()
222 }
223 HigherOrderTypeSignature::VariadicAny => Ok(current_fields.to_vec()),
224 HigherOrderTypeSignature::Any(number) => {
225 if current_fields.len() != number {
226 return plan_err!(
227 "The function '{}' expected {number} arguments but received {}",
228 func.name(),
229 current_fields.len()
230 );
231 }
232
233 Ok(current_fields.to_vec())
234 }
235 HigherOrderTypeSignature::Exact(ref expected) => {
236 if current_fields.len() != expected.len() {
237 let name = func.name();
238 let expected_len = expected.len();
239 let actual_len = current_fields.len();
240 return plan_err!(
241 "The function '{name}' expected {expected_len} argument(s) but received {actual_len}"
242 );
243 }
244
245 for (i, (actual, expected)) in
246 current_fields.iter().zip(expected.iter()).enumerate()
247 {
248 match (actual, expected) {
249 (ValueOrLambda::Value(_), ValueOrLambda::Value(_)) => {}
250 (ValueOrLambda::Lambda(_), ValueOrLambda::Lambda(_)) => {}
251 (ValueOrLambda::Value(_), ValueOrLambda::Lambda(_)) => {
252 let name = func.name();
253 return plan_err!(
254 "The function '{name}' expected a lambda at position {i} but received a value"
255 );
256 }
257 (ValueOrLambda::Lambda(_), ValueOrLambda::Value(_)) => {
258 let name = func.name();
259 return plan_err!(
260 "The function '{name}' expected a value at position {i} but received a lambda"
261 );
262 }
263 }
264 }
265
266 let arg_types = current_fields
267 .iter()
268 .filter_map(|p| match p {
269 ValueOrLambda::Value(field) => Some(field.data_type().clone()),
270 ValueOrLambda::Lambda(_) => None,
271 })
272 .collect::<Vec<_>>();
273
274 let coerced_types = func.coerce_value_types(&arg_types)?;
275
276 if coerced_types.len() != arg_types.len() {
277 return plan_err!(
278 "{} coerce_value_types should have returned {} items but returned {}",
279 func.name(),
280 arg_types.len(),
281 coerced_types.len()
282 );
283 }
284
285 let mut coerced_types = coerced_types.into_iter();
286
287 current_fields
288 .iter()
289 .map(|current_field| match current_field {
290 ValueOrLambda::Value(field) => {
291 let data_type = coerced_types.next().ok_or_else(|| {
292 internal_datafusion_err!(
293 "coerced_types len should have been checked above"
294 )
295 })?;
296
297 Ok(ValueOrLambda::Value(Arc::new(
298 field.as_ref().clone().with_data_type(data_type),
299 )))
300 }
301 ValueOrLambda::Lambda(lambda) => {
302 Ok(ValueOrLambda::Lambda(lambda.clone()))
303 }
304 })
305 .collect()
306 }
307 }
308}
309
310pub fn value_fields_with_higher_order_udf_and_lambdas(
323 current_fields: &[ValueOrLambda<FieldRef, FieldRef>],
324 func: &HigherOrderUDF,
325) -> Result<Vec<ValueOrLambda<FieldRef, FieldRef>>> {
326 let mut new_fields = value_fields_with_higher_order_udf(current_fields, func)?;
327
328 let new_types = new_fields
329 .iter()
330 .map(|f| match f {
331 ValueOrLambda::Value(f) => ValueOrLambda::Value(f.data_type().clone()),
332 ValueOrLambda::Lambda(f) => ValueOrLambda::Lambda(f.data_type().clone()),
333 })
334 .collect::<Vec<_>>();
335
336 if let Some(new_value_types) = func.coerce_values_for_lambdas(&new_types)? {
337 let mut new_value_types = new_value_types.into_iter();
338
339 let value_types_count = new_types
340 .iter()
341 .filter(|e| matches!(e, ValueOrLambda::Value(_)))
342 .count();
343
344 if new_value_types.len() != value_types_count {
345 return plan_err!(
346 "{} coerce_values_for_lambdas returned {} values but {value_types_count} expected",
347 func.name(),
348 new_value_types.len()
349 );
350 }
351
352 for new_field in &mut new_fields {
353 match new_field {
354 ValueOrLambda::Value(value) => {
355 let coerce_to = new_value_types.next().ok_or_else(|| {
356 internal_datafusion_err!(
357 "new_value_types len should have been checked above"
358 )
359 })?;
360
361 if value.data_type() != &coerce_to {
362 Arc::make_mut(value).set_data_type(coerce_to);
363 }
364 }
365 ValueOrLambda::Lambda(_) => {}
366 }
367 }
368 };
369
370 Ok(new_fields)
371}
372
373#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
381pub fn data_types_with_scalar_udf(
382 current_types: &[DataType],
383 func: &ScalarUDF,
384) -> Result<Vec<DataType>> {
385 let current_fields = current_types
386 .iter()
387 .map(|dt| Arc::new(Field::new("f", dt.clone(), true)))
388 .collect::<Vec<_>>();
389 Ok(fields_with_udf(¤t_fields, func)?
390 .iter()
391 .map(|f| f.data_type().clone())
392 .collect())
393}
394
395#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
403pub fn fields_with_aggregate_udf(
404 current_fields: &[FieldRef],
405 func: &AggregateUDF,
406) -> Result<Vec<FieldRef>> {
407 fields_with_udf(current_fields, func)
408}
409
410#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
418pub fn fields_with_window_udf(
419 current_fields: &[FieldRef],
420 func: &WindowUDF,
421) -> Result<Vec<FieldRef>> {
422 fields_with_udf(current_fields, func)
423}
424
425#[deprecated(since = "52.0.0", note = "use fields_with_udf")]
433pub fn data_types(
434 function_name: impl AsRef<str>,
435 current_types: &[DataType],
436 signature: &Signature,
437) -> Result<Vec<DataType>> {
438 let type_signature = &signature.type_signature;
439
440 if current_types.is_empty() && type_signature != &TypeSignature::UserDefined {
441 if type_signature.supports_zero_argument() {
442 return Ok(vec![]);
443 } else if type_signature.used_to_support_zero_arguments() {
444 return plan_err!(
446 "function '{}' has signature {type_signature} which does not support zero arguments. Use TypeSignature::Nullary for zero arguments",
447 function_name.as_ref()
448 );
449 } else {
450 return plan_err!(
451 "Function '{}' has signature {type_signature} which does not support zero arguments",
452 function_name.as_ref()
453 );
454 }
455 }
456
457 let valid_types =
458 get_valid_types(function_name.as_ref(), type_signature, current_types)?;
459 if valid_types
460 .iter()
461 .any(|data_type| data_type == current_types)
462 {
463 return Ok(current_types.to_vec());
464 }
465
466 try_coerce_types(
467 function_name.as_ref(),
468 valid_types,
469 current_types,
470 type_signature,
471 )
472}
473
474fn is_well_supported_signature(type_signature: &TypeSignature) -> bool {
475 match type_signature {
476 TypeSignature::OneOf(type_signatures) => {
477 type_signatures.iter().all(is_well_supported_signature)
478 }
479 TypeSignature::UserDefined
480 | TypeSignature::Numeric(_)
481 | TypeSignature::String(_)
482 | TypeSignature::Coercible(_)
483 | TypeSignature::Any(_)
484 | TypeSignature::Nullary
485 | TypeSignature::Comparable(_) => true,
486 TypeSignature::Variadic(_)
487 | TypeSignature::VariadicAny
488 | TypeSignature::Uniform(_, _)
489 | TypeSignature::Exact(_)
490 | TypeSignature::ArraySignature(_) => false,
491 }
492}
493
494fn try_coerce_types(
495 function_name: &str,
496 valid_types: Vec<Vec<DataType>>,
497 current_types: &[DataType],
498 type_signature: &TypeSignature,
499) -> Result<Vec<DataType>> {
500 let mut valid_types = valid_types;
501
502 if !valid_types.is_empty() && is_well_supported_signature(type_signature) {
504 if !type_signature.is_one_of() {
507 assert_eq!(valid_types.len(), 1);
508 }
509
510 let valid_types = valid_types.swap_remove(0);
511 if let Some(t) = maybe_data_types_without_coercion(&valid_types, current_types) {
512 return Ok(t);
513 }
514 } else {
515 for valid_types in valid_types {
519 if let Some(types) = maybe_data_types(&valid_types, current_types) {
520 return Ok(types);
521 }
522 }
523 }
524
525 plan_err!(
527 "Failed to coerce arguments to satisfy a call to '{function_name}' function: coercion from {} to the signature {type_signature} failed",
528 current_types.iter().join(", ")
529 )
530}
531
532fn get_valid_types_with_udf<F: UDFCoercionExt>(
533 signature: &TypeSignature,
534 current_types: &[DataType],
535 func: &F,
536) -> Result<Vec<Vec<DataType>>> {
537 let valid_types = match signature {
538 TypeSignature::UserDefined => match func.coerce_types(current_types) {
539 Ok(coerced_types) => vec![coerced_types],
540 Err(e) => {
541 return exec_err!(
542 "Function '{}' user-defined coercion failed with: {}",
543 func.name(),
544 e.strip_backtrace()
545 );
546 }
547 },
548 TypeSignature::OneOf(signatures) => {
549 let mut res = vec![];
550 let mut errors = vec![];
551 for sig in signatures {
552 match get_valid_types_with_udf(sig, current_types, func) {
553 Ok(valid_types) => {
554 res.extend(valid_types);
555 }
556 Err(e) => {
557 errors.push(e.to_string());
558 }
559 }
560 }
561
562 if res.is_empty() {
564 return internal_err!(
565 "Function '{}' failed to match any signature, errors: {}",
566 func.name(),
567 errors.join(",")
568 );
569 } else {
570 res
571 }
572 }
573 _ => get_valid_types(func.name(), signature, current_types)?,
574 };
575
576 Ok(valid_types)
577}
578
579fn get_valid_types(
581 function_name: &str,
582 signature: &TypeSignature,
583 current_types: &[DataType],
584) -> Result<Vec<Vec<DataType>>> {
585 fn array_valid_types(
586 function_name: &str,
587 current_types: &[DataType],
588 arguments: &[ArrayFunctionArgument],
589 array_coercion: Option<&ListCoercion>,
590 ) -> Result<Vec<Vec<DataType>>> {
591 fn rebuild_array_type(
592 current_type: &DataType,
593 element_type: &DataType,
594 nullable: bool,
595 large_list: bool,
596 fixed_size: Option<i32>,
597 ) -> DataType {
598 let field = match current_type {
601 DataType::List(field)
602 | DataType::LargeList(field)
603 | DataType::FixedSizeList(field, _) => Some(Arc::new(
604 field
605 .as_ref()
606 .clone()
607 .with_data_type(element_type.clone())
608 .with_nullable(nullable),
609 )),
610 _ => None,
611 };
612
613 if large_list {
614 field.map_or_else(
615 || DataType::new_large_list(element_type.clone(), nullable),
616 DataType::LargeList,
617 )
618 } else if let Some(size) = fixed_size {
619 field.map_or_else(
620 || {
621 DataType::new_fixed_size_list(
622 element_type.clone(),
623 size,
624 nullable,
625 )
626 },
627 |field| DataType::FixedSizeList(field, size),
628 )
629 } else {
630 field.map_or_else(
631 || DataType::new_list(element_type.clone(), nullable),
632 DataType::List,
633 )
634 }
635 }
636
637 if current_types.len() != arguments.len() {
638 return Ok(vec![vec![]]);
639 }
640
641 let mut large_list = false;
642 let mut fixed_size = array_coercion != Some(&ListCoercion::FixedSizedListToList);
643 let mut list_sizes = Vec::with_capacity(arguments.len());
644 let mut element_types = Vec::with_capacity(arguments.len());
645 let mut nested_item_nullability = Vec::with_capacity(arguments.len());
646 for (argument, current_type) in arguments.iter().zip(current_types.iter()) {
647 match argument {
648 ArrayFunctionArgument::Index | ArrayFunctionArgument::String => {
649 nested_item_nullability.push(None);
650 }
651 ArrayFunctionArgument::Element => {
652 element_types.push(current_type.clone());
653 nested_item_nullability.push(None);
654 }
655 ArrayFunctionArgument::Array => match current_type {
656 DataType::Null => {
657 element_types.push(DataType::Null);
658 nested_item_nullability.push(None);
659 }
660 DataType::List(field) | DataType::ListView(field) => {
661 element_types.push(field.data_type().clone());
662 nested_item_nullability.push(Some(field.is_nullable()));
663 fixed_size = false;
664 }
665 DataType::LargeList(field) | DataType::LargeListView(field) => {
666 element_types.push(field.data_type().clone());
667 nested_item_nullability.push(Some(field.is_nullable()));
668 large_list = true;
669 fixed_size = false;
670 }
671 DataType::FixedSizeList(field, size) => {
672 element_types.push(field.data_type().clone());
673 nested_item_nullability.push(Some(field.is_nullable()));
674 list_sizes.push(*size)
675 }
676 arg_type => {
677 plan_err!("{function_name} does not support type {arg_type}")?
678 }
679 },
680 }
681 }
682
683 debug_assert_eq!(nested_item_nullability.len(), arguments.len());
684
685 let Some(element_type) = type_union_resolution(&element_types) else {
686 return Ok(vec![vec![]]);
687 };
688
689 if !fixed_size {
690 list_sizes.clear()
691 };
692
693 let mut list_sizes = list_sizes.into_iter();
694 let valid_types = arguments
695 .iter()
696 .zip(current_types.iter())
697 .zip(nested_item_nullability)
698 .map(|((argument_type, current_type), is_nested_item_nullable)| {
699 match argument_type {
700 ArrayFunctionArgument::Index => DataType::Int64,
701 ArrayFunctionArgument::String => DataType::Utf8,
702 ArrayFunctionArgument::Element => element_type.clone(),
703 ArrayFunctionArgument::Array => {
706 if current_type.is_null() {
707 DataType::Null
708 } else {
709 rebuild_array_type(
710 current_type,
711 &element_type,
712 is_nested_item_nullable.unwrap_or(true),
713 large_list,
714 list_sizes.next(),
715 )
716 }
717 }
718 }
719 });
720
721 Ok(vec![valid_types.collect()])
722 }
723
724 fn recursive_array(array_type: &DataType) -> Option<DataType> {
725 match array_type {
726 DataType::List(_)
727 | DataType::LargeList(_)
728 | DataType::ListView(_)
729 | DataType::LargeListView(_)
730 | DataType::FixedSizeList(_, _) => {
731 let array_type = coerced_fixed_size_list_to_list(array_type);
732 Some(array_type)
733 }
734 _ => None,
735 }
736 }
737
738 fn function_length_check(
739 function_name: &str,
740 length: usize,
741 expected_length: usize,
742 ) -> Result<()> {
743 if length != expected_length {
744 return plan_err!(
745 "Function '{function_name}' expects {expected_length} arguments but received {length}"
746 );
747 }
748 Ok(())
749 }
750
751 let valid_types = match signature {
752 TypeSignature::Variadic(valid_types) => valid_types
753 .iter()
754 .map(|valid_type| vec![valid_type.clone(); current_types.len()])
755 .collect(),
756 TypeSignature::String(number) => {
757 function_length_check(function_name, current_types.len(), *number)?;
758
759 let mut new_types = Vec::with_capacity(current_types.len());
760 for data_type in current_types.iter() {
761 let logical_data_type: NativeType = data_type.into();
762 if logical_data_type == NativeType::String {
763 new_types.push(data_type.to_owned());
764 } else if logical_data_type == NativeType::Null {
765 new_types.push(DataType::Utf8);
767 } else {
768 return plan_err!(
769 "Function '{function_name}' expects String but received {logical_data_type}"
770 );
771 }
772 }
773
774 fn find_common_type(
776 function_name: &str,
777 lhs_type: &DataType,
778 rhs_type: &DataType,
779 ) -> Result<DataType> {
780 match (lhs_type, rhs_type) {
781 (DataType::Dictionary(_, lhs), DataType::Dictionary(_, rhs)) => {
782 find_common_type(function_name, lhs, rhs)
783 }
784 (DataType::Dictionary(_, v), other)
785 | (other, DataType::Dictionary(_, v)) => {
786 find_common_type(function_name, v, other)
787 }
788 _ => {
789 if let Some(coerced_type) = string_coercion(lhs_type, rhs_type) {
790 Ok(coerced_type)
791 } else {
792 plan_err!(
793 "Function '{function_name}' could not coerce {lhs_type} and {rhs_type} to a common string type"
794 )
795 }
796 }
797 }
798 }
799
800 let mut coerced_type = new_types.first().unwrap().to_owned();
802 for t in new_types.iter().skip(1) {
803 coerced_type = find_common_type(function_name, &coerced_type, t)?;
804 }
805
806 fn base_type_or_default_type(data_type: &DataType) -> DataType {
807 if let DataType::Dictionary(_, v) = data_type {
808 base_type_or_default_type(v)
809 } else {
810 data_type.to_owned()
811 }
812 }
813
814 vec![vec![base_type_or_default_type(&coerced_type); *number]]
815 }
816 TypeSignature::Numeric(number) => {
817 function_length_check(function_name, current_types.len(), *number)?;
818
819 let mut valid_type = current_types.first().unwrap().to_owned();
821 for t in current_types.iter().skip(1) {
822 let logical_data_type: NativeType = t.into();
823 if logical_data_type == NativeType::Null {
824 continue;
825 }
826
827 if !logical_data_type.is_numeric() {
828 return plan_err!(
829 "Function '{function_name}' expects Numeric but received {logical_data_type}"
830 );
831 }
832
833 if let Some(coerced_type) = binary_numeric_coercion(&valid_type, t) {
834 valid_type = coerced_type;
835 } else {
836 return plan_err!(
837 "For function '{function_name}' {valid_type} and {t} are not coercible to a common numeric type"
838 );
839 }
840 }
841
842 let logical_data_type: NativeType = valid_type.clone().into();
843 if logical_data_type == NativeType::Null {
847 valid_type = DataType::Float64;
848 } else if !logical_data_type.is_numeric() {
849 return plan_err!(
850 "Function '{function_name}' expects Numeric but received {logical_data_type}"
851 );
852 }
853
854 vec![vec![valid_type; *number]]
855 }
856 TypeSignature::Comparable(num) => {
857 function_length_check(function_name, current_types.len(), *num)?;
858 let mut target_type = current_types[0].to_owned();
859 for data_type in current_types.iter().skip(1) {
860 if let Some(dt) = comparison_coercion(&target_type, data_type) {
861 target_type = dt;
862 } else {
863 return plan_err!(
864 "For function '{function_name}' {target_type} and {data_type} is not comparable"
865 );
866 }
867 }
868 if target_type.is_null() {
870 vec![vec![DataType::Utf8View; *num]]
871 } else {
872 vec![vec![target_type; *num]]
873 }
874 }
875 TypeSignature::Coercible(param_types) => {
876 function_length_check(function_name, current_types.len(), param_types.len())?;
877
878 fn coercion_value_type<'a>(
879 current_type: &'a DataType,
880 desired_type: &TypeSignatureClass,
881 ) -> &'a DataType {
882 if matches!(desired_type, TypeSignatureClass::Any) {
883 return current_type;
884 }
885
886 match current_type {
887 DataType::Dictionary(_, value_type) => {
888 coercion_value_type(value_type, desired_type)
889 }
890 _ => current_type,
891 }
892 }
893
894 fn preserve_encoding(
895 current_type: &DataType,
896 casted_type: DataType,
897 desired_type: &TypeSignatureClass,
898 encoding_preservation: EncodingPreservation,
899 ) -> DataType {
900 if matches!(desired_type, TypeSignatureClass::Any) {
901 return casted_type;
902 }
903
904 match current_type {
905 DataType::Dictionary(key_type, value_type) => {
906 let casted_type = preserve_encoding(
907 value_type,
908 casted_type,
909 desired_type,
910 encoding_preservation,
911 );
912 if encoding_preservation.preserve_dictionary() {
913 DataType::Dictionary(key_type.clone(), Box::new(casted_type))
914 } else {
915 casted_type
916 }
917 }
918 _ => casted_type,
919 }
920 }
921
922 let mut new_types = Vec::with_capacity(current_types.len());
923 for (current_type, param) in current_types.iter().zip(param_types.iter()) {
924 let current_native_type: NativeType = current_type.into();
925 let encoding_preservation = param.encoding_preservation();
926 let coercion_value_type =
927 coercion_value_type(current_type, param.desired_type());
928
929 if param
930 .desired_type()
931 .matches_native_type(¤t_native_type)
932 {
933 let casted_type = param
934 .desired_type()
935 .default_casted_type(¤t_native_type, coercion_value_type)?;
936
937 new_types.push(preserve_encoding(
938 current_type,
939 casted_type,
940 param.desired_type(),
941 encoding_preservation,
942 ));
943 } else if param
944 .allowed_source_types()
945 .iter()
946 .any(|t| t.matches_native_type(¤t_native_type))
947 {
948 let default_casted_type = param.default_casted_type().unwrap();
950 let casted_type =
951 default_casted_type.default_cast_for(coercion_value_type)?;
952 new_types.push(preserve_encoding(
953 current_type,
954 casted_type,
955 param.desired_type(),
956 encoding_preservation,
957 ));
958 } else {
959 let hint = if matches!(current_native_type, NativeType::Binary) {
960 "\n\nHint: Binary types are not automatically coerced to String. Use CAST(column AS VARCHAR) to convert Binary data to String."
961 } else {
962 ""
963 };
964 return plan_err!(
965 "Function '{function_name}' requires {}, but received {} (DataType: {}).{hint}",
966 param.desired_type(),
967 current_native_type,
968 current_type
969 );
970 }
971 }
972
973 vec![new_types]
974 }
975 TypeSignature::Uniform(number, valid_types) => {
976 if *number == 0 {
977 return plan_err!(
978 "The function '{function_name}' expected at least one argument"
979 );
980 }
981
982 valid_types
983 .iter()
984 .map(|valid_type| vec![valid_type.clone(); *number])
985 .collect()
986 }
987 TypeSignature::UserDefined => {
988 return internal_err!(
989 "Function '{function_name}' user-defined signature should be handled by function-specific coerce_types"
990 );
991 }
992 TypeSignature::VariadicAny => {
993 if current_types.is_empty() {
994 return plan_err!(
995 "Function '{function_name}' expected at least one argument but received 0"
996 );
997 }
998 vec![current_types.to_vec()]
999 }
1000 TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
1001 TypeSignature::ArraySignature(function_signature) => match function_signature {
1002 ArrayFunctionSignature::Array {
1003 arguments,
1004 array_coercion,
1005 } => array_valid_types(
1006 function_name,
1007 current_types,
1008 arguments,
1009 array_coercion.as_ref(),
1010 )?,
1011 ArrayFunctionSignature::RecursiveArray => {
1012 if current_types.len() != 1 {
1013 return Ok(vec![vec![]]);
1014 }
1015 recursive_array(¤t_types[0])
1016 .map_or_else(|| vec![vec![]], |array_type| vec![vec![array_type]])
1017 }
1018 ArrayFunctionSignature::MapArray => {
1019 if current_types.len() != 1 {
1020 return Ok(vec![vec![]]);
1021 }
1022
1023 match ¤t_types[0] {
1024 DataType::Map(_, _) => vec![vec![current_types[0].clone()]],
1025 _ => vec![vec![]],
1026 }
1027 }
1028 },
1029 TypeSignature::Nullary => {
1030 if !current_types.is_empty() {
1031 return plan_err!(
1032 "The function '{function_name}' expected zero argument but received {}",
1033 current_types.len()
1034 );
1035 }
1036 vec![vec![]]
1037 }
1038 TypeSignature::Any(number) => {
1039 if current_types.is_empty() {
1040 return plan_err!(
1041 "The function '{function_name}' expected at least one argument but received 0"
1042 );
1043 }
1044
1045 if current_types.len() != *number {
1046 return plan_err!(
1047 "The function '{function_name}' expected {number} arguments but received {}",
1048 current_types.len()
1049 );
1050 }
1051 vec![current_types.to_vec()]
1052 }
1053 TypeSignature::OneOf(types) => types
1054 .iter()
1055 .filter_map(|t| get_valid_types(function_name, t, current_types).ok())
1056 .flatten()
1057 .collect::<Vec<_>>(),
1058 };
1059
1060 Ok(valid_types)
1061}
1062
1063fn maybe_data_types(
1070 valid_types: &[DataType],
1071 current_types: &[DataType],
1072) -> Option<Vec<DataType>> {
1073 if valid_types.len() != current_types.len() {
1074 return None;
1075 }
1076
1077 let mut new_type = Vec::with_capacity(valid_types.len());
1078 for (i, valid_type) in valid_types.iter().enumerate() {
1079 let current_type = ¤t_types[i];
1080
1081 if current_type == valid_type {
1082 new_type.push(current_type.clone())
1083 } else {
1084 let coerced_type = coerced_from(valid_type, current_type)?;
1088 new_type.push(coerced_type)
1089 }
1090 }
1091 Some(new_type)
1092}
1093
1094fn maybe_data_types_without_coercion(
1098 valid_types: &[DataType],
1099 current_types: &[DataType],
1100) -> Option<Vec<DataType>> {
1101 if valid_types.len() != current_types.len() {
1102 return None;
1103 }
1104
1105 let mut new_type = Vec::with_capacity(valid_types.len());
1106 for (i, valid_type) in valid_types.iter().enumerate() {
1107 let current_type = ¤t_types[i];
1108
1109 if current_type == valid_type {
1110 new_type.push(current_type.clone())
1111 } else if can_cast_types(current_type, valid_type) {
1112 new_type.push(valid_type.clone())
1114 } else {
1115 return None;
1116 }
1117 }
1118 Some(new_type)
1119}
1120
1121#[deprecated(since = "53.0.0", note = "Unused internal function")]
1126pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
1127 if type_into == type_from {
1128 return true;
1129 }
1130 if let Some(coerced) = coerced_from(type_into, type_from) {
1131 return coerced == *type_into;
1132 }
1133 false
1134}
1135
1136fn coerced_from<'a>(
1143 type_into: &'a DataType,
1144 type_from: &'a DataType,
1145) -> Option<DataType> {
1146 use self::DataType::*;
1147
1148 match (type_into, type_from) {
1150 (_, Dictionary(_, value_type))
1152 if coerced_from(type_into, value_type).is_some() =>
1153 {
1154 Some(type_into.clone())
1155 }
1156 (Dictionary(_, value_type), _)
1157 if coerced_from(value_type, type_from).is_some() =>
1158 {
1159 Some(type_into.clone())
1160 }
1161 (Int8, Null | Int8) => Some(type_into.clone()),
1163 (Int16, Null | Int8 | Int16 | UInt8) => Some(type_into.clone()),
1164 (Int32, Null | Int8 | Int16 | Int32 | UInt8 | UInt16) => Some(type_into.clone()),
1165 (Int64, Null | Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32) => {
1166 Some(type_into.clone())
1167 }
1168 (UInt8, Null | UInt8) => Some(type_into.clone()),
1169 (UInt16, Null | UInt8 | UInt16) => Some(type_into.clone()),
1170 (UInt32, Null | UInt8 | UInt16 | UInt32) => Some(type_into.clone()),
1171 (UInt64, Null | UInt8 | UInt16 | UInt32 | UInt64) => Some(type_into.clone()),
1172 (Float16, Null | Int8 | Int16 | UInt8 | UInt16 | Float16) => {
1173 Some(type_into.clone())
1174 }
1175 (
1176 Float32,
1177 Null | Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64
1178 | Float16 | Float32,
1179 ) => Some(type_into.clone()),
1180 (
1181 Float64,
1182 Null
1183 | Int8
1184 | Int16
1185 | Int32
1186 | Int64
1187 | UInt8
1188 | UInt16
1189 | UInt32
1190 | UInt64
1191 | Float16
1192 | Float32
1193 | Float64
1194 | Decimal32(_, _)
1195 | Decimal64(_, _)
1196 | Decimal128(_, _)
1197 | Decimal256(_, _),
1198 ) => Some(type_into.clone()),
1199 (
1200 Timestamp(TimeUnit::Nanosecond, None),
1201 Null | Timestamp(_, None) | Date32 | Date64 | Utf8 | LargeUtf8 | Utf8View,
1202 ) => Some(type_into.clone()),
1203 (Interval(_), Null | Utf8 | LargeUtf8 | Utf8View) => Some(type_into.clone()),
1204 (Utf8 | LargeUtf8 | Utf8View, _) => Some(type_into.clone()),
1206 (BinaryView, Binary | LargeBinary | Null) => Some(type_into.clone()),
1208 (Null, _) if can_cast_types(type_from, type_into) => Some(type_into.clone()),
1209
1210 (List(_), FixedSizeList(_, _)) => Some(type_into.clone()),
1211
1212 (List(_) | LargeList(_) | ListView(_) | LargeListView(_), _)
1215 if base_type(type_from).is_null()
1216 || list_ndims(type_from) == list_ndims(type_into) =>
1217 {
1218 Some(type_into.clone())
1219 }
1220 (
1222 FixedSizeList(f_into, FIXED_SIZE_LIST_WILDCARD),
1223 FixedSizeList(f_from, size_from),
1224 ) => match coerced_from(f_into.data_type(), f_from.data_type()) {
1225 Some(data_type) if &data_type != f_into.data_type() => {
1226 let new_field =
1227 Arc::new(f_into.as_ref().clone().with_data_type(data_type));
1228 Some(FixedSizeList(new_field, *size_from))
1229 }
1230 Some(_) => Some(FixedSizeList(Arc::clone(f_into), *size_from)),
1231 _ => None,
1232 },
1233 (Timestamp(unit, Some(tz)), _) if tz.as_ref() == TIMEZONE_WILDCARD => {
1234 match type_from {
1235 Timestamp(_, Some(from_tz)) => {
1236 Some(Timestamp(*unit, Some(Arc::clone(from_tz))))
1237 }
1238 Null | Date32 | Utf8 | LargeUtf8 | Timestamp(_, None) => {
1239 Some(Timestamp(*unit, Some("+00".into())))
1241 }
1242 _ => None,
1243 }
1244 }
1245 (Timestamp(_, Some(_)), Null | Timestamp(_, _) | Date32 | Utf8 | LargeUtf8) => {
1246 Some(type_into.clone())
1247 }
1248 (_, Null) if can_cast_types(type_from, type_into) => Some(type_into.clone()),
1255 _ => None,
1256 }
1257}
1258
1259#[cfg(test)]
1260mod tests {
1261 use crate::{
1262 HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature,
1263 HigherOrderUDFImpl, Volatility,
1264 };
1265
1266 use super::*;
1267 use arrow::datatypes::IntervalUnit;
1268 use datafusion_common::{
1269 assert_contains,
1270 types::{logical_binary, logical_int64, logical_string},
1271 };
1272 use datafusion_expr_common::{
1273 columnar_value::ColumnarValue,
1274 signature::{Coercion, EncodingPreservation, TypeSignatureClass},
1275 };
1276
1277 #[test]
1278 fn test_string_conversion() {
1279 let cases = vec![
1280 (DataType::Utf8View, DataType::Utf8),
1281 (DataType::Utf8View, DataType::LargeUtf8),
1282 (DataType::Utf8View, DataType::Null),
1283 ];
1284
1285 for case in cases {
1286 assert_eq!(coerced_from(&case.0, &case.1), Some(case.0));
1287 }
1288 }
1289
1290 #[test]
1291 fn test_binary_conversion() {
1292 let cases = vec![
1293 (DataType::BinaryView, DataType::Binary),
1294 (DataType::BinaryView, DataType::LargeBinary),
1295 (DataType::BinaryView, DataType::Null),
1296 ];
1297
1298 for case in cases {
1299 assert_eq!(coerced_from(&case.0, &case.1), Some(case.0));
1300 }
1301 }
1302
1303 #[test]
1304 fn test_coerced_from_null() {
1305 assert_eq!(
1307 coerced_from(
1308 &DataType::Interval(IntervalUnit::MonthDayNano),
1309 &DataType::Null
1310 ),
1311 Some(DataType::Interval(IntervalUnit::MonthDayNano))
1312 );
1313
1314 assert_eq!(
1316 coerced_from(&DataType::Date32, &DataType::Null),
1317 Some(DataType::Date32)
1318 );
1319
1320 assert_eq!(
1322 coerced_from(
1323 &DataType::Timestamp(TimeUnit::Microsecond, Some("+00".into())),
1324 &DataType::Null
1325 ),
1326 Some(DataType::Timestamp(
1327 TimeUnit::Microsecond,
1328 Some("+00".into())
1329 ))
1330 );
1331 }
1332
1333 #[test]
1334 fn test_maybe_data_types() {
1335 let cases = vec![
1337 (
1339 vec![DataType::UInt8, DataType::UInt16],
1340 vec![DataType::UInt8, DataType::UInt16],
1341 Some(vec![DataType::UInt8, DataType::UInt16]),
1342 ),
1343 (
1345 vec![DataType::UInt16, DataType::UInt16],
1346 vec![DataType::UInt8, DataType::UInt16],
1347 Some(vec![DataType::UInt16, DataType::UInt16]),
1348 ),
1349 (vec![], vec![], Some(vec![])),
1351 (
1353 vec![DataType::Boolean, DataType::UInt16],
1354 vec![DataType::UInt8, DataType::UInt16],
1355 None,
1356 ),
1357 (
1359 vec![DataType::Boolean, DataType::UInt32],
1360 vec![DataType::Boolean, DataType::UInt16],
1361 Some(vec![DataType::Boolean, DataType::UInt32]),
1362 ),
1363 (
1365 vec![
1366 DataType::Timestamp(TimeUnit::Nanosecond, None),
1367 DataType::Timestamp(TimeUnit::Nanosecond, Some("+TZ".into())),
1368 DataType::Timestamp(TimeUnit::Nanosecond, Some("+01".into())),
1369 ],
1370 vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
1371 Some(vec![
1372 DataType::Timestamp(TimeUnit::Nanosecond, None),
1373 DataType::Timestamp(TimeUnit::Nanosecond, Some("+00".into())),
1374 DataType::Timestamp(TimeUnit::Nanosecond, Some("+01".into())),
1375 ]),
1376 ),
1377 ];
1378
1379 for case in cases {
1380 assert_eq!(maybe_data_types(&case.0, &case.1), case.2)
1381 }
1382 }
1383
1384 #[test]
1385 fn test_get_valid_types_numeric() -> Result<()> {
1386 let get_valid_types_flatten =
1387 |function_name: &str,
1388 signature: &TypeSignature,
1389 current_types: &[DataType]| {
1390 get_valid_types(function_name, signature, current_types)
1391 .unwrap()
1392 .into_iter()
1393 .flatten()
1394 .collect::<Vec<_>>()
1395 };
1396
1397 let got = get_valid_types_flatten(
1399 "test",
1400 &TypeSignature::Numeric(1),
1401 &[DataType::Int32],
1402 );
1403 assert_eq!(got, [DataType::Int32]);
1404
1405 let got = get_valid_types_flatten(
1407 "test",
1408 &TypeSignature::Numeric(2),
1409 &[DataType::Int32, DataType::Int64],
1410 );
1411 assert_eq!(got, [DataType::Int64, DataType::Int64]);
1412
1413 let got = get_valid_types_flatten(
1415 "test",
1416 &TypeSignature::Numeric(3),
1417 &[DataType::Int32, DataType::Int64, DataType::Float64],
1418 );
1419 assert_eq!(
1420 got,
1421 [DataType::Float64, DataType::Float64, DataType::Float64]
1422 );
1423
1424 let got = get_valid_types(
1426 "test",
1427 &TypeSignature::Numeric(2),
1428 &[DataType::Int32, DataType::Utf8],
1429 )
1430 .unwrap_err();
1431 assert_contains!(
1432 got.to_string(),
1433 "Function 'test' expects Numeric but received String"
1434 );
1435
1436 let got = get_valid_types_flatten(
1438 "test",
1439 &TypeSignature::Numeric(1),
1440 &[DataType::Null],
1441 );
1442 assert_eq!(got, [DataType::Float64]);
1443
1444 let got = get_valid_types(
1446 "test",
1447 &TypeSignature::Numeric(1),
1448 &[DataType::Timestamp(TimeUnit::Second, None)],
1449 )
1450 .unwrap_err();
1451 assert_contains!(
1452 got.to_string(),
1453 "Function 'test' expects Numeric but received Timestamp(s)"
1454 );
1455
1456 Ok(())
1457 }
1458
1459 #[test]
1460 fn test_get_valid_types_one_of() -> Result<()> {
1461 let signature =
1462 TypeSignature::OneOf(vec![TypeSignature::Any(1), TypeSignature::Any(2)]);
1463
1464 let invalid_types = get_valid_types(
1465 "test",
1466 &signature,
1467 &[DataType::Int32, DataType::Int32, DataType::Int32],
1468 )?;
1469 assert_eq!(invalid_types.len(), 0);
1470
1471 let args = vec![DataType::Int32, DataType::Int32];
1472 let valid_types = get_valid_types("test", &signature, &args)?;
1473 assert_eq!(valid_types.len(), 1);
1474 assert_eq!(valid_types[0], args);
1475
1476 let args = vec![DataType::Int32];
1477 let valid_types = get_valid_types("test", &signature, &args)?;
1478 assert_eq!(valid_types.len(), 1);
1479 assert_eq!(valid_types[0], args);
1480
1481 Ok(())
1482 }
1483
1484 #[test]
1485 fn test_get_valid_types_length_check() -> Result<()> {
1486 let signature = TypeSignature::Numeric(1);
1487
1488 let err = get_valid_types("test", &signature, &[]).unwrap_err();
1489 assert_contains!(
1490 err.to_string(),
1491 "Function 'test' expects 1 arguments but received 0"
1492 );
1493
1494 let err = get_valid_types(
1495 "test",
1496 &signature,
1497 &[DataType::Int32, DataType::Int32, DataType::Int32],
1498 )
1499 .unwrap_err();
1500 assert_contains!(
1501 err.to_string(),
1502 "Function 'test' expects 1 arguments but received 3"
1503 );
1504
1505 Ok(())
1506 }
1507
1508 struct MockUdf(Signature);
1509
1510 impl UDFCoercionExt for MockUdf {
1511 fn name(&self) -> &str {
1512 "test"
1513 }
1514 fn signature(&self) -> &Signature {
1515 &self.0
1516 }
1517 fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
1518 unimplemented!()
1519 }
1520 }
1521
1522 #[test]
1523 fn test_fixed_list_wildcard_coerce() -> Result<()> {
1524 let inner = Arc::new(Field::new_list_field(DataType::Int32, false));
1525 let current_fields = vec![Arc::new(Field::new(
1527 "t",
1528 DataType::FixedSizeList(Arc::clone(&inner), 2),
1529 true,
1530 ))];
1531
1532 let signature = Signature::exact(
1533 vec![DataType::FixedSizeList(
1534 Arc::clone(&inner),
1535 FIXED_SIZE_LIST_WILDCARD,
1536 )],
1537 Volatility::Stable,
1538 );
1539
1540 let coerced_fields = fields_with_udf(¤t_fields, &MockUdf(signature))?;
1541 assert_eq!(coerced_fields, current_fields);
1542
1543 let signature = Signature::exact(
1545 vec![DataType::FixedSizeList(Arc::clone(&inner), 3)],
1546 Volatility::Stable,
1547 );
1548 let coerced_fields = fields_with_udf(¤t_fields, &MockUdf(signature));
1549 assert!(coerced_fields.is_err());
1550
1551 let signature = Signature::exact(
1553 vec![DataType::FixedSizeList(Arc::clone(&inner), 2)],
1554 Volatility::Stable,
1555 );
1556 let coerced_fields =
1557 fields_with_udf(¤t_fields, &MockUdf(signature)).unwrap();
1558 assert_eq!(coerced_fields, current_fields);
1559
1560 Ok(())
1561 }
1562
1563 #[test]
1564 fn test_nested_wildcard_fixed_size_lists() -> Result<()> {
1565 let type_into = DataType::FixedSizeList(
1566 Arc::new(Field::new_list_field(
1567 DataType::FixedSizeList(
1568 Arc::new(Field::new_list_field(DataType::Int32, false)),
1569 FIXED_SIZE_LIST_WILDCARD,
1570 ),
1571 false,
1572 )),
1573 FIXED_SIZE_LIST_WILDCARD,
1574 );
1575
1576 let type_from = DataType::FixedSizeList(
1577 Arc::new(Field::new_list_field(
1578 DataType::FixedSizeList(
1579 Arc::new(Field::new_list_field(DataType::Int8, false)),
1580 4,
1581 ),
1582 false,
1583 )),
1584 3,
1585 );
1586
1587 assert_eq!(
1588 coerced_from(&type_into, &type_from),
1589 Some(DataType::FixedSizeList(
1590 Arc::new(Field::new_list_field(
1591 DataType::FixedSizeList(
1592 Arc::new(Field::new_list_field(DataType::Int32, false)),
1593 4,
1594 ),
1595 false,
1596 )),
1597 3,
1598 ))
1599 );
1600
1601 Ok(())
1602 }
1603
1604 #[test]
1605 fn test_coerced_from_dictionary() {
1606 let type_into =
1607 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::UInt32));
1608 let type_from = DataType::Int64;
1609 assert_eq!(coerced_from(&type_into, &type_from), None);
1610
1611 let type_from =
1612 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::UInt32));
1613 let type_into = DataType::Int64;
1614 assert_eq!(
1615 coerced_from(&type_into, &type_from),
1616 Some(type_into.clone())
1617 );
1618 }
1619
1620 #[test]
1621 fn test_get_valid_types_array_and_array() -> Result<()> {
1622 let function = "array_and_array";
1623 let signature = Signature::arrays(
1624 2,
1625 Some(ListCoercion::FixedSizedListToList),
1626 Volatility::Immutable,
1627 );
1628
1629 let data_types = vec![
1630 DataType::new_list(DataType::Int32, true),
1631 DataType::new_large_list(DataType::Float64, true),
1632 ];
1633 assert_eq!(
1634 get_valid_types(function, &signature.type_signature, &data_types)?,
1635 vec![vec![
1636 DataType::new_large_list(DataType::Float64, true),
1637 DataType::new_large_list(DataType::Float64, true),
1638 ]]
1639 );
1640
1641 let data_types = vec![
1642 DataType::new_fixed_size_list(DataType::Int64, 3, true),
1643 DataType::new_fixed_size_list(DataType::Int32, 5, true),
1644 ];
1645 assert_eq!(
1646 get_valid_types(function, &signature.type_signature, &data_types)?,
1647 vec![vec![
1648 DataType::new_list(DataType::Int64, true),
1649 DataType::new_list(DataType::Int64, true),
1650 ]]
1651 );
1652
1653 let data_types = vec![
1654 DataType::new_fixed_size_list(DataType::Null, 3, true),
1655 DataType::new_large_list(DataType::Utf8, true),
1656 ];
1657 assert_eq!(
1658 get_valid_types(function, &signature.type_signature, &data_types)?,
1659 vec![vec![
1660 DataType::new_large_list(DataType::Utf8, true),
1661 DataType::new_large_list(DataType::Utf8, true),
1662 ]]
1663 );
1664
1665 let data_types = vec![
1666 DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1667 DataType::new_list(DataType::Int32, true),
1668 ];
1669 assert_eq!(
1670 get_valid_types(function, &signature.type_signature, &data_types)?,
1671 vec![vec![
1672 DataType::new_list(DataType::Int32, true),
1673 DataType::new_list(DataType::Int32, true),
1674 ]]
1675 );
1676
1677 let data_types = vec![
1678 DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1679 DataType::new_list(DataType::Int32, true),
1680 ];
1681 assert_eq!(
1682 get_valid_types(function, &signature.type_signature, &data_types)?,
1683 vec![vec![
1684 DataType::new_large_list(DataType::Int32, true),
1685 DataType::new_large_list(DataType::Int32, true),
1686 ]]
1687 );
1688
1689 let data_types = vec![
1690 DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1691 DataType::ListView(Field::new_list_field(DataType::Int32, true).into()),
1692 ];
1693 assert_eq!(
1694 get_valid_types(function, &signature.type_signature, &data_types)?,
1695 vec![vec![
1696 DataType::new_list(DataType::Int32, true),
1697 DataType::new_list(DataType::Int32, true),
1698 ]]
1699 );
1700
1701 let data_types = vec![
1702 DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1703 DataType::LargeListView(Field::new_list_field(DataType::Int32, true).into()),
1704 ];
1705 assert_eq!(
1706 get_valid_types(function, &signature.type_signature, &data_types)?,
1707 vec![vec![
1708 DataType::new_large_list(DataType::Int32, true),
1709 DataType::new_large_list(DataType::Int32, true),
1710 ]]
1711 );
1712
1713 Ok(())
1714 }
1715
1716 #[test]
1717 fn test_get_valid_types_array_and_element() -> Result<()> {
1718 let function = "array_and_element";
1719 let signature = Signature::array_and_element(Volatility::Immutable);
1720
1721 let data_types =
1722 vec![DataType::new_list(DataType::Int32, true), DataType::Float64];
1723 assert_eq!(
1724 get_valid_types(function, &signature.type_signature, &data_types)?,
1725 vec![vec![
1726 DataType::new_list(DataType::Float64, true),
1727 DataType::Float64,
1728 ]]
1729 );
1730
1731 let data_types = vec![
1732 DataType::new_large_list(DataType::Int32, true),
1733 DataType::Null,
1734 ];
1735 assert_eq!(
1736 get_valid_types(function, &signature.type_signature, &data_types)?,
1737 vec![vec![
1738 DataType::new_large_list(DataType::Int32, true),
1739 DataType::Int32,
1740 ]]
1741 );
1742
1743 let data_types = vec![
1744 DataType::new_fixed_size_list(DataType::Null, 3, true),
1745 DataType::Utf8,
1746 ];
1747 assert_eq!(
1748 get_valid_types(function, &signature.type_signature, &data_types)?,
1749 vec![vec![
1750 DataType::new_list(DataType::Utf8, true),
1751 DataType::Utf8,
1752 ]]
1753 );
1754
1755 Ok(())
1756 }
1757
1758 #[test]
1759 fn test_get_valid_types_array_and_index_preserves_list_field_name() -> Result<()> {
1760 let struct_fields = vec![
1761 Field::new("id", DataType::Utf8, true),
1762 Field::new("prim", DataType::Boolean, true),
1763 ];
1764 let current_type = DataType::List(Arc::new(Field::new(
1765 "element",
1766 DataType::Struct(struct_fields.into()),
1767 true,
1768 )));
1769 let signature = Signature::array_and_index(Volatility::Immutable);
1770
1771 assert_eq!(
1772 get_valid_types(
1773 "array_element",
1774 &signature.type_signature,
1775 &[current_type.clone(), DataType::Int64],
1776 )?,
1777 vec![vec![current_type, DataType::Int64]]
1778 );
1779
1780 Ok(())
1781 }
1782
1783 #[test]
1784 fn test_get_valid_types_element_and_array() -> Result<()> {
1785 let function = "element_and_array";
1786 let signature = Signature::element_and_array(Volatility::Immutable);
1787
1788 let data_types = vec![
1789 DataType::new_large_list(DataType::Null, false),
1790 DataType::new_list(DataType::new_list(DataType::Int64, true), true),
1791 ];
1792 assert_eq!(
1793 get_valid_types(function, &signature.type_signature, &data_types)?,
1794 vec![vec![
1795 DataType::new_large_list(DataType::Int64, true),
1796 DataType::new_list(DataType::new_large_list(DataType::Int64, true), true),
1797 ]]
1798 );
1799
1800 Ok(())
1801 }
1802
1803 #[test]
1804 fn test_coercible_nulls() -> Result<()> {
1805 fn null_input(coercion: Coercion) -> Result<Vec<DataType>> {
1806 fields_with_udf(
1807 &[Field::new("field", DataType::Null, true).into()],
1808 &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1809 )
1810 .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1811 }
1812
1813 let output = null_input(Coercion::new_exact(TypeSignatureClass::Native(
1815 logical_int64(),
1816 )))?;
1817 assert_eq!(vec![DataType::Int64], output);
1818
1819 let output = null_input(Coercion::new_implicit(
1820 TypeSignatureClass::Native(logical_int64()),
1821 vec![],
1822 NativeType::Int64,
1823 ))?;
1824 assert_eq!(vec![DataType::Int64], output);
1825
1826 let output = null_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
1828 assert_eq!(vec![DataType::Null], output);
1829
1830 let output = null_input(Coercion::new_implicit(
1831 TypeSignatureClass::Integer,
1832 vec![],
1833 NativeType::Int64,
1834 ))?;
1835 assert_eq!(vec![DataType::Null], output);
1836
1837 Ok(())
1838 }
1839
1840 #[test]
1841 fn test_coercible_dictionary() -> Result<()> {
1842 let dictionary =
1843 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64));
1844 fn dictionary_input(coercion: Coercion) -> Result<Vec<DataType>> {
1845 fields_with_udf(
1846 &[Field::new(
1847 "field",
1848 DataType::Dictionary(
1849 Box::new(DataType::Int8),
1850 Box::new(DataType::Int64),
1851 ),
1852 true,
1853 )
1854 .into()],
1855 &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1856 )
1857 .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1858 }
1859
1860 let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Native(
1862 logical_int64(),
1863 )))?;
1864 assert_eq!(vec![DataType::Int64], output);
1865
1866 let output = dictionary_input(Coercion::new_implicit(
1867 TypeSignatureClass::Native(logical_int64()),
1868 vec![],
1869 NativeType::Int64,
1870 ))?;
1871 assert_eq!(vec![DataType::Int64], output);
1872
1873 let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Any))?;
1875 assert_eq!(vec![dictionary.clone()], output);
1876
1877 let output = dictionary_input(
1878 Coercion::new_exact(TypeSignatureClass::Any)
1879 .with_encoding_preservation(EncodingPreservation::dictionary()),
1880 )?;
1881 assert_eq!(vec![dictionary.clone()], output);
1882
1883 let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
1885 assert_eq!(vec![DataType::Int64], output);
1886
1887 let output = dictionary_input(Coercion::new_implicit(
1888 TypeSignatureClass::Integer,
1889 vec![],
1890 NativeType::Int64,
1891 ))?;
1892 assert_eq!(vec![DataType::Int64], output);
1893
1894 let output = dictionary_input(
1896 Coercion::new_exact(TypeSignatureClass::Integer)
1897 .with_encoding_preservation(EncodingPreservation::dictionary()),
1898 )?;
1899 assert_eq!(vec![dictionary], output);
1900
1901 Ok(())
1902 }
1903
1904 #[test]
1905 fn test_coercible_dictionary_preserves_encoding() -> Result<()> {
1906 fn dictionary_input(
1907 value_type: DataType,
1908 coercion: Coercion,
1909 ) -> Result<Vec<DataType>> {
1910 fields_with_udf(
1911 &[Field::new(
1912 "field",
1913 DataType::Dictionary(Box::new(DataType::Int8), Box::new(value_type)),
1914 true,
1915 )
1916 .into()],
1917 &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
1918 )
1919 .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
1920 }
1921
1922 let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
1923 .with_encoding_preservation(EncodingPreservation::dictionary());
1924
1925 assert_eq!(
1926 dictionary_input(DataType::LargeUtf8, coercion.clone())?,
1927 vec![DataType::Dictionary(
1928 Box::new(DataType::Int8),
1929 Box::new(DataType::LargeUtf8),
1930 )]
1931 );
1932 assert_eq!(
1933 dictionary_input(
1934 DataType::BinaryView,
1935 Coercion::new_implicit(
1936 TypeSignatureClass::Native(logical_string()),
1937 vec![TypeSignatureClass::Native(logical_binary())],
1938 NativeType::String,
1939 )
1940 .with_encoding_preservation(EncodingPreservation::dictionary()),
1941 )?,
1942 vec![DataType::Dictionary(
1943 Box::new(DataType::Int8),
1944 Box::new(DataType::Utf8View),
1945 )]
1946 );
1947 assert_eq!(
1949 dictionary_input(
1950 DataType::Int32,
1951 Coercion::new_implicit(
1952 TypeSignatureClass::Native(logical_int64()),
1953 vec![TypeSignatureClass::Integer],
1954 NativeType::Int64,
1955 ),
1956 )?,
1957 vec![DataType::Int64]
1958 );
1959 assert_eq!(
1961 dictionary_input(
1962 DataType::Int32,
1963 Coercion::new_implicit(
1964 TypeSignatureClass::Native(logical_int64()),
1965 vec![TypeSignatureClass::Integer],
1966 NativeType::Int64,
1967 )
1968 .with_encoding_preservation(EncodingPreservation::dictionary()),
1969 )?,
1970 vec![DataType::Dictionary(
1971 Box::new(DataType::Int8),
1972 Box::new(DataType::Int64),
1973 )]
1974 );
1975 assert_eq!(
1977 dictionary_input(
1978 DataType::Int32,
1979 Coercion::new_implicit(
1980 TypeSignatureClass::Integer,
1981 vec![],
1982 NativeType::Int64,
1983 ),
1984 )?,
1985 vec![DataType::Int32]
1986 );
1987 assert_eq!(
1989 dictionary_input(
1990 DataType::Int32,
1991 Coercion::new_implicit(
1992 TypeSignatureClass::Integer,
1993 vec![],
1994 NativeType::Int64,
1995 )
1996 .with_encoding_preservation(EncodingPreservation::dictionary()),
1997 )?,
1998 vec![DataType::Dictionary(
1999 Box::new(DataType::Int8),
2000 Box::new(DataType::Int32),
2001 )]
2002 );
2003
2004 Ok(())
2005 }
2006
2007 #[test]
2008 fn test_coercible_nested_dictionary() -> Result<()> {
2009 let nested_dictionary = DataType::Dictionary(
2010 Box::new(DataType::Int8),
2011 Box::new(DataType::Dictionary(
2012 Box::new(DataType::Int16),
2013 Box::new(DataType::Int32),
2014 )),
2015 );
2016 let nested_dictionary_input = |coercion| -> Result<Vec<DataType>> {
2017 fields_with_udf(
2018 &[Field::new("field", nested_dictionary.clone(), true).into()],
2019 &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
2020 )
2021 .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
2022 };
2023
2024 let output =
2026 nested_dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
2027 assert_eq!(vec![DataType::Int32], output);
2028
2029 let output = nested_dictionary_input(
2031 Coercion::new_exact(TypeSignatureClass::Integer)
2032 .with_encoding_preservation(EncodingPreservation::dictionary()),
2033 )?;
2034 assert_eq!(vec![nested_dictionary.clone()], output);
2035
2036 let int64_coercion = || {
2037 Coercion::new_implicit(
2038 TypeSignatureClass::Native(logical_int64()),
2039 vec![TypeSignatureClass::Integer],
2040 NativeType::Int64,
2041 )
2042 };
2043
2044 let output = nested_dictionary_input(int64_coercion())?;
2046 assert_eq!(vec![DataType::Int64], output);
2047
2048 let output = nested_dictionary_input(
2050 int64_coercion()
2051 .with_encoding_preservation(EncodingPreservation::dictionary()),
2052 )?;
2053 assert_eq!(
2054 vec![DataType::Dictionary(
2055 Box::new(DataType::Int8),
2056 Box::new(DataType::Dictionary(
2057 Box::new(DataType::Int16),
2058 Box::new(DataType::Int64),
2059 )),
2060 )],
2061 output
2062 );
2063
2064 Ok(())
2065 }
2066
2067 #[test]
2068 fn test_coercible_run_end_encoded() -> Result<()> {
2069 let run_end_encoded = DataType::RunEndEncoded(
2070 Field::new("run_ends", DataType::Int16, false).into(),
2071 Field::new("values", DataType::Int64, true).into(),
2072 );
2073 fn run_end_encoded_input(coercion: Coercion) -> Result<Vec<DataType>> {
2074 fields_with_udf(
2075 &[Field::new(
2076 "field",
2077 DataType::RunEndEncoded(
2078 Field::new("run_ends", DataType::Int16, false).into(),
2079 Field::new("values", DataType::Int64, true).into(),
2080 ),
2081 true,
2082 )
2083 .into()],
2084 &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)),
2085 )
2086 .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect())
2087 }
2088
2089 let output = run_end_encoded_input(Coercion::new_exact(
2091 TypeSignatureClass::Native(logical_int64()),
2092 ))?;
2093 assert_eq!(vec![DataType::Int64], output);
2094
2095 let output = run_end_encoded_input(Coercion::new_implicit(
2096 TypeSignatureClass::Native(logical_int64()),
2097 vec![],
2098 NativeType::Int64,
2099 ))?;
2100 assert_eq!(vec![DataType::Int64], output);
2101
2102 let output =
2104 run_end_encoded_input(Coercion::new_exact(TypeSignatureClass::Integer))?;
2105 assert_eq!(vec![run_end_encoded.clone()], output);
2106
2107 let output = run_end_encoded_input(Coercion::new_implicit(
2108 TypeSignatureClass::Integer,
2109 vec![],
2110 NativeType::Int64,
2111 ))?;
2112 assert_eq!(vec![run_end_encoded.clone()], output);
2113
2114 Ok(())
2115 }
2116
2117 #[test]
2118 fn test_get_valid_types_coercible_binary() -> Result<()> {
2119 let signature = Signature::coercible(
2120 vec![Coercion::new_exact(TypeSignatureClass::Native(
2121 logical_binary(),
2122 ))],
2123 Volatility::Immutable,
2124 );
2125
2126 for t in [
2128 DataType::Binary,
2129 DataType::BinaryView,
2130 DataType::LargeBinary,
2131 ] {
2132 assert_eq!(
2133 get_valid_types("", &signature.type_signature, std::slice::from_ref(&t))?,
2134 vec![vec![t]]
2135 );
2136 }
2137
2138 Ok(())
2139 }
2140
2141 #[test]
2142 fn test_get_valid_types_fixed_size_arrays() -> Result<()> {
2143 let function = "fixed_size_arrays";
2144 let signature = Signature::arrays(2, None, Volatility::Immutable);
2145
2146 let data_types = vec![
2147 DataType::new_fixed_size_list(DataType::Int64, 3, true),
2148 DataType::new_fixed_size_list(DataType::Int32, 5, true),
2149 ];
2150 assert_eq!(
2151 get_valid_types(function, &signature.type_signature, &data_types)?,
2152 vec![vec![
2153 DataType::new_fixed_size_list(DataType::Int64, 3, true),
2154 DataType::new_fixed_size_list(DataType::Int64, 5, true),
2155 ]]
2156 );
2157
2158 let data_types = vec![
2159 DataType::new_fixed_size_list(DataType::Int64, 3, true),
2160 DataType::new_list(DataType::Int32, true),
2161 ];
2162 assert_eq!(
2163 get_valid_types(function, &signature.type_signature, &data_types)?,
2164 vec![vec![
2165 DataType::new_list(DataType::Int64, true),
2166 DataType::new_list(DataType::Int64, true),
2167 ]]
2168 );
2169
2170 let data_types = vec![
2171 DataType::new_fixed_size_list(DataType::Utf8, 3, true),
2172 DataType::new_list(DataType::new_list(DataType::Int32, true), true),
2173 ];
2174 assert_eq!(
2175 get_valid_types(function, &signature.type_signature, &data_types)?,
2176 vec![vec![]]
2177 );
2178
2179 let data_types = vec![
2180 DataType::new_fixed_size_list(DataType::Int64, 3, false),
2181 DataType::new_list(DataType::Int32, false),
2182 ];
2183 assert_eq!(
2184 get_valid_types(function, &signature.type_signature, &data_types)?,
2185 vec![vec![
2186 DataType::new_list(DataType::Int64, false),
2187 DataType::new_list(DataType::Int64, false),
2188 ]]
2189 );
2190
2191 Ok(())
2192 }
2193
2194 #[derive(Debug, PartialEq, Eq, Hash)]
2195 struct MockHigherOrderUDF {
2196 signature: HigherOrderSignature,
2197 coerced_value_types: Vec<DataType>,
2198 }
2199
2200 impl HigherOrderUDFImpl for MockHigherOrderUDF {
2201 fn name(&self) -> &str {
2202 "mock_higher_order_function"
2203 }
2204
2205 fn signature(&self) -> &HigherOrderSignature {
2206 &self.signature
2207 }
2208
2209 fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
2210 if arg_types.len() != 1 {
2211 return plan_err!(
2212 "mock_higher_order_function expects 1 value arguments, got {}",
2213 arg_types.len()
2214 );
2215 }
2216 Ok(self.coerced_value_types.clone())
2217 }
2218
2219 fn coerce_values_for_lambdas(
2220 &self,
2221 fields: &[ValueOrLambda<DataType, DataType>],
2222 ) -> Result<Option<Vec<DataType>>> {
2223 let [
2225 ValueOrLambda::Value(list),
2226 ValueOrLambda::Value(_initial),
2227 ValueOrLambda::Lambda(merge),
2228 ] = fields
2229 else {
2230 unreachable!()
2231 };
2232
2233 Ok(Some(vec![list.clone(), merge.clone()]))
2234 }
2235
2236 fn lambda_parameters(
2237 &self,
2238 _step: usize,
2239 _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
2240 ) -> Result<crate::LambdaParametersProgress> {
2241 unimplemented!("mock_higher_order_function")
2242 }
2243
2244 fn return_field_from_args(
2245 &self,
2246 _args: HigherOrderReturnFieldArgs,
2247 ) -> Result<FieldRef> {
2248 unimplemented!("mock_higher_order_function")
2249 }
2250
2251 fn invoke_with_args(
2252 &self,
2253 _args: HigherOrderFunctionArgs,
2254 ) -> Result<ColumnarValue> {
2255 unimplemented!("mock_higher_order_function")
2256 }
2257 }
2258
2259 #[test]
2260 fn test_higher_order_function_user_defined_type_coercion() {
2261 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2262 signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2263 coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)],
2264 });
2265
2266 let new_fields = value_fields_with_higher_order_udf(
2267 &[
2268 ValueOrLambda::Value(Arc::new(Field::new_list(
2269 "",
2270 Field::new_list_field(DataType::Int32, false),
2271 false,
2272 ))),
2273 ValueOrLambda::Lambda(()),
2274 ],
2275 &fun,
2276 )
2277 .unwrap();
2278
2279 assert_eq!(
2281 new_fields,
2282 vec![
2283 ValueOrLambda::Value(Arc::new(Field::new_large_list(
2284 "",
2285 Field::new_list_field(DataType::Int32, false),
2286 false
2287 ))),
2288 ValueOrLambda::Lambda(()),
2289 ]
2290 )
2291 }
2292
2293 #[test]
2294 fn test_higher_order_function_coerce_values_for_lambdas() {
2295 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2296 signature: HigherOrderSignature::variadic_any(Volatility::Immutable),
2297 coerced_value_types: vec![],
2298 });
2299
2300 let new_fields = value_fields_with_higher_order_udf_and_lambdas(
2301 &[
2302 ValueOrLambda::Value(Arc::new(Field::new_list(
2303 "",
2304 Field::new_list_field(DataType::Float32, true),
2305 true,
2306 ))),
2307 ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, true))),
2308 ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
2309 ],
2310 &fun,
2311 )
2312 .unwrap();
2313
2314 assert_eq!(
2316 new_fields,
2317 vec![
2318 ValueOrLambda::Value(Arc::new(Field::new_list(
2319 "",
2320 Field::new_list_field(DataType::Float32, true),
2321 true,
2322 ))),
2323 ValueOrLambda::Value(Arc::new(Field::new("", DataType::Float32, true))),
2324 ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
2325 ]
2326 )
2327 }
2328
2329 #[test]
2330 fn test_higher_order_function_user_defined_type_coercion_bad_args() {
2331 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2332 signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2333 coerced_value_types: vec![DataType::Int32],
2334 });
2335
2336 let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
2337
2338 assert_contains!(
2339 err.to_string(),
2340 "mock_higher_order_function expects 1 value arguments, got 0"
2341 );
2342 }
2343
2344 #[test]
2345 fn test_higher_order_function_faulty_user_defined_type_coercion() {
2346 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2347 signature: HigherOrderSignature::user_defined(Volatility::Immutable),
2348 coerced_value_types: vec![DataType::Int32, DataType::Int32],
2349 });
2350
2351 let err = value_fields_with_higher_order_udf::<()>(
2352 &[ValueOrLambda::Value(Arc::new(Field::new(
2353 "",
2354 DataType::Int32,
2355 false,
2356 )))],
2357 &fun,
2358 )
2359 .unwrap_err();
2360
2361 assert_contains!(
2362 err.to_string(),
2363 "mock_higher_order_function coerce_value_types should have returned 1 items but returned 2"
2364 );
2365 }
2366
2367 #[test]
2368 fn test_higher_order_function_any_signature() {
2369 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2370 signature: HigherOrderSignature::any(1, Volatility::Immutable),
2371 coerced_value_types: vec![],
2372 });
2373
2374 let new_fields =
2375 value_fields_with_higher_order_udf(&[ValueOrLambda::Lambda(())], &fun)
2376 .unwrap();
2377
2378 assert_eq!(new_fields, vec![ValueOrLambda::Lambda(())])
2380 }
2381
2382 #[test]
2383 fn test_higher_order_function_any_signature_bad_args() {
2384 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2385 signature: HigherOrderSignature::any(1, Volatility::Immutable),
2386 coerced_value_types: vec![],
2387 });
2388
2389 let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
2390
2391 assert_contains!(
2392 err.to_string(),
2393 "The function 'mock_higher_order_function' expected 1 arguments but received 0"
2394 );
2395 }
2396
2397 #[test]
2398 fn test_higher_order_function_exact_signature() {
2399 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2400 signature: HigherOrderSignature::exact(
2401 vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2402 Volatility::Immutable,
2403 ),
2404 coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)],
2405 });
2406
2407 let new_fields = value_fields_with_higher_order_udf(
2408 &[
2409 ValueOrLambda::Value(Arc::new(Field::new_list(
2410 "",
2411 Field::new_list_field(DataType::Int32, false),
2412 false,
2413 ))),
2414 ValueOrLambda::Lambda(()),
2415 ],
2416 &fun,
2417 )
2418 .unwrap();
2419
2420 assert_eq!(
2422 new_fields,
2423 vec![
2424 ValueOrLambda::Value(Arc::new(Field::new_large_list(
2425 "",
2426 Field::new_list_field(DataType::Int32, false),
2427 false
2428 ))),
2429 ValueOrLambda::Lambda(()),
2430 ]
2431 )
2432 }
2433
2434 #[test]
2435 fn test_higher_order_function_exact_signature_wrong_value_count() {
2436 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2437 signature: HigherOrderSignature::exact(
2438 vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2439 Volatility::Immutable,
2440 ),
2441 coerced_value_types: vec![],
2442 });
2443
2444 let err = value_fields_with_higher_order_udf::<()>(
2445 &[ValueOrLambda::Lambda(()), ValueOrLambda::Lambda(())],
2446 &fun,
2447 )
2448 .unwrap_err();
2449
2450 assert_contains!(
2451 err.to_string(),
2452 "expected a value at position 0 but received a lambda"
2453 );
2454 }
2455
2456 #[test]
2457 fn test_higher_order_function_exact_signature_wrong_lambda_count() {
2458 let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF {
2459 signature: HigherOrderSignature::exact(
2460 vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())],
2461 Volatility::Immutable,
2462 ),
2463 coerced_value_types: vec![],
2464 });
2465
2466 let err = value_fields_with_higher_order_udf::<()>(
2467 &[
2468 ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, false))),
2469 ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, false))),
2470 ],
2471 &fun,
2472 )
2473 .unwrap_err();
2474
2475 assert_contains!(
2476 err.to_string(),
2477 "expected a lambda at position 1 but received a value"
2478 );
2479 }
2480}