Skip to main content

datafusion_functions/datetime/
date_bin.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::temporal_conversions::NANOSECONDS;
21use arrow::array::types::{
22    ArrowTimestampType, IntervalDayTimeType, IntervalMonthDayNanoType,
23    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
24    TimestampSecondType,
25};
26use arrow::array::{ArrayRef, AsArray, PrimitiveArray};
27use arrow::datatypes::DataType::{Time32, Time64, Timestamp};
28use arrow::datatypes::IntervalUnit::{DayTime, MonthDayNano};
29use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
30use arrow::datatypes::{
31    DataType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
32    Time64NanosecondType, TimeUnit,
33};
34use arrow::error::ArrowError;
35use arrow::temporal_conversions::NANOSECONDS_IN_DAY;
36use datafusion_common::cast::as_primitive_array;
37use datafusion_common::{
38    Result, ScalarValue, exec_datafusion_err, exec_err, not_impl_err, plan_err,
39};
40use datafusion_expr::TypeSignature::Exact;
41use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
42use datafusion_expr::{
43    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
44    TIMEZONE_WILDCARD, Volatility,
45};
46use datafusion_macros::user_doc;
47
48use chrono::{DateTime, Datelike, Duration, Months, TimeDelta, Utc};
49
50#[user_doc(
51    doc_section(label = "Time and Date Functions"),
52    description = r#"
53Calculates time intervals and returns the start of the interval nearest to the specified timestamp. Use `date_bin` to downsample time series data by grouping rows into time-based "bins" or "windows" and applying an aggregate or selector function to each window.
54
55For example, if you "bin" or "window" data into 15 minute intervals, an input timestamp of `2023-01-01T18:18:18Z` will be updated to the start time of the 15 minute bin it is in: `2023-01-01T18:15:00Z`.
56"#,
57    syntax_example = "date_bin(interval, expression[, origin_timestamp])",
58    sql_example = r#"```sql
59-- Bin the timestamp into 1 day intervals
60> SELECT date_bin(interval '1 day', time) as bin
61FROM VALUES ('2023-01-01T18:18:18Z'), ('2023-01-03T19:00:03Z')  t(time);
62+---------------------+
63| bin                 |
64+---------------------+
65| 2023-01-01T00:00:00 |
66| 2023-01-03T00:00:00 |
67+---------------------+
682 row(s) fetched.
69
70-- Bin the timestamp into 1 day intervals starting at 3AM on  2023-01-01
71> SELECT date_bin(interval '1 day', time,  '2023-01-01T03:00:00') as bin
72FROM VALUES ('2023-01-01T18:18:18Z'), ('2023-01-03T19:00:03Z')  t(time);
73+---------------------+
74| bin                 |
75+---------------------+
76| 2023-01-01T03:00:00 |
77| 2023-01-03T03:00:00 |
78+---------------------+
792 row(s) fetched.
80
81-- Bin the time into 15 minute intervals starting at 1 min
82>  SELECT date_bin(interval '15 minutes', time, TIME '00:01:00') as bin
83FROM VALUES (TIME '02:18:18'), (TIME '19:00:03')  t(time);
84+----------+
85| bin      |
86+----------+
87| 02:16:00 |
88| 18:46:00 |
89+----------+
902 row(s) fetched.
91```"#,
92    argument(name = "interval", description = "Bin interval."),
93    argument(
94        name = "expression",
95        description = "Time expression to operate on. Can be a constant, column, or function."
96    ),
97    argument(
98        name = "origin_timestamp",
99        description = r#"Optional. Starting point used to determine bin boundaries. If not specified defaults 1970-01-01T00:00:00Z (the UNIX epoch in UTC). The following intervals are supported:
100
101    - nanoseconds
102    - microseconds
103    - milliseconds
104    - seconds
105    - minutes
106    - hours
107    - days
108    - weeks
109    - months
110    - years
111    - century
112"#
113    )
114)]
115#[derive(Debug, PartialEq, Eq, Hash)]
116pub struct DateBinFunc {
117    signature: Signature,
118}
119
120impl Default for DateBinFunc {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126impl DateBinFunc {
127    pub fn new() -> Self {
128        let base_sig = |array_type: TimeUnit| {
129            let mut v = vec![
130                Exact(vec![
131                    DataType::Interval(MonthDayNano),
132                    Timestamp(array_type, None),
133                    Timestamp(Nanosecond, None),
134                ]),
135                Exact(vec![
136                    DataType::Interval(MonthDayNano),
137                    Timestamp(array_type, Some(TIMEZONE_WILDCARD.into())),
138                    Timestamp(Nanosecond, Some(TIMEZONE_WILDCARD.into())),
139                ]),
140                Exact(vec![
141                    DataType::Interval(DayTime),
142                    Timestamp(array_type, None),
143                    Timestamp(Nanosecond, None),
144                ]),
145                Exact(vec![
146                    DataType::Interval(DayTime),
147                    Timestamp(array_type, Some(TIMEZONE_WILDCARD.into())),
148                    Timestamp(Nanosecond, Some(TIMEZONE_WILDCARD.into())),
149                ]),
150                Exact(vec![
151                    DataType::Interval(MonthDayNano),
152                    Timestamp(array_type, None),
153                ]),
154                Exact(vec![
155                    DataType::Interval(MonthDayNano),
156                    Timestamp(array_type, Some(TIMEZONE_WILDCARD.into())),
157                ]),
158                Exact(vec![
159                    DataType::Interval(DayTime),
160                    Timestamp(array_type, None),
161                ]),
162                Exact(vec![
163                    DataType::Interval(DayTime),
164                    Timestamp(array_type, Some(TIMEZONE_WILDCARD.into())),
165                ]),
166            ];
167
168            match array_type {
169                Second | Millisecond => {
170                    v.append(&mut vec![
171                        Exact(vec![
172                            DataType::Interval(MonthDayNano),
173                            Time32(array_type),
174                            Time32(array_type),
175                        ]),
176                        Exact(vec![DataType::Interval(MonthDayNano), Time32(array_type)]),
177                        Exact(vec![
178                            DataType::Interval(DayTime),
179                            Time32(array_type),
180                            Time32(array_type),
181                        ]),
182                        Exact(vec![DataType::Interval(DayTime), Time32(array_type)]),
183                    ]);
184                }
185                Microsecond | Nanosecond => {
186                    v.append(&mut vec![
187                        Exact(vec![
188                            DataType::Interval(DayTime),
189                            Time64(array_type),
190                            Time64(array_type),
191                        ]),
192                        Exact(vec![DataType::Interval(DayTime), Time64(array_type)]),
193                        Exact(vec![
194                            DataType::Interval(MonthDayNano),
195                            Time64(array_type),
196                            Time64(array_type),
197                        ]),
198                        Exact(vec![DataType::Interval(MonthDayNano), Time64(array_type)]),
199                    ]);
200                }
201            }
202
203            v
204        };
205
206        let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
207            .into_iter()
208            .map(base_sig)
209            .collect::<Vec<_>>()
210            .concat();
211
212        Self {
213            signature: Signature::one_of(full_sig, Volatility::Immutable),
214        }
215    }
216}
217
218impl ScalarUDFImpl for DateBinFunc {
219    fn name(&self) -> &str {
220        "date_bin"
221    }
222
223    fn signature(&self) -> &Signature {
224        &self.signature
225    }
226
227    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
228        match &arg_types[1] {
229            Timestamp(tu, tz_opt) => Ok(Timestamp(*tu, tz_opt.clone())),
230            Time32(tu) => Ok(Time32(*tu)),
231            Time64(tu) => Ok(Time64(*tu)),
232            _ => plan_err!(
233                "The date_bin function can only accept timestamp or time as the second arg."
234            ),
235        }
236    }
237
238    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
239        let args = &args.args;
240        if args.len() == 2 {
241            let origin = match args[1].data_type() {
242                Time32(Second) => {
243                    ColumnarValue::Scalar(ScalarValue::Time32Second(Some(0)))
244                }
245                Time32(Millisecond) => {
246                    ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(0)))
247                }
248                Time64(Microsecond) => {
249                    ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0)))
250                }
251                Time64(Nanosecond) => {
252                    ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(0)))
253                }
254                _ => {
255                    // Default to unix EPOCH
256                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
257                        Some(0),
258                        Some("+00:00".into()),
259                    ))
260                }
261            };
262            date_bin_impl(&args[0], &args[1], &origin)
263        } else if args.len() == 3 {
264            date_bin_impl(&args[0], &args[1], &args[2])
265        } else {
266            exec_err!("DATE_BIN expected two or three arguments")
267        }
268    }
269
270    fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
271        // The DATE_BIN function preserves the order of its second argument.
272        let step = &input[0];
273        let date_value = &input[1];
274        let reference = input.get(2);
275
276        if step.sort_properties.eq(&SortProperties::Singleton)
277            && reference
278                .map(|r| r.sort_properties.eq(&SortProperties::Singleton))
279                .unwrap_or(true)
280        {
281            Ok(date_value.sort_properties)
282        } else {
283            Ok(SortProperties::Unordered)
284        }
285    }
286    fn documentation(&self) -> Option<&Documentation> {
287        self.doc()
288    }
289}
290
291const NANOS_PER_MICRO: i64 = 1_000;
292const NANOS_PER_MILLI: i64 = 1_000_000;
293const NANOS_PER_SEC: i64 = NANOSECONDS;
294/// Function type for binning timestamps into intervals
295///
296/// Arguments:
297/// * `stride` - Interval width (nanoseconds for time-based, months for month-based)
298/// * `source` - Timestamp to bin (nanoseconds since epoch)
299/// * `origin` - Origin timestamp (nanoseconds since epoch)
300///
301/// Returns: Binned timestamp in nanoseconds, or error if out of range
302type BinFunction = fn(i64, i64, i64) -> Result<i64>;
303enum Interval {
304    Nanoseconds(i64),
305    Months(i64),
306}
307
308impl Interval {
309    /// Returns (`stride_nanos`, `fn`) where
310    ///
311    /// 1. `stride_nanos` is a width, in nanoseconds
312    /// 2. `fn` is a function that takes (stride_nanos, source, origin)
313    ///
314    /// `source` is the timestamp being binned
315    ///
316    /// `origin`  is the time, in nanoseconds, where windows are measured from
317    fn bin_fn(&self) -> (i64, BinFunction) {
318        match self {
319            Interval::Nanoseconds(nanos) => (*nanos, date_bin_nanos_interval),
320            Interval::Months(months) => (*months, date_bin_months_interval),
321        }
322    }
323}
324
325// return time in nanoseconds that the source timestamp falls into based on the stride and origin
326fn date_bin_nanos_interval(stride_nanos: i64, source: i64, origin: i64) -> Result<i64> {
327    let time_diff = source.checked_sub(origin).ok_or_else(|| {
328        ArrowError::InvalidArgumentError(format!(
329            "date_bin source timestamp {source} - origin {origin} overflows i64"
330        ))
331    })?;
332
333    // distance from origin to bin
334    let time_delta = compute_distance(time_diff, stride_nanos)?;
335
336    origin.checked_add(time_delta).ok_or_else(|| {
337        ArrowError::InvalidArgumentError(format!(
338            "date_bin origin {origin} + delta {time_delta} overflows i64"
339        ))
340        .into()
341    })
342}
343
344// distance from origin to bin
345fn compute_distance(time_diff: i64, stride: i64) -> Result<i64> {
346    let remainder = time_diff.checked_rem(stride).ok_or_else(|| {
347        ArrowError::InvalidArgumentError(format!(
348            "date_bin compute_distance time_diff {time_diff} % stride {stride} overflows i64"
349        ))
350    })?;
351    let time_delta = time_diff.checked_sub(remainder).ok_or_else(|| {
352        ArrowError::InvalidArgumentError(format!(
353            "date_bin compute_distance time_diff {time_diff} - remainder {remainder} overflows i64"
354        ))
355    })?;
356
357    if time_diff < 0 && stride > 1 && time_delta != time_diff {
358        // The origin is later than the source timestamp, round down to the previous bin
359        time_delta.checked_sub(stride).ok_or_else(|| {
360            ArrowError::InvalidArgumentError(format!(
361                "date_bin compute_distance time_delta {time_delta} - stride {stride} overflows i64"
362            ))
363            .into()
364        })
365    } else {
366        Ok(time_delta)
367    }
368}
369
370// Shift `origin_date` by `month_delta` months, mapping an out-of-range result to
371// the same error the binning paths reported when this was written inline.
372fn shift_months(origin_date: DateTime<Utc>, month_delta: i64) -> Result<DateTime<Utc>> {
373    if month_delta < 0 {
374        origin_date
375            .checked_sub_months(Months::new(month_delta.unsigned_abs() as u32))
376            .ok_or_else(|| {
377                exec_datafusion_err!("DATE_BIN month subtraction out of range")
378            })
379    } else {
380        origin_date
381            .checked_add_months(Months::new(month_delta as u32))
382            .ok_or_else(|| exec_datafusion_err!("DATE_BIN month addition out of range"))
383    }
384}
385
386// return time in nanoseconds that the source timestamp falls into based on the stride and origin
387fn date_bin_months_interval(stride_months: i64, source: i64, origin: i64) -> Result<i64> {
388    // convert source and origin to DateTime<Utc>
389    let source_date = to_utc_date_time(source)?;
390    let origin_date = to_utc_date_time(origin)?;
391
392    // calculate the number of months between the source and origin
393    let month_diff = (source_date.year() - origin_date.year()) * 12
394        + source_date.month() as i32
395        - origin_date.month() as i32;
396
397    // distance from origin to bin
398    let month_delta = compute_distance(month_diff as i64, stride_months)?;
399
400    let mut bin_time = shift_months(origin_date, month_delta)?;
401
402    // If origin is not midnight of first date of the month, the bin_time may be larger than the source
403    // In this case, we need to move back to previous bin
404    if bin_time > source_date {
405        let month_delta = month_delta - stride_months;
406        bin_time = shift_months(origin_date, month_delta)?;
407    }
408    match bin_time.timestamp_nanos_opt() {
409        Some(nanos) => Ok(nanos),
410        None => exec_err!("DATE_BIN result timestamp out of range"),
411    }
412}
413
414fn to_utc_date_time(nanos: i64) -> Result<DateTime<Utc>> {
415    // Keep negative sub-second values normalized as seconds + non-negative nanos.
416    let secs = nanos.div_euclid(NANOS_PER_SEC);
417    let nsec = nanos.rem_euclid(NANOS_PER_SEC) as u32;
418    match DateTime::from_timestamp(secs, nsec) {
419        Some(dt) => Ok(dt),
420        None => exec_err!("Invalid timestamp value"),
421    }
422}
423
424fn timestamp_scale<T: ArrowTimestampType>() -> i64 {
425    match T::UNIT {
426        Nanosecond => 1,
427        Microsecond => NANOS_PER_MICRO,
428        Millisecond => NANOS_PER_MILLI,
429        Second => NANOSECONDS,
430    }
431}
432
433// Scale to nanoseconds and report overflow as a normal error.
434fn checked_scale_to_nanos(x: i64, scale: i64) -> Result<i64> {
435    match x.checked_mul(scale) {
436        Some(scaled) => Ok(scaled),
437        None => exec_err!("date_bin timestamp value {x} * scale {scale} overflows i64"),
438    }
439}
440
441// Per-row failures map to NULL, so use Option in the hot path.
442#[inline]
443fn scale_and_bin_to_nanos(
444    value: i64,
445    scale: i64,
446    origin: i64,
447    stride: i64,
448    stride_fn: BinFunction,
449) -> Option<i64> {
450    value
451        .checked_mul(scale)
452        .and_then(|scaled| stride_fn(stride, scaled, origin).ok())
453}
454
455// Per-row timestamp binning shared by scalar and array paths.
456// Source-value failures become None, which callers map to NULL.
457#[inline]
458fn date_bin_timestamp_value<T: ArrowTimestampType>(
459    value: i64,
460    origin: i64,
461    stride: i64,
462    stride_fn: BinFunction,
463) -> Option<i64> {
464    let scale = timestamp_scale::<T>();
465    scale_and_bin_to_nanos(value, scale, origin, stride, stride_fn)
466        .map(|binned| binned / scale)
467}
468
469// Per-row TIME binning shared by scalar and array paths.
470// The modulo keeps the result within a single day before unscaling.
471#[inline]
472fn date_bin_time_value(
473    value: i64,
474    scale: i64,
475    origin: i64,
476    stride: i64,
477    stride_fn: BinFunction,
478) -> Option<i64> {
479    scale_and_bin_to_nanos(value, scale, origin, stride, stride_fn)
480        .map(|binned| (binned % NANOSECONDS_IN_DAY) / scale)
481}
482
483fn validate_time_stride(stride: &Interval) -> Result<()> {
484    match stride {
485        Interval::Months(m) if *m > 0 => {
486            exec_err!("DATE_BIN stride for TIME input must be less than 1 day")
487        }
488        Interval::Nanoseconds(ns) if *ns >= NANOSECONDS_IN_DAY => {
489            exec_err!("DATE_BIN stride for TIME input must be less than 1 day")
490        }
491        _ => Ok(()),
492    }
493}
494
495// Supported intervals:
496//  1. IntervalDayTime: this means that the stride is in days, hours, minutes, seconds and milliseconds
497//     We will assume month interval won't be converted into this type
498//     TODO (my next PR): without `INTERVAL` keyword, the stride was converted into ScalarValue::IntervalDayTime somewhere
499//             for month interval. I need to find that and make it ScalarValue::IntervalMonthDayNano instead
500// 2. IntervalMonthDayNano
501fn date_bin_impl(
502    stride: &ColumnarValue,
503    array: &ColumnarValue,
504    origin: &ColumnarValue,
505) -> Result<ColumnarValue> {
506    let stride = match stride {
507        ColumnarValue::Scalar(s) if s.is_null() => {
508            // NULL stride -> NULL result (standard SQL NULL propagation)
509            return Ok(ColumnarValue::Scalar(ScalarValue::try_from(
510                array.data_type(),
511            )?));
512        }
513        ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(v))) => {
514            let (days, ms) = IntervalDayTimeType::to_parts(*v);
515            let nanos = (TimeDelta::try_days(days as i64).unwrap()
516                + TimeDelta::try_milliseconds(ms as i64).unwrap())
517            .num_nanoseconds();
518
519            match nanos {
520                Some(v) => Interval::Nanoseconds(v),
521                _ => return exec_err!("DATE_BIN stride argument is too large"),
522            }
523        }
524        ColumnarValue::Scalar(ScalarValue::IntervalMonthDayNano(Some(v))) => {
525            let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(*v);
526
527            // If interval is months, its origin must be midnight of first date of the month
528            if months != 0 {
529                // Return error if days or nanos is not zero
530                if days != 0 || nanos != 0 {
531                    return not_impl_err!(
532                        "DATE_BIN stride does not support combination of month, day and nanosecond intervals"
533                    );
534                } else {
535                    Interval::Months(months as i64)
536                }
537            } else {
538                let nanos = (TimeDelta::try_days(days as i64).unwrap()
539                    + Duration::nanoseconds(nanos))
540                .num_nanoseconds();
541                match nanos {
542                    Some(v) => Interval::Nanoseconds(v),
543                    _ => return exec_err!("DATE_BIN stride argument is too large"),
544                }
545            }
546        }
547        ColumnarValue::Scalar(v) => {
548            return exec_err!(
549                "DATE_BIN expects stride argument to be an INTERVAL but got {}",
550                v.data_type()
551            );
552        }
553        ColumnarValue::Array(_) => {
554            return not_impl_err!(
555                "DATE_BIN only supports literal values for the stride argument, not arrays"
556            );
557        }
558    };
559
560    let (origin, is_time) = match origin {
561        ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(v), _)) => {
562            (*v, false)
563        }
564        ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(v))) => {
565            validate_time_stride(&stride)?;
566            // TIME origins can come from reinterpret casts, so scale defensively.
567            (checked_scale_to_nanos(*v as i64, NANOS_PER_MILLI)?, true)
568        }
569        ColumnarValue::Scalar(ScalarValue::Time32Second(Some(v))) => {
570            validate_time_stride(&stride)?;
571            (checked_scale_to_nanos(*v as i64, NANOS_PER_SEC)?, true)
572        }
573        ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(v))) => {
574            validate_time_stride(&stride)?;
575            (checked_scale_to_nanos(*v, NANOS_PER_MICRO)?, true)
576        }
577        ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(v))) => {
578            validate_time_stride(&stride)?;
579            (*v, true)
580        }
581        ColumnarValue::Scalar(v) => {
582            return exec_err!(
583                "DATE_BIN expects origin argument to be a TIMESTAMP with nanosecond precision or a TIME but got {}",
584                v.data_type()
585            );
586        }
587        ColumnarValue::Array(_) => {
588            return not_impl_err!(
589                "DATE_BIN only supports literal values for the origin argument, not arrays"
590            );
591        }
592    };
593
594    let (stride, stride_fn) = stride.bin_fn();
595
596    // Return error if stride is 0
597    if stride == 0 {
598        return exec_err!("DATE_BIN stride must be non-zero");
599    }
600
601    // A TIME source requires a TIME origin. This shared-input check is ordered
602    // after stride/origin parsing and the zero-stride check so error ordering is
603    // unchanged, and replaces the per-arm guards in the TIME branches below.
604    if !is_time {
605        match array.data_type() {
606            Time32(_) => {
607                return exec_err!("DATE_BIN with Time32 source requires Time32 origin");
608            }
609            Time64(_) => {
610                return exec_err!("DATE_BIN with Time64 source requires Time64 origin");
611            }
612            _ => {}
613        }
614    }
615
616    Ok(match array {
617        ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(v, tz_opt)) => {
618            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
619                v.and_then(|x| {
620                    date_bin_timestamp_value::<TimestampNanosecondType>(
621                        x, origin, stride, stride_fn,
622                    )
623                }),
624                tz_opt.clone(),
625            ))
626        }
627        ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(v, tz_opt)) => {
628            ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
629                v.and_then(|x| {
630                    date_bin_timestamp_value::<TimestampMicrosecondType>(
631                        x, origin, stride, stride_fn,
632                    )
633                }),
634                tz_opt.clone(),
635            ))
636        }
637        ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(v, tz_opt)) => {
638            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
639                v.and_then(|x| {
640                    date_bin_timestamp_value::<TimestampMillisecondType>(
641                        x, origin, stride, stride_fn,
642                    )
643                }),
644                tz_opt.clone(),
645            ))
646        }
647        ColumnarValue::Scalar(ScalarValue::TimestampSecond(v, tz_opt)) => {
648            ColumnarValue::Scalar(ScalarValue::TimestampSecond(
649                v.and_then(|x| {
650                    date_bin_timestamp_value::<TimestampSecondType>(
651                        x, origin, stride, stride_fn,
652                    )
653                }),
654                tz_opt.clone(),
655            ))
656        }
657        ColumnarValue::Scalar(ScalarValue::Time32Millisecond(v)) => {
658            let result = v.and_then(|x| {
659                date_bin_time_value(x as i64, NANOS_PER_MILLI, origin, stride, stride_fn)
660                    .map(|binned| binned as i32)
661            });
662            ColumnarValue::Scalar(ScalarValue::Time32Millisecond(result))
663        }
664        ColumnarValue::Scalar(ScalarValue::Time32Second(v)) => {
665            let result = v.and_then(|x| {
666                date_bin_time_value(x as i64, NANOS_PER_SEC, origin, stride, stride_fn)
667                    .map(|binned| binned as i32)
668            });
669            ColumnarValue::Scalar(ScalarValue::Time32Second(result))
670        }
671        ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(v)) => {
672            let result =
673                v.and_then(|x| date_bin_time_value(x, 1, origin, stride, stride_fn));
674            ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(result))
675        }
676        ColumnarValue::Scalar(ScalarValue::Time64Microsecond(v)) => {
677            let result = v.and_then(|x| {
678                date_bin_time_value(x, NANOS_PER_MICRO, origin, stride, stride_fn)
679            });
680            ColumnarValue::Scalar(ScalarValue::Time64Microsecond(result))
681        }
682        ColumnarValue::Array(array) => {
683            fn transform_array_with_stride<T>(
684                origin: i64,
685                stride: i64,
686                stride_fn: BinFunction,
687                array: &ArrayRef,
688                tz_opt: &Option<Arc<str>>,
689            ) -> Result<ColumnarValue>
690            where
691                T: ArrowTimestampType,
692            {
693                let array = as_primitive_array::<T>(array)?;
694
695                // Per-row errors become NULL, matching scalar behavior.
696                let result: PrimitiveArray<T> = array.unary_opt(|val| {
697                    date_bin_timestamp_value::<T>(val, origin, stride, stride_fn)
698                });
699
700                let array = result.with_timezone_opt(tz_opt.clone());
701                Ok(ColumnarValue::Array(Arc::new(array)))
702            }
703
704            match array.data_type() {
705                Timestamp(Nanosecond, tz_opt) => {
706                    transform_array_with_stride::<TimestampNanosecondType>(
707                        origin, stride, stride_fn, array, tz_opt,
708                    )?
709                }
710                Timestamp(Microsecond, tz_opt) => {
711                    transform_array_with_stride::<TimestampMicrosecondType>(
712                        origin, stride, stride_fn, array, tz_opt,
713                    )?
714                }
715                Timestamp(Millisecond, tz_opt) => {
716                    transform_array_with_stride::<TimestampMillisecondType>(
717                        origin, stride, stride_fn, array, tz_opt,
718                    )?
719                }
720                Timestamp(Second, tz_opt) => {
721                    transform_array_with_stride::<TimestampSecondType>(
722                        origin, stride, stride_fn, array, tz_opt,
723                    )?
724                }
725                Time32(Millisecond) => {
726                    let array = array.as_primitive::<Time32MillisecondType>();
727                    let result: PrimitiveArray<Time32MillisecondType> =
728                        array.unary_opt(|x| {
729                            date_bin_time_value(
730                                x as i64,
731                                NANOS_PER_MILLI,
732                                origin,
733                                stride,
734                                stride_fn,
735                            )
736                            .map(|binned| binned as i32)
737                        });
738                    ColumnarValue::Array(Arc::new(result))
739                }
740                Time32(Second) => {
741                    let array = array.as_primitive::<Time32SecondType>();
742                    let result: PrimitiveArray<Time32SecondType> = array.unary_opt(|x| {
743                        date_bin_time_value(
744                            x as i64,
745                            NANOS_PER_SEC,
746                            origin,
747                            stride,
748                            stride_fn,
749                        )
750                        .map(|binned| binned as i32)
751                    });
752                    ColumnarValue::Array(Arc::new(result))
753                }
754                Time64(Microsecond) => {
755                    let array = array.as_primitive::<Time64MicrosecondType>();
756                    let result: PrimitiveArray<Time64MicrosecondType> =
757                        array.unary_opt(|x| {
758                            date_bin_time_value(
759                                x,
760                                NANOS_PER_MICRO,
761                                origin,
762                                stride,
763                                stride_fn,
764                            )
765                        });
766                    ColumnarValue::Array(Arc::new(result))
767                }
768                Time64(Nanosecond) => {
769                    let array = array.as_primitive::<Time64NanosecondType>();
770                    let result: PrimitiveArray<Time64NanosecondType> =
771                        array.unary_opt(|x| {
772                            date_bin_time_value(x, 1, origin, stride, stride_fn)
773                        });
774                    ColumnarValue::Array(Arc::new(result))
775                }
776                _ => {
777                    return exec_err!(
778                        "DATE_BIN expects source argument to be a TIMESTAMP or TIME but got {}",
779                        array.data_type()
780                    );
781                }
782            }
783        }
784        _ => {
785            return exec_err!(
786                "DATE_BIN expects source argument to be a TIMESTAMP or TIME scalar or array"
787            );
788        }
789    })
790}
791
792#[cfg(test)]
793mod tests {
794    use std::sync::Arc;
795
796    use crate::datetime::date_bin::{DateBinFunc, date_bin_nanos_interval};
797    use arrow::array::types::TimestampNanosecondType;
798    use arrow::array::{Array, IntervalDayTimeArray, TimestampNanosecondArray};
799    use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
800    use arrow::datatypes::{DataType, Field, FieldRef, TimeUnit};
801
802    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
803    use datafusion_common::{DataFusionError, ScalarValue};
804    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
805
806    use chrono::TimeDelta;
807    use datafusion_common::config::ConfigOptions;
808
809    fn invoke_date_bin_with_args(
810        args: Vec<ColumnarValue>,
811        number_rows: usize,
812        return_field: &FieldRef,
813    ) -> Result<ColumnarValue, DataFusionError> {
814        let arg_fields = args
815            .iter()
816            .map(|arg| Field::new("a", arg.data_type(), true).into())
817            .collect::<Vec<_>>();
818
819        let args = ScalarFunctionArgs {
820            args,
821            arg_fields,
822            number_rows,
823            return_field: Arc::clone(return_field),
824            config_options: Arc::new(ConfigOptions::default()),
825        };
826        DateBinFunc::new().invoke_with_args(args)
827    }
828
829    fn assert_null_scalar(value: ColumnarValue, expected_type: DataType) {
830        let ColumnarValue::Scalar(value) = value else {
831            panic!("expected scalar, got {value:?}");
832        };
833        assert_eq!(value.data_type(), expected_type);
834        assert!(value.is_null(), "expected NULL, got {value:?}");
835    }
836
837    fn assert_array_null_then_valid(value: ColumnarValue, expected_type: DataType) {
838        let ColumnarValue::Array(array) = value else {
839            panic!("expected array, got {value:?}");
840        };
841        assert_eq!(array.data_type(), &expected_type);
842        assert!(array.is_null(0), "expected NULL at row 0");
843        assert!(array.is_valid(1), "expected valid value at row 1");
844    }
845
846    fn assert_overflow_error(result: Result<ColumnarValue, DataFusionError>) {
847        let err = result.expect_err("expected overflow error");
848        assert!(
849            err.strip_backtrace().contains("overflows i64"),
850            "unexpected error: {err}"
851        );
852    }
853
854    #[test]
855    fn test_date_bin() {
856        let return_field = &Arc::new(Field::new(
857            "f",
858            DataType::Timestamp(TimeUnit::Nanosecond, None),
859            true,
860        ));
861
862        let mut args = vec![
863            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
864                days: 0,
865                milliseconds: 1,
866            }))),
867            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
868            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
869        ];
870        let res = invoke_date_bin_with_args(args, 1, return_field);
871        assert!(res.is_ok());
872
873        let timestamps = Arc::new((1..6).map(Some).collect::<TimestampNanosecondArray>());
874        let batch_len = timestamps.len();
875        args = vec![
876            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
877                days: 0,
878                milliseconds: 1,
879            }))),
880            ColumnarValue::Array(timestamps),
881            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
882        ];
883        let res = invoke_date_bin_with_args(args, batch_len, return_field);
884        assert!(res.is_ok());
885
886        args = vec![
887            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
888                days: 0,
889                milliseconds: 1,
890            }))),
891            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
892        ];
893        let res = invoke_date_bin_with_args(args, 1, return_field);
894        assert!(res.is_ok());
895
896        // stride supports month-day-nano
897        args = vec![
898            ColumnarValue::Scalar(ScalarValue::IntervalMonthDayNano(Some(
899                IntervalMonthDayNano {
900                    months: 0,
901                    days: 0,
902                    nanoseconds: 1,
903                },
904            ))),
905            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
906            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
907        ];
908        let res = invoke_date_bin_with_args(args, 1, return_field);
909        assert!(res.is_ok());
910
911        //
912        // Fallible test cases
913        //
914
915        // invalid number of arguments
916        args = vec![ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(
917            IntervalDayTime {
918                days: 0,
919                milliseconds: 1,
920            },
921        )))];
922        let res = invoke_date_bin_with_args(args, 1, return_field);
923        assert_eq!(
924            res.err().unwrap().strip_backtrace(),
925            "Execution error: DATE_BIN expected two or three arguments"
926        );
927
928        // stride: invalid type
929        args = vec![
930            ColumnarValue::Scalar(ScalarValue::IntervalYearMonth(Some(1))),
931            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
932            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
933        ];
934        let res = invoke_date_bin_with_args(args, 1, return_field);
935        assert_eq!(
936            res.err().unwrap().strip_backtrace(),
937            "Execution error: DATE_BIN expects stride argument to be an INTERVAL but got Interval(YearMonth)"
938        );
939
940        // stride: invalid value
941
942        args = vec![
943            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
944                days: 0,
945                milliseconds: 0,
946            }))),
947            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
948            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
949        ];
950
951        let res = invoke_date_bin_with_args(args, 1, return_field);
952        assert_eq!(
953            res.err().unwrap().strip_backtrace(),
954            "Execution error: DATE_BIN stride must be non-zero"
955        );
956
957        // stride: overflow of day-time interval
958        args = vec![
959            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(
960                IntervalDayTime::MAX,
961            ))),
962            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
963            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
964        ];
965        let res = invoke_date_bin_with_args(args, 1, return_field);
966        assert_eq!(
967            res.err().unwrap().strip_backtrace(),
968            "Execution error: DATE_BIN stride argument is too large"
969        );
970
971        // stride: overflow of month-day-nano interval
972        args = vec![
973            ColumnarValue::Scalar(ScalarValue::new_interval_mdn(0, i32::MAX, 1)),
974            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
975            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
976        ];
977        let res = invoke_date_bin_with_args(args, 1, return_field);
978        assert_eq!(
979            res.err().unwrap().strip_backtrace(),
980            "Execution error: DATE_BIN stride argument is too large"
981        );
982
983        // stride: month intervals
984        args = vec![
985            ColumnarValue::Scalar(ScalarValue::new_interval_mdn(1, 1, 1)),
986            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
987            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
988        ];
989        let res = invoke_date_bin_with_args(args, 1, return_field);
990        assert_eq!(
991            res.err().unwrap().strip_backtrace(),
992            "This feature is not implemented: DATE_BIN stride does not support combination of month, day and nanosecond intervals"
993        );
994
995        // origin: invalid type
996        args = vec![
997            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
998                days: 0,
999                milliseconds: 1,
1000            }))),
1001            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
1002            ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)),
1003        ];
1004        let res = invoke_date_bin_with_args(args, 1, return_field);
1005        assert_eq!(
1006            res.err().unwrap().strip_backtrace(),
1007            "Execution error: DATE_BIN expects origin argument to be a TIMESTAMP with nanosecond precision or a TIME but got Timestamp(µs)"
1008        );
1009
1010        args = vec![
1011            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1012                days: 0,
1013                milliseconds: 1,
1014            }))),
1015            ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(1), None)),
1016            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
1017        ];
1018        let res = invoke_date_bin_with_args(args, 1, return_field);
1019        assert!(res.is_ok());
1020
1021        // unsupported array type for stride
1022        let intervals = Arc::new(
1023            (1..6)
1024                .map(|x| {
1025                    Some(IntervalDayTime {
1026                        days: 0,
1027                        milliseconds: x,
1028                    })
1029                })
1030                .collect::<IntervalDayTimeArray>(),
1031        );
1032        args = vec![
1033            ColumnarValue::Array(intervals),
1034            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
1035            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
1036        ];
1037        let res = invoke_date_bin_with_args(args, 1, return_field);
1038        assert_eq!(
1039            res.err().unwrap().strip_backtrace(),
1040            "This feature is not implemented: DATE_BIN only supports literal values for the stride argument, not arrays"
1041        );
1042
1043        // unsupported array type for origin
1044        let timestamps = Arc::new((1..6).map(Some).collect::<TimestampNanosecondArray>());
1045        let batch_len = timestamps.len();
1046        args = vec![
1047            ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1048                days: 0,
1049                milliseconds: 1,
1050            }))),
1051            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(1), None)),
1052            ColumnarValue::Array(timestamps),
1053        ];
1054        let res = invoke_date_bin_with_args(args, batch_len, return_field);
1055        assert_eq!(
1056            res.err().unwrap().strip_backtrace(),
1057            "This feature is not implemented: DATE_BIN only supports literal values for the origin argument, not arrays"
1058        );
1059    }
1060
1061    #[test]
1062    fn test_date_bin_timezones() {
1063        let cases = [
1064            (
1065                vec![
1066                    "2020-09-08T00:00:00Z",
1067                    "2020-09-08T01:00:00Z",
1068                    "2020-09-08T02:00:00Z",
1069                    "2020-09-08T03:00:00Z",
1070                    "2020-09-08T04:00:00Z",
1071                ],
1072                Some("+00".into()),
1073                "1970-01-01T00:00:00Z",
1074                vec![
1075                    "2020-09-08T00:00:00Z",
1076                    "2020-09-08T00:00:00Z",
1077                    "2020-09-08T00:00:00Z",
1078                    "2020-09-08T00:00:00Z",
1079                    "2020-09-08T00:00:00Z",
1080                ],
1081            ),
1082            (
1083                vec![
1084                    "2020-09-08T00:00:00Z",
1085                    "2020-09-08T01:00:00Z",
1086                    "2020-09-08T02:00:00Z",
1087                    "2020-09-08T03:00:00Z",
1088                    "2020-09-08T04:00:00Z",
1089                ],
1090                None,
1091                "1970-01-01T00:00:00Z",
1092                vec![
1093                    "2020-09-08T00:00:00Z",
1094                    "2020-09-08T00:00:00Z",
1095                    "2020-09-08T00:00:00Z",
1096                    "2020-09-08T00:00:00Z",
1097                    "2020-09-08T00:00:00Z",
1098                ],
1099            ),
1100            (
1101                vec![
1102                    "2020-09-08T00:00:00Z",
1103                    "2020-09-08T01:00:00Z",
1104                    "2020-09-08T02:00:00Z",
1105                    "2020-09-08T03:00:00Z",
1106                    "2020-09-08T04:00:00Z",
1107                ],
1108                Some("-02".into()),
1109                "1970-01-01T00:00:00Z",
1110                vec![
1111                    "2020-09-08T00:00:00Z",
1112                    "2020-09-08T00:00:00Z",
1113                    "2020-09-08T00:00:00Z",
1114                    "2020-09-08T00:00:00Z",
1115                    "2020-09-08T00:00:00Z",
1116                ],
1117            ),
1118            (
1119                vec![
1120                    "2020-09-08T00:00:00+05",
1121                    "2020-09-08T01:00:00+05",
1122                    "2020-09-08T02:00:00+05",
1123                    "2020-09-08T03:00:00+05",
1124                    "2020-09-08T04:00:00+05",
1125                ],
1126                Some("+05".into()),
1127                "1970-01-01T00:00:00+05",
1128                vec![
1129                    "2020-09-08T00:00:00+05",
1130                    "2020-09-08T00:00:00+05",
1131                    "2020-09-08T00:00:00+05",
1132                    "2020-09-08T00:00:00+05",
1133                    "2020-09-08T00:00:00+05",
1134                ],
1135            ),
1136            (
1137                vec![
1138                    "2020-09-08T00:00:00+08",
1139                    "2020-09-08T01:00:00+08",
1140                    "2020-09-08T02:00:00+08",
1141                    "2020-09-08T03:00:00+08",
1142                    "2020-09-08T04:00:00+08",
1143                ],
1144                Some("+08".into()),
1145                "1970-01-01T00:00:00+08",
1146                vec![
1147                    "2020-09-08T00:00:00+08",
1148                    "2020-09-08T00:00:00+08",
1149                    "2020-09-08T00:00:00+08",
1150                    "2020-09-08T00:00:00+08",
1151                    "2020-09-08T00:00:00+08",
1152                ],
1153            ),
1154        ];
1155
1156        cases
1157            .iter()
1158            .for_each(|(original, tz_opt, origin, expected)| {
1159                let input = original
1160                    .iter()
1161                    .map(|s| Some(string_to_timestamp_nanos(s).unwrap()))
1162                    .collect::<TimestampNanosecondArray>()
1163                    .with_timezone_opt(tz_opt.clone());
1164                let right = expected
1165                    .iter()
1166                    .map(|s| Some(string_to_timestamp_nanos(s).unwrap()))
1167                    .collect::<TimestampNanosecondArray>()
1168                    .with_timezone_opt(tz_opt.clone());
1169                let batch_len = input.len();
1170                let args = vec![
1171                    ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)),
1172                    ColumnarValue::Array(Arc::new(input)),
1173                    ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
1174                        Some(string_to_timestamp_nanos(origin).unwrap()),
1175                        tz_opt.clone(),
1176                    )),
1177                ];
1178                let return_field = &Arc::new(Field::new(
1179                    "f",
1180                    DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()),
1181                    true,
1182                ));
1183                let result =
1184                    invoke_date_bin_with_args(args, batch_len, return_field).unwrap();
1185
1186                if let ColumnarValue::Array(result) = result {
1187                    assert_eq!(
1188                        result.data_type(),
1189                        &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone())
1190                    );
1191                    let left = arrow::array::cast::as_primitive_array::<
1192                        TimestampNanosecondType,
1193                    >(&result);
1194                    assert_eq!(left, &right);
1195                } else {
1196                    panic!("unexpected column type");
1197                }
1198            });
1199    }
1200
1201    #[test]
1202    fn test_date_bin_single() {
1203        let cases = [
1204            (
1205                (
1206                    TimeDelta::try_minutes(15),
1207                    "2004-04-09T02:03:04.123456789Z",
1208                    "2001-01-01T00:00:00",
1209                ),
1210                "2004-04-09T02:00:00Z",
1211            ),
1212            (
1213                (
1214                    TimeDelta::try_minutes(15),
1215                    "2004-04-09T02:03:04.123456789Z",
1216                    "2001-01-01T00:02:30",
1217                ),
1218                "2004-04-09T02:02:30Z",
1219            ),
1220            (
1221                (
1222                    TimeDelta::try_minutes(15),
1223                    "2004-04-09T02:03:04.123456789Z",
1224                    "2005-01-01T00:02:30",
1225                ),
1226                "2004-04-09T02:02:30Z",
1227            ),
1228            (
1229                (
1230                    TimeDelta::try_hours(1),
1231                    "2004-04-09T02:03:04.123456789Z",
1232                    "2001-01-01T00:00:00",
1233                ),
1234                "2004-04-09T02:00:00Z",
1235            ),
1236            (
1237                (
1238                    TimeDelta::try_seconds(10),
1239                    "2004-04-09T02:03:11.123456789Z",
1240                    "2001-01-01T00:00:00",
1241                ),
1242                "2004-04-09T02:03:10Z",
1243            ),
1244        ];
1245
1246        cases
1247            .iter()
1248            .for_each(|((stride, source, origin), expected)| {
1249                let stride = stride.unwrap();
1250                let stride1 = stride.num_nanoseconds().unwrap();
1251                let source1 = string_to_timestamp_nanos(source).unwrap();
1252                let origin1 = string_to_timestamp_nanos(origin).unwrap();
1253
1254                let expected1 = string_to_timestamp_nanos(expected).unwrap();
1255                let result = date_bin_nanos_interval(stride1, source1, origin1).unwrap();
1256                assert_eq!(result, expected1, "{source} = {expected}");
1257            })
1258    }
1259
1260    #[test]
1261    fn test_date_bin_before_epoch() {
1262        let cases = [
1263            (
1264                (TimeDelta::try_minutes(15), "1969-12-31T23:44:59.999999999"),
1265                "1969-12-31T23:30:00",
1266            ),
1267            (
1268                (TimeDelta::try_minutes(15), "1969-12-31T23:45:00"),
1269                "1969-12-31T23:45:00",
1270            ),
1271            (
1272                (TimeDelta::try_minutes(15), "1969-12-31T23:45:00.000000001"),
1273                "1969-12-31T23:45:00",
1274            ),
1275        ];
1276
1277        cases.iter().for_each(|((stride, source), expected)| {
1278            let stride = stride.unwrap();
1279            let stride1 = stride.num_nanoseconds().unwrap();
1280            let source1 = string_to_timestamp_nanos(source).unwrap();
1281
1282            let expected1 = string_to_timestamp_nanos(expected).unwrap();
1283            let result = date_bin_nanos_interval(stride1, source1, 0).unwrap();
1284            assert_eq!(result, expected1, "{source} = {expected}");
1285        })
1286    }
1287
1288    #[test]
1289    fn test_date_bin_out_of_range() {
1290        let return_field = &Arc::new(Field::new(
1291            "f",
1292            DataType::Timestamp(TimeUnit::Millisecond, None),
1293            true,
1294        ));
1295        let args = vec![
1296            ColumnarValue::Scalar(ScalarValue::new_interval_mdn(1637426858, 0, 0)),
1297            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
1298                Some(1040292460),
1299                None,
1300            )),
1301            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
1302                Some(string_to_timestamp_nanos("1984-01-07 00:00:00").unwrap()),
1303                None,
1304            )),
1305        ];
1306
1307        let result = invoke_date_bin_with_args(args, 1, return_field);
1308        assert!(result.is_ok());
1309        if let ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(val, _)) =
1310            result.unwrap()
1311        {
1312            assert!(val.is_none(), "Expected None for out of range operation");
1313        }
1314        let args = vec![
1315            ColumnarValue::Scalar(ScalarValue::new_interval_mdn(1637426858, 0, 0)),
1316            ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
1317                Some(-1040292460),
1318                None,
1319            )),
1320            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
1321                Some(string_to_timestamp_nanos("1984-01-07 00:00:00").unwrap()),
1322                None,
1323            )),
1324        ];
1325
1326        let result = invoke_date_bin_with_args(args, 1, return_field);
1327        assert!(result.is_ok());
1328        if let ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(val, _)) =
1329            result.unwrap()
1330        {
1331            assert!(val.is_none(), "Expected None for out of range operation");
1332        }
1333    }
1334
1335    #[test]
1336    fn test_date_bin_compute_distance_i64_min() {
1337        // Regression for #22215: date_bin_nanos_interval on a source near i64::MIN
1338        // previously panicked inside compute_distance with "attempt to subtract with overflow".
1339        // Now it must return a normal Err that the scalar pipeline maps to NULL.
1340        let result = date_bin_nanos_interval(3, i64::MIN, 0);
1341        assert!(
1342            result.is_err(),
1343            "expected Err for source=i64::MIN, got {result:?}"
1344        );
1345
1346        let return_field = &Arc::new(Field::new(
1347            "f",
1348            DataType::Timestamp(TimeUnit::Nanosecond, None),
1349            true,
1350        ));
1351        let args = vec![
1352            ColumnarValue::Scalar(ScalarValue::new_interval_mdn(0, 0, 3)),
1353            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(i64::MIN), None)),
1354            ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)),
1355        ];
1356        let result = invoke_date_bin_with_args(args, 1, return_field);
1357        assert!(result.is_ok(), "expected Ok with NULL, got {result:?}");
1358        if let ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(val, _)) =
1359            result.unwrap()
1360        {
1361            assert!(
1362                val.is_none(),
1363                "Expected None for compute_distance overflow, got {val:?}"
1364            );
1365        } else {
1366            panic!("Expected TimestampNanosecond scalar");
1367        }
1368    }
1369
1370    #[test]
1371    fn test_date_bin_scale_overflow_returns_null() {
1372        // Scaling non-nanosecond timestamps to nanoseconds can overflow.
1373        use arrow::array::{
1374            ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray,
1375            TimestampSecondArray,
1376        };
1377
1378        let scalar_cases = [
1379            ScalarValue::TimestampSecond(Some(i64::MAX), None),
1380            ScalarValue::TimestampMillisecond(Some(i64::MAX), None),
1381            ScalarValue::TimestampMicrosecond(Some(i64::MAX), None),
1382        ];
1383        for source in scalar_cases {
1384            let expected_type = source.data_type();
1385            let return_field = Arc::new(Field::new("f", expected_type.clone(), true));
1386            let args = vec![
1387                ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)),
1388                ColumnarValue::Scalar(source),
1389                ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)),
1390            ];
1391            let result = invoke_date_bin_with_args(args, 1, &return_field)
1392                .unwrap_or_else(|e| panic!("expected Ok for {expected_type}, got {e:?}"));
1393            assert_null_scalar(result, expected_type);
1394        }
1395
1396        let array_cases: Vec<ArrayRef> = vec![
1397            Arc::new(TimestampSecondArray::from(vec![Some(i64::MAX), Some(0)])),
1398            Arc::new(TimestampMillisecondArray::from(vec![
1399                Some(i64::MAX),
1400                Some(0),
1401            ])),
1402            Arc::new(TimestampMicrosecondArray::from(vec![
1403                Some(i64::MAX),
1404                Some(0),
1405            ])),
1406        ];
1407        for array in array_cases {
1408            let dt = array.data_type().clone();
1409            let return_field = Arc::new(Field::new("f", dt.clone(), true));
1410            let args = vec![
1411                ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)),
1412                ColumnarValue::Array(array),
1413                ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)),
1414            ];
1415            let result = invoke_date_bin_with_args(args, 2, &return_field)
1416                .unwrap_or_else(|e| panic!("expected Ok for {dt:?}, got {e:?}"));
1417            assert_array_null_then_valid(result, dt);
1418        }
1419    }
1420
1421    #[test]
1422    fn test_date_bin_time64_micro_overflow_handling() {
1423        // Time64(Microsecond) can hold out-of-range values after reinterpret casts.
1424        use arrow::array::Time64MicrosecondArray;
1425
1426        let data_type = DataType::Time64(TimeUnit::Microsecond);
1427        let return_field = &Arc::new(Field::new("f", data_type.clone(), true));
1428        let stride = || ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1000));
1429        let origin = || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0)));
1430
1431        // Out-of-range source values are per-row data, so they become NULL.
1432        let args = vec![
1433            stride(),
1434            ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX))),
1435            origin(),
1436        ];
1437        let result = invoke_date_bin_with_args(args, 1, return_field).unwrap();
1438        assert_null_scalar(result, data_type.clone());
1439
1440        let array = Arc::new(Time64MicrosecondArray::from(vec![Some(i64::MAX), Some(0)]));
1441        let args = vec![stride(), ColumnarValue::Array(array), origin()];
1442        let result = invoke_date_bin_with_args(args, 2, return_field).unwrap();
1443        assert_array_null_then_valid(result, data_type);
1444
1445        let bad_origin =
1446            || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX)));
1447
1448        // Out-of-range origins are shared inputs, so they return an error.
1449        let args = vec![
1450            stride(),
1451            ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0))),
1452            bad_origin(),
1453        ];
1454        assert_overflow_error(invoke_date_bin_with_args(args, 1, return_field));
1455
1456        let array = Arc::new(Time64MicrosecondArray::from(vec![Some(0), Some(1)]));
1457        let args = vec![stride(), ColumnarValue::Array(array), bad_origin()];
1458        assert_overflow_error(invoke_date_bin_with_args(args, 2, return_field));
1459    }
1460
1461    // Compare scalar execution with a one-row array for the same input.
1462    fn assert_scalar_array_parity(
1463        stride: ScalarValue,
1464        source: ScalarValue,
1465        origin: ScalarValue,
1466    ) {
1467        let return_field = Arc::new(Field::new("f", source.data_type().clone(), true));
1468
1469        let scalar_args = vec![
1470            ColumnarValue::Scalar(stride.clone()),
1471            ColumnarValue::Scalar(source.clone()),
1472            ColumnarValue::Scalar(origin.clone()),
1473        ];
1474        let scalar_result = invoke_date_bin_with_args(scalar_args, 1, &return_field)
1475            .expect("scalar path should not error");
1476        let ColumnarValue::Scalar(scalar_value) = scalar_result else {
1477            panic!("expected scalar result, got {scalar_result:?}");
1478        };
1479
1480        let source_array = source.to_array().expect("source value to array");
1481        let array_args = vec![
1482            ColumnarValue::Scalar(stride),
1483            ColumnarValue::Array(source_array),
1484            ColumnarValue::Scalar(origin),
1485        ];
1486        let array_result = invoke_date_bin_with_args(array_args, 1, &return_field)
1487            .expect("array path should not error");
1488        let ColumnarValue::Array(array) = array_result else {
1489            panic!("expected array result, got {array_result:?}");
1490        };
1491        let array_value =
1492            ScalarValue::try_from_array(&array, 0).expect("array row to scalar");
1493
1494        assert_eq!(
1495            scalar_value, array_value,
1496            "scalar and array results diverged for source {source:?}"
1497        );
1498    }
1499
1500    #[test]
1501    fn test_date_bin_scalar_array_parity() {
1502        // Negative sub-second timestamp with a month interval. This is the case
1503        // that previously diverged (scalar value vs array execution error)
1504        // before #22610; both paths must now agree on the same non-NULL value.
1505        assert_scalar_array_parity(
1506            ScalarValue::new_interval_mdn(1, 0, 0),
1507            ScalarValue::TimestampNanosecond(Some(-1), None),
1508            ScalarValue::TimestampNanosecond(Some(0), None),
1509        );
1510
1511        // Source scaling overflow -> NULL in both paths.
1512        assert_scalar_array_parity(
1513            ScalarValue::new_interval_dt(1, 0),
1514            ScalarValue::TimestampSecond(Some(i64::MAX), None),
1515            ScalarValue::TimestampNanosecond(Some(0), None),
1516        );
1517
1518        // Month interval out-of-range binning -> NULL in both paths.
1519        assert_scalar_array_parity(
1520            ScalarValue::new_interval_mdn(1637426858, 0, 0),
1521            ScalarValue::TimestampMillisecond(Some(1040292460), None),
1522            ScalarValue::TimestampNanosecond(
1523                Some(string_to_timestamp_nanos("1984-01-07 00:00:00").unwrap()),
1524                None,
1525            ),
1526        );
1527    }
1528
1529    #[test]
1530    fn test_date_bin_time_source_requires_time_origin() {
1531        // A TIME source combined with a non-TIME (timestamp) origin is rejected
1532        // with a unit-specific message. This is the shared-input guard that was
1533        // hoisted out of the per-type match arms; cover scalar and array for
1534        // both Time32 and Time64 so the error text stays put.
1535        use arrow::array::{Time32MillisecondArray, Time64NanosecondArray};
1536
1537        let stride = || ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1000));
1538        let ts_origin =
1539            || ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None));
1540
1541        let assert_msg = |args: Vec<ColumnarValue>, dt: DataType, msg: &str| {
1542            let return_field = Arc::new(Field::new("f", dt, true));
1543            assert_eq!(
1544                invoke_date_bin_with_args(args, 1, &return_field)
1545                    .err()
1546                    .unwrap()
1547                    .strip_backtrace(),
1548                msg
1549            );
1550        };
1551
1552        let time32_msg =
1553            "Execution error: DATE_BIN with Time32 source requires Time32 origin";
1554        assert_msg(
1555            vec![
1556                stride(),
1557                ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(0))),
1558                ts_origin(),
1559            ],
1560            DataType::Time32(TimeUnit::Millisecond),
1561            time32_msg,
1562        );
1563        assert_msg(
1564            vec![
1565                stride(),
1566                ColumnarValue::Array(Arc::new(Time32MillisecondArray::from(vec![Some(
1567                    0,
1568                )]))),
1569                ts_origin(),
1570            ],
1571            DataType::Time32(TimeUnit::Millisecond),
1572            time32_msg,
1573        );
1574
1575        let time64_msg =
1576            "Execution error: DATE_BIN with Time64 source requires Time64 origin";
1577        assert_msg(
1578            vec![
1579                stride(),
1580                ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(0))),
1581                ts_origin(),
1582            ],
1583            DataType::Time64(TimeUnit::Nanosecond),
1584            time64_msg,
1585        );
1586        assert_msg(
1587            vec![
1588                stride(),
1589                ColumnarValue::Array(Arc::new(Time64NanosecondArray::from(vec![Some(
1590                    0,
1591                )]))),
1592                ts_origin(),
1593            ],
1594            DataType::Time64(TimeUnit::Nanosecond),
1595            time64_msg,
1596        );
1597    }
1598
1599    #[test]
1600    fn test_date_bin_compute_distance_rem_overflow() {
1601        // Regression for #22215: `time_diff % stride` panics with "attempt to
1602        // calculate the remainder with overflow" when `time_diff == i64::MIN`
1603        // and `stride == -1`. Now it must return a normal Err that the scalar
1604        // pipeline maps to NULL.
1605        let result = date_bin_nanos_interval(-1, i64::MIN, 0);
1606        assert!(
1607            result.is_err(),
1608            "expected Err for time_diff=i64::MIN, stride=-1, got {result:?}"
1609        );
1610    }
1611}