benbellick commented on code in PR #26607:
URL: https://github.com/apache/datafusion/pull/26607#discussion_r4353613487
##########
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() {
Review Comment:
mirrors the date_trunc nullness-preservation tests
https://github.com/apache/datafusion/blob/c7a7f9440ab1e01f7e373c455801c2c9c6903d12/datafusion/functions/src/datetime/date_trunc.rs#L1017-L1063
##########
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:
similar to the date_trunc scalar/array consistency tests
https://github.com/apache/datafusion/blob/2717e25625ff0c51c1805f8df52d317979db74ec/datafusion/functions/src/datetime/date_trunc.rs#L1896-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]
