Skip to content

Commit

Permalink
fix: Fix parse_sql_expr not handling alias
Browse files Browse the repository at this point in the history
  • Loading branch information
Eason0729 committed Sep 27, 2024
1 parent 36affd4 commit 215cd2f
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 5 deletions.
7 changes: 5 additions & 2 deletions datafusion/core/src/dataframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3998,13 +3998,16 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn dev_test() -> Result<()> {
async fn parse_sql_expr_handle_alias() -> Result<()> {
let ctx = SessionContext::new();
let df = ctx
.read_csv("tests/data/example.csv", CsvReadOptions::new())
.await?;
let sql = df.parse_sql_expr("SUM(a) as a_sum")?;
println!("{:?}", sql.to_string());
assert_eq!(
"sum(a) as a_sum".to_string(),
sql.to_string().to_lowercase()
);
Ok(())
}
}
20 changes: 17 additions & 3 deletions datafusion/core/src/execution/session_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ use datafusion_sql::planner::{ContextProvider, ParserOptions, PlannerContext, Sq
use itertools::Itertools;
use log::{debug, info};
use object_store::ObjectStore;
use sqlparser::ast::ExprWithAlias as SQLExprWithAlias;
use sqlparser::ast::{Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias};
use sqlparser::dialect::dialect_from_str;
use std::any::Any;
use std::collections::hash_map::Entry;
Expand Down Expand Up @@ -483,11 +483,25 @@ impl SessionState {

/// parse a sql string into a sqlparser-rs AST [`SQLExpr`].
///
/// discard the alias information.
///
/// See [`Self::create_logical_expr`] for parsing sql to [`Expr`].
pub fn sql_to_expr(
&self,
sql: &str,
dialect: &str,
) -> datafusion_common::Result<SQLExpr> {
self.sql_to_expr_with_alias(sql, dialect)
.map(|expr| expr.expr)
}

/// parse a sql string into a sqlparser-rs AST [`SQLExprWithAlias`].
///
/// keep the alias information if there is any.
pub fn sql_to_expr_with_alias(
&self,
sql: &str,
dialect: &str,
) -> datafusion_common::Result<SQLExprWithAlias> {
let dialect = dialect_from_str(dialect).ok_or_else(|| {
plan_datafusion_err!(
Expand Down Expand Up @@ -591,7 +605,7 @@ impl SessionState {
) -> datafusion_common::Result<Expr> {
let dialect = self.config.options().sql_parser.dialect.as_str();

let sql_expr = self.sql_to_expr(sql, dialect)?;
let sql_expr = self.sql_to_expr_with_alias(sql, dialect)?;

let provider = SessionContextProvider {
state: self,
Expand Down Expand Up @@ -1872,7 +1886,7 @@ mod tests {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let df_schema = DFSchema::try_from(schema)?;
let dialect = state.config.options().sql_parser.dialect.as_str();
let sql_expr = state.sql_to_expr(sql, dialect)?.expr;
let sql_expr = state.sql_to_expr(sql, dialect)?;

let query = SqlToRel::new_with_options(&provider, state.get_parser_options());
query.sql_to_expr(sql_expr, &df_schema, &mut PlannerContext::new())
Expand Down
1 change: 1 addition & 0 deletions datafusion/sql/src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,7 @@ impl<'a> DFParser<'a> {
}
}

/// parse an expression with alias
pub fn parse_expr(&mut self) -> Result<ExprWithAlias, ParserError> {
if let Token::Word(w) = self.parser.peek_token().token {
match w.keyword {
Expand Down

0 comments on commit 215cd2f

Please sign in to comment.