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::any::Any;
19use std::sync::Arc;
20
21use arrow::array::builder::PrimitiveBuilder;
22use arrow::array::cast::AsArray;
23use arrow::array::types::{Date32Type, Int32Type};
24use arrow::array::{Array, PrimitiveArray};
25use arrow::datatypes::DataType;
26use arrow::datatypes::DataType::Date32;
27use chrono::prelude::*;
28
29use datafusion_common::types::{NativeType, logical_int32, logical_string};
30use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
31use datafusion_expr::{
32    Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
33    Volatility,
34};
35use datafusion_macros::user_doc;
36
37#[user_doc(
38    doc_section(label = "Time and Date Functions"),
39    description = "Make a date from year/month/day component parts.",
40    syntax_example = "make_date(year, month, day)",
41    sql_example = r#"```sql
42> select make_date(2023, 1, 31);
43+-------------------------------------------+
44| make_date(Int64(2023),Int64(1),Int64(31)) |
45+-------------------------------------------+
46| 2023-01-31                                |
47+-------------------------------------------+
48> select make_date('2023', '01', '31');
49+-----------------------------------------------+
50| make_date(Utf8("2023"),Utf8("01"),Utf8("31")) |
51+-----------------------------------------------+
52| 2023-01-31                                    |
53+-----------------------------------------------+
54```
55
56Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/date_time.rs)
57"#,
58    argument(
59        name = "year",
60        description = "Year to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
61    ),
62    argument(
63        name = "month",
64        description = "Month to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
65    ),
66    argument(
67        name = "day",
68        description = "Day to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators."
69    )
70)]
71#[derive(Debug, PartialEq, Eq, Hash)]
72pub struct MakeDateFunc {
73    signature: Signature,
74}
75
76impl Default for MakeDateFunc {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl MakeDateFunc {
83    pub fn new() -> Self {
84        let int = Coercion::new_implicit(
85            TypeSignatureClass::Native(logical_int32()),
86            vec![
87                TypeSignatureClass::Integer,
88                TypeSignatureClass::Native(logical_string()),
89            ],
90            NativeType::Int32,
91        );
92        Self {
93            signature: Signature::coercible(vec![int; 3], Volatility::Immutable),
94        }
95    }
96}
97
98impl ScalarUDFImpl for MakeDateFunc {
99    fn as_any(&self) -> &dyn Any {
100        self
101    }
102
103    fn name(&self) -> &str {
104        "make_date"
105    }
106
107    fn signature(&self) -> &Signature {
108        &self.signature
109    }
110
111    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
112        Ok(Date32)
113    }
114
115    fn invoke_with_args(
116        &self,
117        args: datafusion_expr::ScalarFunctionArgs,
118    ) -> Result<ColumnarValue> {
119        let [years, months, days] = take_function_args(self.name(), args.args)?;
120
121        match (years, months, days) {
122            (ColumnarValue::Scalar(y), _, _) if y.is_null() => {
123                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
124            }
125            (_, ColumnarValue::Scalar(m), _) if m.is_null() => {
126                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
127            }
128            (_, _, ColumnarValue::Scalar(d)) if d.is_null() => {
129                Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
130            }
131            (
132                ColumnarValue::Scalar(ScalarValue::Int32(Some(years))),
133                ColumnarValue::Scalar(ScalarValue::Int32(Some(months))),
134                ColumnarValue::Scalar(ScalarValue::Int32(Some(days))),
135            ) => {
136                let mut value = 0;
137                make_date_inner(years, months, days, |days: i32| value = days)?;
138                Ok(ColumnarValue::Scalar(ScalarValue::Date32(Some(value))))
139            }
140            (years, months, days) => {
141                let len = args.number_rows;
142                let years = years.into_array(len)?;
143                let months = months.into_array(len)?;
144                let days = days.into_array(len)?;
145
146                let years = years.as_primitive::<Int32Type>();
147                let months = months.as_primitive::<Int32Type>();
148                let days = days.as_primitive::<Int32Type>();
149
150                let mut builder: PrimitiveBuilder<Date32Type> =
151                    PrimitiveArray::builder(len);
152
153                for i in 0..len {
154                    // match postgresql behaviour which returns null for any null input
155                    if years.is_null(i) || months.is_null(i) || days.is_null(i) {
156                        builder.append_null();
157                    } else {
158                        make_date_inner(
159                            years.value(i),
160                            months.value(i),
161                            days.value(i),
162                            |days: i32| builder.append_value(days),
163                        )?;
164                    }
165                }
166
167                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
168            }
169        }
170    }
171
172    fn documentation(&self) -> Option<&Documentation> {
173        self.doc()
174    }
175}
176
177/// Converts the year/month/day fields to an `i32` representing the days from
178/// the unix epoch and invokes `date_consumer_fn` with the value
179fn make_date_inner<F: FnMut(i32)>(
180    year: i32,
181    month: i32,
182    day: i32,
183    mut date_consumer_fn: F,
184) -> Result<()> {
185    let m = match month {
186        1..=12 => month as u32,
187        _ => return exec_err!("Month value '{month:?}' is out of range"),
188    };
189    let d = match day {
190        1..=31 => day as u32,
191        _ => return exec_err!("Day value '{day:?}' is out of range"),
192    };
193
194    if let Some(date) = NaiveDate::from_ymd_opt(year, m, d) {
195        // The number of days until the start of the unix epoch in the proleptic Gregorian calendar
196        // (with January 1, Year 1 (CE) as day 1). See [Datelike::num_days_from_ce].
197        const UNIX_DAYS_FROM_CE: i32 = 719_163;
198
199        // since the epoch for the date32 datatype is the unix epoch
200        // we need to subtract the unix epoch from the current date
201        // note that this can result in a negative value
202        date_consumer_fn(date.num_days_from_ce() - UNIX_DAYS_FROM_CE);
203        Ok(())
204    } else {
205        exec_err!("Unable to parse date from {year}, {month}, {day}")
206    }
207}