feat(sql): lower typed List element access

crates/frontend-sql/src/sql/map_planning.rs → 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::<Result<Vec<_>>>()?;
-        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::<Result<Vec<_>>>()?;
+            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::<Vec<_>>(),
+            None,
         )
         .map(|output| output.0)
     }
@@ -84,7 +117,8 @@ impl ScalarUDFImpl for MapPlanningFunction {
             .iter()
             .map(|arg| arg.nullable(schema))
             .collect::<Result<Vec<_>>>()?;
-        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::<Result<Vec<_>>>();
         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<ColumnarValue> {
-        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),

crates/frontend-sql/src/sql/expr.rs

@@ -203,7 +203,7 @@ pub(super) fn df_expr_to_unresolved(expr: &Expr) -> Result<Unresolved, LoweringE
             let args: Result<Vec<_>, _> = 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()
                 },

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

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}"
+        );
+    }
+}