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, 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, ReturnInfo, ReturnTypeArgs, 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_type_from_args should be called instead")
146    }
147
148    fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result<ReturnInfo> {
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                            ReturnInfo::new_nullable(DataType::Float64)
159                        } else {
160                            ReturnInfo::new_nullable(DataType::Int32)
161                        }
162                    })
163            })
164            .map_or_else(
165                || exec_err!("{} requires non-empty constant string", self.name()),
166                Ok,
167            )
168    }
169
170    fn invoke_with_args(
171        &self,
172        args: datafusion_expr::ScalarFunctionArgs,
173    ) -> Result<ColumnarValue> {
174        let args = args.args;
175        let [part, array] = take_function_args(self.name(), args)?;
176
177        let part = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = part {
178            v
179        } else if let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(v))) = part {
180            v
181        } else {
182            return exec_err!(
183                "First argument of `DATE_PART` must be non-null scalar Utf8"
184            );
185        };
186
187        let is_scalar = matches!(array, ColumnarValue::Scalar(_));
188
189        let array = match array {
190            ColumnarValue::Array(array) => Arc::clone(&array),
191            ColumnarValue::Scalar(scalar) => scalar.to_array()?,
192        };
193
194        let part_trim = part_normalization(&part);
195
196        // using IntervalUnit here means we hand off all the work of supporting plurals (like "seconds")
197        // and synonyms ( like "ms,msec,msecond,millisecond") to Arrow
198        let arr = if let Ok(interval_unit) = IntervalUnit::from_str(part_trim) {
199            match interval_unit {
200                IntervalUnit::Year => date_part(array.as_ref(), DatePart::Year)?,
201                IntervalUnit::Month => date_part(array.as_ref(), DatePart::Month)?,
202                IntervalUnit::Week => date_part(array.as_ref(), DatePart::Week)?,
203                IntervalUnit::Day => date_part(array.as_ref(), DatePart::Day)?,
204                IntervalUnit::Hour => date_part(array.as_ref(), DatePart::Hour)?,
205                IntervalUnit::Minute => date_part(array.as_ref(), DatePart::Minute)?,
206                IntervalUnit::Second => seconds_as_i32(array.as_ref(), Second)?,
207                IntervalUnit::Millisecond => seconds_as_i32(array.as_ref(), Millisecond)?,
208                IntervalUnit::Microsecond => seconds_as_i32(array.as_ref(), Microsecond)?,
209                IntervalUnit::Nanosecond => seconds_as_i32(array.as_ref(), Nanosecond)?,
210                // century and decade are not supported by `DatePart`, although they are supported in postgres
211                _ => return exec_err!("Date part '{part}' not supported"),
212            }
213        } else {
214            // special cases that can be extracted (in postgres) but are not interval units
215            match part_trim.to_lowercase().as_str() {
216                "qtr" | "quarter" => date_part(array.as_ref(), DatePart::Quarter)?,
217                "doy" => date_part(array.as_ref(), DatePart::DayOfYear)?,
218                "dow" => date_part(array.as_ref(), DatePart::DayOfWeekSunday0)?,
219                "epoch" => epoch(array.as_ref())?,
220                _ => return exec_err!("Date part '{part}' not supported"),
221            }
222        };
223
224        Ok(if is_scalar {
225            ColumnarValue::Scalar(ScalarValue::try_from_array(arr.as_ref(), 0)?)
226        } else {
227            ColumnarValue::Array(arr)
228        })
229    }
230
231    fn aliases(&self) -> &[String] {
232        &self.aliases
233    }
234    fn documentation(&self) -> Option<&Documentation> {
235        self.doc()
236    }
237}
238
239fn is_epoch(part: &str) -> bool {
240    let part = part_normalization(part);
241    matches!(part.to_lowercase().as_str(), "epoch")
242}
243
244// Try to remove quote if exist, if the quote is invalid, return original string and let the downstream function handle the error
245fn part_normalization(part: &str) -> &str {
246    part.strip_prefix(|c| c == '\'' || c == '\"')
247        .and_then(|s| s.strip_suffix(|c| c == '\'' || c == '\"'))
248        .unwrap_or(part)
249}
250
251/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
252/// result to a total number of seconds, milliseconds, microseconds or
253/// nanoseconds
254fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
255    // Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing
256    // with overflow and precision issue we don't support nanosecond
257    if unit == Nanosecond {
258        return not_impl_err!("Date part {unit:?} not supported");
259    }
260
261    let conversion_factor = match unit {
262        Second => 1_000_000_000,
263        Millisecond => 1_000_000,
264        Microsecond => 1_000,
265        Nanosecond => 1,
266    };
267
268    let second_factor = match unit {
269        Second => 1,
270        Millisecond => 1_000,
271        Microsecond => 1_000_000,
272        Nanosecond => 1_000_000_000,
273    };
274
275    let secs = date_part(array, DatePart::Second)?;
276    // This assumes array is primitive and not a dictionary
277    let secs = as_int32_array(secs.as_ref())?;
278    let subsecs = date_part(array, DatePart::Nanosecond)?;
279    let subsecs = as_int32_array(subsecs.as_ref())?;
280
281    // Special case where there are no nulls.
282    if subsecs.null_count() == 0 {
283        let r: Int32Array = binary(secs, subsecs, |secs, subsecs| {
284            secs * second_factor + (subsecs % 1_000_000_000) / conversion_factor
285        })?;
286        Ok(Arc::new(r))
287    } else {
288        // Nulls in secs are preserved, nulls in subsecs are treated as zero to account for the case
289        // where the number of nanoseconds overflows.
290        let r: Int32Array = secs
291            .iter()
292            .zip(subsecs)
293            .map(|(secs, subsecs)| {
294                secs.map(|secs| {
295                    let subsecs = subsecs.unwrap_or(0);
296                    secs * second_factor + (subsecs % 1_000_000_000) / conversion_factor
297                })
298            })
299            .collect();
300        Ok(Arc::new(r))
301    }
302}
303
304/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
305/// result to a total number of seconds, milliseconds, microseconds or
306/// nanoseconds
307///
308/// Given epoch return f64, this is a duplicated function to optimize for f64 type
309fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
310    let sf = match unit {
311        Second => 1_f64,
312        Millisecond => 1_000_f64,
313        Microsecond => 1_000_000_f64,
314        Nanosecond => 1_000_000_000_f64,
315    };
316    let secs = date_part(array, DatePart::Second)?;
317    // This assumes array is primitive and not a dictionary
318    let secs = as_int32_array(secs.as_ref())?;
319    let subsecs = date_part(array, DatePart::Nanosecond)?;
320    let subsecs = as_int32_array(subsecs.as_ref())?;
321
322    // Special case where there are no nulls.
323    if subsecs.null_count() == 0 {
324        let r: Float64Array = binary(secs, subsecs, |secs, subsecs| {
325            (secs as f64 + ((subsecs % 1_000_000_000) as f64 / 1_000_000_000_f64)) * sf
326        })?;
327        Ok(Arc::new(r))
328    } else {
329        // Nulls in secs are preserved, nulls in subsecs are treated as zero to account for the case
330        // where the number of nanoseconds overflows.
331        let r: Float64Array = secs
332            .iter()
333            .zip(subsecs)
334            .map(|(secs, subsecs)| {
335                secs.map(|secs| {
336                    let subsecs = subsecs.unwrap_or(0);
337                    (secs as f64 + ((subsecs % 1_000_000_000) as f64 / 1_000_000_000_f64))
338                        * sf
339                })
340            })
341            .collect();
342        Ok(Arc::new(r))
343    }
344}
345
346fn epoch(array: &dyn Array) -> Result<ArrayRef> {
347    const SECONDS_IN_A_DAY: f64 = 86400_f64;
348
349    let f: Float64Array = match array.data_type() {
350        Timestamp(Second, _) => as_timestamp_second_array(array)?.unary(|x| x as f64),
351        Timestamp(Millisecond, _) => {
352            as_timestamp_millisecond_array(array)?.unary(|x| x as f64 / 1_000_f64)
353        }
354        Timestamp(Microsecond, _) => {
355            as_timestamp_microsecond_array(array)?.unary(|x| x as f64 / 1_000_000_f64)
356        }
357        Timestamp(Nanosecond, _) => {
358            as_timestamp_nanosecond_array(array)?.unary(|x| x as f64 / 1_000_000_000_f64)
359        }
360        Date32 => as_date32_array(array)?.unary(|x| x as f64 * SECONDS_IN_A_DAY),
361        Date64 => as_date64_array(array)?.unary(|x| x as f64 / 1_000_f64),
362        Time32(Second) => as_time32_second_array(array)?.unary(|x| x as f64),
363        Time32(Millisecond) => {
364            as_time32_millisecond_array(array)?.unary(|x| x as f64 / 1_000_f64)
365        }
366        Time64(Microsecond) => {
367            as_time64_microsecond_array(array)?.unary(|x| x as f64 / 1_000_000_f64)
368        }
369        Time64(Nanosecond) => {
370            as_time64_nanosecond_array(array)?.unary(|x| x as f64 / 1_000_000_000_f64)
371        }
372        Interval(_) | Duration(_) => return seconds(array, Second),
373        d => return exec_err!("Cannot convert {d:?} to epoch"),
374    };
375    Ok(Arc::new(f))
376}
datafusion_functions/datetime/date_part.rs

datafusion_functions/datetime/
date_part.rs
