1use 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 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 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;
294type BinFunction = fn(i64, i64, i64) -> Result<i64>;
303enum Interval {
304 Nanoseconds(i64),
305 Months(i64),
306}
307
308impl Interval {
309 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
325fn 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 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
344fn 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 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
370fn 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
386fn date_bin_months_interval(stride_months: i64, source: i64, origin: i64) -> Result<i64> {
388 let source_date = to_utc_date_time(source)?;
390 let origin_date = to_utc_date_time(origin)?;
391
392 let month_diff = (source_date.year() - origin_date.year()) * 12
394 + source_date.month() as i32
395 - origin_date.month() as i32;
396
397 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 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 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
433fn 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#[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#[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#[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
495fn 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 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 months != 0 {
529 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 (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 if stride == 0 {
598 return exec_err!("DATE_BIN stride must be non-zero");
599 }
600
601 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}