diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 10733c48d8b..9432ddb964a 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -337,6 +337,10 @@ harness = false name = "widen" harness = false +[[bench]] +name = "widen_decimal" +harness = false + [[bench]] name = "run_end_null_count" harness = false diff --git a/vortex-array/benches/widen_decimal.rs b/vortex-array/benches/widen_decimal.rs new file mode 100644 index 00000000000..c1e41ffb23a --- /dev/null +++ b/vortex-array/benches/widen_decimal.rs @@ -0,0 +1,333 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compute on decimal `Widen` arrays and canonical arrays with the same values. +//! +//! The cases cover precision 18 over precision 9, precision 10 over precision 5, and precision 38 +//! over precision 18, all at scale 2. The precision 5 values fit in `i16`, but require `i32` to +//! represent their full decimal domain. Each iteration shares value buffers and builds fresh arrays +//! so cached statistics cannot remove the operation's work. `Widen` buffers are selected once with +//! [`WidenArray::encode_decimal`], outside the timed region. + +use std::fmt; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max; +use vortex_array::aggregate_fn::fns::sum; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::WidenArray; +use vortex_array::arrays::widen::WidenArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::Nullability; +use vortex_array::memory::MemorySession; +use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + LazyLock::force(&SESSION); + + // Keep process-wide execution initialization outside the first measured operation. + let case = Case { + input: Input::Widen, + ..CASES[0] + }; + case.values() + .array(case) + .execute::(&mut SESSION.create_execution_ctx()) + .vortex_expect("the benchmark input executes to canonical decimals"); + + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() + .with::() +}); + +const LEN: usize = 32_768; +const SCALE: i8 = 2; + +#[derive(Clone, Copy, Debug)] +struct Case { + precision: u8, + child_precision: u8, + input: Input, +} + +const CASES: &[Case] = &[ + Case { + precision: 18, + child_precision: 9, + input: Input::Canonical, + }, + Case { + precision: 10, + child_precision: 5, + input: Input::Canonical, + }, + Case { + precision: 38, + child_precision: 18, + input: Input::Canonical, + }, +]; + +#[derive(Clone, Copy, Debug)] +enum Input { + Canonical, + Widen, +} + +fn cases() -> Vec { + CASES + .iter() + .flat_map(|&case| { + [ + case, + Case { + input: Input::Widen, + ..case + }, + ] + }) + .collect() +} + +impl fmt::Display for Case { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "p{}_childp{}_{}", + self.precision, + self.child_precision, + match self.input { + Input::Canonical => "canonical", + Input::Widen => "widen", + }, + ) + } +} + +impl Case { + fn decimal_dtype(self) -> DecimalDType { + DecimalDType::new(self.precision, SCALE) + } + + fn values(self) -> Values { + let (minimum, maximum) = match self.child_precision { + 5 => (10_000, 32_767), + 9 => (-999_999_999, 999_999_999), + 18 => (-999_999_999_999_999_999, 999_999_999_999_999_999), + _ => unreachable!("CASES contains only child precisions 5, 9, and 18"), + }; + let mut rng = StdRng::seed_from_u64(0); + + // Fixed extrema make the required child precision independent of the random sample. + let values = [minimum, maximum] + .into_iter() + .chain((2..LEN).map(|_| rng.random_range(minimum..=maximum))); + + let values = if self.precision <= 18 { + Values::I64(values.collect()) + } else { + Values::I128(values.map(i128::from).collect()) + }; + + if matches!(self.input, Input::Canonical) { + return values; + } + + let encoded = WidenArray::encode_decimal( + values.decimal_array(self.decimal_dtype()), + &mut SESSION.create_execution_ctx(), + ) + .vortex_expect("the benchmark values fit a narrower decimal precision"); + let widen = encoded.as_::(); + let child = widen.values().as_::(); + assert_eq!(child.precision(), self.child_precision); + assert_eq!( + child.values_type(), + child.decimal_dtype().canonical_values_type(), + ); + + match child.values_type() { + DecimalType::I32 => Values::I32(child.buffer::()), + DecimalType::I64 => Values::I64(child.buffer::()), + _ => unreachable!("the benchmark child precisions require i32 or i64"), + } + } + + fn constant(self, value: i64) -> ArrayRef { + let scalar = Scalar::decimal(value.into(), self.decimal_dtype(), Nullability::NonNullable); + ConstantArray::new(scalar, LEN).into_array() + } + + fn in_range(self) -> i64 { + if self.child_precision == 5 { 20_000 } else { 0 } + } + + fn out_of_range(self) -> i64 { + 10i64.pow(u32::from(self.child_precision)) + } + + fn cast_target(self) -> DType { + let precision = if self.precision == 38 { 76 } else { 38 }; + DType::Decimal( + DecimalDType::new(precision, SCALE), + Nullability::NonNullable, + ) + } +} + +enum Values { + I32(Buffer), + I64(Buffer), + I128(Buffer), +} + +impl Values { + fn decimal_array(&self, dtype: DecimalDType) -> DecimalArray { + match self { + Self::I32(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + Self::I64(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + Self::I128(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + } + } + + fn array(&self, case: Case) -> ArrayRef { + match case.input { + Input::Canonical => self.decimal_array(case.decimal_dtype()).into_array(), + Input::Widen => { + let child_dtype = DecimalDType::new(case.child_precision, SCALE); + let child = self.decimal_array(child_dtype).into_array(); + let dtype = DType::Decimal(case.decimal_dtype(), Nullability::NonNullable); + WidenArray::try_new(child, dtype) + .vortex_expect("encode_decimal selected a strictly narrower child dtype") + .into_array() + } + } + } +} + +fn bench_case( + bencher: Bencher, + setup: impl Fn() -> T + Sync, + operation: impl Fn(T, &mut ExecutionCtx) -> O + Sync, +) { + bencher + .with_inputs(|| (setup(), SESSION.create_execution_ctx())) + .bench_values(|(operands, mut ctx)| operation(operands, &mut ctx)); +} + +fn execute_canonical(array: VortexResult, ctx: &mut ExecutionCtx) -> Canonical { + array + .vortex_expect("the operation is valid for the benchmark inputs") + .execute::(ctx) + .vortex_expect("the operation executes to canonical for the benchmark inputs") +} + +#[divan::bench(args = cases())] +fn compare_constant(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || (values.array(case), case.constant(case.in_range())), + |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), + ); +} + +#[divan::bench(args = cases())] +fn compare_constant_out_of_range(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || (values.array(case), case.constant(case.out_of_range())), + |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), + ); +} + +#[divan::bench(args = cases())] +fn sum(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| sum::sum(&array, ctx).vortex_expect("the benchmark decimal sum is valid"), + ); +} + +#[divan::bench(args = cases())] +fn min_max(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| { + min_max::min_max(&array, ctx, NumericalAggregateOpts::default()) + .vortex_expect("the benchmark decimal extrema are valid") + }, + ); +} + +#[divan::bench(args = cases())] +fn filter(bencher: Bencher, case: Case) { + let values = case.values(); + let mut rng = StdRng::seed_from_u64(1); + let selection: BitBuffer = (0..LEN).map(|_| rng.random_bool(0.5)).collect(); + bench_case( + bencher, + || (values.array(case), Mask::from_buffer(selection.clone())), + |(array, mask), ctx| execute_canonical(array.filter(mask), ctx), + ); +} + +#[divan::bench(args = cases())] +fn execute(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| execute_canonical(Ok(array), ctx), + ); +} + +#[divan::bench(args = cases())] +fn cast_widen(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| execute_canonical(array.cast(case.cast_target()), ctx), + ); +} diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index c8f561af0b7..af7e2dfd757 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -21,6 +21,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; +use crate::arrays::WidenArray; use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::BigCast; use crate::dtype::DType; @@ -44,8 +45,6 @@ use crate::validity::Validity; impl CastReduce for Decimal { fn cast(array: ArrayView<'_, Decimal>, dtype: &DType) -> VortexResult> { - // Only nullability changes within the same decimal dtype are reducible without execution. - // Precision/scale changes need the kernel. let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { return Ok(None); }; @@ -56,7 +55,10 @@ impl CastReduce for Decimal { ); }; - if from_decimal_dtype != to_decimal_dtype { + // Without execution, only a nullability change and a same-scale cast that grows the + // canonical width reduce. Every other change of the decimal dtype needs the kernel. + let widens_values = widens_canonical_decimal_values(*from_decimal_dtype, *to_decimal_dtype); + if from_decimal_dtype != to_decimal_dtype && !widens_values { return Ok(None); } @@ -67,6 +69,10 @@ impl CastReduce for Decimal { return Ok(None); }; + if widens_values { + return widen_decimal_values(array, dtype, new_validity).map(Some); + } + // SAFETY: validity has the same length, only its nullability tag changes. unsafe { Ok(Some( @@ -121,6 +127,10 @@ impl CastKernel for Decimal { .clone() .cast_nullability(*to_nullability, array.len(), ctx)?; + if widens_canonical_decimal_values(*from_decimal_dtype, *to_decimal_dtype) { + return widen_decimal_values(array, dtype, new_validity).map(Some); + } + // Reuse the values buffer untouched when no rescale is required, the target precision // only widens (so every value still fits), and the current physical type is already wide // enough to hold the target precision. This keeps the common precision-widening cast @@ -169,6 +179,32 @@ impl CastKernel for Decimal { } } +/// Returns whether a cast keeps the scale and grows the canonical width, which implies a higher +/// precision. +fn widens_canonical_decimal_values(from: DecimalDType, to: DecimalDType) -> bool { + from.scale() == to.scale() && from.canonical_values_type() < to.canonical_values_type() +} + +/// Gives the child the target nullability before constructing a decimal `Widen`. +fn widen_decimal_values( + array: ArrayView<'_, Decimal>, + dtype: &DType, + validity: Validity, +) -> VortexResult { + // SAFETY: values and decimal dtype are unchanged, and the caller cast validity without + // changing its length. + let values = unsafe { + DecimalArray::new_unchecked_handle( + array.buffer_handle().clone(), + array.values_type(), + array.decimal_dtype(), + validity, + ) + }; + + Ok(WidenArray::try_new(values.into_array(), dtype.clone())?.into_array()) +} + fn cast_to_integer( array: ArrayView<'_, Decimal>, nullability: Nullability, diff --git a/vortex-array/src/arrays/interleave/execute/decimal.rs b/vortex-array/src/arrays/interleave/execute/decimal.rs new file mode 100644 index 00000000000..793b17e3cb4 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/decimal.rs @@ -0,0 +1,297 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for decimal [`Interleave`] values. +//! +//! Sources can use different stored widths for the same logical dtype. Decimal execution +//! normalizes them to the logical canonical width before gathering rows. + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use super::super::InterleaveSlots; +use super::fixed_width::Source; +use super::fixed_width::gather; +use crate::AnyColumnar; +use crate::ArrayRef; +use crate::array::Array; +use crate::arrays::Constant; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Primitive; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_decimal_value_type; +use crate::require_child; + +pub(super) fn execute( + mut array: Array, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + array = + require_child!(array, array.array_indices(), InterleaveSlots::ARRAY_INDICES => Primitive); + array = require_child!(array, array.row_indices(), InterleaveSlots::ROW_INDICES => Primitive); + for i in 0..num_values { + array = require_child!(array, array.value(i), InterleaveSlots::VALUES_OFFSET + i => AnyColumnar); + } + + let validity = array.as_ref().validity()?; + let dtype = DecimalDType::try_from(array.dtype())?; + let output = match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + DecimalArray::new(gather_values::(&array, dtype, ctx)?, dtype, validity) + }); + + Ok(ExecutionResult::done(output)) +} + +fn gather_values( + array: &Array, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let sources = (0..array.num_values()) + .map(|i| prepare_source::(array.value(i), dtype, ctx)) + .collect::>>()?; + + gather(array, &sources, ctx.allocator()) +} + +/// Normalizes one decimal source while retaining a single payload for a constant. +fn prepare_source( + value: &ArrayRef, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if let Some(constant) = value.as_opt::() { + let payload = match constant.scalar().as_decimal().decimal_value() { + Some(value) => value.cast::().ok_or_else(|| { + vortex_err!("interleave decimal constant must fit {dtype}, got {value}") + })?, + None => T::default(), + }; + + return Ok(Source::constant(payload, value.len(), ctx.allocator())); + } + + let decimal = value.as_::(); + let data = if decimal.values_type() == T::DECIMAL_TYPE { + decimal.buffer::() + } else { + // A source can still store its values at a width other than the canonical width of its + // dtype (#10701). The copy widens such values, or checks their precision and narrows them + // with zeros in the null slots. + let mask = decimal.validity()?.execute_mask(value.len(), ctx)?; + let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { + copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? + }); + Buffer::::from_byte_buffer(copy.unwrap_host()) + }; + + Ok(Source::decoded(data)) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ConstantArray; + use crate::arrays::DecimalArray; + use crate::arrays::InterleaveArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::decimal::DecimalArrayExt; + use crate::assert_arrays_eq; + use crate::dtype::BigCast; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::NativeDecimalType; + use crate::dtype::Nullability; + use crate::dtype::i256; + use crate::match_each_decimal_value_type; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + use crate::validity::Validity; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + /// Interleaves `values` with one `(array_index, row_index)` pair for each output row, and + /// checks that the result has the canonical width of its dtype and equals `expected`. + #[track_caller] + fn assert_interleaved( + values: Vec, + indices: &[(u8, u8)], + expected: DecimalArray, + ) -> VortexResult<()> { + let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); + let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); + let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; + let mut ctx = SESSION.create_execution_ctx(); + + let output = array.into_array().execute::(&mut ctx)?; + assert_eq!( + output.values_type(), + expected.decimal_dtype().canonical_values_type() + ); + assert_eq!(output.dtype(), expected.dtype()); + assert_arrays_eq!(output.into_array(), expected.into_array(), &mut ctx); + + Ok(()) + } + + /// Gathers rows out of order, with repeats and nulls, at each canonical width. The null slot of + /// the nullable source holds `T::MIN`, which can be outside the precision. + #[rstest] + #[case::i8(2)] + #[case::i16(4)] + #[case::i32(9)] + #[case::i64(18)] + #[case::i128(38)] + #[case::i256(76)] + fn test_decimal_interleave(#[case] precision: u8) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 2); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let min = T::MIN_BY_PRECISION[usize::from(precision)]; + let max = T::MAX_BY_PRECISION[usize::from(precision)]; + let zero = T::default(); + let nullable = DecimalArray::new( + Buffer::copy_from([min, T::MIN, max]), + dtype, + Validity::from_iter([true, false, true]), + ); + let non_nullable = DecimalArray::new( + Buffer::copy_from([max, min, zero]), + dtype, + Validity::NonNullable, + ); + let expected = DecimalArray::new( + Buffer::copy_from([zero, min, zero, max, max, min]), + dtype, + Validity::from_iter([true, true, false, true, true, true]), + ); + + assert_interleaved( + vec![nullable.into_array(), non_nullable.into_array()], + &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], + expected, + ) + }) + } + + /// A constant keeps one payload at the canonical width, also when its scalar stores the value + /// at another width. + #[rstest] + #[case::narrower_scalar(76, Some(DecimalValue::I8(-7)))] + #[case::wider_scalar(2, Some(DecimalValue::I256(i256::from_i128(-7))))] + #[case::null(9, None)] + fn test_decimal_interleave_constant( + #[case] precision: u8, + #[case] payload: Option, + ) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 2); + let (scalar, validity) = match payload { + Some(payload) => ( + Scalar::decimal(payload, dtype, Nullability::NonNullable), + Validity::NonNullable, + ), + None => ( + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + Validity::from_iter([false, true, false, false]), + ), + }; + let constant = ConstantArray::new(scalar, 3).into_array(); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let seven = ::from(7i8).vortex_expect("7 fits every decimal type"); + let minus_seven = + ::from(-7i8).vortex_expect("-7 fits every decimal type"); + let column = DecimalArray::new(buffer![seven], dtype, Validity::NonNullable); + let expected = DecimalArray::new( + buffer![minus_seven, seven, minus_seven, minus_seven], + dtype, + validity, + ); + + assert_interleaved( + vec![constant, column.into_array()], + &[(0, 2), (1, 0), (0, 1), (0, 2)], + expected, + ) + }) + } + + /// Sources that store their values at another width are copied to the canonical width first. + /// The `i8` source widens, and the `i256` source narrows with garbage in its null slots. + #[test] + fn test_decimal_interleave_mixed_stored_widths() -> VortexResult<()> { + let dtype = DecimalDType::new(9, 0); + let narrow = DecimalArray::new(buffer![1i8, -2], dtype, Validity::NonNullable); + let wide = DecimalArray::new( + buffer![i256::from_i128(3), i256::MIN, i256::MAX], + dtype, + Validity::from_iter([true, false, false]), + ); + let expected = DecimalArray::new( + buffer![3i32, -2, 0, 1, 0], + dtype, + Validity::from_iter([true, true, false, true, false]), + ); + + assert_interleaved( + vec![narrow.into_array(), wide.into_array()], + &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], + expected, + ) + } + + #[test] + fn test_decimal_interleave_empty() -> VortexResult<()> { + let dtype = DecimalDType::new(9, 0); + let empty = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); + let expected = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); + + assert_interleaved( + vec![empty.clone().into_array(), empty.into_array()], + &[], + expected, + ) + } + + #[test] + fn test_decimal_interleave_rejects_values_outside_precision() -> VortexResult<()> { + let source = DecimalArray::new( + buffer![100i16], + DecimalDType::new(2, 0), + Validity::NonNullable, + ) + .into_array(); + let array = InterleaveArray::try_new( + vec![source.clone(), source], + PrimitiveArray::from_iter([0u8]).into_array(), + PrimitiveArray::from_iter([0u8]).into_array(), + )? + .into_array(); + + let error = array + .execute::(&mut SESSION.create_execution_ctx()) + .expect_err("the narrowing copy to `i8` rejects 100 at precision 2"); + assert!(error.to_string().contains("precision"), "{error}"); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/interleave/execute/fixed_width.rs b/vortex-array/src/arrays/interleave/execute/fixed_width.rs new file mode 100644 index 00000000000..dfd59fd87c7 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/fixed_width.rs @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared gather kernel for fixed-width [`Interleave`] values. +//! +//! Primitive and decimal execution turn each value into a [`Source`] and call [`gather`]. They +//! differ only in how they read a constant payload or a decoded buffer, and in the array that they +//! build from the gathered payloads. + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use crate::array::Array; +use crate::arrays::Primitive; +use crate::match_each_unsigned_integer_ptype; + +/// Physical storage for one fixed-width value of an [`Interleave`] array. +/// +/// A constant keeps one payload instead of one payload for each row. `row_mask` maps each row to +/// its payload, so one gather loop reads constants and decoded buffers without a branch. +pub(super) struct Source { + /// One payload for each row, or a single payload for every row of a constant. + data: Buffer, + /// The number of logical rows. A constant can have more rows than `data` has payloads. + len: usize, + /// `usize::MAX` for a decoded buffer and zero for a constant, so that `row & row_mask` is the + /// position of the payload of `row` in `data`. + row_mask: usize, +} + +impl Source { + /// Returns a source of `len` rows that all hold `payload`. + pub(super) fn constant(payload: T, len: usize, allocator: &BufferAllocatorRef) -> Self { + Self { + data: Buffer::full_in(payload, 1, allocator.clone()), + len, + row_mask: 0, + } + } + + /// Returns a source with one row for each payload in `data`. + pub(super) fn decoded(data: Buffer) -> Self { + Self { + len: data.len(), + data, + row_mask: usize::MAX, + } + } +} + +/// Gathers the payload of each output row of `array` from `sources`. +/// +/// `sources` holds one [`Source`] for each value of `array`, in order. Both selectors of `array` +/// **must** be [`Primitive`] arrays, which the kernels ensure with `require_child!`. +/// +/// # Errors +/// +/// Returns an error if the selectors have different lengths, or if a selector names a value or a +/// row that does not exist. +pub(super) fn gather( + array: &Array, + sources: &[Source], + allocator: &BufferAllocatorRef, +) -> VortexResult> { + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather_rows( + sources, + branches.as_slice::(), + rows.as_slice::(), + allocator, + ) + }) + }) +} + +fn gather_rows( + sources: &[Source], + branches: &[A], + rows: &[R], + allocator: &BufferAllocatorRef, +) -> VortexResult> +where + T: Copy, + A: AsPrimitive, + R: AsPrimitive, +{ + // `zip` truncates to the shorter input. + vortex_ensure_eq!( + rows.len(), + branches.len(), + "interleave row_indices length does not match array_indices length", + ); + + // Write rows through a `&mut` slice. Yielding `Result` per row splits each `[u8; W]` copy + // into narrow moves (the `Ok` payload is misaligned), and pushing through `BufferMut` reloads + // its length on every row. + let len = branches.len(); + let mut output = BufferMut::with_capacity_in(len, allocator.clone()); + let slots = &mut output.spare_capacity_mut()[..len]; + for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) { + let Some(source) = sources.get((*branch).as_()) else { + vortex_bail!("interleave array index out of bounds"); + }; + let row = (*row).as_(); + vortex_ensure!(row < source.len, "interleave row index out of bounds"); + slot.write(source.data[row & source.row_mask]); + } + // SAFETY: the loop returns on any error, so reaching here means it wrote all `len` slots: + // `slots`, `branches`, and `rows` all have length `len`. + unsafe { output.set_len(len) }; + Ok(output.freeze()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::Source; + use super::gather_rows; + use crate::memory::BufferAllocatorRef; + + #[rstest] + #[case::array_index(2, 0, "array index out of bounds")] + #[case::constant_row_index(0, 1, "row index out of bounds")] + #[case::decoded_row_index(1, 1, "row index out of bounds")] + fn rejects_out_of_bounds_selectors( + #[case] branch: u8, + #[case] row: u8, + #[case] message: &str, + ) -> VortexResult<()> { + let allocator = BufferAllocatorRef::statically_allocated(); + let sources = [ + Source::constant(1u32, 1, &allocator), + Source::decoded(buffer![2u32]), + ]; + + let error = gather_rows(&sources, &[branch], &[row], &allocator) + .expect_err("every case names a value or a row that does not exist"); + assert!(error.to_string().contains(message), "{error}"); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 558db9e3309..20133f64cfc 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -10,6 +10,8 @@ //! [`Interleave::check`]: super::Interleave::check mod bool; +mod decimal; +mod fixed_width; mod primitive; use vortex_error::VortexResult; @@ -27,6 +29,8 @@ pub(super) fn execute( ) -> VortexResult { if array.dtype().is_boolean() { bool::execute(array, ctx) + } else if array.dtype().is_decimal() { + decimal::execute(array, ctx) } else if array.dtype().is_primitive() { primitive::execute(array, ctx) } else { diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index f1595045383..c758d49b47b 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -3,19 +3,16 @@ //! Execution for primitive [`Interleave`] values. -use num_traits::AsPrimitive; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; use super::super::Interleave; use super::super::InterleaveArrayExt; use super::super::InterleaveSlots; +use super::fixed_width::Source; +use super::fixed_width::gather; use crate::AnyColumnar; use crate::array::Array; use crate::arrays::Constant; @@ -26,7 +23,6 @@ use crate::dtype::ToBytes; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_ptype_width; -use crate::match_each_unsigned_integer_ptype; use crate::require_child; pub(super) fn execute( @@ -57,22 +53,13 @@ pub(super) fn execute( Ok(ExecutionResult::done(output)) } -/// Physical storage for one primitive source array. -struct PrimitiveSource { - data: Buffer, - len: usize, - // Maps logical rows to `data`: zero for constants and identity for decoded buffers. - row_mask: usize, -} - fn gather_values( array: &Array, allocator: &BufferAllocatorRef, ) -> VortexResult> { - let values = (0..array.num_values()) + let sources = (0..array.num_values()) .map(|i| { let value = array.value(i); - let len = value.len(); if let Some(constant) = value.as_opt::() { // Validity carries nullness; a null constant's payload is never observed. let payload = constant @@ -85,98 +72,14 @@ fn gather_values( bytes }) .unwrap_or([0; W]); - PrimitiveSource { - data: Buffer::full_in(payload, 1, allocator.clone()), - len, - row_mask: 0, - } + Source::constant(payload, value.len(), allocator) } else { - PrimitiveSource { - data: Buffer::<[u8; W]>::from_byte_buffer( - value.as_::().buffer_handle().to_host_sync(), - ), - len, - row_mask: usize::MAX, - } + Source::decoded(Buffer::<[u8; W]>::from_byte_buffer( + value.as_::().buffer_handle().to_host_sync(), + )) } }) .collect::>(); - let branches = array.array_indices().as_::(); - let rows = array.row_indices().as_::(); - - match_each_unsigned_integer_ptype!(branches.ptype(), |A| { - match_each_unsigned_integer_ptype!(rows.ptype(), |R| { - gather( - &values, - branches.as_slice::(), - rows.as_slice::(), - allocator, - ) - }) - }) -} - -fn gather( - values: &[PrimitiveSource], - branches: &[A], - rows: &[R], - allocator: &BufferAllocatorRef, -) -> VortexResult> -where - T: Copy, - A: AsPrimitive, - R: AsPrimitive, -{ - // `zip` truncates to the shorter input. - vortex_ensure_eq!( - rows.len(), - branches.len(), - "interleave row_indices length does not match array_indices length", - ); - - // Write rows through a `&mut` slice. Yielding `Result` per row splits each `[u8; W]` copy - // into narrow moves (the `Ok` payload is misaligned), and pushing through `BufferMut` reloads - // its length on every row. - let len = branches.len(); - let mut output = BufferMut::with_capacity_in(len, allocator.clone()); - let slots = &mut output.spare_capacity_mut()[..len]; - for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) { - let Some(source) = values.get((*branch).as_()) else { - vortex_bail!("interleave array index out of bounds"); - }; - let row = (*row).as_(); - vortex_ensure!(row < source.len, "interleave row index out of bounds"); - slot.write(source.data[row & source.row_mask]); - } - // SAFETY: the loop returns on any error, so reaching here means it wrote all `len` slots: - // `slots`, `branches`, and `rows` all have length `len`. - unsafe { output.set_len(len) }; - Ok(output.freeze()) -} -#[cfg(test)] -mod tests { - use super::*; - use crate::memory::BufferAllocatorRef; - - #[test] - fn rejects_out_of_bounds_selectors() { - let values = [PrimitiveSource { - data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), - len: 1, - row_mask: 0, - }]; - - let allocator = BufferAllocatorRef::statically_allocated(); - let error = gather(&values, &[1u8], &[0u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("array index out of bounds"), - "{error}" - ); - let error = gather(&values, &[0u8], &[1u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("row index out of bounds"), - "{error}" - ); - } + gather(array, &sources, allocator) } diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index 28c40c8d804..b9c5c7374ee 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -322,10 +322,12 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::Bool; use crate::arrays::BoolArray; + use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; + use crate::dtype::DecimalDType; use crate::dtype::PType; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -337,16 +339,25 @@ mod tests { ThirdValue, } + /// The value type, which picks the execution kernel. + #[derive(Clone, Copy)] + enum ValueKind { + Bool, + Primitive, + Decimal, + } + #[rstest] #[case::array_selector(LazyChild::ArraySelector)] #[case::row_selector(LazyChild::RowSelector)] #[case::third_value(LazyChild::ThirdValue)] fn executes_lazy_children( #[case] lazy_child: LazyChild, - #[values(false, true)] boolean: bool, + #[values(ValueKind::Bool, ValueKind::Primitive, ValueKind::Decimal)] value_kind: ValueKind, ) -> VortexResult<()> { - let (mut values, expected) = if boolean { - ( + let decimal_dtype = DecimalDType::new(9, 2); + let (mut values, expected) = match value_kind { + ValueKind::Bool => ( vec![ BoolArray::from_iter([true, false]).into_array(), BoolArray::from_iter([false, true]).into_array(), @@ -354,9 +365,8 @@ mod tests { ], BoolArray::from_iter([None, Some(true), Some(true), Some(false), Some(false)]) .into_array(), - ) - } else { - ( + ), + ValueKind::Primitive => ( vec![ PrimitiveArray::from_iter([10i32, 11]).into_array(), PrimitiveArray::from_iter([20i32, 21]).into_array(), @@ -364,7 +374,19 @@ mod tests { ], PrimitiveArray::from_option_iter([None, Some(10i32), Some(21), Some(30), Some(11)]) .into_array(), - ) + ), + ValueKind::Decimal => ( + vec![ + DecimalArray::from_iter([10i32, 11], decimal_dtype).into_array(), + DecimalArray::from_iter([20i32, 21], decimal_dtype).into_array(), + DecimalArray::from_option_iter([Some(30i32), None], decimal_dtype).into_array(), + ], + DecimalArray::from_option_iter( + [None, Some(10i32), Some(21), Some(30), Some(11)], + decimal_dtype, + ) + .into_array(), + ), }; let mut array_indices = PrimitiveArray::from_iter([2u8, 0, 1, 2, 0]).into_array(); let mut row_indices = PrimitiveArray::from_iter([1u32, 0, 1, 0, 1]).into_array(); @@ -384,10 +406,9 @@ mod tests { !matches!(lazy_child, LazyChild::RowSelector) ); for (i, value) in values.iter().enumerate() { - let matches = if boolean { - value.is::() - } else { - value.is::() + let matches = match value_kind { + ValueKind::Bool => value.is::(), + ValueKind::Primitive | ValueKind::Decimal => value.is::(), }; assert_eq!( matches, diff --git a/vortex-array/src/arrays/widen/array.rs b/vortex-array/src/arrays/widen/array.rs index ed08a50dc47..260a39cf853 100644 --- a/vortex-array/src/arrays/widen/array.rs +++ b/vortex-array/src/arrays/widen/array.rs @@ -16,7 +16,7 @@ use crate::dtype::DType; /// The child slots of a [`WidenArray`]. #[array_slots(Widen)] pub struct WidenSlots { - /// The stored values, in an integer type strictly narrower than the logical dtype. + /// The stored values, with a canonical width strictly narrower than the logical dtype. #[slot(0)] pub values: ArrayRef, } @@ -25,9 +25,10 @@ impl WidenArray { /// Creates a `Widen` that presents `values` with the wider logical `dtype`, without reading /// any values. /// - /// `values` and `dtype` must both be integer dtypes with the same signedness and nullability, - /// and `values` must be strictly narrower. When `values` is itself a `Widen`, the new array - /// stores its values instead, so that `Widen` arrays do not nest. + /// `values` and `dtype` must both be integers with the same signedness, or decimals with the + /// same scale. They must have the same nullability, and the canonical width of `values` must + /// be strictly narrower. This trusts the child dtype and reads no values. When `values` is + /// itself a `Widen`, the new array stores its values instead, so that `Widen` arrays do not nest. 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` @@ -52,23 +53,34 @@ impl WidenArray { /// Checks that the dtypes of `values` and `logical` meet the requirements of /// [`WidenArray::try_new`]. pub(super) fn validate_dtypes(values: &DType, logical: &DType) -> VortexResult<()> { - let (values_ptype, logical_ptype) = match (values, logical) { + match (values, logical) { (DType::Primitive(values_ptype, _), DType::Primitive(logical_ptype, _)) if values_ptype.is_int() && logical_ptype.is_int() => { - (*values_ptype, *logical_ptype) + vortex_ensure!( + values_ptype.is_signed_int() == logical_ptype.is_signed_int(), + "Widen requires values with the signedness of {logical}, got {values}" + ); + vortex_ensure!( + values_ptype.byte_width() < logical_ptype.byte_width(), + "Widen requires values narrower than {logical}, got {values}" + ); } - _ => vortex_bail!("Widen requires integer dtypes, got {values} values for {logical}"), - }; + (DType::Decimal(values_decimal, _), DType::Decimal(logical_decimal, _)) => { + vortex_ensure!( + values_decimal.scale() == logical_decimal.scale(), + "Widen requires values with the scale of {logical}, got {values}" + ); + vortex_ensure!( + values_decimal.canonical_values_type() < logical_decimal.canonical_values_type(), + "Widen requires values narrower than {logical}, got {values}" + ); + } + _ => vortex_bail!( + "Widen requires integer dtypes or decimal dtypes, got {values} values for {logical}" + ), + } - vortex_ensure!( - values_ptype.is_signed_int() == logical_ptype.is_signed_int(), - "Widen requires values with the signedness of {logical}, got {values}" - ); - vortex_ensure!( - values_ptype.byte_width() < logical_ptype.byte_width(), - "Widen requires values narrower than {logical}, got {values}" - ); vortex_ensure!( values.nullability() == logical.nullability(), "Widen requires values with the nullability of {logical}, got {values}" diff --git a/vortex-array/src/arrays/widen/compute/aggregate.rs b/vortex-array/src/arrays/widen/compute/aggregate.rs index aeca44dc042..e00838144f0 100644 --- a/vortex-array/src/arrays/widen/compute/aggregate.rs +++ b/vortex-array/src/arrays/widen/compute/aggregate.rs @@ -7,6 +7,11 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::arrays::Widen; @@ -19,7 +24,10 @@ use crate::scalar::Scalar; /// /// The values run the same accumulator code as the canonical array, so edge cases such as an empty /// batch or a single null match by construction, and the values keep their own cached statistics -/// and kernels. [`Sum`] is served only when [`is_exact_sum`] holds. +/// and kernels. Decimal widening preserves order, equality and validity, so extrema, constancy and +/// sortedness have the same partial after widening. Decimal sums must use the logical precision +/// to select their accumulator, so they fall back to execution. Integer [`Sum`] is served only +/// when [`is_exact_sum`] holds. #[derive(Debug)] pub(crate) struct WidenArrayAggregate; @@ -34,10 +42,19 @@ impl DynAggregateKernel for WidenArrayAggregate { return Ok(None); }; - // NB: Exactness rests on integer decoding. A decimal `Widen` (CD4 in #10733) needs its own - // argument, because `Sum` checks a decimal for overflow against the precision of its input. - if !batch.dtype().is_int() { - return Ok(None); + match batch.dtype() { + DType::Decimal(..) => { + if !(aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::()) + { + return Ok(None); + } + } + dtype if dtype.is_int() => {} + _ => return Ok(None), } let values = widen.values(); @@ -49,8 +66,7 @@ impl DynAggregateKernel for WidenArrayAggregate { accumulator.accumulate(values, ctx)?; let partial_dtype = aggregate_fn.state_dtype(batch.dtype()).vortex_expect( - "an aggregate that accepts the values accepts the logical dtype, which is a wider \ - integer of the same signedness", + "the served aggregate accepts the logical dtype when it accepts the narrower values", ); Ok(Some(accumulator.flush()?.cast(&partial_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/cast.rs b/vortex-array/src/arrays/widen/compute/cast.rs index 6f7131b2f90..771f89316af 100644 --- a/vortex-array/src/arrays/widen/compute/cast.rs +++ b/vortex-array/src/arrays/widen/compute/cast.rs @@ -16,6 +16,13 @@ use crate::scalar_fn::fns::cast::CastReduce; impl CastReduce for Widen { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { let values = array.values(); + if matches!(values.dtype(), DType::Decimal(..)) { + return match dtype { + DType::Decimal(..) | DType::Primitive(..) => values.cast(dtype.clone()).map(Some), + _ => Ok(None), + }; + } + let (DType::Primitive(target_ptype, nullability), DType::Primitive(values_ptype, _)) = (dtype, values.dtype()) else { diff --git a/vortex-array/src/arrays/widen/compute/compare.rs b/vortex-array/src/arrays/widen/compute/compare.rs index 9f3b5c4985d..3ad8967b6d3 100644 --- a/vortex-array/src/arrays/widen/compute/compare.rs +++ b/vortex-array/src/arrays/widen/compute/compare.rs @@ -15,9 +15,11 @@ 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,7 +52,7 @@ 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) } @@ -62,6 +64,31 @@ fn compare_narrow_values( rhs: ArrayView<'_, Widen>, operator: CompareOperator, ) -> VortexResult> { + if let (DType::Decimal(lhs_dtype, _), DType::Decimal(rhs_dtype, _)) = + (lhs.values().dtype(), rhs.values().dtype()) + { + if lhs_dtype.scale() != rhs_dtype.scale() { + return Ok(None); + } + + let decimal_dtype = DecimalDType::new( + lhs_dtype.precision().max(rhs_dtype.precision()), + lhs_dtype.scale(), + ); + let lhs_values = lhs.values().cast(DType::Decimal( + decimal_dtype, + lhs.values().dtype().nullability(), + ))?; + let rhs_values = rhs.values().cast(DType::Decimal( + decimal_dtype, + rhs.values().dtype().nullability(), + ))?; + + return lhs_values + .binary(rhs_values, Operator::from(operator)) + .map(Some); + } + let Some(ptype) = ptype_holding_both_values(lhs, rhs) else { return Ok(None); }; diff --git a/vortex-array/src/arrays/widen/compute/fill_null.rs b/vortex-array/src/arrays/widen/compute/fill_null.rs index e869b4ebc40..644e3ec34c9 100644 --- a/vortex-array/src/arrays/widen/compute/fill_null.rs +++ b/vortex-array/src/arrays/widen/compute/fill_null.rs @@ -8,6 +8,7 @@ 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,6 +21,23 @@ impl FillNullReduce for Widen { array: ArrayView<'_, Self>, fill_value: &Scalar, ) -> VortexResult> { + if matches!(array.values().dtype(), DType::Decimal(..)) { + if !fill_value.dtype().eq_ignore_nullability(array.dtype()) + || constant_fits_values(fill_value, array.values().dtype()) != Some(true) + { + return Ok(None); + } + + let fill_dtype = array + .values() + .dtype() + .with_nullability(fill_value.dtype().nullability()); + let filled = array.values().fill_null(fill_value.cast(&fill_dtype)?)?; + let dtype = array.dtype().with_nullability(filled.dtype().nullability()); + + return wrap_values(filled, dtype).map(Some); + } + let Some(ptype) = ptype_holding_values_and_scalar(array, fill_value)? else { return Ok(None); }; diff --git a/vortex-array/src/arrays/widen/compute/mod.rs b/vortex-array/src/arrays/widen/compute/mod.rs index e409beed88b..5d43ef4fcff 100644 --- a/vortex-array/src/arrays/widen/compute/mod.rs +++ b/vortex-array/src/arrays/widen/compute/mod.rs @@ -100,11 +100,27 @@ 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 { + if let DType::Decimal(decimal_dtype, _) = values { + let DType::Decimal(constant_dtype, _) = constant.dtype() else { + return None; + }; + if constant_dtype.scale() != decimal_dtype.scale() { + return None; + } + + return Some( + constant + .as_decimal() + .decimal_value()? + .fits_in_precision(*decimal_dtype), + ); + } + let DType::Primitive(ptype, _) = values else { return None; }; @@ -126,8 +142,8 @@ pub(super) fn constant_fits_values(constant: &Scalar, values: &DType) -> Option< /// 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 +163,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/encode.rs b/vortex-array/src/arrays/widen/encode.rs index 24d41c334fb..90d3d966ce3 100644 --- a/vortex-array/src/arrays/widen/encode.rs +++ b/vortex-array/src/arrays/widen/encode.rs @@ -4,6 +4,7 @@ use std::ops::RangeInclusive; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::ArrayRef; @@ -11,12 +12,21 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; +use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; +use crate::arrays::decimal::valid_range; use crate::arrays::widen::compute::wrap_values; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; use crate::dtype::PType; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; use crate::scalar::Scalar; impl WidenArray { @@ -30,6 +40,48 @@ impl WidenArray { wrap_values(narrow_values(array, ctx)?.into_array(), dtype) } + /// Stores a decimal array at the smallest precision that its valid values and scale allow. + /// + /// The result keeps the input dtype and nullability. Empty and all-null arrays use the + /// precision required by the scale, or 1 for a non-positive scale. Returns the input unchanged + /// when the selected precision has the canonical width of the input dtype. + /// + /// # Errors + /// + /// Returns an error if a valid value does not fit the input precision. + pub fn encode_decimal(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let parent = array.decimal_dtype(); + let validity = array.validity()?; + let mask = validity.execute_mask(array.len(), ctx)?; + let range = match_each_decimal_value_type!(array.values_type(), |T| { + valid_range(array.buffer::().as_slice(), &mask).map(|(min, max)| { + DecimalValue::from(min).as_i256()..=DecimalValue::from(max).as_i256() + }) + }); + let child_dtype = smallest_decimal_dtype(range, parent)?; + if child_dtype.canonical_values_type() == parent.canonical_values_type() { + return Ok(array.into_array()); + } + + let buffer = match_each_decimal_value_type!(array.values_type(), |T| { + copy_to_canonical_width(array.buffer::().as_slice(), &mask, child_dtype)? + }); + let dtype = array.dtype().clone(); + + // SAFETY: the range proves that every valid value fits the child precision. The copy + // preserves the scale and length, uses the child canonical width, and keeps its validity. + let values = unsafe { + DecimalArray::new_unchecked_handle( + buffer, + child_dtype.canonical_values_type(), + child_dtype, + validity, + ) + }; + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) + } + /// Returns `array` cast to the type that [`Self::encode`] chooses for its values, without a /// `Widen` around them. /// @@ -107,3 +159,37 @@ pub(super) fn integer_value(scalar: &Scalar) -> VortexResult { u64::try_from(scalar).map(i128::from) } } + +/// Returns the smallest decimal dtype that holds an unscaled valid-value range at the parent scale. +/// +/// The precision is at least 1 and at least a positive scale. An absent range represents an empty +/// or all-null array and needs only that minimum precision. The result never exceeds the parent +/// precision. Callers must exclude null payloads from the range, for example with +/// [`valid_range`]. +/// +/// # Errors +/// +/// Returns an error if the minimum exceeds the maximum or an endpoint exceeds the parent precision. +pub fn smallest_decimal_dtype( + range: Option>, + parent: DecimalDType, +) -> VortexResult { + let minimum_precision = parent.scale().max(1) as u8; + let Some(range) = range else { + return Ok(DecimalDType::new(minimum_precision, parent.scale())); + }; + let (min, max) = range.into_inner(); + vortex_ensure!( + min <= max, + "Decimal range requires ordered extrema, got {min} to {max}" + ); + + for precision in minimum_precision..=parent.precision() { + let index = usize::from(precision); + if i256::MIN_BY_PRECISION[index] <= min && max <= i256::MAX_BY_PRECISION[index] { + return Ok(DecimalDType::new(precision, parent.scale())); + } + } + + vortex_bail!("Decimal range must fit {parent}, got {min} to {max}") +} diff --git a/vortex-array/src/arrays/widen/mod.rs b/vortex-array/src/arrays/widen/mod.rs index ac748864ad2..b873655c7ab 100644 --- a/vortex-array/src/arrays/widen/mod.rs +++ b/vortex-array/src/arrays/widen/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! An in-memory encoding that stores an integer array in a narrower integer type. +//! An in-memory encoding that stores integer and decimal arrays at a narrower canonical width. //! //! A [`WidenArray`] keeps the logical dtype of an integer array, and stores its values and their //! validity in a child with a strictly narrower integer type of the same signedness and @@ -9,9 +9,14 @@ //! values, but not `u8` values. [`WidenArray::encode`] picks the narrowest such type for a //! [`PrimitiveArray`]. //! +//! Decimal children keep the logical scale and nullability, with a strictly narrower canonical +//! width. [`WidenArray::encode_decimal`] selects the smallest precision that holds the valid values +//! and the scale. Decimal execution copies the child values into the logical canonical width, +//! without casting, because a widening decimal cast can itself return a `Widen`. +//! //! Narrowing is a physical choice, so it **must not** change a result. Execution casts the values -//! to the logical dtype, and an operation runs on the values only where it gives the canonical -//! result. A `Widen` never wraps another `Widen`. +//! to the logical dtype for integers, and an operation runs on the values only where it gives the +//! canonical result. A `Widen` never wraps another `Widen`. //! //! Take executes its gather instead of leaving a lazy [`Dict`] under a `Widen`, because code such //! as the DuckDB exporter and the dict layout writer looks for a `Dict` at the root of an array. @@ -31,6 +36,7 @@ pub use vtable::WidenArray; pub(crate) mod compute; mod encode; +pub use encode::smallest_decimal_dtype; mod vtable; use vortex_session::VortexSession; diff --git a/vortex-array/src/arrays/widen/tests/aggregates.rs b/vortex-array/src/arrays/widen/tests/aggregates.rs index b91640151b3..46e07b200a2 100644 --- a/vortex-array/src/arrays/widen/tests/aggregates.rs +++ b/vortex-array/src/arrays/widen/tests/aggregates.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt::Debug; use std::num::NonZeroUsize; use rstest::rstest; @@ -99,7 +100,7 @@ fn canonical_batch( /// Runs the kernel that the session registers for `aggregate` on `Widen` over `array`, and returns /// its partial, or `None` when no kernel is registered or the kernel declines. -fn kernel_partial( +pub(super) fn kernel_partial( aggregate: &AggregateFnRef, array: &ArrayRef, ctx: &mut ExecutionCtx, @@ -113,26 +114,28 @@ fn kernel_partial( } } -/// Accumulates `batches` of `dtype` with `aggregate` once as [`WidenArray::encode`] encodes them -/// and once as canonical arrays, and asserts that the running partials and the results are equal. +/// Accumulates `batches` of `dtype` with `aggregate` once as `encode` encodes them and once as +/// canonical arrays, and asserts that the running partials and the results are equal. /// /// On each batch that encodes to a `Widen`, the kernel must return a partial exactly when /// `served`, and the partial must equal the partial of a fresh accumulator on the canonical batch. /// Encoding and each accumulator get their own canonical batch, because both cache statistics on /// the array they read, and an accumulator reads a cached statistic instead of the data. #[track_caller] -fn assert_aggregate_batches( +pub(super) fn assert_aggregate_batches( aggregate: &AggregateFnRef, served: bool, - batches: &[Vec>], + batches: &[T], dtype: &DType, + mut build_batch: impl FnMut(&T, &DType, &mut ExecutionCtx) -> VortexResult, + mut encode: impl FnMut(C, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let mut actual = aggregate.accumulator(dtype)?; let mut expected = aggregate.accumulator(dtype)?; for values in batches { - let encoded = WidenArray::encode(canonical_batch(values, dtype, ctx)?, ctx)?; + let encoded = encode(build_batch(values, dtype, ctx)?, ctx)?; if encoded.is::() { let partial = kernel_partial(aggregate, &encoded, ctx)?; @@ -140,13 +143,13 @@ fn assert_aggregate_batches( if let Some(partial) = partial { let mut reference = aggregate.accumulator(dtype)?; - reference.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?; + reference.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?; assert_eq!(partial, reference.flush()?, "{aggregate} on {values:?}"); } } actual.accumulate(&encoded, ctx)?; - expected.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?; + expected.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?; assert_eq!( actual.partial_scalar()?, expected.partial_scalar()?, @@ -238,7 +241,15 @@ fn test_aggregate_partials( let dtype = DType::Primitive(logical, nullability); for (aggregate, served) in aggregates() { - assert_aggregate_batches(&aggregate, served, &batches, &dtype, &mut ctx)?; + assert_aggregate_batches( + &aggregate, + served, + &batches, + &dtype, + |values, dtype, ctx| canonical_batch(values, dtype, ctx), + WidenArray::encode, + &mut ctx, + )?; } Ok(()) diff --git a/vortex-array/src/arrays/widen/tests/decimal.rs b/vortex-array/src/arrays/widen/tests/decimal.rs new file mode 100644 index 00000000000..dc17440ec22 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Construction, execution, and encoding contracts for decimal `Widen` arrays. + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use super::stored_decimal_dtype; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::Canonical; +use crate::EmptyArrayData; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::arrays::widen::WidenSlots; +use crate::arrays::widen::smallest_decimal_dtype; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::dtype::PType; +use crate::dtype::i256; +use crate::validity::Validity; + +fn dtype(precision: u8, scale: i8, nullability: Nullability) -> DType { + DType::Decimal(DecimalDType::new(precision, scale), nullability) +} + +#[rstest] +#[case::i32_to_i64(dtype(9, 2, NonNullable), dtype(18, 2, NonNullable), true)] +#[case::i128_to_i256(dtype(38, 0, Nullable), dtype(76, 0, Nullable), true)] +#[case::same_i64(dtype(10, 2, NonNullable), dtype(18, 2, NonNullable), false)] +#[case::same_i256(dtype(75, 0, NonNullable), dtype(76, 0, NonNullable), false)] +#[case::scale(dtype(9, 1, NonNullable), dtype(18, 2, NonNullable), false)] +#[case::nullability(dtype(9, 2, Nullable), dtype(18, 2, NonNullable), false)] +#[case::higher_precision(dtype(18, 2, NonNullable), dtype(9, 2, NonNullable), false)] +#[case::integer_child(PType::I32.into(), dtype(18, 2, NonNullable), false)] +#[case::integer_parent(dtype(9, 2, NonNullable), PType::I64.into(), false)] +fn test_try_new_decimal_dtypes( + #[case] child: DType, + #[case] parent: DType, + #[case] accepted: bool, +) -> VortexResult<()> { + let values = Canonical::empty(&child).into_array(); + let parts = ArrayParts::new( + Widen, + parent.clone(), + 0, + EmptyArrayData, + WidenSlots { + values: values.clone(), + } + .into_slots(), + ); + let constructed = WidenArray::try_new(values, parent.clone()); + let validated = WidenArray::try_from_parts(parts); + assert_eq!(constructed.is_ok(), accepted); + assert_eq!(validated.is_ok(), accepted); + if !accepted { + let error = constructed + .expect_err("the dtype pair is invalid") + .to_string(); + assert!(error.contains(&child.to_string()), "{error}"); + assert!(error.contains(&parent.to_string()), "{error}"); + } + + Ok(()) +} + +#[test] +fn test_decimal_widen_flattens() -> VortexResult<()> { + let values = DecimalArray::from_iter([1i16, 99], DecimalDType::new(4, 0)).into_array(); + let inner = WidenArray::try_new(values.clone(), dtype(9, 0, NonNullable))?.into_array(); + let outer = WidenArray::try_new(inner, dtype(18, 0, NonNullable))?; + assert!(ArrayRef::ptr_eq(outer.values(), &values)); + assert!(!outer.values().is::()); + + Ok(()) +} + +#[rstest] +#[case::canonical(DecimalArray::from_iter([-999i32, 0, 999], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::nullable(DecimalArray::from_option_iter([Some(-999i32), None, Some(999)], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::narrow_stored(DecimalArray::from_iter([-99i8, 0, 99], DecimalDType::new(9, 2)), [-99, 0, 99])] +#[case::parent_width_stored(DecimalArray::from_iter([-999i64, 0, 999], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::wider_stored(DecimalArray::from_iter([-999i128, 0, 999], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::null_garbage(DecimalArray::new(buffer![i128::MIN, 999, i128::MAX], DecimalDType::new(9, 2), Validity::from_iter([false, true, false])), [0, 999, 0])] +fn test_decimal_execute_widens_and_terminates( + #[case] child: DecimalArray, + #[case] expected: [i64; 3], +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = child.validity()?; + let parent = dtype(18, 2, validity.nullability()); + let reference = DecimalArray::new( + expected.into_iter().collect::>(), + DecimalDType::new(18, 2), + validity, + ); + let widened = WidenArray::try_new(child.into_array(), parent.clone())?.into_array(); + let executed = widened.clone().execute::(&mut ctx)?.into_array(); + assert_eq!(executed.dtype(), &parent); + let decimal = executed.as_::(); + assert_eq!(decimal.values_type(), DecimalType::I64); + assert_arrays_eq!(executed, reference, &mut ctx); + assert_arrays_eq!(widened, reference, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::scale_raises_precision(vec![Some(1)], 7, 7, true)] +#[case::negative_boundary(vec![Some(-128), Some(99)], 0, 3, true)] +#[case::negative_scale(vec![Some(-99), Some(99)], -3, 2, true)] +#[case::all_null(vec![None, None], 0, 1, true)] +#[case::empty(vec![], 0, 1, true)] +#[case::no_narrower_width(vec![Some(1_000_000_000)], 2, 10, false)] +fn test_encode_decimal_best_precision( + #[case] values: Vec>, + #[case] scale: i8, + #[case] child_precision: u8, + #[case] is_widen: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = DecimalArray::from_option_iter(values, DecimalDType::new(18, scale)); + let encoded = WidenArray::encode_decimal(input.clone(), &mut ctx)?; + assert_eq!(encoded.is::(), is_widen); + if let Some(widen) = encoded.as_opt::() { + assert_eq!(stored_decimal_dtype(&encoded).precision(), child_precision); + let child = widen.values().as_::(); + assert_eq!( + child.values_type(), + child.decimal_dtype().canonical_values_type() + ); + } else { + assert!(ArrayRef::ptr_eq(&encoded, &input.clone().into_array())); + } + assert_arrays_eq!(encoded, input, &mut ctx); + + Ok(()) +} + +#[test] +fn test_encode_decimal_ignores_null_payloads() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = DecimalArray::new( + buffer![i64::MIN, -128, i64::MAX], + DecimalDType::new(18, 0), + Validity::from_iter([false, true, false]), + ); + let encoded = WidenArray::encode_decimal(input.clone(), &mut ctx)?; + assert_eq!( + encoded.as_::().values().dtype(), + &dtype(3, 0, Nullable) + ); + assert_arrays_eq!(encoded, input, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::empty_positive_scale(None, 18, 12, Some(12))] +#[case::empty_negative_scale(None, 18, -4, Some(1))] +#[case::negative_boundary(Some((-128, 0)), 18, 0, Some(3))] +#[case::decimal_boundary(Some((-999, 999)), 18, 0, Some(3))] +#[case::outside_parent(Some((-1_000, 1)), 3, 0, None)] +#[case::reversed(Some((2, 1)), 18, 0, None)] +fn test_smallest_decimal_dtype( + #[case] range: Option<(i64, i64)>, + #[case] parent_precision: u8, + #[case] scale: i8, + #[case] precision: Option, +) -> VortexResult<()> { + let range = range.map(|(min, max)| i256::from_i128(min.into())..=i256::from_i128(max.into())); + let result = smallest_decimal_dtype(range, DecimalDType::new(parent_precision, scale)); + match precision { + Some(precision) => assert_eq!(result?, DecimalDType::new(precision, scale)), + None => assert!(result.is_err()), + } + + Ok(()) +} + +#[test] +fn test_smallest_decimal_dtype_i256_extrema() -> VortexResult<()> { + let parent = DecimalDType::new(76, 0); + let range = Some(i256::MIN_BY_PRECISION[76]..=i256::MAX_BY_PRECISION[76]); + assert_eq!(smallest_decimal_dtype(range, parent)?, parent); + assert!(smallest_decimal_dtype(Some(i256::MIN..=i256::MAX), parent).is_err()); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs b/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs new file mode 100644 index 00000000000..d6e60518275 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs @@ -0,0 +1,248 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal aggregate equivalence, including the precision-dependent sum fallback. +//! +//! Each batch uses fresh canonical arrays so cached statistics cannot hide the kernel path. + +use rstest::rstest; +use vortex_error::VortexResult; + +use super::SESSION; +use super::aggregates::assert_aggregate_batches; +use super::aggregates::kernel_partial; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::first::First; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; +use crate::aggregate_fn::fns::last::Last; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::fns::min_max::MinMax; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::dtype::BigCast; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; + +fn aggregates() -> Vec<(AggregateFnRef, bool)> { + let mut aggregates = vec![ + (IsConstant.bind(EmptyOptions), true), + (IsSorted.bind(IsSortedOptions { strict: false }), true), + (IsSorted.bind(IsSortedOptions { strict: true }), true), + (First.bind(EmptyOptions), false), + (Last.bind(EmptyOptions), false), + ]; + + for skip_nans in [true, false] { + let options = NumericalAggregateOpts { skip_nans }; + aggregates.extend([ + (MinMax.bind(options), true), + (Min.bind(options), true), + (Max.bind(options), true), + (Sum.bind(options), false), + (SumV2.bind(options), false), + ]); + } + + aggregates +} + +fn canonical_batch(values: &[Option], dtype: &DType) -> DecimalArray { + let decimal = *dtype + .as_decimal_opt() + .expect("decimal fixture has a decimal dtype"); + + match_each_decimal_value_type!(decimal.canonical_values_type(), |T| { + let values = values.iter().map(|value| { + value.map(|value| { + ::from(value).expect("fixture values fit the parent canonical width") + }) + }); + match dtype.nullability() { + Nullable => DecimalArray::from_option_iter(values, decimal), + NonNullable => DecimalArray::from_iter( + values.map(|value| value.expect("non-nullable fixture contains only valid values")), + decimal, + ), + } + }) +} + +/// Every fixture must exercise the `Widen` path, including empty and all-null batches. +#[track_caller] +fn encode_batch(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let encoded = WidenArray::encode_decimal(array, ctx)?; + assert!(encoded.is::()); + + Ok(encoded) +} + +#[rstest] +#[case::extrema( + vec![ + vec![Some(-99), None, Some(99)], // + vec![Some(-9999), Some(9999)], // + vec![Some(-999_999_999), Some(999_999_999)], // + vec![Some(-999_999_999_999_999_999), Some(999_999_999_999_999_999)], // + vec![Some(-(10i128.pow(38) - 1)), Some(10i128.pow(38) - 1)], // + ], + Nullable, +)] +#[case::ascending( + vec![vec![Some(-99), Some(-1)], vec![Some(0), Some(99)]], + NonNullable, +)] +#[case::descending_boundary( + vec![vec![Some(0), Some(99)], vec![Some(-99), Some(-1)]], + Nullable, +)] +#[case::unsorted_batch(vec![vec![Some(99), Some(-99)]], NonNullable)] +#[case::equal_boundary(vec![vec![Some(7)], vec![Some(7)]], Nullable)] +#[case::constant(vec![vec![Some(7), Some(7)], vec![Some(7)]], NonNullable)] +#[case::changed_constant(vec![vec![Some(7), Some(7)], vec![Some(8)]], Nullable)] +#[case::constant_with_null(vec![vec![Some(7), None, Some(7)]], Nullable)] +#[case::empty(vec![vec![], vec![]], Nullable)] +#[case::all_null(vec![vec![None], vec![None, None]], Nullable)] +#[case::single_null(vec![vec![None]], Nullable)] +#[case::single_null_then_sorted(vec![vec![None], vec![Some(1), Some(2)]], Nullable)] +#[case::empty_then_constant(vec![vec![], vec![Some(7), Some(7)]], Nullable)] +#[case::leading_and_trailing_nulls( + vec![vec![None, Some(-7)], vec![Some(7), None]], + Nullable, +)] +fn test_decimal_aggregate_partials( + #[case] batches: Vec>>, + #[case] nullability: Nullability, + #[values(-3, 0, 2)] scale: i8, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(76, scale), nullability); + + for (aggregate, served) in aggregates() { + assert_aggregate_batches( + &aggregate, + served, + &batches, + &dtype, + |values, dtype, _ctx| Ok(canonical_batch(values, dtype)), + encode_batch, + &mut ctx, + )?; + } + + Ok(()) +} + +#[rstest] +#[case::sum_skip_nans(Sum.bind(NumericalAggregateOpts { skip_nans: true }))] +#[case::sum_include_nans(Sum.bind(NumericalAggregateOpts { skip_nans: false }))] +#[case::sum_v2_skip_nans(SumV2.bind(NumericalAggregateOpts { skip_nans: true }))] +#[case::sum_v2_include_nans(SumV2.bind(NumericalAggregateOpts { skip_nans: false }))] +fn test_decimal_sum_uses_logical_precision(#[case] aggregate: AggregateFnRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let child_decimal = DecimalDType::new(2, 0); + let parent_decimal = DecimalDType::new(18, 0); + let parent_dtype = DType::Decimal(parent_decimal, NonNullable); + let child_scalar = Scalar::decimal(DecimalValue::I8(99), child_decimal, NonNullable); + let parent_scalar = Scalar::decimal(DecimalValue::I64(99), parent_decimal, NonNullable); + let len = 1usize << 34; + let huge_child = ConstantArray::new(child_scalar, len).into_array(); + let huge_widen = WidenArray::try_new(huge_child.clone(), parent_dtype.clone())?.into_array(); + + // Calling the registered kernel directly must decline without executing this huge array. + let kernel = ctx + .aggregate_kernels() + .find_aggregate_kernel(Widen.id(), aggregate.id()); + if aggregate.is::() { + let kernel = kernel.expect("the session registers the Widen Sum kernel"); + assert!( + kernel + .aggregate(&aggregate, &huge_widen, &mut ctx)? + .is_none() + ); + } else { + assert!(kernel.is_none(), "SumV2 has no Widen kernel"); + } + + let mut parent = aggregate.accumulator(&parent_dtype)?; + parent.accumulate( + &ConstantArray::new(parent_scalar, len).into_array(), + &mut ctx, + )?; + let expected = Scalar::decimal( + DecimalValue::I128(1_700_807_049_216), + DecimalDType::new(28, 0), + Nullable, + ); + assert_eq!(parent.finish()?, expected); + + let mut child = aggregate.accumulator(huge_child.dtype())?; + child.accumulate(&huge_child, &mut ctx)?; + assert_eq!( + child.finish()?, + Scalar::null(DType::Decimal(DecimalDType::new(12, 0), Nullable)) + ); + + let batches = vec![ + vec![], // + vec![Some(99); 256], // + vec![None, Some(-99)], // + ]; + assert_aggregate_batches( + &aggregate, + false, + &batches, + &parent_dtype.as_nullable(), + |values, dtype, _ctx| Ok(canonical_batch(values, dtype)), + encode_batch, + &mut ctx, + )?; + + let values = DecimalArray::from_iter([99i8], child_decimal); + let widen = WidenArray::try_new(values.into_array(), parent_dtype.clone())?.into_array(); + assert!(kernel_partial(&aggregate, &widen, &mut ctx)?.is_none()); + let canonical = DecimalArray::from_iter([99i64], parent_decimal).into_array(); + let mut actual = aggregate.accumulator(&parent_dtype)?; + let mut reference = aggregate.accumulator(&parent_dtype)?; + actual.accumulate(&widen, &mut ctx)?; + reference.accumulate(&canonical, &mut ctx)?; + assert_eq!(actual.partial_scalar()?, reference.partial_scalar()?); + + // Doubling real batch partials reaches the same logical length without allocating its rows. + for doubling in 1..=34 { + let actual_partial = actual.partial_scalar()?; + let reference_partial = reference.partial_scalar()?; + actual.combine_partials(&actual_partial)?; + reference.combine_partials(&reference_partial)?; + assert_eq!( + actual.partial_scalar()?, + reference.partial_scalar()?, + "{aggregate} after {doubling} doublings" + ); + } + + assert_eq!(actual.finish()?, expected); + assert_eq!(reference.finish()?, expected); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/decimal_cast.rs b/vortex-array/src/arrays/widen/tests/decimal_cast.rs new file mode 100644 index 00000000000..dc228aea727 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_cast.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal casts preserve lazy widening, checked precision and nullability. +//! +//! These cases also cover decimal buffers whose stored width differs from their canonical width. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use super::canonical; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::DictArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +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::dtype::PType; +use crate::dtype::i256; +use crate::scalar_fn::fns::cast::CastKernel; +use crate::scalar_fn::fns::cast::CastReduce; +use crate::validity::Validity; + +#[rstest] +#[case::i8(DecimalArray::from_iter([12i8, -12], DecimalDType::new(2, 0)), 4)] +#[case::i16(DecimalArray::from_iter([12i16, -12], DecimalDType::new(4, 0)), 9)] +#[case::i32(DecimalArray::from_iter([12i32, -12], DecimalDType::new(9, 0)), 18)] +#[case::i64(DecimalArray::from_iter([12i64, -12], DecimalDType::new(18, 0)), 38)] +#[case::i128(DecimalArray::from_iter([12i128, -12], DecimalDType::new(38, 0)), 76)] +#[case::stored_target_width(DecimalArray::from_iter([12i16, -12], DecimalDType::new(2, 0)), 4)] +#[case::stored_wider_than_target(DecimalArray::from_iter([12i128, -12], DecimalDType::new(2, 0)), 4)] +#[case::stored_i256(DecimalArray::from_iter([i256::from_i128(12), i256::from_i128(-12)], DecimalDType::new(2, 0)), 4)] +fn test_precision_growth_is_lazy( + #[case] array: DecimalArray, + #[case] precision: u8, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source_dtype = array.dtype().clone(); + let source_values_type = array.values_type(); + let decimal_dtype = DecimalDType::new(precision, 0); + let target = DType::Decimal(decimal_dtype, Nullability::NonNullable); + + let casted = array.into_array().cast(target.clone())?; + + assert!(casted.is::()); + let widen = casted.as_::(); + let values = widen.values(); + assert_eq!(values.dtype(), &source_dtype); + assert_eq!(values.as_::().values_type(), source_values_type); + let executed = casted.execute::(&mut ctx)?; + assert_eq!(executed.dtype(), &target); + assert_eq!( + executed.values_type(), + decimal_dtype.canonical_values_type() + ); + let expected = DecimalArray::from_iter([12i128, -12], decimal_dtype); + assert_arrays_eq!(executed, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::to_nullable(Validity::NonNullable, Nullability::Nullable)] +#[case::to_non_nullable(Validity::AllValid, Nullability::NonNullable)] +#[case::keeps_nullable(Validity::AllValid, Nullability::Nullable)] +fn test_precision_growth_sets_child_nullability( + #[case] validity: Validity, + #[case] nullability: Nullability, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::new(buffer![12i8, -12], DecimalDType::new(2, 0), validity); + let target_decimal = DecimalDType::new(4, 0); + let target = DType::Decimal(target_decimal, nullability); + + let casted = array.into_array().cast(target.clone())?; + + assert!(casted.is::()); + assert_eq!( + casted.as_::().values().dtype().nullability(), + nullability + ); + assert_eq!(casted.dtype(), &target); + let expected_validity = match nullability { + Nullability::Nullable => Validity::AllValid, + Nullability::NonNullable => Validity::NonNullable, + }; + let expected = DecimalArray::new(buffer![12i16, -12], target_decimal, expected_validity); + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::all_valid(true)] +#[case::null_row(false)] +fn test_encoded_validity_falls_back_to_kernel(#[case] second_valid: bool) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = DictArray::try_new( + buffer![0u8, 1].into_array(), + BoolArray::from_iter([true, second_valid]).into_array(), + )?; + let array = DecimalArray::new( + buffer![12i8, -12], + DecimalDType::new(2, 0), + Validity::Array(validity.into_array()), + ); + let target_decimal = DecimalDType::new(4, 0); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + + assert!(::cast(array.as_view(), &target)?.is_none()); + let result = ::cast(array.as_view(), &target, &mut ctx); + if !second_valid { + let error = result.expect_err("the encoded validity contains a null row"); + assert!(error.to_string().contains("invalid values to non-nullable")); + return Ok(()); + } + + let casted = result?.expect("same-scale precision growth has a cast kernel"); + assert!(casted.is::()); + assert_eq!( + casted.as_::().values().dtype().nullability(), + Nullability::NonNullable + ); + let expected = DecimalArray::from_iter([12i16, -12], target_decimal); + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_same_canonical_width_reuses_stored_buffer() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::from_iter([12i128, -12], DecimalDType::new(10, 0)); + let source_ptr = array.buffer::().as_ptr(); + let target_decimal = DecimalDType::new(18, 0); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + + let casted = array + .into_array() + .cast(target)? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ptr(), source_ptr); + assert_arrays_eq!( + casted, + DecimalArray::from_iter([12i128, -12], target_decimal), + &mut ctx + ); + + Ok(()) +} + +#[rstest] +#[case::smaller_decimal(DType::Decimal(DecimalDType::new(3, 2), Nullability::Nullable))] +#[case::wider_decimal(DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable))] +#[case::rescale(DType::Decimal(DecimalDType::new(18, 3), Nullability::Nullable))] +#[case::integer(DType::Primitive(PType::I64, Nullability::Nullable))] +#[case::float(DType::Primitive(PType::F64, Nullability::Nullable))] +fn test_decimal_widen_cast_matches_canonical(#[case] target: DType) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = + DecimalArray::from_option_iter([Some(125i16), None, Some(-250)], DecimalDType::new(4, 2)); + let array = WidenArray::try_new( + values.into_array(), + DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable), + )? + .into_array(); + let expected = canonical(&array, &mut ctx)?.cast(target.clone())?; + + let casted = array.cast(target)?; + + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_precision_narrowing_checks_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::from_iter([1000i16], DecimalDType::new(4, 0)); + let array = WidenArray::try_new( + values.into_array(), + DType::Decimal(DecimalDType::new(18, 0), Nullability::NonNullable), + )? + .into_array(); + let target = DType::Decimal(DecimalDType::new(3, 0), Nullability::NonNullable); + + let error = array + .cast(target)? + .execute::(&mut ctx) + .expect_err("1000 does not fit a precision of three digits"); + + assert!(error.to_string().contains("does not fit in precision")); + + Ok(()) +} + +#[test] +fn test_repeated_casts_terminate_without_nested_widen() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source = DecimalArray::from_iter([12i8, -12], DecimalDType::new(2, 0)).into_array(); + let mut array = source.clone(); + for precision in [4, 9, 18, 38, 76, 4, 2] { + let decimal_dtype = DecimalDType::new(precision, 0); + let target = DType::Decimal(decimal_dtype, Nullability::NonNullable); + array = array.cast(target.clone())?; + assert!(ArrayRef::ptr_eq(&array.cast(target)?, &array)); + if let Some(widen) = array.as_opt::() { + assert!(!widen.values().is::()); + assert_eq!(widen.values().dtype(), source.dtype()); + } + assert_arrays_eq!( + array.clone(), + DecimalArray::from_iter([12i128, -12], decimal_dtype), + &mut ctx + ); + } + + assert_arrays_eq!(array, source, &mut ctx); + + Ok(()) +} + +#[test] +fn test_same_canonical_width_copies_narrow_stored_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::from_iter([125i32, -250], DecimalDType::new(10, 2)); + let target_decimal = DecimalDType::new(18, 2); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + let casted = array.into_array().cast(target)?; + assert!(!casted.is::()); + + let executed = casted.execute::(&mut ctx)?; + assert_eq!( + executed.values_type(), + target_decimal.canonical_values_type() + ); + assert_arrays_eq!( + executed, + DecimalArray::from_iter([125i64, -250], target_decimal), + &mut ctx + ); + + Ok(()) +} 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..56cc208d4c9 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_compare.rs @@ -0,0 +1,312 @@ +// 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 crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::arrays::widen::compute::constant_fits_values; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +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::CompareOperator; +use crate::scalar_fn::fns::operators::Operator; + +/// Checks both operand orders against fully decoded decimal operands. +#[track_caller] +fn assert_compare_matches_canonical( + lhs: &ArrayRef, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical_lhs = lhs.clone().execute::(ctx)?.into_array(); + let canonical_rhs = rhs.clone().execute::(ctx)?.into_array(); + 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(()) +} + +fn widen(values: impl IntoArray) -> VortexResult { + let values = values.into_array(); + let dtype = DType::Decimal(DecimalDType::new(18, 2), values.dtype().nullability()); + + WidenArray::try_new(values, dtype) +} + +fn decimal_scalar(value: Option, nullability: Nullability) -> Scalar { + let dtype = DecimalDType::new(18, 2); + match value { + Some(value) => Scalar::decimal(DecimalValue::I64(value), dtype, nullability), + None => Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + } +} + +#[rstest] +#[case::below_precision(Some(-100_000))] +#[case::minimum(Some(-99_999))] +#[case::negative_inside(Some(-32_768))] +#[case::zero(Some(0))] +#[case::positive_inside(Some(32_768))] +#[case::maximum(Some(99_999))] +#[case::above_precision(Some(100_000))] +#[case::null(None)] +fn test_compare_constant( + #[case] value: Option, + #[values(false, true)] nullable: bool, + #[values( + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::Lte, + Operator::Gt, + Operator::Gte + )] + operator: Operator, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(5, 2); + let values = if nullable { + DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(32_768), Some(99_999)], + decimal_dtype, + ) + } else { + DecimalArray::from_iter([-99_999i32, -32_768, 0, 32_768, 99_999], decimal_dtype) + }; + let array = widen(values)?.into_array(); + let constant = + ConstantArray::new(decimal_scalar(value, Nullability::Nullable), array.len()).into_array(); + + assert_compare_matches_canonical(&array, &constant, operator, &mut ctx) +} + +/// In `different_widths`, the cast of the `i8` child to the common precision returns a `Widen` of +/// its own, so the compare runs through the fallback. +#[rstest] +#[case::same_width( + DecimalArray::from_iter([-99_999i32, -1, 0, 99_999], DecimalDType::new(5, 2)), + DecimalArray::from_option_iter( + [Some(-100_000i32), None, Some(0), Some(999_999_999)], + DecimalDType::new(9, 2), + ), +)] +#[case::different_widths( + DecimalArray::from_iter([-99i8, -1, 0, 99], DecimalDType::new(2, 2)), + DecimalArray::from_option_iter( + [Some(-100i32), None, Some(0), Some(999_999_999)], + 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 lhs = widen(lhs)?.into_array(); + let rhs = widen(rhs)?.into_array(); + + assert_compare_matches_canonical(&lhs, &rhs, operator, &mut ctx) +} + +#[test] +fn test_compare_uses_child_precision_and_nullability() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = widen(DecimalArray::from_iter( + [-99_999i32, 0, 99_999], + DecimalDType::new(5, 2), + ))?; + let rhs = widen(DecimalArray::from_option_iter( + [Some(-100_000i32), None, Some(999_999_999)], + DecimalDType::new(9, 2), + ))? + .into_array(); + + let compared = Widen::compare(lhs.as_view(), &rhs, CompareOperator::Lt, &mut ctx)? + .vortex_expect("Widen compares decimal children at their common precision"); + let lhs_values = compared.nth_child(0).vortex_expect("compare has a lhs"); + let rhs_values = compared.nth_child(1).vortex_expect("compare has a rhs"); + assert_eq!( + lhs_values.dtype(), + &DType::Decimal(DecimalDType::new(9, 2), Nullability::NonNullable) + ); + assert_eq!( + rhs_values.dtype(), + &DType::Decimal(DecimalDType::new(9, 2), Nullability::Nullable) + ); + + Ok(()) +} + +#[test] +fn test_constant_fits_precision_and_scale() { + let dtype = DType::Decimal(DecimalDType::new(75, 2), Nullability::NonNullable); + let max = i256::MAX_BY_PRECISION[75]; + for (value, fits) in [ + (max, true), + (-max, true), + (max + i256::ONE, false), + (-max - i256::ONE, false), + ] { + let scalar = Scalar::decimal( + DecimalValue::I256(value), + DecimalDType::new(76, 2), + Nullability::NonNullable, + ); + assert_eq!(constant_fits_values(&scalar, &dtype), Some(fits)); + } + + let different_scale = Scalar::decimal( + DecimalValue::I8(1), + DecimalDType::new(18, 3), + Nullability::NonNullable, + ); + assert_eq!(constant_fits_values(&different_scale, &dtype), None); + assert_eq!( + constant_fits_values(&Scalar::null(dtype.as_nullable()), &dtype), + None + ); + assert_eq!(constant_fits_values(&Scalar::from(1i64), &dtype), None); +} + +#[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, + #[values(StrictComparison::Strict, StrictComparison::NonStrict)] lower_strict: StrictComparison, + #[values(StrictComparison::Strict, StrictComparison::NonStrict)] upper_strict: StrictComparison, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(99_999)], + DecimalDType::new(5, 2), + ))? + .into_array(); + let lower = + ConstantArray::new(decimal_scalar(lower, Nullability::Nullable), array.len()).into_array(); + let upper = + ConstantArray::new(decimal_scalar(upper, Nullability::Nullable), array.len()).into_array(); + let options = BetweenOptions { + lower_strict, + upper_strict, + }; + + let result = array + .clone() + .between(lower.clone(), upper.clone(), options.clone())?; + let expected = array + .execute::(&mut ctx)? + .into_array() + .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 array = widen(DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(99_999)], + DecimalDType::new(5, 2), + ))?; + let fill = decimal_scalar(Some(value), 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 = array + .execute::(&mut ctx)? + .into_array() + .fill_null(fill)?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_fill_null_rejects_null_scalar() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(DecimalArray::from_option_iter( + [Some(1i32), None], + DecimalDType::new(5, 2), + ))?; + let fill = decimal_scalar(None, Nullability::Nullable); + + assert!(Widen::fill_null(array.as_view(), &fill)?.is_none()); + let error = array + .into_array() + .fill_null(fill) + .and_then(|result| result.execute::(&mut ctx)) + .expect_err("fill_null requires a non-null scalar"); + assert!( + error + .to_string() + .contains("fill_null requires a non-null fill value") + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/decimal_fallback.rs b/vortex-array/src/arrays/widen/tests/decimal_fallback.rs new file mode 100644 index 00000000000..ae9ff486248 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_fallback.rs @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Conformance and canonical fallbacks for decimal widening. +//! +//! Arithmetic, zip, chunked and interleave execution preserve the logical decimal dtype and nulls. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use super::canonical; +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::compute::conformance::binary_numeric::test_binary_numeric_array; +use crate::compute::conformance::cast::test_cast_conformance; +use crate::compute::conformance::consistency::test_array_consistency; +use crate::compute::conformance::filter::test_filter_conformance; +use crate::compute::conformance::mask::test_mask_conformance; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +fn decimal_widen(validity: Validity) -> VortexResult { + let values = DecimalArray::new( + buffer![-125i16, 50, 50, 250, 900], + DecimalDType::new(4, 2), + validity, + ); + let dtype = DType::Decimal(DecimalDType::new(18, 2), values.dtype().nullability()); + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_conformance(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + assert!(array.is::()); + + test_array_consistency(&array, &mut ctx); + test_cast_conformance(&array, &mut ctx); + test_filter_conformance(&array, &mut ctx); + test_take_conformance(&array, &mut ctx); + test_mask_conformance(&array, &mut ctx); + test_binary_numeric_array(&array, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_arithmetic_falls_back_to_canonical( + #[case] validity: Validity, + #[values(Operator::Add, Operator::Sub, Operator::Mul, Operator::Div)] operator: Operator, + #[values(false, true)] widen_on_right: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + let other = DecimalArray::from_iter([25i64, 50, -100, 125, -200], DecimalDType::new(18, 2)) + .into_array(); + let (lhs, rhs) = if widen_on_right { + (other, array) + } else { + (array, other) + }; + let expected = canonical(&lhs, &mut ctx)? + .binary(canonical(&rhs, &mut ctx)?, operator)? + .execute::(&mut ctx)?; + + let actual = lhs.binary(rhs, operator)?.execute::(&mut ctx)?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::true_branch(true, false)] +#[case::false_branch(false, true)] +#[case::both_branches(true, true)] +fn test_decimal_zip_falls_back_to_canonical( + #[case] widen_if_true: bool, + #[case] widen_if_false: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let if_true = decimal_widen(Validity::from_iter([true, false, true, true, false]))?; + let if_false = WidenArray::try_new( + DecimalArray::from_iter([500i16, -600, 700, -800, 900], DecimalDType::new(4, 2)) + .into_array(), + DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable), + )? + .into_array(); + let canonical_true = canonical(&if_true, &mut ctx)?; + let canonical_false = canonical(&if_false, &mut ctx)?; + let mask = + BoolArray::from_iter([Some(true), Some(true), None, Some(false), Some(true)]).into_array(); + let expected = mask + .zip(canonical_true.clone(), canonical_false.clone())? + .execute::(&mut ctx)?; + let if_true = if widen_if_true { + if_true + } else { + canonical_true + }; + let if_false = if widen_if_false { + if_false + } else { + canonical_false + }; + + let actual = mask + .zip(if_true, if_false)? + .execute::(&mut ctx)?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_chunked_execution_matches_canonical( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + let expected = canonical(&array, &mut ctx)?; + let chunks = ChunkedArray::from_iter([ + array.slice(0..2)?, + expected.slice(2..3)?, + array.slice(3..5)?, + ]) + .into_array(); + + let actual = chunks.execute::(&mut ctx)?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_interleave_falls_back_to_canonical(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + let other = DecimalArray::from_iter([500i64, -600, 700, -800, 950], DecimalDType::new(18, 2)) + .into_array(); + let array_indices = buffer![0u8, 1, 0, 1, 0, 0, 1].into_array(); + let row_indices = buffer![4u32, 0, 1, 4, 0, 4, 1].into_array(); + let expected = array_indices + .interleave( + [canonical(&array, &mut ctx)?, other.clone()], + row_indices.clone(), + )? + .execute::(&mut ctx)?; + + let actual = array_indices + .interleave([array, other], row_indices)? + .execute::(&mut ctx)?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, 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 6a959e199d3..7ad3a1ef5cf 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -8,6 +8,11 @@ mod aggregates; mod compute; +mod decimal; +mod decimal_aggregates; +mod decimal_cast; +mod decimal_compare; +mod decimal_fallback; mod encoding; mod numeric; @@ -27,6 +32,7 @@ use crate::arrays::Widen; use crate::arrays::WidenArray; use crate::arrays::widen::WidenArraySlotsExt; use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::dtype::PType; use crate::validity::Validity; @@ -48,6 +54,16 @@ fn stored_ptype(array: &ArrayRef) -> PType { } } +/// Returns the decimal dtype of the values child, or of `array` when it is not a `Widen`. +fn stored_decimal_dtype(array: &ArrayRef) -> DecimalDType { + let dtype = match array.as_opt::() { + Some(widen) => widen.values().dtype().as_decimal_opt().copied(), + None => array.dtype().as_decimal_opt().copied(), + }; + + dtype.vortex_expect("the fixture has a decimal dtype") +} + /// Returns `array` executed to its canonical form. fn canonical(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(array.clone().execute::(ctx)?.into_array()) diff --git a/vortex-array/src/arrays/widen/vtable/mod.rs b/vortex-array/src/arrays/widen/vtable/mod.rs index 3af2178b523..1ee74c6b931 100644 --- a/vortex-array/src/arrays/widen/vtable/mod.rs +++ b/vortex-array/src/arrays/widen/vtable/mod.rs @@ -24,6 +24,9 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; +use crate::arrays::DecimalArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::WidenSlots; use crate::arrays::widen::array::validate_dtypes; @@ -31,13 +34,15 @@ use crate::arrays::widen::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::match_each_decimal_value_type; use crate::serde::ArrayChildren; /// The encoding of a [`WidenArray`]. #[derive(Clone, Debug)] pub struct Widen; -/// A [`Widen`]-encoded integer array. +/// A [`Widen`]-encoded integer or decimal array. pub type WidenArray = Array; impl VTable for Widen { @@ -112,7 +117,14 @@ impl VTable for Widen { WidenSlots::NAMES[idx].to_string() } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + // NB: Decimals copy their values instead of yielding a cast. A same-scale decimal cast that + // grows the canonical width returns a `Widen`, so a yielded cast would execute to itself. + if let DType::Decimal(decimal_dtype, _) = array.dtype() { + let widened = execute_decimal(&array, *decimal_dtype, ctx)?; + return Ok(ExecutionResult::done(widened.into_array())); + } + // Yield a cast, not a canonical array, so that one kernel decodes and widens the values. let widened = array.values().cast(array.dtype().clone())?; Ok(ExecutionResult::done(widened)) @@ -136,3 +148,30 @@ impl VTable for Widen { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Executes the values of `array` and copies them into the canonical width of `decimal_dtype`, the +/// logical dtype of `array`. +fn execute_decimal( + array: &WidenArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let values = array.values().clone().execute::(ctx)?; + let validity = values.validity()?; + let mask = validity.execute_mask(values.len(), ctx)?; + let buffer = match_each_decimal_value_type!(values.values_type(), |T| { + copy_to_canonical_width(values.buffer::().as_slice(), &mask, decimal_dtype)? + }); + + // SAFETY: the copy has the parent canonical width and the child length. Widen construction + // guarantees the same scale and a parent precision that holds the child, and the unchanged + // child validity has the same length. + Ok(unsafe { + DecimalArray::new_unchecked_handle( + buffer, + decimal_dtype.canonical_values_type(), + decimal_dtype, + validity, + ) + }) +} diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index ce4d7e54ba5..7e1becc2a30 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -35,6 +35,8 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::WidenArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; @@ -258,6 +260,37 @@ async fn test_read_simple() { assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); } +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_decimal_widen_default_writer_roundtrip( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(18, 2); + let values = DecimalArray::new( + buffer![-125i16, 50, 0, 250, 900], + DecimalDType::new(4, 2), + validity.clone(), + ); + let dtype = DType::Decimal(decimal_dtype, validity.nullability()); + let widened = WidenArray::try_new(values.into_array(), dtype)?.into_array(); + assert!(widened.is::()); + let array = StructArray::from_fields(&[("decimal", widened)])?.into_array(); + let expected_values = + DecimalArray::new(buffer![-125i64, 50, 0, 250, 900], decimal_dtype, validity); + let expected = StructArray::from_fields(&[("decimal", expected_values.into_array())])?; + + let actual = write_read_roundtrip(array).await?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + /// Wide decimals split into multi-part arrays only when the writer allows the v2 format. The /// default writer allows every registered encoding once editions are disabled; a strategy built /// from the session keeps the enabled editions' restrictions regardless.