fix(core): expand the wildcard before Wren rewrite rules

ibis-server/tests/routers/v3/connector/postgres/test_query.py

@@ -111,16 +111,16 @@ async def test_query(client, manifest_str, connection_info):
     assert len(result["columns"]) == 10
     assert len(result["data"]) == 1
     assert result["data"][0] == [
+        1,
+        370,
+        "O",
+        "172799.49",
+        "1996-01-02 00:00:00.000000",
+        "1_370",
         "2024-01-01 23:59:59.000000",
         "2024-01-01 23:59:59.000000 UTC",
         "2024-01-16 04:00:00.000000 UTC",  # utc-5
         "2024-07-16 03:00:00.000000 UTC",  # utc-4
-        "172799.49",
-        "1_370",
-        370,
-        "1996-01-02 00:00:00.000000",
-        1,
-        "O",
     ]
     assert result["dtypes"] == {
         "o_orderkey": "int32",
@@ -212,7 +212,7 @@ async def test_query_with_connection_url(client, manifest_str, connection_url):
     result = response.json()
     assert len(result["columns"]) == 10
     assert len(result["data"]) == 1
-    assert result["data"][0][0] == "2024-01-01 23:59:59.000000"
+    assert result["data"][0][0] == 1
     assert result["dtypes"] is not None
 
 

wren-core-py/tests/test_modeling_core.py

@@ -90,7 +90,7 @@ def test_session_context():
     rewritten_sql = session_context.transform_sql(sql)
     assert (
         rewritten_sql
-        == "SELECT customer.c_custkey, customer.c_name FROM (SELECT __source.c_custkey AS c_custkey, __source.c_name AS c_name FROM main.customer AS __source) AS customer"
+        == "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 main.customer AS __source) AS customer) AS customer"
     )
 
     session_context = SessionContext(manifest_str, "tests/functions.csv")

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![
+        // To align the lastest change in datafusion, apply this this rule first.
+        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![
+        // To align the lastest change in datafusion, apply this this rule first.
+        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/bin/sqllogictests.rs

@@ -117,12 +117,7 @@ async fn run_tests() -> Result<()> {
             futures::stream::iter(match result {
                 // Tokio panic error
                 Err(e) => Some(DataFusionError::External(Box::new(e))),
-                Ok(thread_result) => match thread_result {
-                    // Test run error
-                    Err(e) => Some(e),
-                    // success
-                    Ok(_) => None,
-                },
+                Ok(thread_result) => thread_result.err(),
             })
         })
         .collect()

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