diff --git a/datafusion/spark/src/function/datetime/date_from_unix_date.rs b/datafusion/spark/src/function/datetime/date_from_unix_date.rs new file mode 100644 index 0000000000000..33b989bec0a37 --- /dev/null +++ b/datafusion/spark/src/function/datetime/date_from_unix_date.rs @@ -0,0 +1,181 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{Array, Date32Array, Int32Array}; +use arrow::datatypes::DataType; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue, exec_err, internal_datafusion_err}; +use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; +use datafusion_expr::{ + ColumnarValue, Expr, ExprSchemable, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; + +/// Spark-compatible `date_from_unix_date` function. +/// Creates a date from the number of days since epoch (1970-01-01). +/// +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkDateFromUnixDate { + signature: Signature, +} + +impl Default for SparkDateFromUnixDate { + fn default() -> Self { + Self::new() + } +} + +impl SparkDateFromUnixDate { + pub fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkDateFromUnixDate { + fn name(&self) -> &str { + "date_from_unix_date" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Date32) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [days] = take_function_args(self.name(), args.args)?; + match days { + ColumnarValue::Array(array) => { + let days = + array.as_any().downcast_ref::().ok_or_else(|| { + internal_datafusion_err!( + "date_from_unix_date expected an Int32 array" + ) + })?; + // Date32 and Int32 both count days since the Unix epoch + // (1970-01-01), so the i32 values transfer directly. + let dates = + Date32Array::new(days.values().clone(), days.nulls().cloned()); + Ok(ColumnarValue::Array(Arc::new(dates))) + } + ColumnarValue::Scalar(ScalarValue::Int32(days)) => { + Ok(ColumnarValue::Scalar(ScalarValue::Date32(days))) + } + ColumnarValue::Scalar(other) => { + exec_err!("date_from_unix_date expected an Int32 argument, got {other:?}") + } + } + } + + fn simplify( + &self, + args: Vec, + info: &SimplifyContext, + ) -> Result { + let [days] = take_function_args(self.name(), args)?; + Ok(ExprSimplifyResult::Simplified( + days.cast_to(&DataType::Date32, info.schema())?, + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::Field; + + /// Exercises the physical `invoke_with_args` path directly, as consumers + /// that build physical `ScalarFunctionExpr` nodes (e.g. DataFusion Comet) + /// call it without ever running the logical `SimplifyExpressions` pass. + #[test] + fn test_date_from_unix_date_array() -> Result<()> { + let func = SparkDateFromUnixDate::new(); + + let input = Int32Array::from(vec![Some(0), Some(1), Some(-1), None, Some(18628)]); + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(input))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Int32, true))], + number_rows: 5, + return_field: Arc::new(Field::new( + "date_from_unix_date", + DataType::Date32, + true, + )), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Array(arr) => { + let expected = Date32Array::from(vec![ + Some(0), + Some(1), + Some(-1), + None, + Some(18628), + ]); + assert_eq!(arr.as_ref(), &expected as &dyn Array); + } + other => panic!("unexpected result: {other:?}"), + } + + Ok(()) + } + + #[test] + fn test_date_from_unix_date_scalar() -> Result<()> { + let func = SparkDateFromUnixDate::new(); + + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Int32(Some(0)))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Int32, true))], + number_rows: 1, + return_field: Arc::new(Field::new( + "date_from_unix_date", + DataType::Date32, + true, + )), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Scalar(ScalarValue::Date32(Some(v))) => assert_eq!(v, 0), + other => panic!("unexpected result: {other:?}"), + } + + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Int32(None))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Int32, true))], + number_rows: 1, + return_field: Arc::new(Field::new( + "date_from_unix_date", + DataType::Date32, + true, + )), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Scalar(ScalarValue::Date32(None)) => {} + other => panic!("unexpected result: {other:?}"), + } + + Ok(()) + } +} diff --git a/datafusion/spark/src/function/datetime/mod.rs b/datafusion/spark/src/function/datetime/mod.rs index 70ab024c329aa..828b3927b7070 100644 --- a/datafusion/spark/src/function/datetime/mod.rs +++ b/datafusion/spark/src/function/datetime/mod.rs @@ -18,6 +18,7 @@ pub mod add_months; pub mod date_add; pub mod date_diff; +pub mod date_from_unix_date; pub mod date_part; pub mod date_sub; pub mod date_trunc; @@ -41,6 +42,10 @@ use std::sync::Arc; make_udf_function!(add_months::SparkAddMonths, add_months); make_udf_function!(date_add::SparkDateAdd, date_add); make_udf_function!(date_diff::SparkDateDiff, date_diff); +make_udf_function!( + date_from_unix_date::SparkDateFromUnixDate, + date_from_unix_date +); make_udf_function!(date_part::SparkDatePart, date_part); make_udf_function!(date_sub::SparkDateSub, date_sub); make_udf_function!(date_trunc::SparkDateTrunc, date_trunc); @@ -138,6 +143,11 @@ pub mod expr_fn { "Returns the number of days from start `start` to end `end`.", end start )); + export_functions!(( + date_from_unix_date, + "Creates a date from the number of days since epoch (1970-01-01) `days`.", + days + )); export_functions!(( date_trunc, "Truncates a timestamp `ts` to the unit specified by the format `fmt`.", @@ -200,6 +210,7 @@ pub fn functions() -> Vec> { add_months(), date_add(), date_diff(), + date_from_unix_date(), date_part(), date_sub(), date_trunc(), diff --git a/datafusion/sqllogictest/test_files/spark/datetime/date_from_unix_date.slt b/datafusion/sqllogictest/test_files/spark/datetime/date_from_unix_date.slt new file mode 100644 index 0000000000000..0f4662f0e7430 --- /dev/null +++ b/datafusion/sqllogictest/test_files/spark/datetime/date_from_unix_date.slt @@ -0,0 +1,135 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# date_from_unix_date tests +# +# The function lowers to an Int32 -> Date32 cast in `simplify()`, so the logical +# optimizer replaces it before physical planning and `invoke_with_args` is never +# reached. Consumers that build physical expressions directly (for example +# DataFusion Comet) do not run `SimplifyExpressions` and call `invoke_with_args` +# instead. Both paths are covered below: the second section disables the logical +# optimizer so that the function survives into the physical plan. + +########## +# simplify() path (default) +########## + +# The plan below confirms this section exercises the rewritten cast. +query TT +EXPLAIN SELECT date_from_unix_date(a) FROM (VALUES (0::INT), (1::INT), (18628::INT)) AS t(a); +---- +logical_plan +01)Projection: CAST(t.a AS Date32) AS date_from_unix_date(t.a) +02)--SubqueryAlias: t +03)----Projection: column1 AS a +04)------Values: (Int32(0) AS Int64(0)), (Int32(1) AS Int64(1)), (Int32(18628) AS Int64(18628)) +physical_plan +01)ProjectionExec: expr=[CAST(column1@0 AS Date32) as date_from_unix_date(t.a)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Scalar cases +query D +SELECT date_from_unix_date(0::INT); +---- +1970-01-01 + +query D +SELECT date_from_unix_date(1::INT); +---- +1970-01-02 + +query D +SELECT date_from_unix_date((-1)::INT); +---- +1969-12-31 + +query D +SELECT date_from_unix_date(18628::INT); +---- +2021-01-01 + +# NULL input +query D +SELECT date_from_unix_date(NULL::INT); +---- +NULL + +# Array input +query D +SELECT date_from_unix_date(a) FROM (VALUES (0::INT), (1::INT), (18628::INT)) AS t(a); +---- +1970-01-01 +1970-01-02 +2021-01-01 + +########## +# invoke_with_args() path +########## + +statement ok +set datafusion.optimizer.max_passes = 0; + +# The plan below confirms this section exercises the function itself. +query TT +EXPLAIN SELECT date_from_unix_date(a) FROM (VALUES (0::INT), (1::INT), (18628::INT)) AS t(a); +---- +logical_plan +01)Projection: date_from_unix_date(t.a) +02)--SubqueryAlias: t +03)----Projection: column1 AS a +04)------Values: (CAST(Int64(0) AS Int32)), (CAST(Int64(1) AS Int32)), (CAST(Int64(18628) AS Int32)) +physical_plan +01)ProjectionExec: expr=[date_from_unix_date(column1@0) as date_from_unix_date(t.a)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Scalar cases +query D +SELECT date_from_unix_date(0::INT); +---- +1970-01-01 + +query D +SELECT date_from_unix_date(1::INT); +---- +1970-01-02 + +query D +SELECT date_from_unix_date((-1)::INT); +---- +1969-12-31 + +query D +SELECT date_from_unix_date(18628::INT); +---- +2021-01-01 + +# NULL input +query D +SELECT date_from_unix_date(NULL::INT); +---- +NULL + +# Array input +query D +SELECT date_from_unix_date(a) FROM (VALUES (0::INT), (1::INT), (18628::INT)) AS t(a); +---- +1970-01-01 +1970-01-02 +2021-01-01 + +statement ok +reset datafusion.optimizer.max_passes;