Skip to main content

datafusion_physical_expr/expressions/
try_cast.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use 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/// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast
36#[derive(Debug, Clone, Eq)]
37pub struct TryCastExpr {
38    /// The expression to cast
39    expr: Arc<dyn PhysicalExpr>,
40    /// The target field.
41    ///
42    /// For a type-only cast (see [`TryCastExpr::new`]) this is a field
43    /// synthesized from the target data type alone and only its data type is
44    /// meaningful. For a cast built from an explicit field (see
45    /// [`TryCastExpr::new_with_target_field`]) its metadata is applied to the
46    /// output field as-is.
47    target_field: FieldRef,
48    /// Whether `target_field` was supplied by the caller (as opposed to being
49    /// synthesized from a `DataType`), and therefore whether its metadata
50    /// describes the output field exactly.
51    explicit_target: bool,
52}
53
54// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808
55impl PartialEq for TryCastExpr {
56    fn eq(&self, other: &Self) -> bool {
57        // Compare the semantically meaningful parts of the target field only:
58        // the field name never affects the output of this expression.
59        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        // Hash the metadata by iterating over sorted keys for deterministic ordering
70        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    /// Create a new `TryCastExpr` using only a `DataType`.
83    ///
84    /// This constructor creates a type-only cast where metadata is passed through
85    /// from the source expression (with extension type keys stripped).
86    /// TRY_CAST results are always nullable since failed casts return NULL.
87    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    /// Create a new `TryCastExpr` with an explicit target `FieldRef`.
96    ///
97    /// The provided `target_field` determines the output characteristics:
98    /// - The field's data type becomes the cast target type
99    /// - The field's metadata is used exactly as provided
100    ///
101    /// TRY_CAST results are always nullable since failed casts return NULL.
102    ///
103    /// See [`TryCastExpr::new`] for type-only casts where source metadata should
104    /// pass through.
105    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    /// The expression to cast
117    pub fn expr(&self) -> &Arc<dyn PhysicalExpr> {
118        &self.expr
119    }
120
121    /// The data type to cast to
122    pub fn cast_type(&self) -> &DataType {
123        self.target_field.data_type()
124    }
125
126    /// Explicit metadata for the output field, or `None` to pass through source metadata.
127    pub fn target_metadata(&self) -> Option<&HashMap<String, String>> {
128        self.explicit_target.then(|| self.target_field.metadata())
129    }
130
131    /// The target field this cast was constructed with.
132    ///
133    /// For a type-only cast this is a field synthesized from the target data
134    /// type alone; only its data type is meaningful. TRY_CAST results are
135    /// always nullable regardless of the target field's nullability.
136    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        // If metadata is explicit, we can build the field without source
167        // (though we still try to get source for the name)
168        let source_result = self.expr.return_field(input_schema);
169
170        if let Some(metadata) = self.target_metadata() {
171            // Explicit metadata: use it exactly, TRY_CAST is always nullable
172            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        // Pass-through metadata from source (stripping extension keys)
183        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) // TRY_CAST is always nullable
194                    .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    /// Reconstruct a [`TryCastExpr`] from its protobuf representation.
242    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
271/// Return a PhysicalExpression representing `expr` casted to
272/// `cast_type`, if any casting is needed.
273///
274/// Note that such casts may lose type information
275pub 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
290/// Return a PhysicalExpression representing `expr` casted to `target_field`,
291/// preserving any explicit field semantics such as metadata.
292///
293/// TRY_CAST results are always nullable since failed casts return NULL.
294///
295/// If the input expression already has the same data type, the target field
296/// has no explicit metadata constraints, and the source has no extension
297/// metadata to strip, the original expression is returned unchanged.
298pub 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    // Check if this is a "default" target field (type-only cast with no explicit
307    // metadata constraints). This is the field created by `into_nullable_field_ref()`
308    // when only a DataType is known.
309    let is_type_only = target_field.name().is_empty()
310        && target_field.is_nullable()
311        && target_field.metadata().is_empty();
312
313    // For same-type casts, we can skip creating a TryCastExpr only if:
314    // 1. The target is type-only (no explicit metadata)
315    // 2. The source has no extension metadata that needs to be stripped
316    // Otherwise we need the TryCastExpr to strip extension metadata from the source.
317    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    // For type-only casts, use TryCastExpr::new which preserves source metadata.
332    // For explicit target fields, use new_with_target_field which applies the target's
333    // metadata exactly.
334    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    // runs an end-to-end test of physical type cast
361    // 1. construct a record batch with a column "a" of type A
362    // 2. construct a physical expression of TRY_CAST(a AS B)
363    // 3. evaluate the expression
364    // 4. verify that the resulting expression is of type B
365    // 5. verify that the resulting values are downcastable and correct
366    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            // verify that we can construct the expression
374            let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
375
376            // verify that its display is correct
377            assert_eq!(
378                format!("TRY_CAST(a@0 AS {})", $TYPE),
379                format!("{}", expression)
380            );
381
382            // verify that the expression's type is correct
383            assert_eq!(expression.data_type(&schema)?, $TYPE);
384
385            // compute
386            let result = expression
387                .evaluate(&batch)?
388                .into_array(batch.num_rows())
389                .expect("Failed to convert to array");
390
391            // verify that the array's data_type is correct
392            assert_eq!(*result.data_type(), $TYPE);
393
394            // verify that the data itself is downcastable
395            let result = result
396                .as_any()
397                .downcast_ref::<$TYPEARRAY>()
398                .expect("failed to downcast");
399
400            // verify that the result itself is correct
401            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    // runs an end-to-end test of physical type cast
411    // 1. construct a record batch with a column "a" of type A
412    // 2. construct a physical expression of TRY_CAST(a AS B)
413    // 3. evaluate the expression
414    // 4. verify that the resulting expression is of type B
415    // 5. verify that the resulting values are downcastable and correct
416    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            // verify that we can construct the expression
425            let expression = try_cast(col("a", &schema)?, &schema, $TYPE)?;
426
427            // verify that its display is correct
428            assert_eq!(
429                format!("TRY_CAST(a@0 AS {})", $TYPE),
430                format!("{}", expression)
431            );
432
433            // verify that the expression's type is correct
434            assert_eq!(expression.data_type(&schema)?, $TYPE);
435
436            // compute
437            let result = expression
438                .evaluate(&batch)?
439                .into_array(batch.num_rows())
440                .expect("Failed to convert to array");
441
442            // verify that the array's data_type is correct
443            assert_eq!(*result.data_type(), $TYPE);
444
445            // verify that the len is correct
446            assert_eq!(result.len(), a_vec_len);
447
448            // verify that the data itself is downcastable
449            let result = result
450                .as_any()
451                .downcast_ref::<$TYPEARRAY>()
452                .expect("failed to downcast");
453
454            // verify that the result itself is correct
455            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        // try cast one decimal data type to another decimal data type
467        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        // TODO we should add function to create Decimal128Array with value and metadata
499        // https://github.com/apache/arrow-rs/issues/1009
500        let array: Vec<i128> = vec![1, 2, 3, 4, 5];
501        let decimal_array = create_decimal_array(&array, 10, 0);
502        // decimal to i8
503        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        // decimal to i16
519        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        // decimal to i32
536        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        // decimal to i64
553        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        // decimal to float32
570        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        // decimal to float64
587        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        // int8
609        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        // int16
619        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        // int32
629        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        // int64
639        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        // int64 to different scale
649        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        // float32
659        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        // float64
669        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        // Ensure a useful error happens at plan time if invalid casts are used
752        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    // create decimal array with the specified precision and scale
763    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        // Test numeric casting
780        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        // Test string casting
787        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        // When using field-aware cast, target's metadata should be used exactly
800        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        // TRY_CAST is always nullable
851        assert!(field.is_nullable());
852
853        Ok(())
854    }
855
856    #[test]
857    fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> {
858        // Target field metadata should be preserved exactly (no merging with source).
859        // TRY_CAST is always nullable regardless of target field's nullability.
860        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) // target says non-nullable
867                    .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            // Field name comes from source
876            assert_eq!(field.name(), "a");
877            assert_eq!(field.data_type(), &DataType::Int64);
878            // TRY_CAST is ALWAYS nullable (ignores target field's nullability)
879            assert!(field.is_nullable(), "TRY_CAST should always be nullable");
880            // Target metadata should be preserved exactly
881            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        // Type-only cast should strip extension keys but preserve other source metadata
898        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        // Extension keys should be stripped
917        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        // Non-extension metadata should pass through
926        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        // Field name preserved, type changed, always nullable
932        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        // TRY_CAST is always nullable even when source is non-nullable
942        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}