diff --git a/crates/frontend-sql/src/sql/map_planning.rs b/crates/frontend-sql/src/sql/collection_planning.rs similarity index 61% rename from crates/frontend-sql/src/sql/map_planning.rs rename to crates/frontend-sql/src/sql/collection_planning.rs index be99d4a8..d4582afd 100644 --- a/crates/frontend-sql/src/sql/map_planning.rs +++ b/crates/frontend-sql/src/sql/collection_planning.rs @@ -1,7 +1,8 @@ //! DataFusion planning adapters. Types come from the canonical signature rules; //! physical evaluation deliberately remains the query engine's responsibility. -use super::types::{arrow_to_dtype, dtype_to_arrow}; -use asap_types::pre_asap::scalar_signature::MapScalarFunction; +use super::types::{arrow_to_dtype, dtype_to_arrow, scalar_value_to_asap}; +use asap_types::pre_asap::scalar_signature::{element_access_type, MapScalarFunction}; +use asap_types::pre_asap::{Column, QueryExpr, Schema}; use datafusion::arrow::datatypes::DataType; use datafusion::common::{DataFusionError, ExprSchema, Result}; use datafusion::logical_expr::{ @@ -16,7 +17,7 @@ pub(super) fn register(context: &SessionContext) { ("mapconcat", MapScalarFunction::Concat), ("arrayelement", MapScalarFunction::Access), ] { - context.register_udf(ScalarUDF::from(MapPlanningFunction { + context.register_udf(ScalarUDF::from(CollectionPlanningFunction { name, function, signature: match function { @@ -31,13 +32,18 @@ pub(super) fn register(context: &SessionContext) { } } #[derive(Debug)] -struct MapPlanningFunction { +struct CollectionPlanningFunction { name: &'static str, function: MapScalarFunction, signature: Signature, } -impl MapPlanningFunction { - fn output(&self, args: &[DataType], nullable: &[bool]) -> Result<(DataType, bool)> { +impl CollectionPlanningFunction { + fn output( + &self, + args: &[DataType], + nullable: &[bool], + expressions: Option<&[Expr]>, + ) -> Result<(DataType, bool)> { let inputs = args .iter() .zip(nullable) @@ -47,14 +53,40 @@ impl MapPlanningFunction { .map_err(|e| DataFusionError::Plan(e.to_string())) }) .collect::>>()?; - let (dtype, nullable) = self - .function - .output_type(&inputs) - .map_err(DataFusionError::Plan)?; + let (dtype, nullable) = if self.name == "arrayelement" { + // DataFusion asks for argument-dependent types before canonical + // expression binding. Reuse the shared resolver over typed argument + // slots; final canonical binding also validates literal selectors. + let schema = Schema::new( + inputs + .into_iter() + .enumerate() + .map(|(index, (dtype, nullable))| { + Column::new(format!("argument_{index}"), dtype, nullable) + }) + .collect(), + ); + let args = (0..schema.columns.len()) + .map(|index| { + if let Some(Expr::Literal(value)) = expressions.and_then(|args| args.get(index)) + { + scalar_value_to_asap(value) + .map(QueryExpr::Literal) + .map_err(|error| DataFusionError::Plan(error.to_string())) + } else { + Ok(QueryExpr::Column(index)) + } + }) + .collect::>>()?; + element_access_type(&args, &schema) + } else { + self.function.output_type(&inputs) + } + .map_err(DataFusionError::Plan)?; Ok((dtype_to_arrow(&dtype), nullable)) } } -impl ScalarUDFImpl for MapPlanningFunction { +impl ScalarUDFImpl for CollectionPlanningFunction { fn as_any(&self) -> &dyn std::any::Any { self } @@ -71,6 +103,7 @@ impl ScalarUDFImpl for MapPlanningFunction { .iter() .map(|dtype| *dtype == DataType::Null) .collect::>(), + None, ) .map(|output| output.0) } @@ -84,7 +117,8 @@ impl ScalarUDFImpl for MapPlanningFunction { .iter() .map(|arg| arg.nullable(schema)) .collect::>>()?; - self.output(types, &nullable).map(|output| output.0) + self.output(types, &nullable, Some(args)) + .map(|output| output.0) } fn is_nullable(&self, args: &[Expr], schema: &dyn ExprSchema) -> bool { let types = args @@ -97,14 +131,14 @@ impl ScalarUDFImpl for MapPlanningFunction { .collect::>>(); match (types, nullable) { (Ok(types), Ok(nullable)) => self - .output(&types, &nullable) + .output(&types, &nullable, Some(args)) .map(|out| out.1) .unwrap_or(true), _ => true, } } fn invoke_batch(&self, _args: &[ColumnarValue], _number_rows: usize) -> Result { - Err(DataFusionError::NotImplemented("map planning adapter cannot execute; use a capable query engine or external exact subtree".into())) + Err(DataFusionError::NotImplemented("collection planning adapter cannot execute; use a capable query engine or external exact subtree".into())) } } @@ -113,7 +147,7 @@ mod tests { use super::*; #[test] fn planning_adapter_explicitly_refuses_physical_execution() { - let adapter = MapPlanningFunction { + let adapter = CollectionPlanningFunction { name: "map", function: MapScalarFunction::Construct, signature: Signature::any(0, Volatility::Immutable), diff --git a/crates/frontend-sql/src/sql/expr.rs b/crates/frontend-sql/src/sql/expr.rs index 029842d5..cd57079f 100644 --- a/crates/frontend-sql/src/sql/expr.rs +++ b/crates/frontend-sql/src/sql/expr.rs @@ -203,7 +203,7 @@ pub(super) fn df_expr_to_unresolved(expr: &Expr) -> Result, _> = sf.args.iter().map(df_expr_to_unresolved).collect(); Ok(Unresolved::FunctionCall { name: if sf.func.name().eq_ignore_ascii_case("arrayelement") { - "asap_map_access".into() + "asap_element_access".into() } else { sf.func.name().to_string() }, diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 4cf9fad5..dc9116eb 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -68,8 +68,8 @@ use asap_types::workload::SqlDialect; use crate::error::SqlError as LoweringError; mod clickhouse_ast; +mod collection_planning; mod expr; -mod map_planning; mod types; pub use types::SqlCatalog; @@ -201,7 +201,7 @@ impl<'a> SqlLowerer<'a> { let config = SessionConfig::new().set_str("datafusion.sql_parser.dialect", dialect_name); let ctx = SessionContext::new_with_config(config); if matches!(self.dialect, SqlDialect::ClickhouseSQL) { - map_planning::register(&ctx); + collection_planning::register(&ctx); } // A catalog key like "bgp.bgp_updates" schema-qualifies the table // (e.g. a ClickHouse database name). DataFusion requires the parent diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index bb9d4318..556c0e5b 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -2344,3 +2344,52 @@ async fn arg_selector_result_schema_tracks_selected_argument() { assert_eq!(schema.columns[0].nullable, nullable); } } + +#[tokio::test] +async fn clickhouse_list_element_uses_canonical_typed_access() { + let catalog = SqlCatalog::new().with_table( + "t", + Schema::new(vec![ + Column::new( + "samples", + DataType::List { + element: Box::new(Column::new("item", DataType::Int64, false)), + }, + false, + ), + Column::new("index", DataType::Int64, true), + ]), + ); + for (sql, nullable) in [ + ("SELECT samples[1] AS selected FROM t", false), + ("SELECT arrayElement(samples, -1) AS selected FROM t", false), + ("SELECT samples[index] AS selected FROM t", true), + ] { + let query = lower_sql_dialect( + sql, + &catalog, + SqlDialect::ClickhouseSQL, + AccuracyTarget::Exact, + ) + .await + .unwrap(); + let output = query.output_schema().unwrap(); + assert_eq!(output.columns[0].dtype, DataType::Int64); + assert_eq!(output.columns[0].nullable, nullable); + let serialized = serde_json::to_string(&query).unwrap(); + assert!(serialized.contains("asap_element_access"), "{serialized}"); + } + for sql in ["SELECT samples[0] FROM t", "SELECT samples['bad'] FROM t"] { + assert!( + lower_sql_dialect( + sql, + &catalog, + SqlDialect::ClickhouseSQL, + AccuracyTarget::Exact + ) + .await + .is_err(), + "{sql}" + ); + } +}