Preserve and validate native ARRAY INSERT VALUES

datafusion/sql/src/statement.rs

@@ -2960,27 +2960,47 @@ impl<S: ContextProvider> 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
                         };

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<S: ContextProvider> SqlToRel<'_, S> {
@@ -40,6 +44,7 @@ impl<S: ContextProvider> 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<S: ContextProvider> SqlToRel<'_, S> {
                     .map(|v| self.sql_to_expr(v, &empty_schema, planner_context))
                     .collect::<Result<Vec<_>>>()
             })
-            .collect::<Result<Vec<_>>>()?;
+            .collect::<Result<Vec<_>>>();
+        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::<Vec<_>>();
+        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<ScalarUDF>> = 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()
         }
     }
 }

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),