From a43d4ba4ac884ed0a02931a52314261e7c167d16 Mon Sep 17 00:00:00 2001 From: nuno-faria Date: Mon, 21 Apr 2025 15:24:51 +0100 Subject: [PATCH] Fix errors in the sql feature --- datafusion-federation/src/sql/analyzer.rs | 5 +- datafusion-federation/src/sql/ast_analyzer.rs | 17 ++++-- datafusion-federation/src/sql/mod.rs | 4 +- datafusion-federation/src/sql/table.rs | 2 +- .../src/sql/table_reference.rs | 58 +++++++++++-------- 5 files changed, 54 insertions(+), 32 deletions(-) diff --git a/datafusion-federation/src/sql/analyzer.rs b/datafusion-federation/src/sql/analyzer.rs index 1d8498d..27cd9fe 100644 --- a/datafusion-federation/src/sql/analyzer.rs +++ b/datafusion-federation/src/sql/analyzer.rs @@ -1,7 +1,7 @@ use std::{collections::HashMap, sync::Arc}; use datafusion::{ - common::Column, + common::{Column, Spans}, logical_expr::{ expr::{ AggregateFunction, AggregateFunctionParams, Alias, Exists, InList, InSubquery, @@ -192,6 +192,7 @@ fn rewrite_table_scans_in_expr( Ok(Expr::ScalarSubquery(Subquery { subquery: Arc::new(new_subquery), outer_ref_columns, + spans: Spans::new(), })) } Expr::BinaryExpr(binary_expr) => { @@ -465,6 +466,7 @@ fn rewrite_table_scans_in_expr( let subquery = Subquery { subquery: Arc::new(subquery_plan), outer_ref_columns, + spans: Spans::new(), }; Ok(Expr::Exists(Exists::new(subquery, exists.negated))) } @@ -480,6 +482,7 @@ fn rewrite_table_scans_in_expr( let subquery = Subquery { subquery: Arc::new(subquery_plan), outer_ref_columns, + spans: Spans::new(), }; Ok(Expr::InSubquery(InSubquery::new( Box::new(expr), diff --git a/datafusion-federation/src/sql/ast_analyzer.rs b/datafusion-federation/src/sql/ast_analyzer.rs index 5242553..e974d24 100644 --- a/datafusion-federation/src/sql/ast_analyzer.rs +++ b/datafusion-federation/src/sql/ast_analyzer.rs @@ -28,7 +28,7 @@ pub fn replace_table_args_analyzer(mut visitor: TableArgReplace) -> AstAnalyzer /// let mut analyzer = TableArgReplace::default().with( /// TableReference::parse_str("table1"), /// vec![FunctionArg::Unnamed( -/// Expr::Value( +/// Expr::value( /// Value::Number("1".to_string(), false), /// ) /// .into(), @@ -108,9 +108,18 @@ impl VisitorMut for TableArgReplace { } fn name_to_table_reference(name: &ObjectName) -> TableReference { - let first = name.0.first().map(|n| n.value.to_string()); - let second = name.0.get(1).map(|n| n.value.to_string()); - let third = name.0.get(2).map(|n| n.value.to_string()); + let first = name + .0 + .first() + .map(|n| n.as_ident().expect("expected Ident").value.to_string()); + let second = name + .0 + .get(1) + .map(|n| n.as_ident().expect("expected Ident").value.to_string()); + let third = name + .0 + .get(2) + .map(|n| n.as_ident().expect("expected Ident").value.to_string()); match (first, second, third) { (Some(first), Some(second), Some(third)) => TableReference::full(first, second, third), diff --git a/datafusion-federation/src/sql/mod.rs b/datafusion-federation/src/sql/mod.rs index 7d74f19..08da182 100644 --- a/datafusion-federation/src/sql/mod.rs +++ b/datafusion-federation/src/sql/mod.rs @@ -505,7 +505,7 @@ mod tests { let expected = vec![ "SELECT table_a1.a, table_a1.b, table_a1.c FROM table_a1", "SELECT table_a2.a, table_a2.b, table_a2.c FROM table_a2", - "SELECT table_b1.a, table_b1.b, table_b1.c FROM table_b1(1)", + "SELECT table_b1.a, table_b1.b, table_b1.c FROM table_b1(1) AS table_b1", ]; assert_eq!( @@ -593,7 +593,7 @@ mod tests { }); let expected = vec![ - r#"SELECT "table".a, "table".b, "table".c FROM "default"."table" UNION ALL SELECT "Table".a, "Table".b, "Table".c FROM "default"."Table"(1)"#, + r#"SELECT "table".a, "table".b, "table".c FROM "default"."table" UNION ALL SELECT "Table".a, "Table".b, "Table".c FROM "default"."Table"(1) AS Table"#, ]; assert_eq!( diff --git a/datafusion-federation/src/sql/table.rs b/datafusion-federation/src/sql/table.rs index 1626bad..42d433c 100644 --- a/datafusion-federation/src/sql/table.rs +++ b/datafusion-federation/src/sql/table.rs @@ -50,7 +50,7 @@ impl RemoteTable { /// Creates a new `RemoteTable` instance. /// /// Examples: - /// ```rust + /// ```ignore /// use datafusion::sql::TableReference; /// /// RemoteTable::new("myschema.table".try_into()?, schema); diff --git a/datafusion-federation/src/sql/table_reference.rs b/datafusion-federation/src/sql/table_reference.rs index bb2cc2f..0bee170 100644 --- a/datafusion-federation/src/sql/table_reference.rs +++ b/datafusion-federation/src/sql/table_reference.rs @@ -5,7 +5,7 @@ use datafusion::{ sql::{ sqlparser::{ self, - ast::FunctionArg, + ast::{FunctionArg, ObjectNamePart}, dialect::{Dialect, GenericDialect}, tokenizer::Token, }, @@ -15,14 +15,16 @@ use datafusion::{ /// A multipart identifier to a remote table, view or parameterized view. /// -/// RemoteTableRef can be created by parsing from a string represeting a table obbject with optional +/// RemoteTableRef can be created by parsing from a string representing a table object with optional /// ```rust +/// use datafusion_federation::sql::RemoteTableRef; +/// use datafusion::sql::sqlparser::dialect::PostgreSqlDialect; /// /// RemoteTableRef::try_from("myschema.table"); /// RemoteTableRef::try_from(r#"myschema."Table""#); /// RemoteTableRef::try_from("myschema.view('obj')"); /// -/// RemoteTableRef::parse_with_dialect("myschema.view(name = 'obj')", &PostgresSqlDialect {}); +/// RemoteTableRef::parse_with_dialect("myschema.view(name = 'obj')", &PostgreSqlDialect {}); /// ``` #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct RemoteTableRef { @@ -41,7 +43,7 @@ impl RemoteTableRef { Self::parse_with_dialect(s, &GenericDialect {}) } - /// Create new using a specfic instance of dialect. + /// Create new using a specific instance of dialect. pub fn parse_with_dialect(s: &str, dialect: &dyn Dialect) -> Result { let mut parser = sqlparser::parser::Parser::new(dialect).try_with_sql(s)?; let name = parser.parse_object_name(true)?; @@ -52,15 +54,23 @@ impl RemoteTableRef { }; let table_ref = match (name.0.first(), name.0.get(1), name.0.get(2)) { - (Some(catalog), Some(schema), Some(table)) => TableReference::full( + ( + Some(ObjectNamePart::Identifier(catalog)), + Some(ObjectNamePart::Identifier(schema)), + Some(ObjectNamePart::Identifier(table)), + ) => TableReference::full( catalog.value.clone(), schema.value.clone(), table.value.clone(), ), - (Some(schema), Some(table), None) => { - TableReference::partial(schema.value.clone(), table.value.clone()) + ( + Some(ObjectNamePart::Identifier(schema)), + Some(ObjectNamePart::Identifier(table)), + None, + ) => TableReference::partial(schema.value.clone(), table.value.clone()), + (Some(ObjectNamePart::Identifier(table)), None, None) => { + TableReference::bare(table.value.clone()) } - (Some(table), None, None) => TableReference::bare(table.value.clone()), _ => { return Err(DataFusionError::NotImplemented( "Unable to parse string into TableReference".to_string(), @@ -166,8 +176,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::bare("table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -176,8 +186,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::bare("Table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -189,8 +199,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::bare("table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -199,8 +209,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::bare("Table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -223,8 +233,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::partial("schema", "table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -233,8 +243,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::partial("schema", "Table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -246,8 +256,8 @@ mod tests { let expected = RemoteTableRef::from(( TableReference::partial("schema", "table"), vec![ - FunctionArg::Unnamed(Expr::Value(Value::Number("1".to_string(), false)).into()), - FunctionArg::Unnamed(Expr::Value(Value::Number("2".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("1".to_string(), false)).into()), + FunctionArg::Unnamed(Expr::value(Value::Number("2".to_string(), false)).into()), ], )); assert_eq!(table_ref, expected); @@ -265,12 +275,12 @@ mod tests { vec![ FunctionArg::ExprNamed { name: ast::Expr::Identifier(Ident::new("user_id")), - arg: Expr::Value(Value::Number("1".to_string(), false)).into(), + arg: Expr::value(Value::Number("1".to_string(), false)).into(), operator: FunctionArgOperator::RightArrow, }, FunctionArg::ExprNamed { name: ast::Expr::Identifier(Ident::new("age")), - arg: Expr::Value(Value::Number("2".to_string(), false)).into(), + arg: Expr::value(Value::Number("2".to_string(), false)).into(), operator: FunctionArgOperator::RightArrow, }, ],