diff --git a/vortex-array/src/arrays/widen/array.rs b/vortex-array/src/arrays/widen/array.rs index 803fd8a0ce0..c59446d0210 100644 --- a/vortex-array/src/arrays/widen/array.rs +++ b/vortex-array/src/arrays/widen/array.rs @@ -30,7 +30,9 @@ impl WidenArray { /// **must** be strictly narrower. When `values` is itself a `Widen`, the new array stores its /// values instead, so that `Widen` arrays do not nest. /// - /// A decimal child **must** be a valid decimal array for its declared dtype. + /// A decimal child **must** be a valid decimal array for its declared dtype. Comparisons, + /// bounds checks and fills rely on this requirement. A violation gives incorrect results, + /// but no undefined behavior. pub fn try_new(mut values: ArrayRef, dtype: DType) -> VortexResult { // NB: Check `values` before the flatten. The values of a `Widen` child are narrower than // the child, so the check in `try_from_parts` alone would accept a child as wide as `dtype` diff --git a/vortex-array/src/arrays/widen/compute/between.rs b/vortex-array/src/arrays/widen/compute/between.rs index 9f6068651ce..33d73ce5f9a 100644 --- a/vortex-array/src/arrays/widen/compute/between.rs +++ b/vortex-array/src/arrays/widen/compute/between.rs @@ -9,7 +9,7 @@ use crate::arrays::Constant; use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::compute::constant_fits_values; -use crate::arrays::widen::compute::constant_in_values_ptype; +use crate::arrays::widen::compute::constant_in_values_dtype; use crate::builtins::ArrayBuiltins; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::BetweenReduce; @@ -28,8 +28,8 @@ impl BetweenReduce for Widen { && constant_fits_values(lower, values.dtype()) == Some(true) && constant_fits_values(upper, values.dtype()) == Some(true) { - let lower = constant_in_values_ptype(lower, values, array.len())?; - let upper = constant_in_values_ptype(upper, values, array.len())?; + let lower = constant_in_values_dtype(lower, values, array.len())?; + let upper = constant_in_values_dtype(upper, values, array.len())?; return values .clone() .between(lower, upper, options.clone()) diff --git a/vortex-array/src/arrays/widen/compute/compare.rs b/vortex-array/src/arrays/widen/compute/compare.rs index 9f3b5c4985d..0c181dae080 100644 --- a/vortex-array/src/arrays/widen/compute/compare.rs +++ b/vortex-array/src/arrays/widen/compute/compare.rs @@ -13,11 +13,12 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; -use crate::arrays::widen::compute::cast_keeping_nullability; use crate::arrays::widen::compute::constant_fits_values; -use crate::arrays::widen::compute::constant_in_values_ptype; +use crate::arrays::widen::compute::constant_in_values_dtype; use crate::arrays::widen::compute::ptype_holding_both_values; use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::scalar::Scalar; use crate::scalar_fn::fns::binary::CompareKernel; use crate::scalar_fn::fns::operators::CompareOperator; @@ -50,24 +51,47 @@ impl CompareKernel for Widen { return compare_out_of_range(lhs, constant, operator, ctx).map(Some); } - let constant = constant_in_values_ptype(constant, values, lhs.len())?; + let constant = constant_in_values_dtype(constant, values, lhs.len())?; values.binary(constant, Operator::from(operator)).map(Some) } } /// Compares the values of two `Widen` arrays at the wider of the two value types. +/// +/// Decimal values with the same scale compare at the larger of the two child precisions, which +/// holds every value of both children. fn compare_narrow_values( lhs: ArrayView<'_, Widen>, rhs: ArrayView<'_, Widen>, operator: CompareOperator, ) -> VortexResult> { - let Some(ptype) = ptype_holding_both_values(lhs, rhs) else { - return Ok(None); - }; + let lhs_values = lhs.values(); + let rhs_values = rhs.values(); + + let dtype = match (lhs_values.dtype(), rhs_values.dtype()) { + (DType::Primitive(..), DType::Primitive(..)) => { + let Some(ptype) = ptype_holding_both_values(lhs, rhs) else { + return Ok(None); + }; - let lhs_values = cast_keeping_nullability(lhs.values(), ptype)?; - let rhs_values = cast_keeping_nullability(rhs.values(), ptype)?; + DType::from(ptype) + } + (DType::Decimal(lhs_dtype, _), DType::Decimal(rhs_dtype, _)) + if lhs_dtype.scale() == rhs_dtype.scale() => + { + DType::Decimal( + DecimalDType::new( + lhs_dtype.precision().max(rhs_dtype.precision()), + lhs_dtype.scale(), + ), + lhs_values.dtype().nullability(), + ) + } + _ => return Ok(None), + }; + let lhs_values = lhs_values.cast(dtype.with_nullability(lhs_values.dtype().nullability()))?; + let rhs_values = rhs_values.cast(dtype.with_nullability(rhs_values.dtype().nullability()))?; lhs_values .binary(rhs_values, Operator::from(operator)) diff --git a/vortex-array/src/arrays/widen/compute/fill_null.rs b/vortex-array/src/arrays/widen/compute/fill_null.rs index e869b4ebc40..d0b36961911 100644 --- a/vortex-array/src/arrays/widen/compute/fill_null.rs +++ b/vortex-array/src/arrays/widen/compute/fill_null.rs @@ -7,7 +7,7 @@ use crate::ArrayRef; use crate::array::ArrayView; use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; -use crate::arrays::widen::compute::cast_keeping_nullability; +use crate::arrays::widen::compute::constant_fits_values; use crate::arrays::widen::compute::ptype_holding_values_and_scalar; use crate::arrays::widen::compute::wrap_values; use crate::builtins::ArrayBuiltins; @@ -20,15 +20,30 @@ impl FillNullReduce for Widen { array: ArrayView<'_, Self>, fill_value: &Scalar, ) -> VortexResult> { - let Some(ptype) = ptype_holding_values_and_scalar(array, fill_value)? else { + if !fill_value.dtype().eq_ignore_nullability(array.dtype()) { return Ok(None); - }; + } + + let values = array.values(); + let dtype = match values.dtype() { + DType::Primitive(..) => { + let Some(ptype) = ptype_holding_values_and_scalar(array, fill_value)? else { + return Ok(None); + }; - // The values widen to a type that holds the fill value. When only the logical type holds - // it, the result is a plain array of the logical type. - let widened = cast_keeping_nullability(array.values(), ptype)?; + DType::Primitive(ptype, values.dtype().nullability()) + } + DType::Decimal(..) + if constant_fits_values(fill_value, values.dtype()) == Some(true) => + { + values.dtype().clone() + } + // Decimal fills outside the child precision use the canonical fill kernel. + _ => return Ok(None), + }; + let widened = values.cast(dtype.clone())?; let fill_value = - fill_value.cast(&DType::Primitive(ptype, fill_value.dtype().nullability()))?; + fill_value.cast(&dtype.with_nullability(fill_value.dtype().nullability()))?; let filled = widened.fill_null(fill_value)?; let dtype = array.dtype().with_nullability(filled.dtype().nullability()); diff --git a/vortex-array/src/arrays/widen/compute/mod.rs b/vortex-array/src/arrays/widen/compute/mod.rs index e409beed88b..b930fdc90d1 100644 --- a/vortex-array/src/arrays/widen/compute/mod.rs +++ b/vortex-array/src/arrays/widen/compute/mod.rs @@ -100,34 +100,41 @@ pub(super) fn cast_keeping_nullability(array: &ArrayRef, ptype: PType) -> Vortex array.cast(DType::Primitive(ptype, array.dtype().nullability())) } -/// Returns whether the range of the integer dtype `values` holds `constant`. +/// Returns whether the dtype range of `values` holds `constant`. /// -/// Returns `None` when `values` is not an integer dtype, when `constant` is null, and when -/// `constant` is not an integer with the signedness of `values`. +/// Integers must have the same signedness, and decimals must have the same scale. Returns `None` +/// for null constants, unsupported dtypes, or incompatible constants. pub(super) fn constant_fits_values(constant: &Scalar, values: &DType) -> Option { - let DType::Primitive(ptype, _) = values else { - return None; - }; - let value = constant.as_primitive_opt()?.pvalue()?; - - // `PValue` compares an integer with another value through `i64` or `u64`, and it panics when - // the other value does not convert. So only integers of the same signedness compare here. - let value_ptype = value.ptype(); - if !ptype.is_int() || !value_ptype.is_int() { - return None; + match (values, constant.dtype()) { + (DType::Decimal(values_dtype, _), DType::Decimal(constant_dtype, _)) + if values_dtype.scale() == constant_dtype.scale() => + { + let value = constant.as_decimal().decimal_value()?; + + Some(value.fits_in_precision(*values_dtype)) + } + (DType::Primitive(values_ptype, _), DType::Primitive(constant_ptype, _)) + if values_ptype.is_int() + && constant_ptype.is_int() + && values_ptype.is_signed_int() == constant_ptype.is_signed_int() => + { + // `PValue` compares through `i64` or `u64`, so the signedness must match. + let value = constant.as_primitive().pvalue()?; + + Some( + PValue::min_value(*values_ptype) <= value + && value <= PValue::max_value(*values_ptype), + ) + } + _ => None, } - if ptype.is_signed_int() != value_ptype.is_signed_int() { - return None; - } - - Some(PValue::min_value(*ptype) <= value && value <= PValue::max_value(*ptype)) } /// Returns the wider of the value types of two [`Widen`] arrays, or `None` when a value type is /// not primitive. /// -/// The logical dtypes **must** have the same signedness, which a debug build asserts. The value -/// types then have that signedness too, so the wider one holds every value of the narrower one. +/// The logical dtypes **must** have the same signedness. The value types then have that signedness +/// too, so the wider one holds every value of the narrower one. pub(super) fn ptype_holding_both_values( lhs: ArrayView<'_, Widen>, rhs: ArrayView<'_, Widen>, @@ -147,7 +154,7 @@ pub(super) fn ptype_holding_both_values( /// Returns `constant`, which **must** fit the type of `values`, as a constant array of that type /// with the nullability of `constant`. -pub(super) fn constant_in_values_ptype( +pub(super) fn constant_in_values_dtype( constant: &Scalar, values: &ArrayRef, len: usize, diff --git a/vortex-array/src/arrays/widen/mod.rs b/vortex-array/src/arrays/widen/mod.rs index 7007a309f8e..f387802df0b 100644 --- a/vortex-array/src/arrays/widen/mod.rs +++ b/vortex-array/src/arrays/widen/mod.rs @@ -11,6 +11,7 @@ //! [`WidenArray::encode`] picks the narrowest such type for a [`PrimitiveArray`]. //! //! A decimal child has the logical scale and a lower precision, and **must** be valid for its dtype. +//! Comparisons, bounds checks and fills use the child dtype to bound its valid values. //! [`WidenArray::encode_decimal`] picks the smallest precision that holds the valid values and the //! scale. Decimal execution copies the child values into the logical canonical width. //! diff --git a/vortex-array/src/arrays/widen/tests/compute.rs b/vortex-array/src/arrays/widen/tests/compute.rs index e809914a62c..f6fe2c3ea57 100644 --- a/vortex-array/src/arrays/widen/tests/compute.rs +++ b/vortex-array/src/arrays/widen/tests/compute.rs @@ -9,6 +9,7 @@ use vortex_mask::Mask; use super::SESSION; use super::all_invalid; +use super::assert_compare_matches_canonical; use super::canonical; use super::non_nullable; use super::one_row; @@ -18,7 +19,6 @@ use super::stored_ptype; use super::unsigned; use super::widen; use crate::ArrayRef; -use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::BoolArray; @@ -50,26 +50,6 @@ use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// Asserts that `lhs op rhs` and `rhs op lhs` equal the same compares of the canonical operands. -#[track_caller] -fn assert_compare_matches_canonical( - lhs: &ArrayRef, - rhs: &ArrayRef, - operator: Operator, - ctx: &mut ExecutionCtx, -) -> VortexResult<()> { - let canonical_lhs = canonical(lhs, ctx)?; - let canonical_rhs = canonical(rhs, ctx)?; - - let expected = canonical_lhs.binary(canonical_rhs.clone(), operator)?; - assert_arrays_eq!(lhs.binary(rhs.clone(), operator)?, expected, ctx); - - let expected = canonical_rhs.binary(canonical_lhs, operator)?; - assert_arrays_eq!(rhs.binary(lhs.clone(), operator)?, expected, ctx); - - Ok(()) -} - /// Slice, filter and mask run on the values, so each result is a `Widen`. The conformance suites /// check the values. #[test] diff --git a/vortex-array/src/arrays/widen/tests/decimal_compare.rs b/vortex-array/src/arrays/widen/tests/decimal_compare.rs new file mode 100644 index 00000000000..cd71ebc1851 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_compare.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal comparison, bounds, and fill-null pushdown tests. +//! +//! Canonical decimal operations provide the reference results, including nullable operands. + +use rstest::rstest; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use super::SESSION; +use super::assert_compare_matches_canonical; +use super::canonical; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::Bool; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::between::BetweenOptions; +use crate::scalar_fn::fns::between::StrictComparison; +use crate::scalar_fn::fns::binary::CompareKernel; +use crate::scalar_fn::fns::fill_null::FillNullReduce; +use crate::scalar_fn::fns::operators::Operator; + +/// Returns a `Widen` of the `parent` dtype over `values`, with the nullability of `values`. +fn widen(values: impl IntoArray, parent: DecimalDType) -> VortexResult { + let values = values.into_array(); + let dtype = DType::Decimal(parent, values.dtype().nullability()); + + WidenArray::try_new(values, dtype) +} + +/// Returns partially valid values at both ends of the child precision. +fn partially_valid(parent: DecimalDType) -> VortexResult { + widen( + DecimalArray::from_option_iter( + [ + Some(-99_999i32), // Minimum. + None, // Null. + Some(0), // Zero. + Some(99_999), // Maximum. + ], + DecimalDType::new(5, 2), + ), + parent, + ) +} + +/// Returns `value` as a scalar of the `parent` dtype, or a null scalar. +fn decimal_scalar(value: Option, parent: DecimalDType, nullability: Nullability) -> Scalar { + match value { + Some(value) => Scalar::decimal(DecimalValue::I64(value), parent, nullability), + None => Scalar::null(DType::Decimal(parent, Nullability::Nullable)), + } +} + +#[rstest] +#[case::below_precision(Some(-100_000))] +#[case::minimum(Some(-99_999))] +#[case::zero(Some(0))] +#[case::maximum(Some(99_999))] +#[case::above_precision(Some(100_000))] +#[case::null(None)] +fn test_compare_constant( + #[case] value: Option, + #[values(18, 76)] parent_precision: u8, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, + #[values( + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::Lte, + Operator::Gt, + Operator::Gte + )] + operator: Operator, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let parent = DecimalDType::new(parent_precision, 2); + let array = match nullability { + Nullability::NonNullable => widen( + DecimalArray::from_iter( + [ + -99_999i32, // Minimum. + 0, // Zero. + 99_999, // Maximum. + ], + DecimalDType::new(5, 2), + ), + parent, + )?, + Nullability::Nullable => partially_valid(parent)?, + } + .into_array(); + let constant = decimal_scalar(value, parent, Nullability::Nullable); + let constant = ConstantArray::new(constant, array.len()).into_array(); + + let compared = array + .binary(constant.clone(), operator)? + .execute::(&mut ctx)?; + assert_eq!(compared.dtype(), &DType::Bool(Nullability::Nullable)); + assert_eq!(compared.len(), array.len()); + + match value { + Some(-99_999..=99_999) => { + for (index, nullability) in [nullability, Nullability::Nullable].into_iter().enumerate() + { + let operand = compared + .nth_child(index) + .vortex_expect("the narrow comparison has two operands"); + assert_eq!( + operand.dtype(), + &DType::Decimal(DecimalDType::new(5, 2), nullability) + ); + } + } + Some(_) => match nullability { + Nullability::NonNullable => assert!(compared.is::()), + Nullability::Nullable => assert!(compared.is::()), + }, + None => assert!(compared.as_::().scalar().is_null()), + } + + assert_compare_matches_canonical(&array, &constant, operator, &mut ctx) +} + +#[rstest] +#[case::same_width( + DecimalArray::from_iter( + [ + -99_999i32, // Minimum. + -1, // Negative. + 0, // Zero. + 99_999, // Maximum. + ], + DecimalDType::new(5, 2), + ), + DecimalArray::from_option_iter( + [ + Some(-100_000i32), // Below the left child precision. + None, // Null. + Some(0), // Equal. + Some(999_999_999), // Above the left child precision. + ], + DecimalDType::new(9, 2), + ), +)] +#[case::different_widths( + DecimalArray::from_iter( + [ + -99i8, // Minimum. + -1, // Negative. + 0, // Zero. + 99, // Maximum. + ], + DecimalDType::new(2, 2), + ), + DecimalArray::from_option_iter( + [ + Some(-100i32), // Below the left child precision. + None, // Null. + Some(0), // Equal. + Some(999_999_999), // Above the left child precision. + ], + DecimalDType::new(9, 2), + ), +)] +fn test_compare_two_widen( + #[case] lhs: DecimalArray, + #[case] rhs: DecimalArray, + #[values( + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::Lte, + Operator::Gt, + Operator::Gte + )] + operator: Operator, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let parent = DecimalDType::new(18, 2); + let lhs = widen(lhs, parent)?; + let rhs = widen(rhs, parent)?.into_array(); + let compared = Widen::compare(lhs.as_view(), &rhs, operator.try_into()?, &mut ctx)? + .vortex_expect("Widen compares decimal children at their common precision"); + + for (index, nullability) in [Nullability::NonNullable, Nullability::Nullable] + .into_iter() + .enumerate() + { + let values = compared + .nth_child(index) + .vortex_expect("compare has two operands"); + assert_eq!( + values.dtype(), + &DType::Decimal(DecimalDType::new(9, 2), nullability) + ); + } + + assert_compare_matches_canonical(&lhs.into_array(), &rhs, operator, &mut ctx) +} + +#[rstest] +#[case::inside(Some(-99_999), Some(99_999))] +#[case::outside(Some(-100_000), Some(100_000))] +#[case::below(Some(-100_001), Some(-100_000))] +#[case::above(Some(100_000), Some(100_001))] +#[case::reversed(Some(99_999), Some(-99_999))] +#[case::null_lower(None, Some(99_999))] +#[case::null_upper(Some(-99_999), None)] +fn test_between(#[case] lower: Option, #[case] upper: Option) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let parent = DecimalDType::new(18, 2); + let array = partially_valid(parent)?.into_array(); + let lower = ConstantArray::new( + decimal_scalar(lower, parent, Nullability::Nullable), + array.len(), + ) + .into_array(); + let upper = ConstantArray::new( + decimal_scalar(upper, parent, Nullability::Nullable), + array.len(), + ) + .into_array(); + let options = BetweenOptions { + lower_strict: StrictComparison::Strict, + upper_strict: StrictComparison::NonStrict, + }; + + let result = array + .clone() + .between(lower.clone(), upper.clone(), options.clone())?; + let expected = canonical(&array, &mut ctx)?.between(lower, upper, options)?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::negative_fits(-99_999, true)] +#[case::positive_fits(99_999, true)] +#[case::negative_outside(-100_000, false)] +#[case::positive_outside(100_000, false)] +fn test_fill_null( + #[case] value: i64, + #[case] fits: bool, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let parent = DecimalDType::new(18, 2); + let array = partially_valid(parent)?; + let fill = decimal_scalar(Some(value), parent, nullability); + + let reduced = Widen::fill_null(array.as_view(), &fill)?; + assert_eq!(reduced.is_some(), fits); + if let Some(reduced) = reduced { + let widen = reduced.as_::(); + let child = widen.values(); + assert_eq!( + child.dtype(), + &DType::Decimal(DecimalDType::new(5, 2), nullability) + ); + } + + let array = array.into_array(); + let result = array.fill_null(fill.clone())?; + let expected = canonical(&array, &mut ctx)?.fill_null(fill)?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index 56c15dd4803..df67b13f0fa 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -10,6 +10,7 @@ mod aggregates; mod compute; mod decimal; +mod decimal_compare; mod decimal_fallback; mod encoding; mod numeric; @@ -30,8 +31,11 @@ use crate::arrays::PrimitiveArray; use crate::arrays::Widen; use crate::arrays::WidenArray; use crate::arrays::widen::WidenArraySlotsExt; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::PType; +use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -57,6 +61,26 @@ fn canonical(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult Ok(array.clone().execute::(ctx)?.into_array()) } +/// Asserts that `lhs op rhs` and `rhs op lhs` equal the same compares of the canonical operands. +#[track_caller] +fn assert_compare_matches_canonical( + lhs: &ArrayRef, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical_lhs = canonical(lhs, ctx)?; + let canonical_rhs = canonical(rhs, ctx)?; + + let expected = canonical_lhs.binary(canonical_rhs.clone(), operator)?; + assert_arrays_eq!(lhs.binary(rhs.clone(), operator)?, expected, ctx); + + let expected = canonical_rhs.binary(canonical_lhs, operator)?; + assert_arrays_eq!(rhs.binary(lhs.clone(), operator)?, expected, ctx); + + Ok(()) +} + /// Nullable `i64` values, all valid, at both ends of the `i8` range. fn signed() -> WidenArray { let values = [Some(i8::MIN), Some(-1), Some(0), Some(1), Some(i8::MAX)];