From b2577eb62356b2b018bd504a38bf8dc5f69b9102 Mon Sep 17 00:00:00 2001
From: Jia-Xuan Liu <liugs963@gmail.com>
Date: Fri, 26 Jul 2024 23:00:37 +0800
Subject: [PATCH] implement `preserve_wildcard_expression` in ExprPlanner

---
 .../core/tests/user_defined/expr_planner.rs   | 29 ++++++++++++++++++-
 datafusion/expr/src/planner.rs                | 11 +++++++
 datafusion/sql/src/select.rs                  | 23 +++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs
index 1b23bf9ab2ef..a8c5c8a97fb0 100644
--- a/datafusion/core/tests/user_defined/expr_planner.rs
+++ b/datafusion/core/tests/user_defined/expr_planner.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_array::RecordBatch;
+use arrow_array::{ArrayRef, RecordBatch, StringArray};
 use std::sync::Arc;
 
 use datafusion::common::{assert_batches_eq, DFSchema};
@@ -62,6 +62,10 @@ impl ExprPlanner for MyCustomPlanner {
             _ => Ok(PlannerResult::Original(expr)),
         }
     }
+
+    fn preserve_wildcard_expression(&self) -> bool {
+        true
+    }
 }
 
 async fn plan_and_collect(sql: &str) -> Result<Vec<RecordBatch>> {
@@ -121,3 +125,26 @@ async fn test_question_filter() {
     let expected = ["+---+", "| a |", "+---+", "| 1 |", "+---+"];
     assert_batches_eq!(&expected, &actual);
 }
+
+#[tokio::test]
+async fn test_keep_wildcard() -> Result<()> {
+    let mut ctx = SessionContext::new();
+    ctx.register_expr_planner(Arc::new(MyCustomPlanner))?;
+    ctx.register_batch("t", t_batch())?;
+
+    let actual = ctx.sql("select * from t").await?.into_optimized_plan()?;
+    let actual = format!("{:?}", actual);
+    let expected = "Projection: *\n  TableScan: t projection=[]";
+    assert_eq!(&expected, &actual);
+
+    let actual = ctx.sql("select t.* from t").await?.into_optimized_plan()?;
+    let actual = format!("{:?}", actual);
+    let expected = "Projection: t.*\n  TableScan: t projection=[]";
+    assert_eq!(&expected, &actual);
+    Ok(())
+}
+
+fn t_batch() -> RecordBatch {
+    let c1: ArrayRef = Arc::new(StringArray::from_iter_values(["a", "b", "c"]));
+    RecordBatch::try_from_iter(vec![("c1", c1)]).unwrap()
+}
diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index c775427df138..47800a3eebc4 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -197,6 +197,17 @@ pub trait ExprPlanner: Send + Sync {
             "Default planner compound identifier hasn't been implemented for ExprPlanner"
         )
     }
+
+    /// Disable the expansion of wildcard expressions when selecting items.
+    /// This is used to keep the wildcard expression in the logical plan.
+    ///
+    /// Note:
+    /// Physical Planner doesn't support plan the wildcard expression.
+    /// If the wildcard expression is kept in the logical plan,
+    /// it should be transformed to the list of columns before physical planning.
+    fn preserve_wildcard_expression(&self) -> bool {
+        false
+    }
 }
 
 /// An operator with two arguments to plan
diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs
index 4de3952dc7ea..9c74914df5cc 100644
--- a/datafusion/sql/src/select.rs
+++ b/datafusion/sql/src/select.rs
@@ -590,6 +590,16 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 if empty_from {
                     return plan_err!("SELECT * with no tables specified is not valid");
                 }
+                let preserve_wildcard_expression = self
+                    .context_provider
+                    .get_expr_planners()
+                    .iter()
+                    .any(|planner| planner.preserve_wildcard_expression());
+
+                if preserve_wildcard_expression {
+                    return Ok(vec![Expr::Wildcard { qualifier: None }]);
+                }
+
                 // do not expand from outer schema
                 let expanded_exprs =
                     expand_wildcard(plan.schema().as_ref(), plan, Some(&options))?;
@@ -610,6 +620,19 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             SelectItem::QualifiedWildcard(object_name, options) => {
                 Self::check_wildcard_options(&options)?;
                 let qualifier = idents_to_table_reference(object_name.0, false)?;
+
+                let preserve_wildcard_expression = self
+                    .context_provider
+                    .get_expr_planners()
+                    .iter()
+                    .any(|planner| planner.preserve_wildcard_expression());
+
+                if preserve_wildcard_expression {
+                    return Ok(vec![Expr::Wildcard {
+                        qualifier: Some(qualifier),
+                    }]);
+                }
+
                 // do not expand from outer schema
                 let expanded_exprs = expand_qualified_wildcard(
                     &qualifier,