From 114cb801e64d6b07b7fbddf966b6a42b1eb639d3 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 11:07:43 +0800 Subject: [PATCH 1/7] fix(core): modified the argument ordering of DATE_DIFF --- .../v3/connector/bigquery/test_functions.py | 82 ++++++++++++++++++ wren-core-py/Cargo.lock | 58 ++++++------- .../core/src/mdl/dialect/inner_dialect.rs | 53 +++++++++++- wren-core/core/src/mdl/mod.rs | 84 +++++++++++++++++++ 4 files changed, 246 insertions(+), 31 deletions(-) diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py index 336f8bf7b..7af20d7b5 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py @@ -160,3 +160,85 @@ async def test_datetime_function(client, manifest_str: str, connection_info): "data": [["2001-01-01 00:11:11.000000"]], "dtypes": {"col": "timestamp[us]"}, } + + +async def test_date_diff_function(client, manifest_str: str, connection_info): + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT DATE_DIFF(DAY, CURRENT_DATE(), o_orderdate) AS col FROM orders LIMIT 1", + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result == { + "columns": ["col"], + "data": [[result["data"][0][0]]], + "dtypes": {"col": "int64"}, + } + + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT DATEDIFF(DAY, CURRENT_DATE(), o_orderdate) AS col FROM orders LIMIT 1", + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result == { + "columns": ["col"], + "data": [[result["data"][0][0]]], + "dtypes": {"col": "int64"}, + } + + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT DATEDIFF('DAY', CURRENT_DATE(), o_orderdate) AS col FROM orders LIMIT 1", + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result == { + "columns": ["col"], + "data": [[result["data"][0][0]]], + "dtypes": {"col": "int64"}, + } + + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT DATEDIFF('DAYS', CURRENT_DATE(), o_orderdate) AS col FROM orders LIMIT 1", + }, + ) + + assert response.status_code == 422 + assert "Unsupported date part 'DAYS' for BIGQUERY." in response.text + + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT DATEDIFF('HOUR', CURRENT_TIMESTAMP(), TIMESTAMP WITH TIME ZONE '2025-01-01 00:00:00') AS col FROM orders LIMIT 1", + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result == { + "columns": ["col"], + "data": [[result["data"][0][0]]], + "dtypes": {"col": "int64"}, + } diff --git a/wren-core-py/Cargo.lock b/wren-core-py/Cargo.lock index 0ffb9648c..0300fb3b7 100644 --- a/wren-core-py/Cargo.lock +++ b/wren-core-py/Cargo.lock @@ -735,7 +735,7 @@ dependencies = [ [[package]] name = "datafusion" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "arrow-ipc", @@ -789,7 +789,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -814,7 +814,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -836,7 +836,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -861,7 +861,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "futures", "log", @@ -871,7 +871,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-compression", @@ -906,7 +906,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -930,7 +930,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -954,7 +954,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -986,12 +986,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" [[package]] name = "datafusion-execution" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "dashmap", @@ -1009,7 +1009,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -1030,7 +1030,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "datafusion-common", @@ -1042,7 +1042,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "arrow-buffer", @@ -1070,7 +1070,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -1090,7 +1090,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -1102,7 +1102,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "arrow-ord", @@ -1123,7 +1123,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -1138,7 +1138,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "datafusion-common", @@ -1155,7 +1155,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1164,7 +1164,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "datafusion-expr", "quote", @@ -1174,7 +1174,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "chrono", @@ -1193,7 +1193,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -1214,7 +1214,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -1227,7 +1227,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "datafusion-common", @@ -1246,7 +1246,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "ahash", "arrow", @@ -1275,7 +1275,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "arrow-schema", @@ -1292,7 +1292,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "async-trait", @@ -1315,7 +1315,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "49.0.1" -source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#732dd555f485863647113565d7aaa22188db8bcc" +source = "git+https://github.com/Canner/datafusion.git?branch=canner%2Fv49.0.1#4f9bb6d92f99a041977c49573800a7f72a48a2d5" dependencies = [ "arrow", "bigdecimal", diff --git a/wren-core/core/src/mdl/dialect/inner_dialect.rs b/wren-core/core/src/mdl/dialect/inner_dialect.rs index 26ad1bcfe..a2cabcab6 100644 --- a/wren-core/core/src/mdl/dialect/inner_dialect.rs +++ b/wren-core/core/src/mdl/dialect/inner_dialect.rs @@ -24,7 +24,7 @@ use datafusion::logical_expr::sqlparser::keywords::ALL_KEYWORDS; use datafusion::logical_expr::Expr; use datafusion::scalar::ScalarValue; -use datafusion::sql::sqlparser::ast::{self, ExtractSyntax, Ident, WindowFrameBound}; +use datafusion::sql::sqlparser::ast::{self, ExtractSyntax, Function, Ident, ObjectName, ObjectNamePart, WindowFrameBound}; use datafusion::sql::unparser::Unparser; use regex::Regex; @@ -147,6 +147,50 @@ impl InnerDialect for BigQueryDialect { expr: Box::new(unparser.expr_to_sql(&args[1])?), })) } + "date_diff" => { + if args.len() != 3 { + return plan_err!( + "date_diff requires exactly 3 arguments, found {}", + args.len() + ); + } + let Expr::Literal(ScalarValue::Utf8(Some(s)), _) = args[0].clone() else { + return plan_err!( + "date_diff requires a string literal as the third argument" + ); + }; + let granularity = ast::Expr::Identifier(Ident::new( + self.datetime_field_from_str(&s)?.to_string(), + )); + // DATE_DIFF(end_date, start_date, granularity) + // https://cloud.google.com/bigquery/docs/reference/standard-sql/date_functions#date_diff + Ok(Some(ast::Expr::Function(Function { + name: ObjectName(vec![ObjectNamePart::Identifier(Ident::new( + "DATE_DIFF", + ))]), + args: ast::FunctionArguments::List(ast::FunctionArgumentList { + duplicate_treatment: None, + args: vec![ + unparser.expr_to_sql(&args[2]).map(|e| { + ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e)) + })?, + unparser.expr_to_sql(&args[1]).map(|e| { + ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e)) + })?, + ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr( + granularity, + )), + ], + clauses: vec![], + }), + filter: None, + null_treatment: None, + over: None, + within_group: vec![], + parameters: ast::FunctionArguments::None, + uses_odbc_syntax: false, + }))) + } "now" => { scalar_function_to_sql_internal(unparser, None, "CURRENT_TIMESTAMP", args) } @@ -227,8 +271,13 @@ impl BigQueryDialect { "QUARTER" => Ok(ast::DateTimeField::Quarter), "YEAR" => Ok(ast::DateTimeField::Year), "ISOYEAR" => Ok(ast::DateTimeField::Isoyear), + "MICROSECOND" => Ok(ast::DateTimeField::Microsecond), + "MILLISECOND" => Ok(ast::DateTimeField::Millisecond), + "SECOND" => Ok(ast::DateTimeField::Second), + "MINUTE" => Ok(ast::DateTimeField::Minute), + "HOUR" => Ok(ast::DateTimeField::Hour), _ => { - plan_err!("Unsupported date part '{}' for BigQuery", s) + plan_err!("Unsupported date part '{}' for BIGQUERY. Valid values are: WEEK, DAYOFWEEK, DAY, DAYOFYEAR, ISOWEEK, MONTH, QUARTER, YEAR, ISOYEAR, MICROSECOND, MILLISECOND, SECOND, MINUTE, HOUR", s) } } } diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs index abc3c693d..3c58ba48e 100644 --- a/wren-core/core/src/mdl/mod.rs +++ b/wren-core/core/src/mdl/mod.rs @@ -3088,6 +3088,90 @@ mod test { Ok(()) } + #[tokio::test] + async fn test_date_diff_bigquery() -> Result<()> { + let ctx = SessionContext::new(); + let manifest = ManifestBuilder::new() + .catalog("wren") + .schema("test") + .model( + ModelBuilder::new("date_table") + .table_reference("date_table") + .column(ColumnBuilder::new("date_1", "date").build()) + .column(ColumnBuilder::new("date_2", "date").build()) + .build(), + ) + .model( + ModelBuilder::new("timestamp_table") + .table_reference("timestamp_table") + .column(ColumnBuilder::new("ts_1", "timestamp").build()) + .column(ColumnBuilder::new("ts_2", "timestamp").build()) + .build(), + ) + .data_source(DataSource::BigQuery) + .build(); + let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze( + manifest, + Arc::new(HashMap::default()), + Mode::Unparse, + )?); + + let sql = "select date_diff(DAY, date_1, date_2) from date_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT DATE_DIFF(date_table.date_2, date_table.date_1, DAY) FROM (SELECT date_table.date_1, date_table.date_2 FROM (SELECT __source.date_1 AS date_1, __source.date_2 AS date_2 FROM date_table AS __source) AS date_table) AS date_table" + ); + + let sql = "select datediff(DAY, date_1, date_2) from date_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT DATE_DIFF(date_table.date_2, date_table.date_1, DAY) FROM (SELECT date_table.date_1, date_table.date_2 FROM (SELECT __source.date_1 AS date_1, __source.date_2 AS date_2 FROM date_table AS __source) AS date_table) AS date_table" + ); + let sql = "select datediff('DAY', date_1, date_2) from date_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT DATE_DIFF(date_table.date_2, date_table.date_1, DAY) FROM (SELECT date_table.date_1, date_table.date_2 FROM (SELECT __source.date_1 AS date_1, __source.date_2 AS date_2 FROM date_table AS __source) AS date_table) AS date_table" + ); + + let sql = "select datediff(DAY, date_1, date_2) from date_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT DATE_DIFF(date_table.date_2, date_table.date_1, DAY) FROM (SELECT date_table.date_1, date_table.date_2 FROM (SELECT __source.date_1 AS date_1, __source.date_2 AS date_2 FROM date_table AS __source) AS date_table) AS date_table" + ); + + let sql = "select datediff('DAYS', date_1, date_2) from date_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + match transform_sql_with_ctx( + &ctx, + Arc::clone(&analyzed_mdl), + &[], + Arc::clone(&headers), + sql, + ) + .await + { + Ok(_) => { + panic!("Expected error, but got SQL"); + } + Err(e) => assert_snapshot!( + e.to_string(), + @"Error during planning: Unsupported date part 'DAYS' for BIGQUERY. Valid values are: WEEK, DAYOFWEEK, DAY, DAYOFYEAR, ISOWEEK, MONTH, QUARTER, YEAR, ISOYEAR, MICROSECOND, MILLISECOND, SECOND, MINUTE, HOUR" + ), + } + + let sql = "select datediff(HOUR, ts_1, ts_2) from timestamp_table"; + let headers: Arc>> = Arc::new(HashMap::default()); + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT DATE_DIFF(timestamp_table.ts_2, timestamp_table.ts_1, HOUR) FROM (SELECT timestamp_table.ts_1, timestamp_table.ts_2 FROM (SELECT __source.ts_1 AS ts_1, __source.ts_2 AS ts_2 FROM timestamp_table AS __source) AS timestamp_table) AS timestamp_table" + ); + Ok(()) + } + #[tokio::test] async fn test_window_functions_without_frame_bigquery() -> Result<()> { let ctx = SessionContext::new(); From 16598e4e00c14531470a74210a96075d5d1a07e7 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 11:15:48 +0800 Subject: [PATCH 2/7] fix test --- wren-core/core/src/mdl/function.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/wren-core/core/src/mdl/function.rs b/wren-core/core/src/mdl/function.rs index 1c46c622c..5c430f44f 100644 --- a/wren-core/core/src/mdl/function.rs +++ b/wren-core/core/src/mdl/function.rs @@ -387,15 +387,15 @@ mod test { #[tokio::test] async fn test_by_pass_scalar_udf() -> Result<()> { - let udf = ByPassScalarUDF::new("date_diff", DataType::Int64); + let udf = ByPassScalarUDF::new("date_test", DataType::Int64); let ctx = SessionContext::new(); ctx.register_udf(ScalarUDF::new_from_impl(udf)); let plan = ctx - .sql("SELECT date_diff(1, 2)") + .sql("SELECT date_test(1, 2)") .await? .into_unoptimized_plan(); - let expected = "Projection: date_diff(Int64(1), Int64(2))\n EmptyRelation"; + let expected = "Projection: date_test(Int64(1), Int64(2))\n EmptyRelation"; assert_eq!(format!("{plan}"), expected); ctx.register_udf(ScalarUDF::new_from_impl(ByPassScalarUDF::new( From 55127094abda220fed27fb2aa23fc9d51c356d44 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 11:23:53 +0800 Subject: [PATCH 3/7] fix fmt --- wren-core/core/src/mdl/dialect/inner_dialect.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wren-core/core/src/mdl/dialect/inner_dialect.rs b/wren-core/core/src/mdl/dialect/inner_dialect.rs index a2cabcab6..584e27f69 100644 --- a/wren-core/core/src/mdl/dialect/inner_dialect.rs +++ b/wren-core/core/src/mdl/dialect/inner_dialect.rs @@ -24,7 +24,9 @@ use datafusion::logical_expr::sqlparser::keywords::ALL_KEYWORDS; use datafusion::logical_expr::Expr; use datafusion::scalar::ScalarValue; -use datafusion::sql::sqlparser::ast::{self, ExtractSyntax, Function, Ident, ObjectName, ObjectNamePart, WindowFrameBound}; +use datafusion::sql::sqlparser::ast::{ + self, ExtractSyntax, Function, Ident, ObjectName, ObjectNamePart, WindowFrameBound, +}; use datafusion::sql::unparser::Unparser; use regex::Regex; From ab1e15f7d6ecbbff95b02338bc72dc24f92922ed Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 12:05:48 +0800 Subject: [PATCH 4/7] fix test --- ibis-server/tests/conftest.py | 2 +- wren-core-py/tests/test_modeling_core.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ibis-server/tests/conftest.py b/ibis-server/tests/conftest.py index 96455334c..eb379ebe1 100644 --- a/ibis-server/tests/conftest.py +++ b/ibis-server/tests/conftest.py @@ -11,7 +11,7 @@ def file_path(path: str) -> str: return os.path.join(os.path.dirname(__file__), path) -DATAFUSION_FUNCTION_COUNT = 283 +DATAFUSION_FUNCTION_COUNT = 285 @pytest.fixture(scope="session") diff --git a/wren-core-py/tests/test_modeling_core.py b/wren-core-py/tests/test_modeling_core.py index 46cc95100..8bd214e2b 100644 --- a/wren-core-py/tests/test_modeling_core.py +++ b/wren-core-py/tests/test_modeling_core.py @@ -137,7 +137,7 @@ def test_read_function_list(): path = "tests/functions.csv" session_context = SessionContext(manifest_str, path) functions = session_context.get_available_functions() - assert len(functions) == 290 + assert len(functions) == 292 rewritten_sql = session_context.transform_sql( "SELECT add_two(c_custkey, c_custkey) FROM my_catalog.my_schema.customer" @@ -149,7 +149,7 @@ def test_read_function_list(): session_context = SessionContext(manifest_str, None) functions = session_context.get_available_functions() - assert len(functions) == 283 + assert len(functions) == 285 def test_get_available_functions(): From 93d57032c452d492d716c739bd8a1f0ea71a3e90 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 12:59:45 +0800 Subject: [PATCH 5/7] fix test --- ibis-server/resources/function_list/bigquery.csv | 1 + ibis-server/resources/white_function_list/bigquery.csv | 1 + .../tests/routers/v3/connector/bigquery/test_functions.py | 4 ++-- .../tests/routers/v3/connector/local_file/test_functions.py | 2 +- 4 files changed, 5 insertions(+), 3 deletions(-) diff --git a/ibis-server/resources/function_list/bigquery.csv b/ibis-server/resources/function_list/bigquery.csv index c8ddf404d..669c4d21f 100644 --- a/ibis-server/resources/function_list/bigquery.csv +++ b/ibis-server/resources/function_list/bigquery.csv @@ -11,6 +11,7 @@ scalar,current_timestamp,timestamptz,,"","Returns current timestamp." scalar,date_add,date,,"date,int64","Adds a number of day to a date." scalar,date_sub,date,,"date,interval","Subtracts a specified interval from a date." scalar,date_diff,int64,,"date,date,granularity","Returns the difference between two dates." +scalar,datediff,int64,,"date,date,granularity","Returns the difference between two dates." scalar,timestamp_add,timestamp,,"timestamp,granularity","Adds a specified interval to a timestamp." scalar,timestamp_sub,timestamp,,"timestamp,granularity","Subtracts a specified interval from a timestamp." scalar,timestamp_diff,int64,,"timestamp,timestamp,granularity","Returns the difference between two timestamps." diff --git a/ibis-server/resources/white_function_list/bigquery.csv b/ibis-server/resources/white_function_list/bigquery.csv index 727cff78a..4aead35c6 100644 --- a/ibis-server/resources/white_function_list/bigquery.csv +++ b/ibis-server/resources/white_function_list/bigquery.csv @@ -10,6 +10,7 @@ scalar,current_datetime,timestamp,,"","Returns current date and time." scalar,date_add,date,,"date,int64","Adds a number of day to a date." scalar,date_sub,date,,"date,interval","Subtracts a specified interval from a date." scalar,date_diff,int64,,"date,date,granularity","Returns the difference between two dates." +scalar,datediff,int64,,"date,date,granularity","Returns the difference between two dates." scalar,timestamp_add,timestamp,,"timestamp,granularity","Adds a specified interval to a timestamp." scalar,timestamp_sub,timestamp,,"timestamp,granularity","Subtracts a specified interval from a timestamp." scalar,timestamp_diff,int64,,"timestamp,timestamp,granularity","Returns the difference between two timestamps." diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py index 7af20d7b5..1a161b106 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py @@ -25,6 +25,7 @@ }, "columns": [ {"name": "o_orderkey", "type": "integer"}, + {"name": "o_orderdate", "type": "date"}, ], }, ], @@ -52,7 +53,7 @@ async def test_function_list(client): response = await client.get(url=f"{base_url}/functions") assert response.status_code == 200 result = response.json() - assert len(result) == 175 + assert len(result) == 176 the_func = next( filter( lambda x: x["name"] == "string_agg", @@ -161,7 +162,6 @@ async def test_datetime_function(client, manifest_str: str, connection_info): "dtypes": {"col": "timestamp[us]"}, } - async def test_date_diff_function(client, manifest_str: str, connection_info): response = await client.post( url=f"{base_url}/query", diff --git a/ibis-server/tests/routers/v3/connector/local_file/test_functions.py b/ibis-server/tests/routers/v3/connector/local_file/test_functions.py index 9e3aeee4a..374d07fca 100644 --- a/ibis-server/tests/routers/v3/connector/local_file/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/local_file/test_functions.py @@ -54,7 +54,7 @@ async def test_function_list(client): assert response.status_code == 200 result = response.json() # 429 is the number of functions in `resources/function_list/duckdb.csv` file excluded the default functions in DataFusion. - assert len(result) == DATAFUSION_FUNCTION_COUNT + 427 + assert len(result) == DATAFUSION_FUNCTION_COUNT + 425 the_func = next(filter(lambda x: x["name"] == "regexp_escape", result)) assert the_func == { "name": "regexp_escape", From 2faec4fa52637b1212043825bacca00fecf0d288 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 13:01:04 +0800 Subject: [PATCH 6/7] fix fmt --- .../tests/routers/v3/connector/bigquery/test_functions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py index 1a161b106..8a9a1ec3d 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_functions.py @@ -162,6 +162,7 @@ async def test_datetime_function(client, manifest_str: str, connection_info): "dtypes": {"col": "timestamp[us]"}, } + async def test_date_diff_function(client, manifest_str: str, connection_info): response = await client.post( url=f"{base_url}/query", From 4a8aa3c91bfce632d96b27afef7a196594d9b63e Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Tue, 2 Sep 2025 13:26:15 +0800 Subject: [PATCH 7/7] fix test --- ibis-server/tests/routers/v3/connector/mysql/test_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ibis-server/tests/routers/v3/connector/mysql/test_functions.py b/ibis-server/tests/routers/v3/connector/mysql/test_functions.py index e219a1488..3b83ed337 100644 --- a/ibis-server/tests/routers/v3/connector/mysql/test_functions.py +++ b/ibis-server/tests/routers/v3/connector/mysql/test_functions.py @@ -54,7 +54,7 @@ async def test_function_list(client): response = await client.get(url=f"{base_url}/functions") assert response.status_code == 200 result = response.json() - assert len(result) == DATAFUSION_FUNCTION_COUNT + 25 + assert len(result) == DATAFUSION_FUNCTION_COUNT + 24 the_func = next(filter(lambda x: x["name"] == "lcase", result)) assert the_func == { "name": "lcase",