From 215cd2f8d5dfed36c3b6f05fd49fa2f90412e549 Mon Sep 17 00:00:00 2001
From: "Eason(G Ray)" <30045503+Eason0729@users.noreply.github.com>
Date: Fri, 27 Sep 2024 12:15:45 +0800
Subject: [PATCH] fix: Fix parse_sql_expr not handling alias

---
 datafusion/core/src/dataframe/mod.rs          |  7 +++++--
 .../core/src/execution/session_state.rs       | 20 ++++++++++++++++---
 datafusion/sql/src/parser.rs                  |  1 +
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
index 228860f803f7a..ee44f089d41ea 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -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(())
     }
 }
diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index 10ccd466f19ba..fc6656383d1b1 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -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;
@@ -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!(
@@ -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,
@@ -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())
diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs
index 8e6ed1de61dc3..cb383bfc7f149 100644
--- a/datafusion/sql/src/parser.rs
+++ b/datafusion/sql/src/parser.rs
@@ -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 {