datafusion_functions/datetime/
date_part.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::str::FromStr;
20use std::sync::Arc;
21
22use arrow::array::{Array, ArrayRef, Float64Array, Int32Array};
23use arrow::compute::kernels::cast_utils::IntervalUnit;
24use arrow::compute::{binary, date_part, DatePart};
25use arrow::datatypes::DataType::{
26    Date32, Date64, Duration, Interval, Time32, Time64, Timestamp,
27};
28use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
29use arrow::datatypes::{DataType, Field, FieldRef, TimeUnit};
30use datafusion_common::types::{logical_date, NativeType};
31
32use datafusion_common::{
33    cast::{
34        as_date32_array, as_date64_array, as_int32_array, as_time32_millisecond_array,
35        as_time32_second_array, as_time64_microsecond_array, as_time64_nanosecond_array,
36        as_timestamp_microsecond_array, as_timestamp_millisecond_array,
37        as_timestamp_nanosecond_array, as_timestamp_second_array,
38    },
39    exec_err, internal_err, not_impl_err,
40    types::logical_string,
41    utils::take_function_args,
42    Result, ScalarValue,
43};
44use datafusion_expr::{
45    ColumnarValue, Documentation, ReturnFieldArgs, ScalarUDFImpl, Signature,
46    TypeSignature, Volatility,
47};
48use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
49use datafusion_macros::user_doc;
50
51#[user_doc(
52    doc_section(label = "Time and Date Functions"),
53    description = "Returns the specified part of the date as an integer.",
54    syntax_example = "date_part(part, expression)",
55    alternative_syntax = "extract(field FROM source)",
56    argument(
57        name = "part",
58        description = r#"Part of the date to return. The following date parts are supported:
59        
60    - year
61    - quarter (emits value in inclusive range [1, 4] based on which quartile of the year the date is in)
62    - month
63    - week (week of the year)
64    - day (day of the month)
65    - hour
66    - minute
67    - second
68    - millisecond
69    - microsecond
70    - nanosecond
71    - dow (day of the week)
72    - doy (day of the year)
73    - epoch (seconds since Unix epoch)
74"#
75    ),
76    argument(
77        name = "expression",
78        description = "Time expression to operate on. Can be a constant, column, or function."
79    )
80)]
81#[derive(Debug)]
82pub struct DatePartFunc {
83    signature: Signature,
84    aliases: Vec<String>,
85}
86
87impl Default for DatePartFunc {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93impl DatePartFunc {
94    pub fn new() -> Self {
95        Self {
96            signature: Signature::one_of(
97                vec![
98                    TypeSignature::Coercible(vec![
99                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
100                        Coercion::new_implicit(
101                            TypeSignatureClass::Timestamp,
102                            // Not consistent with Postgres and DuckDB but to avoid regression we implicit cast string to timestamp
103                            vec![TypeSignatureClass::Native(logical_string())],
104                            NativeType::Timestamp(Nanosecond, None),
105                        ),
106                    ]),
107                    TypeSignature::Coercible(vec![
108                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
109                        Coercion::new_exact(TypeSignatureClass::Native(logical_date())),
110                    ]),
111                    TypeSignature::Coercible(vec![
112                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
113                        Coercion::new_exact(TypeSignatureClass::Time),
114                    ]),
115                    TypeSignature::Coercible(vec![
116                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
117                        Coercion::new_exact(TypeSignatureClass::Interval),
118                    ]),
119                    TypeSignature::Coercible(vec![
120                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
121                        Coercion::new_exact(TypeSignatureClass::Duration),
122                    ]),
123                ],
124                Volatility::Immutable,
125            ),
126            aliases: vec![String::from("datepart")],
127        }
128    }
129}
130
131impl ScalarUDFImpl for DatePartFunc {
132    fn as_any(&self) -> &dyn Any {
133        self
134    }
135
136    fn name(&self) -> &str {
137        "date_part"
138    }
139
140    fn signature(&self) -> &Signature {
141        &self.signature
142    }
143
144    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
145        internal_err!("return_field_from_args should be called instead")
146    }
147
148    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
149        let [field, _] = take_function_args(self.name(), args.scalar_arguments)?;
150
151        field
152            .and_then(|sv| {
153                sv.try_as_str()
154                    .flatten()
155                    .filter(|s| !s.is_empty())
156                    .map(|part| {
157                        if is_epoch(part) {
158                            Field::new(self.name(), DataType::Float64, true)
159                        } else {
160                            Field::new(self.name(), DataType::Int32, true)
161                        }
162                    })
163            })
164            .map(Arc::new)
165            .map_or_else(
166                || exec_err!("{} requires non-empty constant string", self.name()),
167                Ok,
168            )
169    }
170
171    fn invoke_with_args(
172        &self,
173        args: datafusion_expr::ScalarFunctionArgs,
174    ) -> Result<ColumnarValue> {
175        let args = args.args;
176        let [part, array] = take_function_args(self.name(), args)?;
177
178        let part = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = part {
179            v
180        } else if let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(v))) = part {
181            v
182        } else {
183            return exec_err!(
184                "First argument of `DATE_PART` must be non-null scalar Utf8"
185            );
186        };
187
188        let is_scalar = matches!(array, ColumnarValue::Scalar(_));
189
190        let array = match array {
191            ColumnarValue::Array(array) => Arc::clone(&array),
192            ColumnarValue::Scalar(scalar) => scalar.to_array()?,
193        };
194
195        let part_trim = part_normalization(&part);
196
197        // using IntervalUnit here means we hand off all the work of supporting plurals (like "seconds")
198        // and synonyms ( like "ms,msec,msecond,millisecond") to Arrow
199        let arr = if let Ok(interval_unit) = IntervalUnit::from_str(part_trim) {
200            match interval_unit {
201                IntervalUnit::Year => date_part(array.as_ref(), DatePart::Year)?,
202                IntervalUnit::Month => date_part(array.as_ref(), DatePart::Month)?,
203                IntervalUnit::Week => date_part(array.as_ref(), DatePart::Week)?,
204                IntervalUnit::Day => date_part(array.as_ref(), DatePart::Day)?,
205                IntervalUnit::Hour => date_part(array.as_ref(), DatePart::Hour)?,
206                IntervalUnit::Minute => date_part(array.as_ref(), DatePart::Minute)?,
207                IntervalUnit::Second => seconds_as_i32(array.as_ref(), Second)?,
208                IntervalUnit::Millisecond => seconds_as_i32(array.as_ref(), Millisecond)?,
209                IntervalUnit::Microsecond => seconds_as_i32(array.as_ref(), Microsecond)?,
210                IntervalUnit::Nanosecond => seconds_as_i32(array.as_ref(), Nanosecond)?,
211                // century and decade are not supported by `DatePart`, although they are supported in postgres
212                _ => return exec_err!("Date part '{part}' not supported"),
213            }
214        } else {
215            // special cases that can be extracted (in postgres) but are not interval units
216            match part_trim.to_lowercase().as_str() {
217                "qtr" | "quarter" => date_part(array.as_ref(), DatePart::Quarter)?,
218                "doy" => date_part(array.as_ref(), DatePart::DayOfYear)?,
219                "dow" => date_part(array.as_ref(), DatePart::DayOfWeekSunday0)?,
220                "epoch" => epoch(array.as_ref())?,
221                _ => return exec_err!("Date part '{part}' not supported"),
222            }
223        };
224
225        Ok(if is_scalar {
226            ColumnarValue::Scalar(ScalarValue::try_from_array(arr.as_ref(), 0)?)
227        } else {
228            ColumnarValue::Array(arr)
229        })
230    }
231
232    fn aliases(&self) -> &[String] {
233        &self.aliases
234    }
235    fn documentation(&self) -> Option<&Documentation> {
236        self.doc()
237    }
238}
239
240fn is_epoch(part: &str) -> bool {
241    let part = part_normalization(part);
242    matches!(part.to_lowercase().as_str(), "epoch")
243}
244
245// Try to remove quote if exist, if the quote is invalid, return original string and let the downstream function handle the error
246fn part_normalization(part: &str) -> &str {
247    part.strip_prefix(|c| c == '\'' || c == '\"')
248        .and_then(|s| s.strip_suffix(|c| c == '\'' || c == '\"'))
249        .unwrap_or(part)
250}
251
252/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
253/// result to a total number of seconds, milliseconds, microseconds or
254/// nanoseconds
255fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
256    // Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing
257    // with overflow and precision issue we don't support nanosecond
258    if unit == Nanosecond {
259        return not_impl_err!("Date part {unit:?} not supported");
260    }
261
262    let conversion_factor = match unit {
263        Second => 1_000_000_000,
264        Millisecond => 1_000_000,
265        Microsecond => 1_000,
266        Nanosecond => 1,
267    };
268
269    let second_factor = match unit {
270        Second => 1,
271        Millisecond => 1_000,
272        Microsecond => 1_000_000,
273        Nanosecond => 1_000_000_000,
274    };
275
276    let secs = date_part(array, DatePart::Second)?;
277    // This assumes array is primitive and not a dictionary
278    let secs = as_int32_array(secs.as_ref())?;
279    let subsecs = date_part(array, DatePart::Nanosecond)?;
280    let subsecs = as_int32_array(subsecs.as_ref())?;
281
282    // Special case where there are no nulls.
283    if subsecs.null_count() == 0 {
284        let r: Int32Array = binary(secs, subsecs, |secs, subsecs| {
285            secs * second_factor + (subsecs % 1_000_000_000) / conversion_factor
286        })?;
287        Ok(Arc::new(r))
288    } else {
289        // Nulls in secs are preserved, nulls in subsecs are treated as zero to account for the case
290        // where the number of nanoseconds overflows.
291        let r: Int32Array = secs
292            .iter()
293            .zip(subsecs)
294            .map(|(secs, subsecs)| {
295                secs.map(|secs| {
296                    let subsecs = subsecs.unwrap_or(0);
297                    secs * second_factor + (subsecs % 1_000_000_000) / conversion_factor
298                })
299            })
300            .collect();
301        Ok(Arc::new(r))
302    }
303}
304
305/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
306/// result to a total number of seconds, milliseconds, microseconds or
307/// nanoseconds
308///
309/// Given epoch return f64, this is a duplicated function to optimize for f64 type
310fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
311    let sf = match unit {
312        Second => 1_f64,
313        Millisecond => 1_000_f64,
314        Microsecond => 1_000_000_f64,
315        Nanosecond => 1_000_000_000_f64,
316    };
317    let secs = date_part(array, DatePart::Second)?;
318    // This assumes array is primitive and not a dictionary
319    let secs = as_int32_array(secs.as_ref())?;
320    let subsecs = date_part(array, DatePart::Nanosecond)?;
321    let subsecs = as_int32_array(subsecs.as_ref())?;
322
323    // Special case where there are no nulls.
324    if subsecs.null_count() == 0 {
325        let r: Float64Array = binary(secs, subsecs, |secs, subsecs| {
326            (secs as f64 + ((subsecs % 1_000_000_000) as f64 / 1_000_000_000_f64)) * sf
327        })?;
328        Ok(Arc::new(r))
329    } else {
330        // Nulls in secs are preserved, nulls in subsecs are treated as zero to account for the case
331        // where the number of nanoseconds overflows.
332        let r: Float64Array = secs
333            .iter()
334            .zip(subsecs)
335            .map(|(secs, subsecs)| {
336                secs.map(|secs| {
337                    let subsecs = subsecs.unwrap_or(0);
338                    (secs as f64 + ((subsecs % 1_000_000_000) as f64 / 1_000_000_000_f64))
339                        * sf
340                })
341            })
342            .collect();
343        Ok(Arc::new(r))
344    }
345}
346
347fn epoch(array: &dyn Array) -> Result<ArrayRef> {
348    const SECONDS_IN_A_DAY: f64 = 86400_f64;
349
350    let f: Float64Array = match array.data_type() {
351        Timestamp(Second, _) => as_timestamp_second_array(array)?.unary(|x| x as f64),
352        Timestamp(Millisecond, _) => {
353            as_timestamp_millisecond_array(array)?.unary(|x| x as f64 / 1_000_f64)
354        }
355        Timestamp(Microsecond, _) => {
356            as_timestamp_microsecond_array(array)?.unary(|x| x as f64 / 1_000_000_f64)
357        }
358        Timestamp(Nanosecond, _) => {
359            as_timestamp_nanosecond_array(array)?.unary(|x| x as f64 / 1_000_000_000_f64)
360        }
361        Date32 => as_date32_array(array)?.unary(|x| x as f64 * SECONDS_IN_A_DAY),
362        Date64 => as_date64_array(array)?.unary(|x| x as f64 / 1_000_f64),
363        Time32(Second) => as_time32_second_array(array)?.unary(|x| x as f64),
364        Time32(Millisecond) => {
365            as_time32_millisecond_array(array)?.unary(|x| x as f64 / 1_000_f64)
366        }
367        Time64(Microsecond) => {
368            as_time64_microsecond_array(array)?.unary(|x| x as f64 / 1_000_000_f64)
369        }
370        Time64(Nanosecond) => {
371            as_time64_nanosecond_array(array)?.unary(|x| x as f64 / 1_000_000_000_f64)
372        }
373        Interval(_) | Duration(_) => return seconds(array, Second),
374        d => return exec_err!("Cannot convert {d:?} to epoch"),
375    };
376    Ok(Arc::new(f))
377}
datafusion_functions/datetime/date_part.rs

datafusion_functions/datetime/
date_part.rs
