[HSTACK] fix: compound_field_access doesn't identifier qualifier. apache#15153

Author: adragomir

datafusion/core/tests/sql/select.rs

@@ -350,3 +350,44 @@ async fn test_version_function() {
 
     assert_eq!(version.value(0), expected_version);
 }
+#[tokio::test]
+async fn test_subscript() -> Result<()> {
+    let ctx = SessionContext::new();
+    ctx.sql(
+        r#"CREATE TABLE test
+        (struct_field STRUCT<substruct INT>)
+    "#,
+    )
+    .await?;
+    ctx.sql(r#"INSERT INTO test VALUES (STRUCT(1))"#).await?;
+
+    let df = ctx
+        .sql(
+            r#"SELECT *
+        FROM test AS test1, test AS test2 WHERE
+        test1.struct_field['substruct'] = test2.struct_field['substruct']"#,
+        )
+        .await?;
+    let _ = df.collect().await?;
+
+    ctx.sql(
+        r#"CREATE TABLE testx
+        (struct_field STRUCT<substruct STRUCT<subsubstruct INT>>)
+    "#,
+    )
+    .await?;
+    ctx.sql(r#"INSERT INTO testx VALUES (STRUCT(STRUCT(1)))"#)
+        .await?;
+
+    let df = ctx.sql(r#"SELECT *
+        FROM testx AS test1, testx AS test2 WHERE
+        test1.struct_field.substruct['subsubstruct'] = test2.struct_field.substruct['subsubstruct']"#).await?;
+    let _ = df.collect().await?;
+
+    let df = ctx.sql(r#"SELECT *
+        FROM testx AS test1, testx AS test2 WHERE
+        test1.struct_field['substruct']['subsubstruct'] = test2.struct_field['substruct']['subsubstruct']"#).await?;
+    let _ = df.collect().await?;
+
+    Ok(())
+}

datafusion/sql/src/expr/identifier.rs

@@ -290,7 +290,7 @@ fn form_identifier(idents: &[String]) -> Result<(Option<TableReference>, &String
     }
 }
 
-fn search_dfschema<'ids, 'schema>(
+pub(crate) fn search_dfschema<'ids, 'schema>(
     ids: &'ids [String],
     schema: &'schema DFSchema,
 ) -> Option<(

datafusion/sql/src/expr/mod.rs

@@ -21,13 +21,13 @@ use datafusion_expr::planner::{
 };
 use sqlparser::ast::{
     AccessExpr, BinaryOperator, CastFormat, CastKind, DataType as SQLDataType,
-    DictionaryField, Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, MapEntry,
-    StructField, Subscript, TrimWhereField, Value,
+    DictionaryField, Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, Ident, MapEntry,
+    Spanned, StructField, Subscript, TrimWhereField, Value,
 };
 
 use datafusion_common::{
     internal_datafusion_err, internal_err, not_impl_err, plan_err, Column, DFSchema,
-    Result, ScalarValue,
+    Result, ScalarValue, Span,
 };
 
 use datafusion_expr::expr::ScalarFunction;
@@ -983,14 +983,100 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         Ok(Expr::Cast(Cast::new(Box::new(expr), dt)))
     }
 
+    /// Extracts the root expression and access chain from a compound expression.
+    ///
+    /// This function attempts to identify if a compound expression (like `a.b.c`) should be treated
+    /// as a column reference with a qualifier (like `table.column`) or as a field access expression.
+    ///
+    /// # Arguments
+    ///
+    /// * `root` - The root SQL expression (e.g., the first part of `a.b.c`)
+    /// * `access_chain` - Vector of access expressions (e.g., `.b` and `.c` parts)
+    /// * `schema` - The schema to resolve column references against
+    /// * `planner_context` - Context for planning expressions
+    ///
+    /// # Returns
+    ///
+    /// A tuple containing:
+    /// * The resolved root expression
+    /// * The remaining access chain that should be processed as field accesses
+    fn extract_root_and_access_chain(
+        &self,
+        root: SQLExpr,
+        access_chain: Vec<AccessExpr>,
+        schema: &DFSchema,
+        planner_context: &mut PlannerContext,
+    ) -> Result<(Expr, Vec<AccessExpr>)> {
+        if let SQLExpr::Identifier(Ident { value: id, .. }) = &root {
+            let mut ids = vec![id.clone()];
+            for access in &access_chain {
+                if let AccessExpr::Dot(SQLExpr::Identifier(Ident { value: id, .. })) =
+                    access
+                {
+                    ids.push(id.clone());
+                } else {
+                    break;
+                }
+            }
+
+            if ids.len() > 1 {
+                // maybe it's a compound identifier
+                if let Some((field, Some(qualifier), nested_names)) =
+                    identifier::search_dfschema(&ids, schema)
+                {
+                    let span_num = ids.len() - nested_names.len();
+                    let mut idx = 0;
+                    let mut spans = vec![];
+                    if let Some(s) = Span::try_from_sqlparser_span(root.span()) {
+                        spans.push(s);
+                    }
+                    idx += 1;
+
+                    let mut nested_access_chain = vec![];
+                    for access in &access_chain {
+                        if idx < span_num {
+                            idx += 1;
+                            if let AccessExpr::Dot(expr) = access {
+                                if let Some(s) =
+                                    Span::try_from_sqlparser_span(expr.span())
+                                {
+                                    spans.push(s);
+                                }
+                            } else {
+                                unreachable!();
+                            }
+                        } else {
+                            nested_access_chain.push(access.clone());
+                        }
+                    }
+
+                    let root = Expr::Column(Column {
+                        name: field.name().clone(),
+                        relation: Some(qualifier.clone()),
+                        spans: datafusion_common::Spans(spans),
+                    });
+                    let access_chain = nested_access_chain;
+                    return Ok((root, access_chain));
+                }
+            }
+        }
+        let root = self.sql_expr_to_logical_expr(root, schema, planner_context)?;
+        Ok((root, access_chain))
+    }
+
     fn sql_compound_field_access_to_expr(
         &self,
         root: SQLExpr,
         access_chain: Vec<AccessExpr>,
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
     ) -> Result<Expr> {
-        let mut root = self.sql_expr_to_logical_expr(root, schema, planner_context)?;
+        let (mut root, access_chain) = self.extract_root_and_access_chain(
+            root,
+            access_chain,
+            schema,
+            planner_context,
+        )?;
         let fields = access_chain
             .into_iter()
             .map(|field| match field {