datafusion_spark/function/datetime/
weekday.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow::array::AsArray;
21use arrow::compute::{DatePart, date_part};
22use arrow::datatypes::{DataType, Field, FieldRef, Int32Type};
23use datafusion_common::types::{NativeType, logical_date};
24use datafusion_common::utils::take_function_args;
25use datafusion_common::{Result, ScalarValue, internal_err};
26use datafusion_expr::{
27    Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl,
28    Signature, TypeSignatureClass, Volatility,
29};
30
31/// Spark-compatible `weekday` expression.
32/// Returns the day of the week for a date or timestamp as an integer index where
33/// Monday = 0, Tuesday = 1, ..., Sunday = 6.
34///
35/// Note: this differs from `dayofweek`, which is 1-indexed with Sunday = 1.
36///
37/// <https://spark.apache.org/docs/latest/api/sql/index.html#weekday>
38#[derive(Debug, PartialEq, Eq, Hash)]
39pub struct SparkWeekDay {
40    signature: Signature,
41}
42
43impl Default for SparkWeekDay {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl SparkWeekDay {
50    pub fn new() -> Self {
51        Self {
52            signature: Signature::coercible(
53                vec![Coercion::new_implicit(
54                    TypeSignatureClass::Native(logical_date()),
55                    vec![TypeSignatureClass::Timestamp],
56                    NativeType::Date,
57                )],
58                Volatility::Immutable,
59            ),
60        }
61    }
62}
63
64impl ScalarUDFImpl for SparkWeekDay {
65    fn name(&self) -> &str {
66        "weekday"
67    }
68
69    fn signature(&self) -> &Signature {
70        &self.signature
71    }
72
73    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
74        internal_err!("return_field_from_args should be used instead")
75    }
76
77    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
78        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
79        Ok(Arc::new(Field::new(self.name(), DataType::Int32, nullable)))
80    }
81
82    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
83        let [arg] = take_function_args(self.name(), args.args)?;
84        match arg {
85            ColumnarValue::Scalar(scalar) => {
86                if scalar.is_null() {
87                    return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None)));
88                }
89                let arr = scalar.to_array_of_size(1)?;
90                // `DayOfWeekMonday0` returns 0..=6 with Monday = 0, which
91                // matches Spark `weekday` semantics exactly.
92                let weekday_arr = date_part(&arr, DatePart::DayOfWeekMonday0)?;
93                let value = weekday_arr.as_primitive::<Int32Type>().value(0);
94                Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(value))))
95            }
96            ColumnarValue::Array(arr) => {
97                let weekday_arr = date_part(&arr, DatePart::DayOfWeekMonday0)?;
98                Ok(ColumnarValue::Array(weekday_arr))
99            }
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use arrow::array::{Date32Array, Int32Array};
108
109    #[test]
110    fn test_weekday_return_field_nullability_matches_input() {
111        let func = SparkWeekDay::new();
112
113        let non_nullable_arg = Arc::new(Field::new("arg", DataType::Date32, false));
114        let nullable_arg = Arc::new(Field::new("arg", DataType::Date32, true));
115
116        let non_nullable_out = func
117            .return_field_from_args(ReturnFieldArgs {
118                arg_fields: &[Arc::clone(&non_nullable_arg)],
119                scalar_arguments: &[None],
120            })
121            .expect("non-nullable arg should succeed");
122        assert_eq!(non_nullable_out.data_type(), &DataType::Int32);
123        assert!(!non_nullable_out.is_nullable());
124
125        let nullable_out = func
126            .return_field_from_args(ReturnFieldArgs {
127                arg_fields: &[Arc::clone(&nullable_arg)],
128                scalar_arguments: &[None],
129            })
130            .expect("nullable arg should succeed");
131        assert_eq!(nullable_out.data_type(), &DataType::Int32);
132        assert!(nullable_out.is_nullable());
133    }
134
135    #[test]
136    fn test_weekday_scalar() -> Result<()> {
137        let func = SparkWeekDay::new();
138
139        // 2024-03-15 is a Friday -> Spark weekday = 4 (Mon=0).
140        let result = func.invoke_with_args(ScalarFunctionArgs {
141            args: vec![ColumnarValue::Scalar(ScalarValue::Date32(Some(19797)))],
142            arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))],
143            number_rows: 1,
144            return_field: Arc::new(Field::new("weekday", DataType::Int32, true)),
145            config_options: Arc::new(Default::default()),
146        })?;
147        match result {
148            ColumnarValue::Scalar(ScalarValue::Int32(Some(v))) => assert_eq!(v, 4),
149            other => panic!("unexpected result: {other:?}"),
150        }
151
152        // NULL input -> NULL output.
153        let result = func.invoke_with_args(ScalarFunctionArgs {
154            args: vec![ColumnarValue::Scalar(ScalarValue::Date32(None))],
155            arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))],
156            number_rows: 1,
157            return_field: Arc::new(Field::new("weekday", DataType::Int32, true)),
158            config_options: Arc::new(Default::default()),
159        })?;
160        match result {
161            ColumnarValue::Scalar(ScalarValue::Int32(None)) => {}
162            other => panic!("unexpected result: {other:?}"),
163        }
164
165        Ok(())
166    }
167
168    #[test]
169    fn test_weekday_array() -> Result<()> {
170        let func = SparkWeekDay::new();
171
172        // 2024-01-01 Mon(0), 2024-01-06 Sat(5), 2024-01-07 Sun(6), NULL.
173        let input = Date32Array::from(vec![Some(19723), Some(19728), Some(19729), None]);
174        let result = func.invoke_with_args(ScalarFunctionArgs {
175            args: vec![ColumnarValue::Array(Arc::new(input))],
176            arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))],
177            number_rows: 4,
178            return_field: Arc::new(Field::new("weekday", DataType::Int32, true)),
179            config_options: Arc::new(Default::default()),
180        })?;
181        match result {
182            ColumnarValue::Array(arr) => {
183                let expected = Int32Array::from(vec![Some(0), Some(5), Some(6), None]);
184                assert_eq!(arr.as_primitive::<Int32Type>(), &expected);
185            }
186            other => panic!("unexpected result: {other:?}"),
187        }
188
189        Ok(())
190    }
191}
datafusion_spark/function/datetime/weekday.rs

datafusion_spark/function/datetime/
weekday.rs
