expand wildcard first to fix missing column when filter

Author: goldmedal

wren-core/core/src/mdl/context.rs

@@ -91,6 +91,8 @@ fn analyze_rule_for_local_runtime(
     session_state_ref: SessionStateRef,
 ) -> Vec<Arc<dyn AnalyzerRule + Send + Sync>> {
     vec![
+        // Every rule that will generate [Expr::Wildcard] should be placed in front of [ExpandWildcardRule].
+        Arc::new(ExpandWildcardRule::new()),
         // expand the view should be the first rule
         Arc::new(ExpandWrenViewRule::new(
             Arc::clone(&analyzed_mdl),
@@ -118,6 +120,8 @@ fn analyze_rule_for_unparsing(
     session_state_ref: SessionStateRef,
 ) -> Vec<Arc<dyn AnalyzerRule + Send + Sync>> {
     vec![
+        // Every rule that will generate [Expr::Wildcard] should be placed in front of [ExpandWildcardRule].
+        Arc::new(ExpandWildcardRule::new()),
         // expand the view should be the first rule
         Arc::new(ExpandWrenViewRule::new(
             Arc::clone(&analyzed_mdl),

wren-core/core/src/mdl/mod.rs

@@ -618,8 +618,8 @@ mod test {
         .await?;
         assert_eq!(actual,
             "SELECT \"Customer\".\"Custkey\", \"Customer\".\"Name\" FROM \
-            (SELECT __source.\"Custkey\" AS \"Custkey\", __source.\"Name\" AS \"Name\" FROM \
-            datafusion.\"public\".customer AS __source) AS \"Customer\"");
+            (SELECT \"Customer\".\"Custkey\", \"Customer\".\"Name\" FROM \
+            (SELECT __source.\"Custkey\" AS \"Custkey\", __source.\"Name\" AS \"Name\" FROM datafusion.\"public\".customer AS __source) AS \"Customer\") AS \"Customer\"");
         Ok(())
     }
 
@@ -725,9 +725,9 @@ mod test {
         )
         .await?;
         assert_eq!(actual,
-                   "SELECT artist.\"名字\", artist.name_append, artist.\"group\", artist.subscribe_plus, artist.subscribe \
-                   FROM (SELECT __source.\"名字\" AS \"名字\", __source.\"名字\" || __source.\"名字\" AS name_append, __source.\"組別\" AS \"group\", \
-                   CAST(__source.\"訂閱數\" AS BIGINT) + 1 AS subscribe_plus, __source.\"訂閱數\" AS subscribe FROM artist AS __source) AS artist"
+                   "SELECT artist.\"名字\", artist.name_append, artist.\"group\", artist.subscribe, artist.subscribe_plus FROM \
+                   (SELECT artist.\"group\", artist.name_append, artist.subscribe, artist.subscribe_plus, artist.\"名字\" FROM \
+                   (SELECT __source.\"名字\" AS \"名字\", __source.\"名字\" || __source.\"名字\" AS name_append, __source.\"組別\" AS \"group\", CAST(__source.\"訂閱數\" AS BIGINT) + 1 AS subscribe_plus, __source.\"訂閱數\" AS subscribe FROM artist AS __source) AS artist) AS artist"
 );
         ctx.sql(&actual).await?.show().await?;
 
@@ -843,9 +843,7 @@ mod test {
         )
         .await?;
         assert_eq!(actual,
-                   "SELECT artist.\"串接名字\" FROM (SELECT artist.\"串接名字\" FROM \
-                   (SELECT __source.\"名字\" || __source.\"名字\" AS \"串接名字\" FROM artist AS __source) AS artist) AS artist");
-
+                   "SELECT artist.\"串接名字\" FROM (SELECT artist.\"串接名字\" FROM (SELECT __source.\"名字\" || __source.\"名字\" AS \"串接名字\" FROM artist AS __source) AS artist) AS artist");
         let sql = r#"select * from wren.test.artist"#;
         let actual = transform_sql_with_ctx(
             &SessionContext::new(),
@@ -855,7 +853,7 @@ mod test {
         )
         .await?;
         assert_eq!(actual,
-                   "SELECT artist.\"串接名字\" FROM (SELECT __source.\"名字\" || __source.\"名字\" AS \"串接名字\" FROM artist AS __source) AS artist");
+                   "SELECT artist.\"串接名字\" FROM (SELECT artist.\"串接名字\" FROM (SELECT __source.\"名字\" || __source.\"名字\" AS \"串接名字\" FROM artist AS __source) AS artist) AS artist");
 
         let sql = r#"select "名字" from wren.test.artist"#;
         let _ = transform_sql_with_ctx(
@@ -1414,6 +1412,33 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_wildcard_where() -> Result<()> {
+        let ctx = SessionContext::new();
+        let manifest = ManifestBuilder::new()
+            .catalog("wren")
+            .schema("test")
+            .model(
+                ModelBuilder::new("customer")
+                    .table_reference("customer")
+                    .column(ColumnBuilder::new("c_custkey", "int").build())
+                    .column(ColumnBuilder::new("c_name", "string").build())
+                    .build(),
+            )
+            .build();
+        let sql = r#"SELECT * FROM customer WHERE c_custkey = 1"#;
+        let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(manifest)?);
+        let result =
+            transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], sql).await?;
+        assert_eq!(
+            result,
+            "SELECT customer.c_custkey, customer.c_name FROM (SELECT customer.c_custkey, customer.c_name FROM \
+            (SELECT __source.c_custkey AS c_custkey, __source.c_name AS c_name FROM customer AS __source) AS customer) AS customer \
+            WHERE customer.c_custkey = 1"
+        );
+        Ok(())
+    }
+
     /// Return a RecordBatch with made up data about customer
     fn customer() -> RecordBatch {
         let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));

wren-core/sqllogictest/test_files/model.slt

@@ -71,3 +71,14 @@ select "Customer_id" from wrenai.public."Orders" where exists (select 1 from wre
 02d1b5b8831241174c6ef13efd35abbd
 04eafb40a16989307464f27f1fed8907
 0732c0881c70ebcda536a4b14e9db106
+
+query RIITRTTT
+select * from "Order_items" where "Order_id" in ('03c83b31dbc387f83f1b5579b53182fb', '08cbb1d4cd574b126569b208fd4b26ea')
+----
+14.68 1 1 03c83b31dbc387f83f1b5579b53182fb 119.8 a04087ab6a96ffa041f8a2701a72b616 2023/1/15 7:26 CA
+6.9 4 1 08cbb1d4cd574b126569b208fd4b26ea 287.4 588531f8ec37e7d5ff5b7b22ea0488f8 2022/10/19 19:35 CA
+
+query RIITRTTT
+select * from "Order_items" where "Order_id" = '03c83b31dbc387f83f1b5579b53182fb'
+----
+14.68 1 1 03c83b31dbc387f83f1b5579b53182fb 119.8 a04087ab6a96ffa041f8a2701a72b616 2023/1/15 7:26 CA