Parsing SQL strings to Exprs wtih the qualified schema

datafusion/core/tests/expr_api/parse_sql_expr.rs

@@ -17,10 +17,12 @@
 
 use arrow_schema::{DataType, Field, Schema};
 use datafusion::prelude::{CsvReadOptions, SessionContext};
+use datafusion_common::DFSchema;
 use datafusion_common::{DFSchemaRef, Result, ToDFSchema};
+use datafusion_expr::col;
+use datafusion_expr::lit;
 use datafusion_expr::Expr;
 use datafusion_sql::unparser::Unparser;
-
 /// A schema like:
 ///
 /// a: Int32 (possibly with nulls)
@@ -85,6 +87,18 @@ async fn round_trip_dataframe(sql: &str) -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn roundtrip_qualified_schema() -> Result<()> {
+    let sql = "a < 5 OR a = 8";
+    let expr = col("t.a").lt(lit(5_i64)).or(col("t.a").eq(lit(8_i64)));
+    let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+    let df_schema = DFSchema::try_from_qualified_schema("t", &schema).unwrap();
+    let ctx = SessionContext::new();
+    let parsed_expr = ctx.parse_sql_expr(sql, &df_schema)?;
+    assert_eq!(parsed_expr, expr);
+    Ok(())
+}
+
 fn unparse_sql_expr(expr: &Expr) -> Result<String> {
     let unparser = Unparser::default();
 

datafusion/sql/src/expr/identifier.rs

@@ -47,38 +47,48 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             // compound identifiers, but this is not a compound
             // identifier. (e.g. it is "foo.bar" not foo.bar)
             let normalize_ident = self.normalizer.normalize(id);
-            match schema.field_with_unqualified_name(normalize_ident.as_str()) {
-                Ok(_) => {
-                    // found a match without a qualified name, this is a inner table column
-                    Ok(Expr::Column(Column {
-                        relation: None,
-                        name: normalize_ident,
-                    }))
-                }
-                Err(_) => {
-                    // check the outer_query_schema and try to find a match
-                    if let Some(outer) = planner_context.outer_query_schema() {
-                        match outer.qualified_field_with_unqualified_name(
-                            normalize_ident.as_str(),
-                        ) {
-                            Ok((qualifier, field)) => {
-                                // found an exact match on a qualified name in the outer plan schema, so this is an outer reference column
-                                Ok(Expr::OuterReferenceColumn(
-                                    field.data_type().clone(),
-                                    Column::from((qualifier, field)),
-                                ))
-                            }
-                            Err(_) => Ok(Expr::Column(Column {
-                                relation: None,
-                                name: normalize_ident,
-                            })),
-                        }
-                    } else {
+            if let Ok((Some(qualifier), _)) =
+                schema.qualified_field_with_unqualified_name(normalize_ident.as_str())
+            {
+                // Found a match with a qualified name, return it with the qualifier
+                Ok(Expr::Column(Column {
+                    relation: Some(qualifier.clone()),
+                    name: normalize_ident,
+                }))
+            } else {
+                match schema.field_with_unqualified_name(normalize_ident.as_str()) {
+                    Ok(_) => {
+                        // found a match without a qualified name, this is a inner table column
                         Ok(Expr::Column(Column {
                             relation: None,
                             name: normalize_ident,
                         }))
                     }
+                    Err(_) => {
+                        // check the outer_query_schema and try to find a match
+                        if let Some(outer) = planner_context.outer_query_schema() {
+                            match outer.qualified_field_with_unqualified_name(
+                                normalize_ident.as_str(),
+                            ) {
+                                Ok((qualifier, field)) => {
+                                    // found an exact match on a qualified name in the outer plan schema, so this is an outer reference column
+                                    Ok(Expr::OuterReferenceColumn(
+                                        field.data_type().clone(),
+                                        Column::from((qualifier, field)),
+                                    ))
+                                }
+                                Err(_) => Ok(Expr::Column(Column {
+                                    relation: None,
+                                    name: normalize_ident,
+                                })),
+                            }
+                        } else {
+                            Ok(Expr::Column(Column {
+                                relation: None,
+                                name: normalize_ident,
+                            }))
+                        }
+                    }
                 }
             }
         }