diff --git a/wren-core/core/src/mdl/dialect/inner_dialect.rs b/wren-core/core/src/mdl/dialect/inner_dialect.rs index ccad0a476..34c53b416 100644 --- a/wren-core/core/src/mdl/dialect/inner_dialect.rs +++ b/wren-core/core/src/mdl/dialect/inner_dialect.rs @@ -36,12 +36,18 @@ pub trait InnerDialect: Send + Sync { ) -> Result> { 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 { 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 + } +} diff --git a/wren-core/core/src/mdl/dialect/wren_dialect.rs b/wren-core/core/src/mdl/dialect/wren_dialect.rs index 15036bb33..c6df175d4 100644 --- a/wren-core/core/src/mdl/dialect/wren_dialect.rs +++ b/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 { diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs index f669a2fd6..b3f070abe 100644 --- a/wren-core/core/src/mdl/mod.rs +++ b/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();