fix(core): handle the unnest table factor for the specific dialect (#1104)

Author: goldmedal

wren-core/core/src/mdl/dialect/inner_dialect.rs

@@ -36,12 +36,18 @@ pub trait InnerDialect: Send + Sync {
     ) -> Result<Option<ast::Expr>> {
         Ok(None)
     }
+
+    /// A wrapper for [datafusion::sql::unparser::dialect::Dialect::unnest_as_table_factor].
+    fn unnest_as_table_factor(&self) -> bool {
+        false
+    }
 }
 
 /// [get_inner_dialect] returns the suitable InnerDialect for the given data source.
 pub fn get_inner_dialect(data_source: &DataSource) -> Box<dyn InnerDialect> {
     match data_source {
         DataSource::MySQL => Box::new(MySQLDialect {}),
+        DataSource::BigQuery => Box::new(BigQueryDialect {}),
         _ => Box::new(GenericDialect {}),
     }
 }
@@ -68,3 +74,11 @@ impl InnerDialect for MySQLDialect {
         }
     }
 }
+
+pub struct BigQueryDialect {}
+
+impl InnerDialect for BigQueryDialect {
+    fn unnest_as_table_factor(&self) -> bool {
+        true
+    }
+}

wren-core/core/src/mdl/dialect/wren_dialect.rs

@@ -83,6 +83,10 @@ impl Dialect for WrenDialect {
             _ => Ok(None),
         }
     }
+
+    fn unnest_as_table_factor(&self) -> bool {
+        self.inner_dialect.unnest_as_table_factor()
+    }
 }
 
 impl Default for WrenDialect {

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

@@ -449,6 +449,7 @@ mod test {
     use datafusion::config::ConfigOptions;
     use datafusion::prelude::{SessionConfig, SessionContext};
     use datafusion::sql::unparser::plan_to_sql;
+    use wren_core_base::mdl::DataSource;
 
     #[test]
     fn test_sync_transform() -> Result<()> {
@@ -976,6 +977,25 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_unnest_as_table_factor() -> Result<()> {
+        let ctx = SessionContext::new();
+        let manifest = ManifestBuilder::new().build();
+        let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(manifest)?);
+        let sql = "select * from unnest([1, 2, 3])";
+        let actual = transform_sql_with_ctx(&ctx, analyzed_mdl, &[], sql).await?;
+        assert_eq!(actual, "SELECT \"UNNEST(make_array(Int64(1),Int64(2),Int64(3)))\" FROM (SELECT UNNEST([1, 2, 3]) AS \"UNNEST(make_array(Int64(1),Int64(2),Int64(3)))\")");
+
+        let manifest = ManifestBuilder::new()
+            .data_source(DataSource::BigQuery)
+            .build();
+        let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(manifest)?);
+        let sql = "select * from unnest([1, 2, 3])";
+        let actual = transform_sql_with_ctx(&ctx, analyzed_mdl, &[], sql).await?;
+        assert_eq!(actual, "SELECT \"UNNEST(make_array(Int64(1),Int64(2),Int64(3)))\" FROM UNNEST([1, 2, 3])");
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_simplify_timestamp() -> Result<()> {
         let ctx = SessionContext::new();