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,63 +2586,108 @@ 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()) | ||
| } | ||
|
|
||
| /// Helper to convert a rounded float distance to u64, returning None if it exceeds u64::MAX, is negative, or is not finite. | ||
| fn rounded_float_distance_u64(diff: f64) -> Option<u64> { | ||
| if diff.is_finite() && diff >= 0.0 && diff < u64::MAX as f64 { | ||
| Some(diff as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
|
|
||
| /// 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(); | ||
| Self::rounded_float_distance_u64(diff as f64) | ||
| } | ||
| (Self::Float32(Some(l)), Self::Float32(Some(r))) => { | ||
| let diff = (l - r).abs().round(); | ||
| Self::rounded_float_distance_u64(diff as f64) | ||
| } | ||
| (Self::Float64(Some(l)), Self::Float64(Some(r))) => { | ||
| let diff = (l - r).abs().round(); | ||
| Self::rounded_float_distance_u64(diff) | ||
| } | ||
| (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), _, lscale), | ||
| Self::Decimal32(Some(r), _, rscale), | ||
| ) => { | ||
| // In order to be aligned with PartialOrd we only | ||
| // check for equal scale, ignoring precision | ||
| if lscale == rscale { | ||
| Some(l.abs_diff(*r) as u64) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| ( | ||
| Self::Decimal64(Some(l), _, lscale), | ||
| Self::Decimal64(Some(r), _, rscale), | ||
| ) => { | ||
| if lscale == rscale { | ||
| Some(l.abs_diff(*r)) | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| ( | ||
| Self::Decimal128(Some(l), _, lscale), | ||
| Self::Decimal128(Some(r), _, rscale), | ||
| ) => { | ||
| if lscale == rscale { | ||
| l.checked_sub(*r)?.checked_abs()?.to_u64() | ||
| } else { | ||
| None | ||
| } | ||
| } | ||
| ( | ||
| Self::Decimal256(Some(l), _, lscale), | ||
| Self::Decimal256(Some(r), _, rscale), | ||
| ) => { | ||
| if lscale == rscale { | ||
| l.checked_sub(*r)?.checked_abs()?.to_u64() | ||
| } else { | ||
| None | ||
| } | ||
|
|
@@ -9437,8 +9483,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 +9501,7 @@ mod tests { | |
| ), | ||
| ]; | ||
| for (lhs, rhs) in cases.iter() { | ||
| let distance = lhs.distance_u64(rhs); | ||
| assert!(distance.is_none(), "{lhs} vs {rhs}"); | ||
| } | ||
| } | ||
|
|
@@ -9501,13 +9547,9 @@ mod tests { | |
| ScalarValue::Decimal128(Some(123), 5, 5), | ||
| ScalarValue::Decimal128(Some(120), 5, 3), | ||
| ), | ||
| ( | ||
| ScalarValue::Decimal256(Some(123.into()), 5, 5), | ||
| ScalarValue::Decimal256(Some(120.into()), 5, 3), | ||
| ), | ||
| // Distance 2 * 2^50 is larger than usize | ||
| ( | ||
|
|
@@ -9529,11 +9571,124 @@ 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); | ||
|
|
||
| // exact 2^64 boundary (18446744073709551616.0) is greater than u64::MAX, so it should return None | ||
| let exact_2_64_f64 = ScalarValue::Float64(Some(18446744073709551616.0)); | ||
| assert_eq!(lhs.distance_u64(&exact_2_64_f64), None); | ||
|
|
||
| // exact 2^64 boundary as Float32 should also return None | ||
| let lhs_f32 = ScalarValue::Float32(Some(0.0)); | ||
| let exact_2_64_f32 = ScalarValue::Float32(Some(18446744073709551616.0)); | ||
| assert_eq!(lhs_f32.distance_u64(&exact_2_64_f32), None); | ||
|
|
||
| // largest float32 value below 2^64 (2^64 - 2^41 = 18446741874686296064.0) should fit | ||
| let below_2_64_f32 = ScalarValue::Float32(Some(18446741874686296064.0)); | ||
| assert_eq!( | ||
| lhs_f32.distance_u64(&below_2_64_f32), | ||
| Some(18446741874686296064) | ||
| ); | ||
|
|
||
| // 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); | ||
|
|
||
| let inf_f32 = ScalarValue::Float32(Some(f32::INFINITY)); | ||
| let neg_inf_f32 = ScalarValue::Float32(Some(f32::NEG_INFINITY)); | ||
| let nan_f32 = ScalarValue::Float32(Some(f32::NAN)); | ||
| assert_eq!(lhs_f32.distance_u64(&inf_f32), None); | ||
| assert_eq!(lhs_f32.distance_u64(&neg_inf_f32), None); | ||
| assert_eq!(lhs_f32.distance_u64(&nan_f32), None); | ||
|
|
||
| let lhs_f16 = ScalarValue::Float16(Some(f16::ZERO)); | ||
| let inf_f16 = ScalarValue::Float16(Some(f16::INFINITY)); | ||
| let neg_inf_f16 = ScalarValue::Float16(Some(f16::NEG_INFINITY)); | ||
| let nan_f16 = ScalarValue::Float16(Some(f16::NAN)); | ||
| assert_eq!(lhs_f16.distance_u64(&inf_f16), None); | ||
| assert_eq!(lhs_f16.distance_u64(&neg_inf_f16), None); | ||
| assert_eq!(lhs_f16.distance_u64(&nan_f16), 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)); | ||
|
|
||
| // 6. Decimal scale matching (ignoring precision) | ||
| let lhs = ScalarValue::Decimal128(Some(100), 10, 2); | ||
| let rhs = ScalarValue::Decimal128(Some(150), 15, 2); | ||
| assert_eq!(lhs.distance_u64(&rhs), Some(50)); | ||
| assert_eq!(rhs.distance_u64(&lhs), Some(50)); | ||
|
|
||
| let lhs = ScalarValue::Decimal128(Some(100), 10, 2); | ||
| let rhs = ScalarValue::Decimal128(Some(150), 10, 3); | ||
| assert_eq!(lhs.distance_u64(&rhs), None); | ||
| } | ||
|
|
||
| #[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.