diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 800fb7e1306e..404299ac3f83 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -2960,27 +2960,47 @@ impl SqlToRel<'_, S> { let source_field = source.schema().field(v); let expr = Expr::Column(Column::from(source.schema().qualified_field(v))); - let expr = if source_field + let source_is_array = source_field .metadata() .get("embucket.snowflake_logical_type") - .is_some_and(|value| value == "ARRAY") - && target_field.extension_type_name() - == Some("arrow.parquet.variant") + .is_some_and(|value| value == "ARRAY"); + let target_is_array = target_field + .metadata() + .get("embucket.snowflake_logical_type") + .is_some_and(|value| value == "ARRAY"); + let source_is_text = matches!( + source_field.data_type(), + arrow::datatypes::DataType::Utf8 + | arrow::datatypes::DataType::Utf8View + | arrow::datatypes::DataType::LargeUtf8 + ); + let source_is_list = matches!( + source_field.data_type(), + arrow::datatypes::DataType::List(_) + | arrow::datatypes::DataType::LargeList(_) + | arrow::datatypes::DataType::FixedSizeList(_, _) + ); + let expr = if target_field.extension_type_name() + == Some("arrow.parquet.variant") && source_field.data_type() != target_field.data_type() + && (source_is_array + || source_is_list + || (target_is_array && source_is_text)) { + let function_name = if target_is_array && source_is_text { + "__rustice_array_to_variant" + } else { + "to_variant" + }; let to_variant = self .context_provider - .get_function_meta("to_variant") + .get_function_meta(function_name) .ok_or_else(|| { - DataFusionError::Plan( - "TO_VARIANT is required to insert ARRAY into VARIANT storage" - .to_string(), - ) + DataFusionError::Plan(format!( + "{function_name} is required to insert ARRAY into VARIANT storage" + )) })?; - Expr::ScalarFunction(ScalarFunction::new_udf( - to_variant, - vec![expr], - )) + Expr::ScalarFunction(ScalarFunction::new_udf(to_variant, vec![expr])) } else { expr }; diff --git a/datafusion/sql/src/values.rs b/datafusion/sql/src/values.rs index a1df1fe1b18d..cc466be7f2e1 100644 --- a/datafusion/sql/src/values.rs +++ b/datafusion/sql/src/values.rs @@ -15,11 +15,15 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{DFSchema, Result, not_impl_err}; -use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; +use arrow::datatypes::DataType; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion_common::{DFSchema, DataFusionError, Result, not_impl_err, plan_err}; +use datafusion_expr::expr::ScalarFunction; +use datafusion_expr::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder, ScalarUDF}; use sqlparser::ast::Values as SQLValues; impl SqlToRel<'_, S> { @@ -40,6 +44,7 @@ impl SqlToRel<'_, S> { } let empty_schema = Arc::new(DFSchema::empty()); + let table_schema = planner_context.set_table_schema(None); let values = rows .into_iter() .map(|row| { @@ -48,13 +53,169 @@ impl SqlToRel<'_, S> { .map(|v| self.sql_to_expr(v, &empty_schema, planner_context)) .collect::>>() }) - .collect::>>()?; + .collect::>>(); + planner_context.set_table_schema(table_schema.clone()); + let mut values = values?; - let schema = planner_context.table_schema().unwrap_or(empty_schema); - if schema.fields().is_empty() { + let schema = table_schema.unwrap_or(empty_schema); + let array_columns = schema + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + (field.extension_type_name() == Some("arrow.parquet.variant") + && field + .metadata() + .get("embucket.snowflake_logical_type") + .is_some_and(|value| value == "ARRAY")) + .then_some(index) + }) + .collect::>(); + if array_columns.is_empty() { + return if schema.fields().is_empty() { + LogicalPlanBuilder::values(values)?.build() + } else { + LogicalPlanBuilder::values_with_schema(values, &schema)?.build() + }; + } + for row in &values { + for expr in row { + let mut has_subquery = false; + expr.apply(|node| { + if matches!( + node, + Expr::ScalarSubquery(_) | Expr::Exists(_) | Expr::InSubquery(_) + ) { + has_subquery = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + })?; + if has_subquery { + return plan_err!( + "Subqueries in ARRAY INSERT VALUES are not supported; use INSERT ... SELECT" + ); + } + } + } + let mut value_fields = schema.fields().to_vec(); + let mut per_row_columns = Vec::new(); + for index in array_columns.iter().copied() { + let mut common_type = None; + let mut mixed = false; + for row in &values { + let Some(expr) = row.get(index) else { + mixed = true; + break; + }; + let (_, source_field) = expr.to_field(&schema)?; + if source_field + .metadata() + .get("snowflake.structured_array") + .is_some_and(|value| value == "true") + { + mixed = true; + break; + } + let data_type = source_field.data_type().clone(); + if data_type == DataType::Null { + continue; + } + if common_type + .as_ref() + .is_some_and(|previous| previous != &data_type) + { + mixed = true; + break; + } + common_type = Some(data_type); + } + let batch_type = + (!mixed) + .then_some(common_type) + .flatten() + .filter(|data_type| { + matches!( + data_type, + DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) + | DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + ) + }); + if let Some(data_type) = batch_type { + // Let the INSERT projection convert homogeneous rows as a batch. + value_fields[index] = Arc::new( + schema + .field(index) + .as_ref() + .clone() + .with_data_type(data_type) + .with_metadata(HashMap::new()), + ); + } else { + per_row_columns.push(index); + } + } + if !per_row_columns.is_empty() { + let mut conversions: HashMap<&'static str, Arc> = HashMap::new(); + for row in &mut values { + for &index in &per_row_columns { + let Some(expr) = row.get_mut(index) else { + continue; + }; + let function_name = if matches!(expr, Expr::Placeholder(_)) { + "__rustice_bound_array_to_variant" + } else { + match expr.get_type(&schema)? { + DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) => "to_variant", + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => { + "__rustice_array_to_variant" + } + _ => continue, + } + }; + let conversion = if let Some(conversion) = + conversions.get(function_name) + { + Arc::clone(conversion) + } else { + let conversion = self + .context_provider + .get_function_meta(function_name) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "{function_name} is required to insert ARRAY into VARIANT storage" + )) + })?; + conversions.insert(function_name, Arc::clone(&conversion)); + conversion + }; + let input = std::mem::take(expr); + *expr = Expr::ScalarFunction(ScalarFunction::new_udf( + conversion, + vec![input], + )); + } + } + } + let values_schema = if per_row_columns.len() == array_columns.len() { + schema + } else { + Arc::new(DFSchema::from_unqualified_fields( + value_fields.into(), + Default::default(), + )?) + }; + if values_schema.fields().is_empty() { LogicalPlanBuilder::values(values)?.build() } else { - LogicalPlanBuilder::values_with_schema(values, &schema)?.build() + LogicalPlanBuilder::values_with_schema(values, &values_schema)?.build() } } } diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index a2c7065ea525..481d3d14070c 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -230,6 +230,31 @@ impl ContextProvider for MockContextProvider { .into(), ), ])), + "mixed_variant_array_target" => Ok(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + 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(), + ), + ])), + "plain_variant_target" => Ok(Schema::new(vec![ + Field::new("items", variant_storage_type(), true).with_metadata( + [( + "ARROW:extension:name".to_string(), + "arrow.parquet.variant".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 1951afc3fc60..7f67c800bbf3 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -845,7 +845,7 @@ fn plan_insert_does_not_promise_ordinary_target_field_metadata() { #[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", + "__rustice_array_to_variant", vec![DataType::Utf8], variant_storage_type(), ))); @@ -855,7 +855,10 @@ fn plan_insert_converts_tagged_array_to_variant_storage() -> Result<()> { ParserOptions::default(), state, )?; - assert_contains!(plan.display_indent().to_string(), "to_variant("); + assert_contains!( + plan.display_indent().to_string(), + "__rustice_array_to_variant(untyped_array_source.items)" + ); let LogicalPlan::Dml(dml) = &plan else { panic!("expected DML plan"); }; @@ -867,11 +870,212 @@ fn plan_insert_converts_tagged_array_to_variant_storage() -> Result<()> { } #[test] -fn plan_insert_does_not_convert_plain_string_to_array() { - let error = logical_plan_from_state( +fn plan_insert_uses_explicit_array_conversion_for_untagged_text() -> Result<()> { + let state = mock_session_state().with_scalar_function(Arc::new(make_udf( + "__rustice_array_to_variant", + vec![DataType::Utf8], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( "INSERT INTO variant_array_target SELECT items FROM plain_string_source", &SnowflakeDialect {}, ParserOptions::default(), + state, + )?; + assert_contains!( + plan.display_indent().to_string(), + "__rustice_array_to_variant(plain_string_source.items)" + ); + Ok(()) +} + +#[test] +fn plan_insert_values_converts_array_before_variant_storage() -> Result<()> { + let state = mock_session_state() + .with_scalar_function(Arc::new(make_udf( + "array_construct", + vec![DataType::Int64], + DataType::Utf8, + ))) + .with_scalar_function(Arc::new(make_udf( + "__rustice_array_to_variant", + vec![DataType::Utf8], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( + "INSERT INTO variant_array_target VALUES (ARRAY_CONSTRUCT(1))", + &SnowflakeDialect {}, + ParserOptions::default(), + state, + )?; + assert_contains!( + plan.display_indent().to_string(), + "__rustice_array_to_variant(" + ); + Ok(()) +} + +#[test] +fn plan_insert_values_keeps_list_elements_typed() -> Result<()> { + let state = mock_session_state() + .with_scalar_function(make_array_udf()) + .with_scalar_function(Arc::new(make_udf( + "to_variant", + vec![DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( + "INSERT INTO variant_array_target VALUES (make_array('1'))", + &SnowflakeDialect {}, + ParserOptions::default(), + state, + )?; + assert_contains!(plan.display_indent().to_string(), "to_variant("); + let LogicalPlan::Dml(dml) = &plan else { + panic!("expected DML plan"); + }; + let LogicalPlan::Projection(projection) = dml.input.as_ref() else { + panic!("expected INSERT projection"); + }; + assert!(matches!( + projection.input.schema().field(0).data_type(), + DataType::List(_) + )); + assert_contains!(plan.display_indent().to_string(), "to_variant("); + Ok(()) +} + +#[test] +fn plan_insert_values_keeps_mixed_array_rows_independent() -> Result<()> { + let state = mock_session_state() + .with_scalar_function(make_array_udf()) + .with_scalar_function(Arc::new(make_udf( + "array_construct", + vec![DataType::Int64], + DataType::Utf8, + ))) + .with_scalar_function(Arc::new(make_udf( + "to_variant", + vec![DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))], + variant_storage_type(), + ))) + .with_scalar_function(Arc::new(make_udf( + "__rustice_array_to_variant", + vec![DataType::Utf8], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( + "INSERT INTO variant_array_target VALUES (make_array('1')), (make_array(2)), (ARRAY_CONSTRUCT(3))", + &SnowflakeDialect {}, + ParserOptions::default(), + state, + )?; + let formatted = plan.display_indent().to_string(); + assert!(formatted.contains("to_variant(make_array("), "{formatted}"); + assert!( + formatted.contains("__rustice_array_to_variant(array_construct("), + "{formatted}" + ); + assert!(!formatted.contains("CAST(make_array(")); + Ok(()) +} + +#[test] +fn plan_insert_values_preserves_other_target_column_types() -> Result<()> { + let state = mock_session_state() + .with_scalar_function(make_array_udf()) + .with_scalar_function(Arc::new(make_udf( + "to_variant", + vec![DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( + "INSERT INTO mixed_variant_array_target VALUES ('7', make_array('x'))", + &SnowflakeDialect {}, + ParserOptions::default(), + state, + )?; + let LogicalPlan::Dml(dml) = &plan else { + panic!("expected DML plan"); + }; + let LogicalPlan::Projection(projection) = dml.input.as_ref() else { + panic!("expected INSERT projection"); + }; + assert_eq!( + projection.input.schema().field(0).data_type(), + &DataType::Int64 + ); + assert!(matches!( + projection.input.schema().field(1).data_type(), + DataType::List(_) + )); + Ok(()) +} + +#[test] +fn plan_insert_values_validates_array_placeholder() -> Result<()> { + let state = mock_session_state().with_scalar_function(Arc::new(make_udf( + "__rustice_bound_array_to_variant", + vec![variant_storage_type()], + variant_storage_type(), + ))); + let plan = logical_plan_from_state( + "INSERT INTO variant_array_target VALUES ($1)", + &SnowflakeDialect {}, + ParserOptions::default(), + state, + )?; + assert_contains!( + plan.display_indent().to_string(), + "__rustice_bound_array_to_variant($1)" + ); + Ok(()) +} + +#[test] +fn plan_insert_values_does_not_type_nested_values_as_insert_target() -> Result<()> { + let error = logical_plan_from_state( + "INSERT INTO mixed_variant_array_target VALUES (7, (SELECT '[]' FROM (VALUES (1)) AS v))", + &SnowflakeDialect {}, + ParserOptions::default(), + mock_session_state(), + ) + .unwrap_err(); + assert_contains!(error.to_string(), "use INSERT ... SELECT"); + Ok(()) +} + +#[test] +fn plan_insert_values_rejects_native_array_scalar_subquery() -> Result<()> { + let error = logical_plan_from_state( + "INSERT INTO variant_array_target VALUES ((SELECT items FROM variant_array_target LIMIT 1))", + &SnowflakeDialect {}, + ParserOptions::default(), + mock_session_state(), + ) + .unwrap_err(); + assert_contains!(error.to_string(), "use INSERT ... SELECT"); + Ok(()) +} + +#[test] +fn plan_insert_does_not_convert_plain_string_to_unmarked_variant() { + let error = logical_plan_from_state( + "INSERT INTO plain_variant_target SELECT items FROM plain_string_source", + &SnowflakeDialect {}, + ParserOptions::default(), mock_session_state(), ) .unwrap_err();