datafusion_functions/datetime/
make_date.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::cast::AsArray;
21use arrow::array::types::{Date32Type, Int32Type};
22use arrow::array::{Array, PrimitiveArray};
23use arrow::buffer::NullBuffer;
24use arrow::datatypes::DataType;
25use arrow::datatypes::DataType::Date32;
26use chrono::prelude::*;
27
28use datafusion_common::types::{NativeType, logical_int32, logical_string};
29use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
30use datafusion_expr::{
31    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    TypeSignatureClass, Volatility,
33};
34use datafusion_macros::user_doc;
35
36#[user_doc(
37    doc_section(label = "Time and Date Functions"),
38    description = "Make a date from year/month/day component parts.",
39    syntax_example = "make_date(year, month, day)",
40    sql_example = r#"```sql
41> select make_date(2023, 1, 31);
42+-------------------------------------------+
43| make_date(Int64(2023),Int64(1),Int64(31)) |
44+-------------------------------------------+
45| 2023-01-31                                |
46+-------------------------------------------+
47> select make_date('2023', '01', '31');
48+-----------------------------------------------+
49| make_date(Utf8("2023"),Utf8("01"),Utf8("31")) |
50+-----------------------------------------------+
51| 2023-01-31                                    |
52+-----------------------------------------------+
53```
54
55Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/date_time.rs)
56"#,
57    argument(
58        name = "year",
59        description = "Year to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
60    ),
61    argument(
62        name = "month",
63        description = "Month to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
64    ),
65    argument(
66        name = "day",
67        description = "Day to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
68    )
69)]
70#[derive(Debug, PartialEq, Eq, Hash)]
71pub struct MakeDateFunc {
72    signature: Signature,
73}
74
75impl Default for MakeDateFunc {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81impl MakeDateFunc {
82    pub fn new() -> Self {
83        let int = Coercion::new_implicit(
84            TypeSignatureClass::Native(logical_int32()),
85            vec![
86                TypeSignatureClass::Integer,
87                TypeSignatureClass::Native(logical_string()),
88            ],
89            NativeType::Int32,
90        );
91        Self {
92            signature: Signature::coercible(vec![int; 3], Volatility::Immutable),
93        }
94    }
95}
96
97impl ScalarUDFImpl for MakeDateFunc {
98    fn name(&self) -> &str {
99        "make_date"
100    }
101
102    fn signature(&self) -> &Signature {
103        &self.signature
104    }
105
106    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
107        Ok(Date32)
108    }
109
110    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
111        let [years, months, days] = take_function_args(self.name(), args.args)?;
112
113        match (years, months, days) {
114            (ColumnarValue::Scalar(y), _, _) if y.is_null() => {
115                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
116            }
117            (_, ColumnarValue::Scalar(m), _) if m.is_null() => {
118                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
119            }
120            (_, _, ColumnarValue::Scalar(d)) if d.is_null() => {
121                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
122            }
123            (
124                ColumnarValue::Scalar(ScalarValue::Int32(Some(years))),
125                ColumnarValue::Scalar(ScalarValue::Int32(Some(months))),
126                ColumnarValue::Scalar(ScalarValue::Int32(Some(days))),
127            ) => {
128                let mut value = 0;
129                make_date_inner(years, months, days, |days: i32| value = days)?;
130                Ok(ColumnarValue::Scalar(ScalarValue::Date32(Some(value))))
131            }
132            (years, months, days) => {
133                let len = args.number_rows;
134                let years = years.into_array(len)?;
135                let months = months.into_array(len)?;
136                let days = days.into_array(len)?;
137
138                let years = years.as_primitive::<Int32Type>();
139                let months = months.as_primitive::<Int32Type>();
140                let days = days.as_primitive::<Int32Type>();
141
142                let nulls =
143                    NullBuffer::union_many([years.nulls(), months.nulls(), days.nulls()]);
144
145                let mut values = Vec::with_capacity(len);
146                for i in 0..len {
147                    // match postgresql behaviour which returns null for any null input
148                    if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
149                        values.push(0);
150                    } else {
151                        make_date_inner(
152                            years.value(i),
153                            months.value(i),
154                            days.value(i),
155                            |days: i32| values.push(days),
156                        )?;
157                    }
158                }
159
160                Ok(ColumnarValue::Array(Arc::new(
161                    PrimitiveArray::<Date32Type>::new(values.into(), nulls),
162                )))
163            }
164        }
165    }
166
167    fn documentation(&self) -> Option<&Documentation> {
168        self.doc()
169    }
170}
171
172/// Converts the year/month/day fields to an `i32` representing the days from
173/// the unix epoch and invokes `date_consumer_fn` with the value
174fn make_date_inner<F: FnMut(i32)>(
175    year: i32,
176    month: i32,
177    day: i32,
178    mut date_consumer_fn: F,
179) -> Result<()> {
180    let m = match month {
181        1..=12 => month as u32,
182        _ => return exec_err!("Month value '{month:?}' is out of range"),
183    };
184    let d = match day {
185        1..=31 => day as u32,
186        _ => return exec_err!("Day value '{day:?}' is out of range"),
187    };
188
189    if let Some(date) = NaiveDate::from_ymd_opt(year, m, d) {
190        // The number of days until the start of the unix epoch in the proleptic Gregorian calendar
191        // (with January 1, Year 1 (CE) as day 1). See [Datelike::num_days_from_ce].
192        const UNIX_DAYS_FROM_CE: i32 = 719_163;
193
194        // since the epoch for the date32 datatype is the unix epoch
195        // we need to subtract the unix epoch from the current date
196        // note that this can result in a negative value
197        date_consumer_fn(date.num_days_from_ce() - UNIX_DAYS_FROM_CE);
198        Ok(())
199    } else {
200        exec_err!("Unable to parse date from {year}, {month}, {day}")
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use arrow::array::Int32Array;
208    use arrow::datatypes::Field;
209    use datafusion_common::config::ConfigOptions;
210
211    fn invoke(args: Vec<ColumnarValue>, number_rows: usize) -> Result<ColumnarValue> {
212        let arg_fields = args
213            .iter()
214            .map(|a| Field::new("a", a.data_type(), true).into())
215            .collect::<Vec<_>>();
216        MakeDateFunc::new().invoke_with_args(ScalarFunctionArgs {
217            args,
218            arg_fields,
219            number_rows,
220            return_field: Field::new("f", Date32, true).into(),
221            config_options: Arc::new(ConfigOptions::default()),
222        })
223    }
224
225    #[test]
226    fn test_make_date_array() {
227        let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![
228            Some(1970),
229            Some(1970),
230        ])));
231        let months =
232            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(1)])));
233        let days =
234            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(2)])));
235
236        let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap()
237        else {
238            panic!("expected array result");
239        };
240        let arr = arr.as_primitive::<Date32Type>();
241        // Days since the unix epoch.
242        assert_eq!(arr.value(0), 0);
243        assert_eq!(arr.value(1), 1);
244    }
245
246    #[test]
247    fn test_make_date_null_propagation() {
248        // A NULL in any component column yields a NULL row.
249        let years =
250            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000), None])));
251        let months =
252            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(6), Some(6)])));
253        let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![None, Some(15)])));
254
255        let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap()
256        else {
257            panic!("expected array result");
258        };
259        let arr = arr.as_primitive::<Date32Type>();
260        assert!(arr.is_null(0));
261        assert!(arr.is_null(1));
262    }
263
264    #[test]
265    fn test_make_date_scalar_array_mix() {
266        let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(1970)));
267        let month = ColumnarValue::Scalar(ScalarValue::Int32(Some(1)));
268        let days =
269            ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(3)])));
270
271        let ColumnarValue::Array(arr) = invoke(vec![year, month, days], 2).unwrap()
272        else {
273            panic!("expected array result");
274        };
275        let arr = arr.as_primitive::<Date32Type>();
276        assert_eq!(arr.value(0), 0);
277        assert_eq!(arr.value(1), 2);
278    }
279
280    #[test]
281    fn test_make_date_out_of_range_errors() {
282        let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000)])));
283        let months = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(13)])));
284        let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1)])));
285        assert!(invoke(vec![years, months, days], 1).is_err());
286    }
287}
datafusion_functions/datetime/make_date.rs

datafusion_functions/datetime/
make_date.rs
