From 5d2ac6a26277f970b36b2d01d95b09af917797fa Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sun, 18 Oct 2026 17:00:25 +0300 Subject: [PATCH] Handle tagged ARRAY insert into Variant storage --- datafusion/sql/src/statement.rs | 31 +++++++++++++++++-- datafusion/sql/tests/common/mod.rs | 34 ++++++++++++++++++++ datafusion/sql/tests/sql_integration.rs | 41 ++++++++++++++++++++++++- 3 files changed, 103 insertions(+), 3 deletions(-) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 99ed385393295..800fb7e1306ea 100644 --- a/datafusion/sql/src/statement.rs +++ b/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 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 diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index 6f40e2f1efb8a..a2c7065ea5253 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/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), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 31611f8f47300..1951afc3fc605 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/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)",