refactor: make scalar distance u64 and overflow aware #22892
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -92,6 +92,7 @@ use arrow::util::display::{ArrayFormatter, FormatOptions, array_value_to_string} | |
| use cache::{get_or_create_cached_key_array, get_or_create_cached_null_array}; | ||
| use chrono::{Duration, NaiveDate}; | ||
| use half::f16; | ||
| use num_traits::ToPrimitive; | ||
| pub use struct_builder::ScalarStructBuilder; | ||
|
|
||
| const SECONDS_PER_DAY: i64 = 86_400; | ||
|
|
@@ -2585,53 +2586,99 @@ impl ScalarValue { | |
| /// distance is greater than [`usize::MAX`]. If the type is a float, then the distance will be | ||
| /// rounded to the nearest integer. | ||
| /// | ||
| /// Note: the datatype itself must support subtraction. | ||
| #[deprecated(since = "54.0.0", note = "Use distance_u64 instead")] | ||
| pub fn distance(&self, other: &ScalarValue) -> Option<usize> { | ||
| self.distance_u64(other) | ||
| .and_then(|d| usize::try_from(d).ok()) | ||
| } | ||
|
|
||
| /// Absolute distance between two numeric values (of the same type). This method will return | ||
| /// None if either one of the arguments are null. It might also return None if the resulting | ||
| /// distance is greater than [`u64::MAX`]. If the type is a float, then the distance will be | ||
| /// rounded to the nearest integer. | ||
| /// | ||
| /// Note: the datatype itself must support subtraction. | ||
| pub fn distance_u64(&self, other: &ScalarValue) -> Option<u64> { | ||
| match (self, other) { | ||
| (Self::Int8(Some(l)), Self::Int8(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::Int16(Some(l)), Self::Int16(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::Int32(Some(l)), Self::Int32(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::Int64(Some(l)), Self::Int64(Some(r))) => Some(l.abs_diff(*r)), | ||
| (Self::UInt8(Some(l)), Self::UInt8(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::UInt16(Some(l)), Self::UInt16(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::UInt32(Some(l)), Self::UInt32(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::UInt64(Some(l)), Self::UInt64(Some(r))) => Some(l.abs_diff(*r)), | ||
| // TODO: we might want to look into supporting ceil/floor here for floats. | ||
| (Self::Float16(Some(l)), Self::Float16(Some(r))) => { | ||
|
Contributor
There was a problem hiding this comment. Non-blocking thought: the float arms all repeat the same pattern of taking a rounded absolute diff and then doing the checked conversion. Once the boundary check is fixed, it might be worth using a small helper like
Contributor
Author
There was a problem hiding this comment. Addressed in c356c99 Since this is now a helper, I am also checking for >= 0 |
||
| let diff = (f16::to_f32(*l) - f16::to_f32(*r)).abs().round(); | ||
| if diff <= u64::MAX as f32 { | ||
| Some(diff as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| (Self::Float32(Some(l)), Self::Float32(Some(r))) => { | ||
| let diff = (l - r).abs().round(); | ||
| if diff <= u64::MAX as f32 { | ||
|
Contributor
There was a problem hiding this comment. Same issue here for Could we reject the rounded value at the first float representation above the |
||
| Some(diff as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| (Self::Float64(Some(l)), Self::Float64(Some(r))) => { | ||
| let diff = (l - r).abs().round(); | ||
| if diff <= u64::MAX as f64 { | ||
|
Contributor
There was a problem hiding this comment.
For example, this shape currently returns ScalarValue::Float64(Some(0.0))
.distance_u64(&ScalarValue::Float64(Some(18446744073709551616.0)))Could we make the upper bound strict, for example
Contributor
Author
There was a problem hiding this comment. Thank you very much for your review! Addressed in c356c99 |
||
| Some(diff as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| (Self::Date32(Some(l)), Self::Date32(Some(r))) => Some(l.abs_diff(*r) as u64), | ||
| (Self::Date64(Some(l)), Self::Date64(Some(r))) => Some(l.abs_diff(*r)), | ||
| // Timestamp values are stored as epoch ticks regardless of timezone | ||
| // annotation, so the distance is tz-independent (tz is display metadata). | ||
| (Self::TimestampSecond(Some(l), _), Self::TimestampSecond(Some(r), _)) => { | ||
| Some(l.abs_diff(*r)) | ||
| } | ||
| ( | ||
| Self::TimestampMillisecond(Some(l), _), | ||
| Self::TimestampMillisecond(Some(r), _), | ||
| ) => Some(l.abs_diff(*r)), | ||
| ( | ||
| Self::TimestampMicrosecond(Some(l), _), | ||
| Self::TimestampMicrosecond(Some(r), _), | ||
| ) => Some(l.abs_diff(*r)), | ||
| ( | ||
| Self::TimestampNanosecond(Some(l), _), | ||
| Self::TimestampNanosecond(Some(r), _), | ||
| ) => Some(l.abs_diff(*r)), | ||
| ( | ||
| Self::Decimal32(Some(l), lprecision, lscale), | ||
| Self::Decimal32(Some(r), rprecision, rscale), | ||
| ) => { | ||
| if lprecision == rprecision && lscale == rscale { | ||
| Some(l.abs_diff(*r) as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| ( | ||
| Self::Decimal64(Some(l), lprecision, lscale), | ||
| Self::Decimal64(Some(r), rprecision, rscale), | ||
| ) => { | ||
| if lprecision == rprecision && lscale == rscale { | ||
| Some(l.abs_diff(*r)) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| ( | ||
| Self::Decimal128(Some(l), lprecision, lscale), | ||
| Self::Decimal128(Some(r), rprecision, rscale), | ||
| ) => { | ||
| if lprecision == rprecision && lscale == rscale { | ||
| l.checked_sub(*r)?.checked_abs()?.to_u64() | ||
| } else { | ||
| None | ||
| } | ||
|
|
@@ -2641,7 +2688,7 @@ impl ScalarValue { | |
| Self::Decimal256(Some(r), rprecision, rscale), | ||
| ) => { | ||
| if lprecision == rprecision && lscale == rscale { | ||
| l.checked_sub(*r)?.checked_abs()?.to_u64() | ||
| } else { | ||
| None | ||
| } | ||
|
|
@@ -9437,8 +9484,8 @@ mod tests { | |
| ), | ||
| ]; | ||
| for (lhs, rhs, expected) in cases.iter() { | ||
| let distance = lhs.distance_u64(rhs).unwrap(); | ||
| assert_eq!(distance, *expected as u64); | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -9455,7 +9502,7 @@ mod tests { | |
| ), | ||
| ]; | ||
| for (lhs, rhs) in cases.iter() { | ||
| let distance = lhs.distance_u64(rhs); | ||
| assert!(distance.is_none(), "{lhs} vs {rhs}"); | ||
| } | ||
| } | ||
|
|
@@ -9529,11 +9576,83 @@ mod tests { | |
| ), | ||
| ]; | ||
| for (lhs, rhs) in cases { | ||
| let distance = lhs.distance_u64(&rhs); | ||
| assert!(distance.is_none()); | ||
| } | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_scalar_distance_u64_boundaries() { | ||
| // 1. Full-domain integer ranges | ||
| // i64::MIN to i64::MAX -> distance is u64::MAX | ||
| let lhs = ScalarValue::Int64(Some(i64::MIN)); | ||
| let rhs = ScalarValue::Int64(Some(i64::MAX)); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); | ||
| assert_eq!(rhs.distance_u64(&lhs), Some(u64::MAX)); | ||
|
|
||
| // u64::MIN to u64::MAX -> distance is u64::MAX | ||
| let lhs = ScalarValue::UInt64(Some(u64::MIN)); | ||
| let rhs = ScalarValue::UInt64(Some(u64::MAX)); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); | ||
| assert_eq!(rhs.distance_u64(&lhs), Some(u64::MAX)); | ||
|
|
||
| // 2. Decimal128 overflow edges (around u64::MAX) | ||
| // distance equal to u64::MAX fits | ||
| let lhs = ScalarValue::Decimal128(Some(0), 20, 0); | ||
| let rhs = ScalarValue::Decimal128(Some(u64::MAX as i128), 20, 0); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); | ||
|
|
||
| // distance greater than u64::MAX overflows | ||
| let lhs = ScalarValue::Decimal128(Some(0), 20, 0); | ||
| let rhs = ScalarValue::Decimal128(Some(u64::MAX as i128 + 1), 20, 0); | ||
| assert_eq!(lhs.distance_u64(&rhs), None); | ||
|
|
||
| // 3. Decimal256 overflow edges (around u64::MAX) | ||
| // distance equal to u64::MAX fits | ||
| let lhs = ScalarValue::Decimal256(Some(i256::from_parts(0, 0)), 20, 0); | ||
| let rhs = | ||
| ScalarValue::Decimal256(Some(i256::from_parts(u64::MAX as u128, 0)), 20, 0); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); | ||
|
|
||
| // distance greater than u64::MAX overflows | ||
| let lhs = ScalarValue::Decimal256(Some(i256::from_parts(0, 0)), 20, 0); | ||
| let rhs = ScalarValue::Decimal256( | ||
| Some(i256::from_parts(u64::MAX as u128 + 1, 0)), | ||
| 20, | ||
| 0, | ||
| ); | ||
| assert_eq!(lhs.distance_u64(&rhs), None); | ||
|
|
||
| // 4. Float64 overflow edges (around u64::MAX) | ||
| let lhs = ScalarValue::Float64(Some(0.0)); | ||
| let val: f64 = 18446744073709500000.0; | ||
| let rhs = ScalarValue::Float64(Some(val)); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(18446744073709500416)); | ||
|
|
||
| // float value > u64::MAX overflows | ||
| let rhs = ScalarValue::Float64(Some(1.9e19)); | ||
| assert_eq!(lhs.distance_u64(&rhs), None); | ||
|
|
||
| // Inf, NegInf, NaN | ||
| let inf = ScalarValue::Float64(Some(f64::INFINITY)); | ||
| let neg_inf = ScalarValue::Float64(Some(f64::NEG_INFINITY)); | ||
| let nan = ScalarValue::Float64(Some(f64::NAN)); | ||
| assert_eq!(lhs.distance_u64(&inf), None); | ||
| assert_eq!(lhs.distance_u64(&neg_inf), None); | ||
| assert_eq!(lhs.distance_u64(&nan), None); | ||
|
|
||
| // 5. Date and Timestamp boundaries | ||
| // Date32: i32::MIN to i32::MAX | ||
| let lhs = ScalarValue::Date32(Some(i32::MIN)); | ||
| let rhs = ScalarValue::Date32(Some(i32::MAX)); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u32::MAX as u64)); | ||
|
|
||
| // TimestampSecond: i64::MIN to i64::MAX | ||
| let lhs = ScalarValue::TimestampSecond(Some(i64::MIN), None); | ||
| let rhs = ScalarValue::TimestampSecond(Some(i64::MAX), None); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_scalar_interval_negate() { | ||
| let cases = [ | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the updates here. One remaining concern:
ScalarValue::distanceis a public method and it is now marked as#[deprecated], while the workspace version is still54.0.0.The semver bot is reporting this as a minor-version API change, so could we either remove or defer the deprecation in this PR, or make sure the semver and versioning policy is explicitly satisfied before merging?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@kosiew I removed the deprecation note as I am not sure what the right next version would be - I also updated the PR description mentioning that deprecation of
distanceis not part of this PR.