diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index de5e6b97c1af9..3a72f9d9fff16 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -30,7 +30,9 @@ use crate::datasource::provider_as_source; use crate::execution::SessionStateDefaults; use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner}; use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; -use arrow_schema::{DataType, FieldRef}; +#[cfg(feature = "sql")] +use arrow_schema::DataType; +use arrow_schema::FieldRef; use datafusion_catalog::MemoryCatalogProviderList; use datafusion_catalog::information_schema::{ INFORMATION_SCHEMA, InformationSchemaProvider, diff --git a/datafusion/spark/src/function/datetime/mod.rs b/datafusion/spark/src/function/datetime/mod.rs index 98afa91ddc834..bfd4c01810db7 100644 --- a/datafusion/spark/src/function/datetime/mod.rs +++ b/datafusion/spark/src/function/datetime/mod.rs @@ -32,6 +32,7 @@ pub mod time_trunc; pub mod to_utc_timestamp; pub mod trunc; pub mod unix; +pub mod weekday; use datafusion_expr::ScalarUDF; use datafusion_functions::make_udf_function; @@ -74,6 +75,7 @@ make_udf_function!( unix_seconds, unix::SparkUnixTimestamp::seconds ); +make_udf_function!(weekday::SparkWeekday, weekday); pub mod expr_fn { use datafusion_functions::export_functions; @@ -186,6 +188,11 @@ pub mod expr_fn { "Returns the number of seconds since epoch (1970-01-01 00:00:00 UTC) for the given timestamp `ts`.", ts )); + export_functions!(( + weekday, + "Returns the day of the week for date or timestamp with Monday as 0 and Sunday as 6.", + dt + )); } pub fn functions() -> Vec> { @@ -212,5 +219,6 @@ pub fn functions() -> Vec> { unix_micros(), unix_millis(), unix_seconds(), + weekday(), ] } diff --git a/datafusion/spark/src/function/datetime/weekday.rs b/datafusion/spark/src/function/datetime/weekday.rs new file mode 100644 index 0000000000000..b86e7c7df9296 --- /dev/null +++ b/datafusion/spark/src/function/datetime/weekday.rs @@ -0,0 +1,141 @@ +// 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::ArrayRef; +use arrow::compute::{DatePart, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::types::{NativeType, logical_date}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, internal_err}; +use datafusion_expr::{ + Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignatureClass, Volatility, +}; +use datafusion_functions::utils::make_scalar_function; + +/// Spark-compatible `weekday` expression. +/// +/// +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkWeekday { + signature: Signature, +} + +impl Default for SparkWeekday { + fn default() -> Self { + Self::new() + } +} + +impl SparkWeekday { + pub fn new() -> Self { + Self { + signature: Signature::coercible( + vec![Coercion::new_implicit( + TypeSignatureClass::Native(logical_date()), + vec![TypeSignatureClass::Timestamp], + NativeType::Date, + )], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkWeekday { + fn name(&self) -> &str { + "weekday" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(Field::new(self.name(), DataType::Int32, nullable))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(spark_weekday, vec![])(&args.args) + } +} + +fn spark_weekday(args: &[ArrayRef]) -> Result { + let [date_arg] = take_function_args("weekday", args)?; + Ok(date_part(date_arg.as_ref(), DatePart::DayOfWeekMonday0)?) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Date32Array, Int32Array}; + + #[test] + fn test_weekday_array() { + let date_array = Date32Array::from(vec![ + Some(4), // 1970-01-05, Monday + Some(5), // 1970-01-06, Tuesday + Some(6), // 1970-01-07, Wednesday + Some(7), // 1970-01-08, Thursday + Some(8), // 1970-01-09, Friday + Some(9), // 1970-01-10, Saturday + Some(10), // 1970-01-11, Sunday + None, + ]); + + let result = spark_weekday(&[Arc::new(date_array)]).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + + for (idx, expected) in (0..=6).enumerate() { + assert_eq!(result.value(idx), expected); + } + assert!(result.is_null(7)); + } + + #[test] + fn test_weekday_nullability_matches_input() { + let func = SparkWeekday::new(); + + let non_nullable_arg = Arc::new(Field::new("arg", DataType::Date32, false)); + let nullable_arg = Arc::new(Field::new("arg", DataType::Date32, true)); + + let non_nullable_out = func + .return_field_from_args(ReturnFieldArgs { + arg_fields: &[Arc::clone(&non_nullable_arg)], + scalar_arguments: &[None], + }) + .unwrap(); + assert_eq!(non_nullable_out.data_type(), &DataType::Int32); + assert!(!non_nullable_out.is_nullable()); + + let nullable_out = func + .return_field_from_args(ReturnFieldArgs { + arg_fields: &[Arc::clone(&nullable_arg)], + scalar_arguments: &[None], + }) + .unwrap(); + assert_eq!(nullable_out.data_type(), &DataType::Int32); + assert!(nullable_out.is_nullable()); + } +} diff --git a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt index b4f5444e8a2da..7d2ee7ea87efa 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt @@ -23,5 +23,65 @@ ## Original Query: SELECT weekday('2009-07-30'); ## PySpark 3.5.5 Result: {'weekday(2009-07-30)': 3, 'typeof(weekday(2009-07-30))': 'int', 'typeof(2009-07-30)': 'string'} -#query -#SELECT weekday('2009-07-30'::string); + +# Scalar date input +query I +SELECT weekday('2009-07-30'::DATE); +---- +3 + +# Monday = 0, ..., Sunday = 6 +query I +SELECT weekday(d) FROM (VALUES + (0, '2024-01-01'::DATE), + (1, '2024-01-02'::DATE), + (2, '2024-01-03'::DATE), + (3, '2024-01-04'::DATE), + (4, '2024-01-05'::DATE), + (5, '2024-01-06'::DATE), + (6, '2024-01-07'::DATE) +) AS t(i, d) ORDER BY i; +---- +0 +1 +2 +3 +4 +5 +6 + +# NULL handling +query I +SELECT weekday(NULL::DATE); +---- +NULL + +# Timestamp input: Spark coerces TIMESTAMP/TIMESTAMP_NTZ to DATE before evaluation +query I +SELECT weekday('2024-01-07 23:59:59'::TIMESTAMP); +---- +6 + +query I +SELECT weekday(arrow_cast('2024-01-03 12:34:56', 'Timestamp(Microsecond, None)')); +---- +2 + +query I +SELECT weekday(arrow_cast(NULL, 'Timestamp(Microsecond, None)')); +---- +NULL + +# Return type +query T +SELECT arrow_typeof(weekday('2024-01-01'::DATE)); +---- +Int32 + +# Error: wrong argument type +statement error Function 'weekday' requires Date, but received Int64 +SELECT weekday(123); + +# Error: no arguments +statement error 'weekday' does not support zero arguments +SELECT weekday();