1use std::collections::HashMap;
19use std::fmt;
20use std::hash::Hash;
21use std::sync::Arc;
22
23use crate::PhysicalExpr;
24use arrow::compute;
25use arrow::compute::CastOptions;
26use arrow::datatypes::{DataType, Field, FieldRef, Schema};
27use arrow::record_batch::RecordBatch;
28use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
29use compute::can_cast_types;
30use datafusion_common::datatype::DataTypeExt;
31use datafusion_common::format::DEFAULT_FORMAT_OPTIONS;
32use datafusion_common::{Result, not_impl_err};
33use datafusion_expr::ColumnarValue;
34
35#[derive(Debug, Clone, Eq)]
37pub struct TryCastExpr {
38 expr: Arc<dyn PhysicalExpr>,
40 target_field: FieldRef,
48 explicit_target: bool,
52}
53
54impl PartialEq for TryCastExpr {
56 fn eq(&self, other: &Self) -> bool {
57 self.expr.eq(&other.expr)
60 && self.cast_type() == other.cast_type()
61 && self.target_metadata() == other.target_metadata()
62 }
63}
64
65impl Hash for TryCastExpr {
66 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
67 self.expr.hash(state);
68 self.cast_type().hash(state);
69 if let Some(metadata) = self.target_metadata() {
71 let mut entries: Vec<_> = metadata.iter().collect();
72 entries.sort_by_key(|(k, _)| *k);
73 for (k, v) in entries {
74 k.hash(state);
75 v.hash(state);
76 }
77 }
78 }
79}
80
81impl TryCastExpr {
82 pub fn new(expr: Arc<dyn PhysicalExpr>, cast_type: DataType) -> Self {
88 Self {
89 expr,
90 target_field: cast_type.into_nullable_field_ref(),
91 explicit_target: false,
92 }
93 }
94
95 pub fn new_with_target_field(
106 expr: Arc<dyn PhysicalExpr>,
107 target_field: FieldRef,
108 ) -> Self {
109 Self {
110 expr,
111 target_field,
112 explicit_target: true,
113 }
114 }
115
116 pub fn expr(&self) -> &Arc<dyn PhysicalExpr> {
118 &self.expr
119 }
120
121 pub fn cast_type(&self) -> &DataType {
123 self.target_field.data_type()
124 }
125
126 pub fn target_metadata(&self) -> Option<&HashMap<String, String>> {
128 self.explicit_target.then(|| self.target_field.metadata())
129 }
130
131 pub fn target_field(&self) -> &FieldRef {
137 &self.target_field
138 }
139}
140
141impl fmt::Display for TryCastExpr {
142 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
143 write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type())
144 }
145}
146
147impl PhysicalExpr for TryCastExpr {
148 fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
149 Ok(self.cast_type().clone())
150 }
151
152 fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
153 Ok(true)
154 }
155
156 fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
157 let value = self.expr.evaluate(batch)?;
158 let options = CastOptions {
159 safe: true,
160 format_options: DEFAULT_FORMAT_OPTIONS,
161 };
162 value.cast_to(self.cast_type(), Some(&options))
163 }
164
165 fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> {
166 let source_result = self.expr.return_field(input_schema);
169
170 if let Some(metadata) = self.target_metadata() {
171 let name = source_result
173 .as_ref()
174 .map(|f| f.name().to_string())
175 .unwrap_or_default();
176 return Ok(Arc::new(
177 Field::new(name, self.cast_type().clone(), true)
178 .with_metadata(metadata.clone()),
179 ));
180 }
181
182 source_result.map(|source_field| {
184 let mut metadata = source_field.metadata().clone();
185 metadata.remove(EXTENSION_TYPE_NAME_KEY);
186 metadata.remove(EXTENSION_TYPE_METADATA_KEY);
187
188 Arc::new(
189 source_field
190 .as_ref()
191 .clone()
192 .with_data_type(self.cast_type().clone())
193 .with_nullable(true) .with_metadata(metadata),
195 )
196 })
197 }
198
199 fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
200 vec![&self.expr]
201 }
202
203 fn with_new_children(
204 self: Arc<Self>,
205 children: Vec<Arc<dyn PhysicalExpr>>,
206 ) -> Result<Arc<dyn PhysicalExpr>> {
207 Ok(Arc::new(TryCastExpr {
208 expr: Arc::clone(&children[0]),
209 target_field: Arc::clone(&self.target_field),
210 explicit_target: self.explicit_target,
211 }))
212 }
213
214 fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
215 write!(f, "TRY_CAST(")?;
216 self.expr.fmt_sql(f)?;
217 write!(f, " AS {:?})", self.cast_type())
218 }
219
220 #[cfg(feature = "proto")]
221 fn try_to_proto(
222 &self,
223 ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
224 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
225 use datafusion_proto_models::protobuf;
226
227 Ok(Some(protobuf::PhysicalExprNode {
228 expr_id: None,
229 expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new(
230 protobuf::PhysicalTryCastNode {
231 expr: Some(Box::new(ctx.encode_child(&self.expr)?)),
232 arrow_type: Some(self.cast_type().try_into()?),
233 },
234 ))),
235 }))
236 }
237}
238
239#[cfg(feature = "proto")]
240impl TryCastExpr {
241 pub fn try_from_proto(
243 node: &datafusion_proto_models::protobuf::PhysicalExprNode,
244 ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
245 ) -> Result<Arc<dyn PhysicalExpr>> {
246 use datafusion_physical_expr_common::expect_expr_variant;
247 use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field;
248 use datafusion_proto_models::protobuf;
249
250 let try_cast = expect_expr_variant!(
251 node,
252 protobuf::physical_expr_node::ExprType::TryCast,
253 "TryCastExpr",
254 );
255 let expr = ctx.decode_required_expression(
256 try_cast.expr.as_deref(),
257 "TryCastExpr",
258 "expr",
259 )?;
260 let arrow_type = require_proto_field(
261 try_cast.arrow_type.as_ref(),
262 "TryCastExpr",
263 "arrow_type",
264 )?;
265 let cast_type: DataType = arrow_type.try_into()?;
266
267 Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
268 }
269}
270
271pub fn try_cast(
276 expr: Arc<dyn PhysicalExpr>,
277 input_schema: &Schema,
278 cast_type: DataType,
279) -> Result<Arc<dyn PhysicalExpr>> {
280 let expr_type = expr.data_type(input_schema)?;
281 if expr_type == cast_type {
282 Ok(Arc::clone(&expr))
283 } else if can_cast_types(&expr_type, &cast_type) {
284 Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
285 } else {
286 not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}")
287 }
288}
289
290pub fn try_cast_with_target_field(
299 expr: Arc<dyn PhysicalExpr>,
300 input_schema: &Schema,
301 target_field: &FieldRef,
302) -> Result<Arc<dyn PhysicalExpr>> {
303 let expr_type = expr.data_type(input_schema)?;
304 let cast_type = target_field.data_type();
305
306 let is_type_only = target_field.name().is_empty()
310 && target_field.is_nullable()
311 && target_field.metadata().is_empty();
312
313 if expr_type == *cast_type && is_type_only {
318 let source_field = expr.return_field(input_schema)?;
319 let has_extension_metadata = source_field
320 .metadata()
321 .contains_key(EXTENSION_TYPE_NAME_KEY);
322 if !has_extension_metadata {
323 return Ok(Arc::clone(&expr));
324 }
325 }
326
327 if !can_cast_types(&expr_type, cast_type) {
328 return not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}");
329 }
330
331 if is_type_only {
335 Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone())))
336 } else {
337 Ok(Arc::new(TryCastExpr::new_with_target_field(
338 expr,
339 Arc::clone(target_field),
340 )))
341 }
342}
343
344#[cfg(test)]
345mod tests {
346 use super::*;
347 use crate::expressions::col;
348 use arrow::array::{
349 Decimal128Array, Decimal128Builder, StringArray, Time64NanosecondArray,
350 };
351 use arrow::{
352 array::{
353 Array, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
354 Int64Array, TimestampNanosecondArray, UInt32Array,
355 },
356 datatypes::*,
357 };
358 use datafusion_physical_expr_common::physical_expr::fmt_sql;
359
360 macro_rules! generic_decimal_to_other_test_cast {
367 ($DECIMAL_ARRAY:ident, $A_TYPE:expr, $TYPEARRAY:ident, $TYPE:expr, $VEC:expr) => {{
368 let schema = Schema::new(vec![Field::new("a", $A_TYPE, true)]);
369 let batch = RecordBatch::try_new(
370 Arc::new(schema.clone()),
371 vec![Arc::new($DECIMAL_ARRAY)],
372 )?;
373 let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
375
376 assert_eq!(
378 format!("TRY_CAST(a@0 AS {})", $TYPE),
379 format!("{}", expression)
380 );
381
382 assert_eq!(expression.data_type(&schema)?, $TYPE);
384
385 let result = expression
387 .evaluate(&batch)?
388 .into_array(batch.num_rows())
389 .expect("Failed to convert to array");
390
391 assert_eq!(*result.data_type(), $TYPE);
393
394 let result = result
396 .as_any()
397 .downcast_ref::<$TYPEARRAY>()
398 .expect("failed to downcast");
399
400 for (i, x) in $VEC.iter().enumerate() {
402 match x {
403 Some(x) => assert_eq!(result.value(i), *x),
404 None => assert!(result.is_null(i)),
405 }
406 }
407 }};
408 }
409
410 macro_rules! generic_test_cast {
417 ($A_ARRAY:ident, $A_TYPE:expr, $A_VEC:expr, $TYPEARRAY:ident, $TYPE:expr, $VEC:expr) => {{
418 let schema = Schema::new(vec![Field::new("a", $A_TYPE, true)]);
419 let a_vec_len = $A_VEC.len();
420 let a = $A_ARRAY::from($A_VEC);
421 let batch =
422 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;
423
424 let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
426
427 assert_eq!(
429 format!("TRY_CAST(a@0 AS {})", $TYPE),
430 format!("{}", expression)
431 );
432
433 assert_eq!(expression.data_type(&schema)?, $TYPE);
435
436 let result = expression
438 .evaluate(&batch)?
439 .into_array(batch.num_rows())
440 .expect("Failed to convert to array");
441
442 assert_eq!(*result.data_type(), $TYPE);
444
445 assert_eq!(result.len(), a_vec_len);
447
448 let result = result
450 .as_any()
451 .downcast_ref::<$TYPEARRAY>()
452 .expect("failed to downcast");
453
454 for (i, x) in $VEC.iter().enumerate() {
456 match x {
457 Some(x) => assert_eq!(result.value(i), *x),
458 None => assert!(result.is_null(i)),
459 }
460 }
461 }};
462 }
463
464 #[test]
465 fn test_try_cast_decimal_to_decimal() -> Result<()> {
466 let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
468 let decimal_array = create_decimal_array(&array, 10, 3);
469 generic_decimal_to_other_test_cast!(
470 decimal_array,
471 DataType::Decimal128(10, 3),
472 Decimal128Array,
473 DataType::Decimal128(20, 6),
474 [
475 Some(1_234_000),
476 Some(2_222_000),
477 Some(3_000),
478 Some(4_000_000),
479 Some(5_000_000),
480 None
481 ]
482 );
483
484 let decimal_array = create_decimal_array(&array, 10, 3);
485 generic_decimal_to_other_test_cast!(
486 decimal_array,
487 DataType::Decimal128(10, 3),
488 Decimal128Array,
489 DataType::Decimal128(10, 2),
490 [Some(123), Some(222), Some(0), Some(400), Some(500), None]
491 );
492
493 Ok(())
494 }
495
496 #[test]
497 fn test_try_cast_decimal_to_numeric() -> Result<()> {
498 let array: Vec<i128> = vec![1, 2, 3, 4, 5];
501 let decimal_array = create_decimal_array(&array, 10, 0);
502 generic_decimal_to_other_test_cast!(
504 decimal_array,
505 DataType::Decimal128(10, 0),
506 Int8Array,
507 DataType::Int8,
508 [
509 Some(1_i8),
510 Some(2_i8),
511 Some(3_i8),
512 Some(4_i8),
513 Some(5_i8),
514 None
515 ]
516 );
517
518 let decimal_array = create_decimal_array(&array, 10, 0);
520 generic_decimal_to_other_test_cast!(
521 decimal_array,
522 DataType::Decimal128(10, 0),
523 Int16Array,
524 DataType::Int16,
525 [
526 Some(1_i16),
527 Some(2_i16),
528 Some(3_i16),
529 Some(4_i16),
530 Some(5_i16),
531 None
532 ]
533 );
534
535 let decimal_array = create_decimal_array(&array, 10, 0);
537 generic_decimal_to_other_test_cast!(
538 decimal_array,
539 DataType::Decimal128(10, 0),
540 Int32Array,
541 DataType::Int32,
542 [
543 Some(1_i32),
544 Some(2_i32),
545 Some(3_i32),
546 Some(4_i32),
547 Some(5_i32),
548 None
549 ]
550 );
551
552 let decimal_array = create_decimal_array(&array, 10, 0);
554 generic_decimal_to_other_test_cast!(
555 decimal_array,
556 DataType::Decimal128(10, 0),
557 Int64Array,
558 DataType::Int64,
559 [
560 Some(1_i64),
561 Some(2_i64),
562 Some(3_i64),
563 Some(4_i64),
564 Some(5_i64),
565 None
566 ]
567 );
568
569 let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
571 let decimal_array = create_decimal_array(&array, 10, 3);
572 generic_decimal_to_other_test_cast!(
573 decimal_array,
574 DataType::Decimal128(10, 3),
575 Float32Array,
576 DataType::Float32,
577 [
578 Some(1.234_f32),
579 Some(2.222_f32),
580 Some(0.003_f32),
581 Some(4.0_f32),
582 Some(5.0_f32),
583 None
584 ]
585 );
586 let decimal_array = create_decimal_array(&array, 20, 6);
588 generic_decimal_to_other_test_cast!(
589 decimal_array,
590 DataType::Decimal128(20, 6),
591 Float64Array,
592 DataType::Float64,
593 [
594 Some(0.001234_f64),
595 Some(0.002222_f64),
596 Some(0.000003_f64),
597 Some(0.004_f64),
598 Some(0.005_f64),
599 None
600 ]
601 );
602
603 Ok(())
604 }
605
606 #[test]
607 fn test_try_cast_numeric_to_decimal() -> Result<()> {
608 generic_test_cast!(
610 Int8Array,
611 DataType::Int8,
612 vec![1, 2, 3, 4, 5],
613 Decimal128Array,
614 DataType::Decimal128(3, 0),
615 [Some(1), Some(2), Some(3), Some(4), Some(5)]
616 );
617
618 generic_test_cast!(
620 Int16Array,
621 DataType::Int16,
622 vec![1, 2, 3, 4, 5],
623 Decimal128Array,
624 DataType::Decimal128(5, 0),
625 [Some(1), Some(2), Some(3), Some(4), Some(5)]
626 );
627
628 generic_test_cast!(
630 Int32Array,
631 DataType::Int32,
632 vec![1, 2, 3, 4, 5],
633 Decimal128Array,
634 DataType::Decimal128(10, 0),
635 [Some(1), Some(2), Some(3), Some(4), Some(5)]
636 );
637
638 generic_test_cast!(
640 Int64Array,
641 DataType::Int64,
642 vec![1, 2, 3, 4, 5],
643 Decimal128Array,
644 DataType::Decimal128(20, 0),
645 [Some(1), Some(2), Some(3), Some(4), Some(5)]
646 );
647
648 generic_test_cast!(
650 Int64Array,
651 DataType::Int64,
652 vec![1, 2, 3, 4, 5],
653 Decimal128Array,
654 DataType::Decimal128(20, 2),
655 [Some(100), Some(200), Some(300), Some(400), Some(500)]
656 );
657
658 generic_test_cast!(
660 Float32Array,
661 DataType::Float32,
662 vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50],
663 Decimal128Array,
664 DataType::Decimal128(10, 2),
665 [Some(150), Some(250), Some(300), Some(112), Some(550)]
666 );
667
668 generic_test_cast!(
670 Float64Array,
671 DataType::Float64,
672 vec![1.5, 2.5, 3.0, 1.123_456_8, 5.50],
673 Decimal128Array,
674 DataType::Decimal128(20, 4),
675 [
676 Some(15000),
677 Some(25000),
678 Some(30000),
679 Some(11235),
680 Some(55000)
681 ]
682 );
683 Ok(())
684 }
685
686 #[test]
687 fn test_cast_i32_u32() -> Result<()> {
688 generic_test_cast!(
689 Int32Array,
690 DataType::Int32,
691 vec![1, 2, 3, 4, 5],
692 UInt32Array,
693 DataType::UInt32,
694 [
695 Some(1_u32),
696 Some(2_u32),
697 Some(3_u32),
698 Some(4_u32),
699 Some(5_u32)
700 ]
701 );
702 Ok(())
703 }
704
705 #[test]
706 fn test_cast_i32_utf8() -> Result<()> {
707 generic_test_cast!(
708 Int32Array,
709 DataType::Int32,
710 vec![1, 2, 3, 4, 5],
711 StringArray,
712 DataType::Utf8,
713 [Some("1"), Some("2"), Some("3"), Some("4"), Some("5")]
714 );
715 Ok(())
716 }
717
718 #[test]
719 fn test_try_cast_utf8_i32() -> Result<()> {
720 generic_test_cast!(
721 StringArray,
722 DataType::Utf8,
723 vec!["a", "2", "3", "b", "5"],
724 Int32Array,
725 DataType::Int32,
726 [None, Some(2), Some(3), None, Some(5)]
727 );
728 Ok(())
729 }
730
731 #[test]
732 fn test_cast_i64_t64() -> Result<()> {
733 let original = vec![1, 2, 3, 4, 5];
734 let expected: Vec<Option<i64>> = original
735 .iter()
736 .map(|i| Some(Time64NanosecondArray::from(vec![*i]).value(0)))
737 .collect();
738 generic_test_cast!(
739 Int64Array,
740 DataType::Int64,
741 original,
742 TimestampNanosecondArray,
743 DataType::Timestamp(TimeUnit::Nanosecond, None),
744 expected
745 );
746 Ok(())
747 }
748
749 #[test]
750 fn invalid_cast() {
751 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
753
754 let result = try_cast(
755 col("a", &schema).unwrap(),
756 &schema,
757 DataType::Interval(IntervalUnit::MonthDayNano),
758 );
759 result.expect_err("expected Invalid TRY_CAST");
760 }
761
762 fn create_decimal_array(array: &[i128], precision: u8, scale: i8) -> Decimal128Array {
764 let mut decimal_builder = Decimal128Builder::with_capacity(array.len());
765 for value in array {
766 decimal_builder.append_value(*value);
767 }
768 decimal_builder.append_null();
769 decimal_builder
770 .finish()
771 .with_precision_and_scale(precision, scale)
772 .unwrap()
773 }
774
775 #[test]
776 fn test_fmt_sql() -> Result<()> {
777 let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
778
779 let expr = try_cast(col("a", &schema)?, &schema, DataType::Int64)?;
781 let display_string = expr.to_string();
782 assert_eq!(display_string, "TRY_CAST(a@0 AS Int64)");
783 let sql_string = fmt_sql(expr.as_ref()).to_string();
784 assert_eq!(sql_string, "TRY_CAST(a AS Int64)");
785
786 let schema = Schema::new(vec![Field::new("b", DataType::Utf8, true)]);
788 let expr = try_cast(col("b", &schema)?, &schema, DataType::Int32)?;
789 let display_string = expr.to_string();
790 assert_eq!(display_string, "TRY_CAST(b@0 AS Int32)");
791 let sql_string = fmt_sql(expr.as_ref()).to_string();
792 assert_eq!(sql_string, "TRY_CAST(b AS Int32)");
793
794 Ok(())
795 }
796
797 #[test]
798 fn field_aware_try_cast_uses_exact_target_metadata() -> Result<()> {
799 let source_meta = HashMap::from([
801 (
802 EXTENSION_TYPE_NAME_KEY.to_string(),
803 "source.type".to_string(),
804 ),
805 ("source_key".to_string(), "source_value".to_string()),
806 ]);
807 let target_meta = HashMap::from([
808 (
809 EXTENSION_TYPE_NAME_KEY.to_string(),
810 "target.type".to_string(),
811 ),
812 (
813 EXTENSION_TYPE_METADATA_KEY.to_string(),
814 "target_ext_meta".to_string(),
815 ),
816 ("target_key".to_string(), "target_value".to_string()),
817 ]);
818 let schema = Schema::new(vec![
819 Field::new("a", DataType::FixedSizeBinary(16), false)
820 .with_metadata(source_meta),
821 ]);
822
823 let target_field =
824 Arc::new(Field::new("b", DataType::Utf8, true).with_metadata(target_meta));
825 let expr = TryCastExpr::new_with_target_field(
826 col("a", &schema)?,
827 Arc::clone(&target_field),
828 );
829
830 let field = expr.return_field(&schema)?;
831 assert_eq!(
832 field.metadata().get(EXTENSION_TYPE_NAME_KEY),
833 Some(&"target.type".to_string()),
834 "Field-aware try_cast should use target's extension type name"
835 );
836 assert_eq!(
837 field.metadata().get(EXTENSION_TYPE_METADATA_KEY),
838 Some(&"target_ext_meta".to_string()),
839 "Field-aware try_cast should use target's extension type metadata"
840 );
841 assert!(
842 field.metadata().get("source_key").is_none(),
843 "Field-aware try_cast should NOT preserve source metadata"
844 );
845 assert_eq!(
846 field.metadata().get("target_key"),
847 Some(&"target_value".to_string()),
848 "Field-aware try_cast should preserve target's non-extension metadata"
849 );
850 assert!(field.is_nullable());
852
853 Ok(())
854 }
855
856 #[test]
857 fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> {
858 let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]);
861
862 for child_nullable in [true, false] {
863 let schema =
864 Schema::new(vec![Field::new("a", DataType::Int32, child_nullable)]);
865 let target_field = Arc::new(
866 Field::new("cast_target", DataType::Int64, false) .with_metadata(metadata.clone()),
868 );
869 let expr = TryCastExpr::new_with_target_field(
870 col("a", &schema)?,
871 Arc::clone(&target_field),
872 );
873
874 let field = expr.return_field(&schema)?;
875 assert_eq!(field.name(), "a");
877 assert_eq!(field.data_type(), &DataType::Int64);
878 assert!(field.is_nullable(), "TRY_CAST should always be nullable");
880 assert_eq!(
882 field.metadata().get("target_meta"),
883 Some(&"1".to_string()),
884 "Target metadata should be preserved exactly"
885 );
886 assert!(
887 expr.nullable(&schema)?,
888 "TRY_CAST should always be nullable"
889 );
890 }
891
892 Ok(())
893 }
894
895 #[test]
896 fn type_only_try_cast_strips_extension_keys() -> Result<()> {
897 let source_meta = HashMap::from([
899 (
900 EXTENSION_TYPE_NAME_KEY.to_string(),
901 "source.extension".to_string(),
902 ),
903 (
904 EXTENSION_TYPE_METADATA_KEY.to_string(),
905 "ext_meta".to_string(),
906 ),
907 ("custom_key".to_string(), "custom_value".to_string()),
908 ]);
909 let schema = Schema::new(vec![
910 Field::new("a", DataType::Int32, false).with_metadata(source_meta),
911 ]);
912
913 let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64);
914 let field = expr.return_field(&schema)?;
915
916 assert!(
918 field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(),
919 "Type-only try_cast should strip extension type name"
920 );
921 assert!(
922 field.metadata().get(EXTENSION_TYPE_METADATA_KEY).is_none(),
923 "Type-only try_cast should strip extension type metadata"
924 );
925 assert_eq!(
927 field.metadata().get("custom_key"),
928 Some(&"custom_value".to_string()),
929 "Type-only try_cast should preserve non-extension metadata"
930 );
931 assert_eq!(field.name(), "a");
933 assert_eq!(field.data_type(), &DataType::Int64);
934 assert!(field.is_nullable());
935
936 Ok(())
937 }
938
939 #[test]
940 fn type_only_try_cast_is_always_nullable() -> Result<()> {
941 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
943 let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64);
944
945 let field = expr.return_field(&schema)?;
946
947 assert_eq!(field.name(), "a");
948 assert_eq!(field.data_type(), &DataType::Int64);
949 assert!(field.is_nullable(), "TRY_CAST should always be nullable");
950 assert!(
951 expr.nullable(&schema)?,
952 "TRY_CAST should always be nullable"
953 );
954
955 Ok(())
956 }
957}
958
959#[cfg(all(test, feature = "proto"))]
960mod proto_tests {
961 use super::*;
962 use crate::expressions::{Column, col};
963 use crate::proto_test_util::{
964 StubDecoder, StubEncoder, UnreachableDecoder, column_node,
965 };
966 use arrow::datatypes::Field;
967 use datafusion_common::DataFusionError;
968 use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
969 use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
970 use datafusion_proto_models::datafusion_common::ArrowType;
971 use datafusion_proto_models::protobuf::{
972 PhysicalExprNode, PhysicalTryCastNode, physical_expr_node,
973 };
974
975 fn try_cast_fixture() -> TryCastExpr {
976 let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
977 TryCastExpr::new(col("a", &schema).unwrap(), DataType::Int32)
978 }
979
980 fn int32_arrow_type() -> ArrowType {
981 (&DataType::Int32).try_into().unwrap()
982 }
983
984 fn try_cast_node(
985 expr: Option<Box<PhysicalExprNode>>,
986 arrow_type: Option<ArrowType>,
987 ) -> PhysicalExprNode {
988 PhysicalExprNode {
989 expr_id: None,
990 expr_type: Some(physical_expr_node::ExprType::TryCast(Box::new(
991 PhysicalTryCastNode { expr, arrow_type },
992 ))),
993 }
994 }
995
996 #[test]
997 fn try_to_proto_encodes_try_cast_expr() {
998 let try_cast = try_cast_fixture();
999 let encoder = StubEncoder::ok();
1000 let ctx = PhysicalExprEncodeCtx::new(&encoder);
1001
1002 let node = try_cast
1003 .try_to_proto(&ctx)
1004 .unwrap()
1005 .expect("TryCastExpr should encode to Some(node)");
1006
1007 assert!(node.expr_id.is_none());
1008 let try_cast_node = match node.expr_type {
1009 Some(physical_expr_node::ExprType::TryCast(boxed)) => *boxed,
1010 other => panic!("expected a TryCastExpr node, got {other:?}"),
1011 };
1012 assert!(try_cast_node.expr.is_some());
1013
1014 let arrow_type = try_cast_node
1015 .arrow_type
1016 .as_ref()
1017 .expect("try cast type should be encoded");
1018 let data_type: DataType = arrow_type.try_into().unwrap();
1019 assert_eq!(data_type, DataType::Int32);
1020 }
1021
1022 #[test]
1023 fn try_to_proto_propagates_child_encode_error() {
1024 let try_cast = try_cast_fixture();
1025 let encoder = StubEncoder::failing_on(1);
1026 let ctx = PhysicalExprEncodeCtx::new(&encoder);
1027 let err = try_cast.try_to_proto(&ctx).unwrap_err();
1028 assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
1029 }
1030
1031 #[test]
1032 fn try_from_proto_decodes_try_cast_expr() {
1033 let node =
1034 try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type()));
1035 let schema = Schema::empty();
1036 let decoder = StubDecoder::ok();
1037 let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
1038
1039 let decoded = TryCastExpr::try_from_proto(&node, &ctx).unwrap();
1040 let try_cast = decoded
1041 .downcast_ref::<TryCastExpr>()
1042 .expect("decoded expr should be a TryCastExpr");
1043
1044 assert_eq!(try_cast.cast_type(), &DataType::Int32);
1045 assert!(try_cast.expr().downcast_ref::<Column>().is_some());
1046 }
1047
1048 #[test]
1049 fn try_from_proto_rejects_non_try_cast_node() {
1050 let node = column_node("a");
1051 let schema = Schema::empty();
1052 let decoder = UnreachableDecoder;
1053 let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
1054
1055 let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
1056 assert!(
1057 matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a TryCastExpr"))
1058 );
1059 }
1060
1061 #[test]
1062 fn try_from_proto_rejects_missing_expr() {
1063 let node = try_cast_node(None, Some(int32_arrow_type()));
1064 let schema = Schema::empty();
1065 let decoder = UnreachableDecoder;
1066 let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
1067
1068 let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
1069 assert!(
1070 matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'expr'"))
1071 );
1072 }
1073
1074 #[test]
1075 fn try_from_proto_rejects_missing_arrow_type() {
1076 let node = try_cast_node(Some(Box::new(column_node("a"))), None);
1077 let schema = Schema::empty();
1078 let decoder = StubDecoder::ok();
1079 let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
1080
1081 let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
1082 assert!(
1083 matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'arrow_type'"))
1084 );
1085 }
1086
1087 #[test]
1088 fn try_from_proto_propagates_child_decode_error() {
1089 let node =
1090 try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type()));
1091 let schema = Schema::empty();
1092 let decoder = StubDecoder::failing_on(1);
1093 let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
1094 let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err();
1095 assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
1096 }
1097}