support union comparison for try-cast

friendlymatthew · friendlymatthew · commit 923d749bbb06 · 2026-04-01T11:50:26.000-04:00
diff --git a/datafusion/core/tests/sql/union_comparison.rs b/datafusion/core/tests/sql/union_comparison.rs
@@ -433,6 +433,68 @@ async fn test_union_cast_compatible_variant() -> Result<()> {
     Ok(())
 }
 
+/// Tests TRY_CAST from union type to scalar type.
+/// TRY_CAST should extract the matching variant and return NULL for non-matching ones,
+/// just like CAST, but also return NULL (instead of erroring) for cast failures.
+#[tokio::test]
+async fn test_union_try_cast() -> Result<()> {
+    let union_array = create_sparse_union_array(vec![
+        UnionValue::Int(Some(10)),
+        UnionValue::Str(Some("hello")),
+        UnionValue::Int(Some(30)),
+        UnionValue::Str(Some("42")),
+    ]);
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, false),
+        Field::new(
+            "val",
+            DataType::Union(
+                UnionFields::new(
+                    vec![0, 1],
+                    vec![
+                        Field::new("int", DataType::Int32, true),
+                        Field::new("str", DataType::Utf8, true),
+                    ],
+                ),
+                UnionMode::Sparse,
+            ),
+            true,
+        ),
+    ]));
+
+    let batch = RecordBatch::try_new(
+        schema,
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
+            Arc::new(union_array),
+        ],
+    )?;
+
+    let ctx = SessionContext::new();
+    ctx.register_batch("test", batch)?;
+
+    // TRY_CAST to INT should extract int variants, NULL for string variants
+    let df = ctx
+        .sql("SELECT id, TRY_CAST(val AS INT) as val_int FROM test")
+        .await?;
+    let results = df.collect().await?;
+
+    let expected = [
+        "+----+---------+",
+        "| id | val_int |",
+        "+----+---------+",
+        "| 1  | 10      |",
+        "| 2  |         |",
+        "| 3  | 30      |",
+        "| 4  |         |",
+        "+----+---------+",
+    ];
+    assert_batches_eq!(expected, &results);
+
+    Ok(())
+}
+
 /// Tests union-to-union equality comparison (supported via arrow-ord).
 #[tokio::test]
 async fn test_union_eq_same_union() -> Result<()> {
@@ -481,9 +543,7 @@ async fn test_union_eq_same_union() -> Result<()> {
     ctx.register_batch("test", batch)?;
 
     // Row 1: Int(10) = Int(10) -> true; Row 2: Str("hello") = Str("world") -> false
-    let df = ctx
-        .sql("SELECT id FROM test WHERE val1 = val2")
-        .await?;
+    let df = ctx.sql("SELECT id FROM test WHERE val1 = val2").await?;
     let results = df.collect().await?;
 
     let expected = ["+----+", "| id |", "+----+", "| 1  |", "+----+"];
diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs
@@ -78,9 +78,9 @@ use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter};
 use parquet::file::metadata::ParquetMetaData;
 use parquet::schema::types::SchemaDescriptor;
 
-use datafusion_common::{Result, ScalarValue};
 use datafusion_common::cast::as_boolean_array;
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
+use datafusion_common::{Result, ScalarValue};
 use datafusion_physical_expr::ScalarFunctionExpr;
 use datafusion_physical_expr::expressions::{Column, Literal};
 use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns};
@@ -398,7 +398,7 @@ impl<'schema> PushdownChecker<'schema> {
 /// See <https://github.com/datafusion-contrib/datafusion-variant> for the
 /// `datafusion-variant` crate that defines these UDFs.
 const VARIANT_UDF_NAMES: &[&str] = &[
-    "variant_get",       // variant_get, variant_get_str, variant_get_int, etc.
+    "variant_get", // variant_get, variant_get_str, variant_get_int, etc.
     "is_variant_null",
 ];
 
@@ -513,76 +513,71 @@ impl TreeNodeVisitor<'_> for PushdownChecker<'_> {
         //   - `metadata` — always needed (variant metadata dictionary)
         //   - `value` — always needed (fallback for non-shredded values)
         //   - `typed_value.<path...>` — the specific shredded field(s)
-        if let Some(func_expr) = node.as_any().downcast_ref::<ScalarFunctionExpr>() {
-            if is_variant_udf_name(func_expr.name()) {
-                if let Some(column) = func_expr
-                    .args()
-                    .first()
-                    .and_then(|a| a.as_any().downcast_ref::<Column>())
-                {
-                    let Ok(root_idx) = self.file_schema.index_of(column.name()) else {
-                        self.projected_columns = true;
-                        return Ok(TreeNodeRecursion::Jump);
-                    };
-
-                    // Extract the variant path from the second argument.
-                    // It can be a string literal or a list of string literals.
-                    let variant_path: Option<Vec<String>> =
-                        func_expr.args().get(1).and_then(|arg| {
-                            let lit = arg.as_any().downcast_ref::<Literal>()?;
-                            match lit.value() {
-                                ScalarValue::Utf8(Some(s))
-                                | ScalarValue::Utf8View(Some(s))
-                                | ScalarValue::LargeUtf8(Some(s)) => {
-                                    Some(vec![s.to_string()])
-                                }
-                                ScalarValue::List(arr) if !arr.is_null(0) => {
-                                    let values = arr.value(0);
-                                    let strings =
-                                        values.as_any().downcast_ref::<StringArray>()?;
-                                    let path: Vec<String> = (0..strings.len())
-                                        .filter_map(|i| {
-                                            strings
-                                                .is_valid(i)
-                                                .then(|| strings.value(i).to_string())
-                                        })
-                                        .collect();
-                                    Some(path)
-                                }
-                                _ => None,
-                            }
-                        });
-
-                    // Record struct field accesses for the variant sub-fields:
-                    // metadata, value, and typed_value.<path>
-                    self.struct_field_accesses.push(StructFieldAccess {
-                        root_index: root_idx,
-                        field_path: vec!["metadata".to_string()],
-                    });
-                    self.struct_field_accesses.push(StructFieldAccess {
-                        root_index: root_idx,
-                        field_path: vec!["value".to_string()],
-                    });
-
-                    if let Some(path) = variant_path {
-                        // typed_value.<field1>.<field2>...
-                        let mut typed_value_path = vec!["typed_value".to_string()];
-                        typed_value_path.extend(path);
-                        self.struct_field_accesses.push(StructFieldAccess {
-                            root_index: root_idx,
-                            field_path: typed_value_path,
-                        });
-                    } else {
-                        // Can't determine path statically — read entire typed_value
-                        self.struct_field_accesses.push(StructFieldAccess {
-                            root_index: root_idx,
-                            field_path: vec!["typed_value".to_string()],
-                        });
-                    }
+        if let Some(func_expr) = node.as_any().downcast_ref::<ScalarFunctionExpr>()
+            && is_variant_udf_name(func_expr.name())
+            && let Some(column) = func_expr
+                .args()
+                .first()
+                .and_then(|a| a.as_any().downcast_ref::<Column>())
+        {
+            let Ok(root_idx) = self.file_schema.index_of(column.name()) else {
+                self.projected_columns = true;
+                return Ok(TreeNodeRecursion::Jump);
+            };
 
-                    return Ok(TreeNodeRecursion::Jump);
-                }
+            // Extract the variant path from the second argument.
+            // It can be a string literal or a list of string literals.
+            let variant_path: Option<Vec<String>> =
+                func_expr.args().get(1).and_then(|arg| {
+                    let lit = arg.as_any().downcast_ref::<Literal>()?;
+                    match lit.value() {
+                        ScalarValue::Utf8(Some(s))
+                        | ScalarValue::Utf8View(Some(s))
+                        | ScalarValue::LargeUtf8(Some(s)) => {
+                            Some(vec![s.to_string()])
+                        }
+                        ScalarValue::List(arr) if !arr.is_null(0) => {
+                            let values = arr.value(0);
+                            let strings =
+                                values.as_any().downcast_ref::<StringArray>()?;
+                            let path: Vec<String> = (0..strings.len())
+                                .filter(|&i| strings.is_valid(i))
+                                .map(|i| strings.value(i).to_string())
+                                .collect();
+                            Some(path)
+                        }
+                        _ => None,
+                    }
+                });
+
+            // Record struct field accesses for the variant sub-fields:
+            // metadata, value, and typed_value.<path>
+            self.struct_field_accesses.push(StructFieldAccess {
+                root_index: root_idx,
+                field_path: vec!["metadata".to_string()],
+            });
+            self.struct_field_accesses.push(StructFieldAccess {
+                root_index: root_idx,
+                field_path: vec!["value".to_string()],
+            });
+
+            if let Some(path) = variant_path {
+                // typed_value.<field1>.<field2>...
+                let mut typed_value_path = vec!["typed_value".to_string()];
+                typed_value_path.extend(path);
+                self.struct_field_accesses.push(StructFieldAccess {
+                    root_index: root_idx,
+                    field_path: typed_value_path,
+                });
+            } else {
+                // Can't determine path statically — read entire typed_value
+                self.struct_field_accesses.push(StructFieldAccess {
+                    root_index: root_idx,
+                    field_path: vec!["typed_value".to_string()],
+                });
             }
+
+            return Ok(TreeNodeRecursion::Jump);
         }
 
         if let Some(column) = node.as_any().downcast_ref::<Column>()
diff --git a/datafusion/expr-common/src/columnar_value.rs b/datafusion/expr-common/src/columnar_value.rs
@@ -20,7 +20,7 @@
 use arrow::buffer::NullBuffer;
 use arrow::{
     array::{Array, ArrayRef, Date32Array, Date64Array, NullArray},
-    compute::{can_cast_types, CastOptions, kernels, max, min},
+    compute::{CastOptions, can_cast_types, kernels, max, min},
     datatypes::DataType,
     error::ArrowError,
     util::pretty::pretty_format_columns,
@@ -322,7 +322,8 @@ fn cast_array_by_name(
         _ if datafusion_common::nested_struct::requires_nested_struct_cast(
             array.data_type(),
             cast_type,
-        ) => {
+        ) =>
+        {
             datafusion_common::nested_struct::cast_column(array, cast_type, cast_options)
         }
         _ => {
@@ -346,7 +347,7 @@ fn cast_union_array(
     to_type: &DataType,
     cast_options: &CastOptions,
 ) -> Result<ArrayRef, ArrowError> {
-    use arrow::array::{make_array, new_null_array, UInt32Array, UnionArray};
+    use arrow::array::{UInt32Array, UnionArray, make_array, new_null_array};
     use arrow::compute::cast_with_options;
     use arrow::datatypes::UnionMode;
 
diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs
@@ -138,8 +138,18 @@ pub fn try_cast(
 ) -> Result<Arc<dyn PhysicalExpr>> {
     let expr_type = expr.data_type(input_schema)?;
     if expr_type == cast_type {
-        Ok(Arc::clone(&expr))
-    } else if can_cast_types(&expr_type, &cast_type) {
+        return Ok(Arc::clone(&expr));
+    }
+
+    // Check if cast is supported, with special handling for union types.
+    let can_cast = match &expr_type {
+        DataType::Union(fields, _) => fields
+            .iter()
+            .any(|(_, f)| can_cast_types(f.data_type(), &cast_type)),
+        _ => can_cast_types(&expr_type, &cast_type),
+    };
+
+    if can_cast {
         Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
     } else {
         not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}")
