|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! `Widen` operations over a `BitPacked` child. |
| 5 | +//! |
| 6 | +//! Selection and comparison keep the stored width of the child, aggregates read the child, and |
| 7 | +//! canonical execution widens the values to the logical dtype of the `Widen`. |
| 8 | +
|
| 9 | +#![cfg(test)] |
| 10 | + |
| 11 | +use std::sync::LazyLock; |
| 12 | + |
| 13 | +use rstest::rstest; |
| 14 | +use vortex_array::ArrayRef; |
| 15 | +use vortex_array::IntoArray; |
| 16 | +use vortex_array::VortexSessionExecute; |
| 17 | +use vortex_array::aggregate_fn::NumericalAggregateOpts; |
| 18 | +use vortex_array::aggregate_fn::fns::is_constant::is_constant; |
| 19 | +use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; |
| 20 | +use vortex_array::aggregate_fn::fns::is_sorted::is_strict_sorted; |
| 21 | +use vortex_array::aggregate_fn::fns::min_max::min_max; |
| 22 | +use vortex_array::aggregate_fn::fns::sum::sum; |
| 23 | +use vortex_array::arrays::BoolArray; |
| 24 | +use vortex_array::arrays::ConstantArray; |
| 25 | +use vortex_array::arrays::PrimitiveArray; |
| 26 | +use vortex_array::arrays::Widen; |
| 27 | +use vortex_array::arrays::WidenArray; |
| 28 | +use vortex_array::arrays::widen::WidenArraySlotsExt; |
| 29 | +use vortex_array::assert_arrays_eq; |
| 30 | +use vortex_array::builtins::ArrayBuiltins; |
| 31 | +use vortex_array::dtype::DType; |
| 32 | +use vortex_array::dtype::PType; |
| 33 | +use vortex_array::scalar::Scalar; |
| 34 | +use vortex_array::scalar_fn::fns::operators::Operator; |
| 35 | +use vortex_buffer::buffer; |
| 36 | +use vortex_error::VortexResult; |
| 37 | +use vortex_fastlanes::BitPacked; |
| 38 | +use vortex_mask::Mask; |
| 39 | +use vortex_session::VortexSession; |
| 40 | + |
| 41 | +static SESSION: LazyLock<VortexSession> = LazyLock::new(|| { |
| 42 | + let session = VortexSession::empty(); |
| 43 | + vortex_fastlanes::initialize(&session); |
| 44 | + session |
| 45 | +}); |
| 46 | + |
| 47 | +/// Returns a `u64` `Widen` over `values` packed at `bit_width` bits. |
| 48 | +fn widen_bitpacked(values: PrimitiveArray, bit_width: u8) -> VortexResult<ArrayRef> { |
| 49 | + let mut ctx = SESSION.create_execution_ctx(); |
| 50 | + let packed = BitPacked::encode(&values.into_array(), bit_width, &mut ctx)?.into_array(); |
| 51 | + let dtype = DType::Primitive(PType::U64, packed.dtype().nullability()); |
| 52 | + |
| 53 | + Ok(WidenArray::try_new(packed, dtype)?.into_array()) |
| 54 | +} |
| 55 | + |
| 56 | +#[test] |
| 57 | +fn test_operations() -> VortexResult<()> { |
| 58 | + let mut ctx = SESSION.create_execution_ctx(); |
| 59 | + let array = widen_bitpacked(PrimitiveArray::from_iter([0u8, 1, 3, 7]), 3)?; |
| 60 | + |
| 61 | + let selected = array.filter(Mask::from_iter([true, false, true, true]))?; |
| 62 | + assert!(selected.is::<Widen>()); |
| 63 | + assert_eq!(selected.as_::<Widen>().values().dtype(), &PType::U8.into()); |
| 64 | + assert_arrays_eq!(selected, buffer![0u64, 3, 7].into_array(), &mut ctx); |
| 65 | + |
| 66 | + assert_arrays_eq!( |
| 67 | + array.take(buffer![3u32, 1].into_array())?, |
| 68 | + buffer![7u64, 1].into_array(), |
| 69 | + &mut ctx |
| 70 | + ); |
| 71 | + assert_arrays_eq!( |
| 72 | + array.binary( |
| 73 | + ConstantArray::new(Scalar::from(3u64), 4).into_array(), |
| 74 | + Operator::Lt |
| 75 | + )?, |
| 76 | + BoolArray::from_iter([true, true, false, false]), |
| 77 | + &mut ctx |
| 78 | + ); |
| 79 | + assert_eq!( |
| 80 | + array.execute::<PrimitiveArray>(&mut ctx)?.as_slice::<u64>(), |
| 81 | + &[0, 1, 3, 7] |
| 82 | + ); |
| 83 | + |
| 84 | + Ok(()) |
| 85 | +} |
| 86 | + |
| 87 | +#[rstest] |
| 88 | +#[case::sorted(PrimitiveArray::from_iter([0u8, 1, 3, 7]))] |
| 89 | +#[case::constant(PrimitiveArray::from_iter([5u8, 5, 5, 5]))] |
| 90 | +#[case::partially_valid(PrimitiveArray::from_option_iter([None, Some(6u8), Some(2), None]))] |
| 91 | +fn test_aggregates_match_canonical(#[case] values: PrimitiveArray) -> VortexResult<()> { |
| 92 | + let mut ctx = SESSION.create_execution_ctx(); |
| 93 | + let canonical = widen_bitpacked(values.clone(), 3)? |
| 94 | + .execute::<PrimitiveArray>(&mut ctx)? |
| 95 | + .into_array(); |
| 96 | + let options = NumericalAggregateOpts::default(); |
| 97 | + |
| 98 | + // Each aggregate gets a fresh `Widen`. A statistic cached by an earlier aggregate, such as the |
| 99 | + // min and max that `min_max` caches, would otherwise answer a later one before its kernel runs. |
| 100 | + let widen = || widen_bitpacked(values.clone(), 3); |
| 101 | + |
| 102 | + assert_eq!( |
| 103 | + min_max(&widen()?, &mut ctx, options)?, |
| 104 | + min_max(&canonical, &mut ctx, options)? |
| 105 | + ); |
| 106 | + assert_eq!( |
| 107 | + is_constant(&widen()?, &mut ctx)?, |
| 108 | + is_constant(&canonical, &mut ctx)? |
| 109 | + ); |
| 110 | + assert_eq!( |
| 111 | + is_sorted(&widen()?, &mut ctx)?, |
| 112 | + is_sorted(&canonical, &mut ctx)? |
| 113 | + ); |
| 114 | + assert_eq!( |
| 115 | + is_strict_sorted(&widen()?, &mut ctx)?, |
| 116 | + is_strict_sorted(&canonical, &mut ctx)? |
| 117 | + ); |
| 118 | + assert_eq!(sum(&widen()?, &mut ctx)?, sum(&canonical, &mut ctx)?); |
| 119 | + |
| 120 | + Ok(()) |
| 121 | +} |
0 commit comments