diff --git a/ibis-server/tests/routers/v3/connector/postgres/test_query.py b/ibis-server/tests/routers/v3/connector/postgres/test_query.py index 5a9436469..78638546f 100644 --- a/ibis-server/tests/routers/v3/connector/postgres/test_query.py +++ b/ibis-server/tests/routers/v3/connector/postgres/test_query.py @@ -1279,3 +1279,23 @@ async def test_case_sensitive_without_quote(client, connection_info): "O_orderkey": "int32", "O_custkey": "int32", } + + +async def test_to_char_numeric(client, manifest_str, connection_info): + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT to_char(1234, '9999') AS formatted_number, to_char(1234.567, '0000.00') AS formatted_double", + }, + headers={X_WREN_FALLBACK_DISABLE: "true"}, + ) + assert response.status_code == 200 + result = response.json() + assert result["data"][0][0] == " 1234" + assert result["data"][0][1] == " 1234.57" + assert result["dtypes"] == { + "formatted_number": "string", + "formatted_double": "string", + } diff --git a/wren-core/core/src/mdl/function/macros.rs b/wren-core/core/src/mdl/function/macros.rs new file mode 100644 index 000000000..fdc7f5076 --- /dev/null +++ b/wren-core/core/src/mdl/function/macros.rs @@ -0,0 +1,21 @@ +/// Creates a singleton `ScalarUDF` of the `$UDF` function and a function +/// named `$NAME` which returns that singleton. +/// +/// This is used to ensure creating the list of `ScalarUDF` only happens once. +#[macro_export] +macro_rules! make_udf_function { + ($UDF:ty, $NAME:ident) => { + #[doc = concat!("Return a [`ScalarUDF`](datafusion_expr::ScalarUDF) implementation of ", stringify!($NAME))] + pub fn $NAME() -> std::sync::Arc { + // Singleton instance of the function + static INSTANCE: std::sync::LazyLock< + std::sync::Arc, + > = std::sync::LazyLock::new(|| { + std::sync::Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl( + <$UDF>::new(), + )) + }); + std::sync::Arc::clone(&INSTANCE) + } + }; +} diff --git a/wren-core/core/src/mdl/function/mod.rs b/wren-core/core/src/mdl/function/mod.rs index 54526b1ac..1cae051d7 100644 --- a/wren-core/core/src/mdl/function/mod.rs +++ b/wren-core/core/src/mdl/function/mod.rs @@ -1,4 +1,5 @@ mod aggregate; +mod macros; mod remote_function; mod scalar; mod table; diff --git a/wren-core/core/src/mdl/function/scalar/mod.rs b/wren-core/core/src/mdl/function/scalar/mod.rs index 7219efacf..ab3388082 100644 --- a/wren-core/core/src/mdl/function/scalar/mod.rs +++ b/wren-core/core/src/mdl/function/scalar/mod.rs @@ -9,6 +9,12 @@ use datafusion::{ logical_expr::ScalarUDF, }; +use crate::make_udf_function; + +mod to_char; + +make_udf_function!(to_char::ToCharFunc, to_char); + pub fn scalar_functions() -> Vec> { vec![ // datefusion core diff --git a/wren-core/core/src/mdl/function/scalar/to_char.rs b/wren-core/core/src/mdl/function/scalar/to_char.rs new file mode 100644 index 000000000..7fd733bd6 --- /dev/null +++ b/wren-core/core/src/mdl/function/scalar/to_char.rs @@ -0,0 +1,83 @@ +use std::any::Any; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::not_impl_err; +use datafusion::common::types::{ + logical_date, logical_float32, logical_float64, logical_string, +}; +use datafusion::error::Result; +use datafusion::logical_expr::TypeSignatureClass; +use datafusion::logical_expr::{ + Coercion, ColumnarValue, DocSection, Documentation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; + +#[derive(Debug)] +pub struct ToCharFunc { + signature: Signature, + aliases: Vec, + doc: Option, +} + +impl Default for ToCharFunc { + fn default() -> Self { + Self::new() + } +} + +impl ToCharFunc { + pub fn new() -> Self { + Self { + signature: Signature::one_of(vec![ + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Duration), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Interval), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Native(logical_date())), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Time), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Timestamp), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Integer), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Native(logical_float64())), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Coercible(vec![Coercion::new_exact(TypeSignatureClass::Native(logical_float32())), Coercion::new_exact(TypeSignatureClass::Native(logical_string()))]), + TypeSignature::Exact(vec![DataType::Decimal128(38, 10), DataType::Utf8]), + TypeSignature::Exact(vec![DataType::Decimal256(38, 10), DataType::Utf8]), + ], Volatility::Immutable,), + aliases: vec!["date_format".to_string()], + doc: Some(Documentation::builder( + DocSection { include: false, label: "Time and Date Functions", description: None }, + "Returns a string representation of a date, time, timestamp or duration based on a [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html).", + "to_char(expression, format)") + .with_argument("expression", "The date, time, timestamp or duration to format.") + .with_argument("format", "The [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) string to use for formatting.") + .build()), + } + } +} + +impl ScalarUDFImpl for ToCharFunc { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "to_char" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("to_char is not implemented") + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc.as_ref() + } +} diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs index 3d42e62e3..850902e81 100644 --- a/wren-core/core/src/mdl/mod.rs +++ b/wren-core/core/src/mdl/mod.rs @@ -3732,6 +3732,48 @@ mod test { Ok(()) } + #[tokio::test] + async fn test_to_char() -> Result<()> { + let ctx = create_wren_ctx(None); + let manifest = ManifestBuilder::new() + .catalog("wren") + .schema("test") + .model( + ModelBuilder::new("customer") + .table_reference("customer") + .column(ColumnBuilder::new("c_date", "date").build()) + .column(ColumnBuilder::new("c_timestamp", "timestamp").build()) + .column(ColumnBuilder::new("c_timestamptz", "timestamptz").build()) + .column(ColumnBuilder::new("c_int", "int").build()) + .column(ColumnBuilder::new("c_bigint", "bigint").build()) + .column(ColumnBuilder::new("c_float", "float").build()) + .column(ColumnBuilder::new("c_double", "double").build()) + .column(ColumnBuilder::new("c_decimal", "decimal").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 to_char(c_date, '%Y-%m-%d'), to_char(c_timestamp, '%Y-%m-%d'), to_char(c_timestamptz, '%Y-%m-%d') FROM customer"; + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT to_char(customer.c_date, '%Y-%m-%d'), to_char(customer.c_timestamp, '%Y-%m-%d'), to_char(customer.c_timestamptz, '%Y-%m-%d') FROM (SELECT customer.c_date, customer.c_timestamp, customer.c_timestamptz FROM (SELECT __source.c_date AS c_date, __source.c_timestamp AS c_timestamp, __source.c_timestamptz AS c_timestamptz FROM customer AS __source) AS customer) AS customer" + ); + + let sql = "SELECT to_char(c_int, '999'), to_char(c_bigint, '999'), to_char(c_float, '999.99'), to_char(c_double, '999.99'), to_char(c_decimal, '999.99') FROM customer"; + assert_snapshot!( + transform_sql_with_ctx(&ctx, Arc::clone(&analyzed_mdl), &[], Arc::clone(&headers), sql).await?, + @"SELECT to_char(customer.c_int, '999'), to_char(customer.c_bigint, '999'), to_char(customer.c_float, '999.99'), to_char(customer.c_double, '999.99'), to_char(customer.c_decimal, '999.99') FROM (SELECT customer.c_bigint, customer.c_decimal, customer.c_double, customer.c_float, customer.c_int FROM (SELECT __source.c_bigint AS c_bigint, __source.c_decimal AS c_decimal, __source.c_double AS c_double, __source.c_float AS c_float, __source.c_int AS c_int FROM customer AS __source) AS customer) AS customer" + ); + + Ok(()) + } + /// Return a RecordBatch with made up data about customer fn customer() -> RecordBatch { let custkey: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3]));