Handle tagged ARRAY insert into Variant storage

Author: osipovartem

datafusion/sql/src/statement.rs

@@ -43,6 +43,7 @@ use datafusion_common::{
 use datafusion_expr::dml::{
     CopyTo, InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp,
 };
+use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check;
 use datafusion_expr::logical_plan::DdlStatement;
 use datafusion_expr::logical_plan::builder::project;
@@ -2956,8 +2957,34 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 let target_field = table_schema.field(i);
                 let expr = match value_index {
                     Some(v) => {
-                        Expr::Column(Column::from(source.schema().qualified_field(v)))
-                            .cast_to(target_field.data_type(), source.schema())?
+                        let source_field = source.schema().field(v);
+                        let expr =
+                            Expr::Column(Column::from(source.schema().qualified_field(v)));
+                        let expr = if source_field
+                            .metadata()
+                            .get("embucket.snowflake_logical_type")
+                            .is_some_and(|value| value == "ARRAY")
+                            && target_field.extension_type_name()
+                                == Some("arrow.parquet.variant")
+                            && source_field.data_type() != target_field.data_type()
+                        {
+                            let to_variant = self
+                                .context_provider
+                                .get_function_meta("to_variant")
+                                .ok_or_else(|| {
+                                    DataFusionError::Plan(
+                                        "TO_VARIANT is required to insert ARRAY into VARIANT storage"
+                                            .to_string(),
+                                    )
+                                })?;
+                            Expr::ScalarFunction(ScalarFunction::new_udf(
+                                to_variant,
+                                vec![expr],
+                            ))
+                        } else {
+                            expr
+                        };
+                        expr.cast_to(target_field.data_type(), source.schema())?
                     }
                     // The value is not specified. Fill in the default value for the column.
                     None => table_source

datafusion/sql/tests/common/mod.rs

@@ -114,6 +114,13 @@ impl MockSessionState {
     }
 }
 
+pub(crate) fn variant_storage_type() -> DataType {
+    DataType::Struct(Fields::from(vec![
+        Field::new("metadata", DataType::Binary, false),
+        Field::new("value", DataType::Binary, false),
+    ]))
+}
+
 pub(crate) struct MockContextProvider {
     pub(crate) state: MockSessionState,
 }
@@ -196,6 +203,33 @@ impl ContextProvider for MockContextProvider {
                     .into(),
                 ),
             ])),
+            "untyped_array_source" => Ok(Schema::new(vec![
+                Field::new("items", DataType::Utf8, true).with_metadata(
+                    [(
+                        "embucket.snowflake_logical_type".to_string(),
+                        "ARRAY".to_string(),
+                    )]
+                    .into(),
+                ),
+            ])),
+            "plain_string_source" => {
+                Ok(Schema::new(vec![Field::new("items", DataType::Utf8, true)]))
+            }
+            "variant_array_target" => Ok(Schema::new(vec![
+                Field::new("items", variant_storage_type(), true).with_metadata(
+                    [
+                        (
+                            "ARROW:extension:name".to_string(),
+                            "arrow.parquet.variant".to_string(),
+                        ),
+                        (
+                            "embucket.snowflake_logical_type".to_string(),
+                            "ARRAY".to_string(),
+                        ),
+                    ]
+                    .into(),
+                ),
+            ])),
             "orders" => Ok(Schema::new(vec![
                 Field::new("order_id", DataType::UInt32, false),
                 Field::new("o_orderkey", DataType::UInt32, false),

datafusion/sql/tests/sql_integration.rs

@@ -25,7 +25,7 @@ use std::sync::Arc;
 use std::vec;
 
 use arrow::datatypes::{TimeUnit::Nanosecond, *};
-use common::MockContextProvider;
+use common::{MockContextProvider, variant_storage_type};
 use datafusion_common::{
     DFSchema, DataFusionError, Result, ScalarValue, assert_contains,
 };
@@ -842,6 +842,45 @@ fn plan_insert_does_not_promise_ordinary_target_field_metadata() {
     assert!(dml.input.schema().field(0).metadata().is_empty());
 }
 
+#[test]
+fn plan_insert_converts_tagged_array_to_variant_storage() -> Result<()> {
+    let state = mock_session_state().with_scalar_function(Arc::new(make_udf(
+        "to_variant",
+        vec![DataType::Utf8],
+        variant_storage_type(),
+    )));
+    let plan = logical_plan_from_state(
+        "INSERT INTO variant_array_target SELECT items FROM untyped_array_source",
+        &SnowflakeDialect {},
+        ParserOptions::default(),
+        state,
+    )?;
+    assert_contains!(plan.display_indent().to_string(), "to_variant(");
+    let LogicalPlan::Dml(dml) = &plan else {
+        panic!("expected DML plan");
+    };
+    assert_eq!(
+        dml.input.schema().field(0).extension_type_name(),
+        Some("arrow.parquet.variant")
+    );
+    Ok(())
+}
+
+#[test]
+fn plan_insert_does_not_convert_plain_string_to_array() {
+    let error = logical_plan_from_state(
+        "INSERT INTO variant_array_target SELECT items FROM plain_string_source",
+        &SnowflakeDialect {},
+        ParserOptions::default(),
+        mock_session_state(),
+    )
+    .unwrap_err();
+    assert_contains!(
+        error.to_string(),
+        "Cannot automatically convert Utf8 to Struct"
+    );
+}
+
 #[rstest]
 #[case::duplicate_columns(
     "INSERT INTO test_decimal (id, price, price) VALUES (1, 2, 3), (4, 5, 6)",