feat(core): add BigQuery EXTRACT function support and window frame validation

wren-core-base/src/mdl/utils.rs

@@ -43,7 +43,7 @@ pub(crate) fn parse_identifiers_normalized(
     })
 }
 
-pub fn quote_identifier(s: &str) -> Cow<str> {
+pub fn quote_identifier(s: &str) -> Cow<'_, str> {
     if needs_quotes(s) {
         Cow::Owned(format!("\"{}\"", s.replace('"', "\"\"")))
     } else {

wren-core/core/src/logical_plan/analyze/model_generation.rs

@@ -97,12 +97,14 @@ impl ModelGenerationRule {
                     let rls_filter = filters
                         .into_iter()
                         .reduce(|acc, filter| {
-                            if acc.is_none() {
-                                filter
-                            } else if let Some(filter) = filter {
-                                Some(acc.unwrap().and(filter))
+                            if let Some(acc) = acc {
+                                if let Some(filter) = filter {
+                                    Some(acc.and(filter))
+                                } else {
+                                    Some(acc)
+                                }
                             } else {
-                                acc
+                                filter
                             }
                         })
                         .flatten();
@@ -231,7 +233,7 @@ impl ModelGenerationRule {
                             .build()?;
                         Ok(Transformed::yes(alias))
                     } else {
-                        return plan_err!("measures should have an alias");
+                        plan_err!("measures should have an alias")
                     }
                 } else if let Some(partial_model) = extension
                     .node

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

@@ -19,11 +19,12 @@
 
 use crate::mdl::dialect::utils::scalar_function_to_sql_internal;
 use crate::mdl::manifest::DataSource;
-use datafusion::common::Result;
+use datafusion::common::{plan_err, Result};
 use datafusion::logical_expr::sqlparser::keywords::ALL_KEYWORDS;
 use datafusion::logical_expr::Expr;
 
-use datafusion::sql::sqlparser::ast;
+use datafusion::scalar::ScalarValue;
+use datafusion::sql::sqlparser::ast::{self, ExtractSyntax};
 use datafusion::sql::unparser::Unparser;
 use regex::Regex;
 
@@ -116,6 +117,81 @@ impl InnerDialect for BigQueryDialect {
             Ok(Some(alias.to_string()))
         }
     }
+
+    fn scalar_function_to_sql_overrides(
+        &self,
+        unparser: &Unparser,
+        function_name: &str,
+        args: &[Expr],
+    ) -> Result<Option<ast::Expr>> {
+        match function_name {
+            "date_part" => Ok(Some(ast::Expr::Extract {
+                field: datetime_field_from_expr(&args[0])?,
+                syntax: ExtractSyntax::From,
+                expr: Box::new(unparser.expr_to_sql(&args[1])?),
+            })),
+            _ => Ok(None),
+        }
+    }
+}
+
+fn datetime_field_from_expr(expr: &Expr) -> Result<ast::DateTimeField> {
+    match expr {
+        Expr::Literal(ScalarValue::Utf8(Some(s))) => Ok(datetime_field_from_str(s)?),
+        _ => plan_err!("Invalid argument type for datetime field. Expected UTF8 string."),
+    }
+}
+
+fn datetime_field_from_str(s: &str) -> Result<ast::DateTimeField> {
+    match s.to_uppercase().as_str() {
+        "YEAR" => Ok(ast::DateTimeField::Year),
+        "YEARS" => Ok(ast::DateTimeField::Years),
+        "QUARTER" => Ok(ast::DateTimeField::Quarter),
+        "MONTH" => Ok(ast::DateTimeField::Month),
+        "MONTHS" => Ok(ast::DateTimeField::Months),
+        "WEEK" => Ok(ast::DateTimeField::Week(None)),
+        "WEEKS" => Ok(ast::DateTimeField::Weeks),
+        "DAY" => Ok(ast::DateTimeField::Day),
+        "DAYOFWEEK" => Ok(ast::DateTimeField::DayOfWeek),
+        "DAYOFYEAR" => Ok(ast::DateTimeField::DayOfYear),
+        "DAYS" => Ok(ast::DateTimeField::Days),
+        "DATE" => Ok(ast::DateTimeField::Date),
+        "DATETIME" => Ok(ast::DateTimeField::Datetime),
+        "HOUR" => Ok(ast::DateTimeField::Hour),
+        "HOURS" => Ok(ast::DateTimeField::Hours),
+        "MINUTE" => Ok(ast::DateTimeField::Minute),
+        "MINUTES" => Ok(ast::DateTimeField::Minutes),
+        "SECOND" => Ok(ast::DateTimeField::Second),
+        "SECONDS" => Ok(ast::DateTimeField::Seconds),
+        "CENTURY" => Ok(ast::DateTimeField::Century),
+        "DECADE" => Ok(ast::DateTimeField::Decade),
+        "DOW" => Ok(ast::DateTimeField::Dow),
+        "DOY" => Ok(ast::DateTimeField::Doy),
+        "EPOCH" => Ok(ast::DateTimeField::Epoch),
+        "ISODOW" => Ok(ast::DateTimeField::Isodow),
+        "ISOWEEK" => Ok(ast::DateTimeField::IsoWeek),
+        "ISOYEAR" => Ok(ast::DateTimeField::Isoyear),
+        "JULIAN" => Ok(ast::DateTimeField::Julian),
+        "MICROSECOND" => Ok(ast::DateTimeField::Microsecond),
+        "MICROSECONDS" => Ok(ast::DateTimeField::Microseconds),
+        "MILLENIUM" => Ok(ast::DateTimeField::Millenium),
+        "MILLENNIUM" => Ok(ast::DateTimeField::Millennium),
+        "MILLISECOND" => Ok(ast::DateTimeField::Millisecond),
+        "MILLISECONDS" => Ok(ast::DateTimeField::Milliseconds),
+        "NANOSECOND" => Ok(ast::DateTimeField::Nanosecond),
+        "NANOSECONDS" => Ok(ast::DateTimeField::Nanoseconds),
+        "TIME" => Ok(ast::DateTimeField::Time),
+        "TIMEZONE" => Ok(ast::DateTimeField::Timezone),
+        "TIMEZONE_ABBR" => Ok(ast::DateTimeField::TimezoneAbbr),
+        "TIMEZONE_HOUR" => Ok(ast::DateTimeField::TimezoneHour),
+        "TIMEZONE_MINUTE" => Ok(ast::DateTimeField::TimezoneMinute),
+        "TIMEZONE_REGION" => Ok(ast::DateTimeField::TimezoneRegion),
+        "NODATETIME" => Ok(ast::DateTimeField::NoDateTime),
+        _ => {
+            let ident = ast::Ident::new(s);
+            Ok(ast::DateTimeField::Custom(ident))
+        }
+    }
 }
 
 pub struct OracleDialect {}

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

@@ -2930,6 +2930,34 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_extract_roundtrip_bigquery() -> Result<()> {
+        let ctx = SessionContext::new();
+        let manifest = ManifestBuilder::new()
+            .catalog("wren")
+            .schema("test")
+            .model(
+                ModelBuilder::new("orders")
+                    .table_reference("orders")
+                    .column(ColumnBuilder::new("o_orderdate", "date").build())
+                    .build(),
+            )
+            .data_source(DataSource::BigQuery)
+            .build();
+        let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze(
+            manifest,
+            Arc::new(HashMap::default()),
+            Mode::Unparse,
+        )?);
+        let headers = Arc::new(HashMap::default());
+        let sql = "SELECT EXTRACT(YEAR FROM o_orderdate) FROM orders";
+        assert_snapshot!(
+            transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?,
+            @"SELECT EXTRACT(YEAR FROM orders.o_orderdate) FROM (SELECT orders.o_orderdate FROM (SELECT __source.o_orderdate AS o_orderdate FROM orders AS __source) AS orders) AS orders"
+        );
+        Ok(())
+    }
+
     /// Return a RecordBatch with made up data about customer
     fn customer() -> RecordBatch {
         let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));