benbellick commented on code in PR #26607:
URL: https://github.com/apache/datafusion/pull/26607#discussion_r4353613678
##########
datafusion/functions/src/datetime/date_bin.rs:
##########
@@ -1629,11 +1635,121 @@ mod tests {
),
];
for (stride, source) in cases {
- let expected = ScalarValue::try_from(source.data_type()).unwrap();
- assert_date_bin(stride, source, epoch(), expected);
+ let return_field = Arc::new(Field::new("f", source.data_type(),
true));
+ for source in [
+ ColumnarValue::Scalar(source.clone()),
+ ColumnarValue::Array(source.to_array().unwrap()),
+ ] {
+ let args = vec![
+ ColumnarValue::Scalar(stride.clone()),
+ source,
+ ColumnarValue::Scalar(epoch()),
+ ];
+ let error = invoke_date_bin_with_args(args, 1, &return_field)
+ .expect_err("out-of-range timestamp binning must fail");
+ assert_eq!(
+ error.strip_backtrace(),
+ "Execution error: DATE_BIN timestamp out of range"
+ );
+ }
+ }
+ }
+
+ fn timestamp_binning_strides() -> [(&'static str, ScalarValue); 5] {
+ [
+ ("three nanoseconds", ScalarValue::new_interval_mdn(0, 0, 3)),
+ ("one second", ScalarValue::new_interval_dt(0, 1000)),
+ ("one hour", ScalarValue::new_interval_dt(0, 3_600_000)),
+ ("one day", ScalarValue::new_interval_dt(1, 0)),
+ ("one month", ScalarValue::new_interval_mdn(1, 0, 0)),
+ ]
+ }
+
+ fn assert_timestamp_nullness_preserved<T: ArrowTimestampType>() {
+ let input = PrimitiveArray::<T>::from_iter([Some(-1), None, Some(0),
Some(1)]);
+ let return_field = Arc::new(Field::new("f", input.data_type().clone(),
true));
+
+ for (stride_name, stride) in timestamp_binning_strides() {
+ for origin in [0, 1_234_567] {
+ let origin = ScalarValue::TimestampNanosecond(Some(origin),
None);
+ let args = vec![
+ ColumnarValue::Scalar(stride.clone()),
+ ColumnarValue::Array(Arc::new(input.clone())),
+ ColumnarValue::Scalar(origin.clone()),
+ ];
+ let ColumnarValue::Array(output) =
+ invoke_date_bin_with_args(args, input.len(),
&return_field).unwrap()
+ else {
+ panic!("expected array output");
+ };
+ assert_eq!(
+ input.nulls(),
+ output.nulls(),
+ "stride {stride_name}, origin {origin:?}, unit {:?}",
+ T::UNIT,
+ );
+ }
+ }
+ }
+
+ fn assert_timestamp_scalar_and_array_agree<T: ArrowTimestampType>() {
+ let input = PrimitiveArray::<T>::from_iter([Some(-1), None, Some(0),
Some(1)]);
+ let return_field = Arc::new(Field::new("f", input.data_type().clone(),
true));
+
+ for (stride_name, stride) in timestamp_binning_strides() {
+ for origin in [0, 1_234_567] {
+ let origin = ScalarValue::TimestampNanosecond(Some(origin),
None);
+ let array_args = vec![
+ ColumnarValue::Scalar(stride.clone()),
+ ColumnarValue::Array(Arc::new(input.clone())),
+ ColumnarValue::Scalar(origin.clone()),
+ ];
+ let ColumnarValue::Array(array_output) =
+ invoke_date_bin_with_args(array_args, input.len(),
&return_field)
+ .unwrap()
+ else {
+ panic!("expected array output");
+ };
+
+ for index in 0..input.len() {
+ let source = ScalarValue::try_from_array(&input,
index).unwrap();
+ let scalar_args = vec![
+ ColumnarValue::Scalar(stride.clone()),
+ ColumnarValue::Scalar(source),
+ ColumnarValue::Scalar(origin.clone()),
+ ];
+ let ColumnarValue::Scalar(scalar_output) =
+ invoke_date_bin_with_args(scalar_args, 1,
&return_field).unwrap()
+ else {
+ panic!("expected scalar output");
+ };
+ assert_eq!(
+ scalar_output,
+ ScalarValue::try_from_array(&array_output,
index).unwrap(),
+ "stride {stride_name}, origin {origin:?}, unit {:?},
index {index}",
+ T::UNIT,
+ );
+ }
+ }
}
}
+ #[test]
+ fn test_date_bin_preserves_timestamp_nullness() {
+ assert_timestamp_nullness_preserved::<TimestampSecondType>();
+ assert_timestamp_nullness_preserved::<TimestampMillisecondType>();
+ assert_timestamp_nullness_preserved::<TimestampMicrosecondType>();
+ assert_timestamp_nullness_preserved::<TimestampNanosecondType>();
+ }
+
+ #[test]
+ fn test_date_bin_timestamp_scalar_and_array_agree() {
Review Comment:
related date_trunc scalar/array consistency test
https://github.com/apache/datafusion/blob/2717e25625ff0c51c1805f8df52d317979db74ec/datafusion/functions/src/datetime/date_trunc.rs#L1930-L1991
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
