|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +use rstest::rstest; |
| 5 | +use vortex_buffer::Buffer; |
| 6 | +use vortex_error::VortexExpect; |
| 7 | +use vortex_error::VortexResult; |
| 8 | +use vortex_session::VortexSession; |
| 9 | + |
| 10 | +use crate::ArrayRef; |
| 11 | +use crate::IntoArray; |
| 12 | +use crate::VortexSessionExecute; |
| 13 | +use crate::arrays::ConstantArray; |
| 14 | +use crate::arrays::DecimalArray; |
| 15 | +use crate::arrays::InterleaveArray; |
| 16 | +use crate::arrays::PrimitiveArray; |
| 17 | +use crate::arrays::SliceArray; |
| 18 | +use crate::arrays::decimal::DecimalArrayExt; |
| 19 | +use crate::assert_arrays_eq; |
| 20 | +use crate::dtype::BigCast; |
| 21 | +use crate::dtype::DType; |
| 22 | +use crate::dtype::DecimalDType; |
| 23 | +use crate::dtype::DecimalType; |
| 24 | +use crate::dtype::NativeDecimalType; |
| 25 | +use crate::dtype::Nullability; |
| 26 | +use crate::dtype::i256; |
| 27 | +use crate::match_each_decimal_value_type; |
| 28 | +use crate::scalar::DecimalValue; |
| 29 | +use crate::scalar::Scalar; |
| 30 | +use crate::validity::Validity; |
| 31 | + |
| 32 | +/// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle. |
| 33 | +fn check(values: Vec<ArrayRef>, indices: &[(u8, u64)], expected: DecimalArray) -> VortexResult<()> { |
| 34 | + let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); |
| 35 | + let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); |
| 36 | + let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; |
| 37 | + let mut ctx = VortexSession::empty().create_execution_ctx(); |
| 38 | + let output = array.into_array().execute::<DecimalArray>(&mut ctx)?; |
| 39 | + assert_eq!( |
| 40 | + output.values_type(), |
| 41 | + expected.decimal_dtype().canonical_values_type() |
| 42 | + ); |
| 43 | + assert_eq!(output.dtype(), expected.dtype()); |
| 44 | + assert_arrays_eq!(output.into_array(), expected.into_array(), &mut ctx); |
| 45 | + |
| 46 | + Ok(()) |
| 47 | +} |
| 48 | + |
| 49 | +#[rstest] |
| 50 | +#[case::i8(2)] |
| 51 | +#[case::i16(4)] |
| 52 | +#[case::i32(9)] |
| 53 | +#[case::i64(18)] |
| 54 | +#[case::i128(38)] |
| 55 | +#[case::i256(76)] |
| 56 | +fn test_decimal_interleave( |
| 57 | + #[case] precision: u8, |
| 58 | + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, |
| 59 | + #[values(-3, 0, 2)] scale: i8, |
| 60 | +) -> VortexResult<()> { |
| 61 | + let dtype = DecimalDType::new(precision, scale); |
| 62 | + let (first_validity, output_validity) = if nullability.is_nullable() { |
| 63 | + ( |
| 64 | + Validity::from_iter([true, false, true]), |
| 65 | + Validity::from_iter([true, true, false, true, true, true]), |
| 66 | + ) |
| 67 | + } else { |
| 68 | + (Validity::NonNullable, Validity::NonNullable) |
| 69 | + }; |
| 70 | + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { |
| 71 | + let min = T::MIN_BY_PRECISION[usize::from(precision)]; |
| 72 | + let max = T::MAX_BY_PRECISION[usize::from(precision)]; |
| 73 | + let zero = T::default(); |
| 74 | + let middle = if nullability.is_nullable() { |
| 75 | + T::MIN |
| 76 | + } else { |
| 77 | + zero |
| 78 | + }; |
| 79 | + let first = DecimalArray::new(Buffer::copy_from([min, middle, max]), dtype, first_validity); |
| 80 | + let second = DecimalArray::new( |
| 81 | + Buffer::copy_from([max, min, zero]), |
| 82 | + dtype, |
| 83 | + Validity::NonNullable, |
| 84 | + ); |
| 85 | + let expected = DecimalArray::new( |
| 86 | + Buffer::copy_from([zero, min, zero, max, max, min]), |
| 87 | + dtype, |
| 88 | + output_validity, |
| 89 | + ); |
| 90 | + check( |
| 91 | + vec![first.into_array(), second.into_array()], |
| 92 | + &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], |
| 93 | + expected, |
| 94 | + ) |
| 95 | + }) |
| 96 | +} |
| 97 | + |
| 98 | +#[rstest] |
| 99 | +#[case::i8(2)] |
| 100 | +#[case::i16(4)] |
| 101 | +#[case::i32(9)] |
| 102 | +#[case::i64(18)] |
| 103 | +#[case::i128(38)] |
| 104 | +#[case::i256(76)] |
| 105 | +fn test_decimal_interleave_constants( |
| 106 | + #[case] precision: u8, |
| 107 | + #[values(false, true)] null_constant: bool, |
| 108 | + #[values(DecimalType::I8, DecimalType::I256)] scalar_storage: DecimalType, |
| 109 | +) -> VortexResult<()> { |
| 110 | + let dtype = DecimalDType::new(precision, 2); |
| 111 | + let scalar = if null_constant { |
| 112 | + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) |
| 113 | + } else { |
| 114 | + let payload = match_each_decimal_value_type!(scalar_storage, |S| { |
| 115 | + DecimalValue::from( |
| 116 | + <S as BigCast>::from(-7i8).vortex_expect("-7 fits every decimal type"), |
| 117 | + ) |
| 118 | + }); |
| 119 | + Scalar::decimal(payload, dtype, Nullability::NonNullable) |
| 120 | + }; |
| 121 | + let constant = ConstantArray::new(scalar, 3).into_array(); |
| 122 | + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { |
| 123 | + let seven = <T as BigCast>::from(7i8).vortex_expect("7 fits every decimal type"); |
| 124 | + let minus_seven = <T as BigCast>::from(-7i8).vortex_expect("-7 fits every decimal type"); |
| 125 | + let column = DecimalArray::new(Buffer::copy_from([seven]), dtype, Validity::NonNullable) |
| 126 | + .into_array(); |
| 127 | + let expected = DecimalArray::new( |
| 128 | + Buffer::copy_from([minus_seven, seven, minus_seven, minus_seven]), |
| 129 | + dtype, |
| 130 | + if null_constant { |
| 131 | + Validity::from_iter([false, true, false, false]) |
| 132 | + } else { |
| 133 | + Validity::NonNullable |
| 134 | + }, |
| 135 | + ); |
| 136 | + check( |
| 137 | + vec![constant, column], |
| 138 | + &[(0, 2), (1, 0), (0, 1), (0, 2)], |
| 139 | + expected, |
| 140 | + ) |
| 141 | + }) |
| 142 | +} |
| 143 | + |
| 144 | +#[rstest] |
| 145 | +#[case::i8(2)] |
| 146 | +#[case::i16(4)] |
| 147 | +#[case::i32(9)] |
| 148 | +#[case::i64(18)] |
| 149 | +#[case::i128(38)] |
| 150 | +#[case::i256(76)] |
| 151 | +fn test_decimal_interleave_mixed_stored_widths(#[case] precision: u8) -> VortexResult<()> { |
| 152 | + let dtype = DecimalDType::new(precision, 0); |
| 153 | + let first = |
| 154 | + DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); |
| 155 | + let second = DecimalArray::new( |
| 156 | + Buffer::copy_from([i256::from_i128(3), i256::MIN, i256::MAX]), |
| 157 | + dtype, |
| 158 | + Validity::from_iter([true, false, false]), |
| 159 | + ) |
| 160 | + .into_array(); |
| 161 | + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { |
| 162 | + let values = [3i8, -2, 0, 1, 0].map(|v| { |
| 163 | + <T as BigCast>::from(v).vortex_expect("small test values fit every decimal type") |
| 164 | + }); |
| 165 | + let expected = DecimalArray::new( |
| 166 | + Buffer::copy_from(values), |
| 167 | + dtype, |
| 168 | + Validity::from_iter([true, true, false, true, false]), |
| 169 | + ); |
| 170 | + check( |
| 171 | + vec![first, second], |
| 172 | + &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], |
| 173 | + expected, |
| 174 | + ) |
| 175 | + }) |
| 176 | +} |
| 177 | + |
| 178 | +#[rstest] |
| 179 | +#[case::empty(&[])] |
| 180 | +#[case::all_null(&[(0, 0), (1, 1), (0, 1)])] |
| 181 | +fn test_decimal_interleave_without_valid_values(#[case] indices: &[(u8, u64)]) -> VortexResult<()> { |
| 182 | + let dtype = DecimalDType::new(9, -2); |
| 183 | + let first = DecimalArray::new( |
| 184 | + Buffer::copy_from([i256::MIN, i256::MAX]), |
| 185 | + dtype, |
| 186 | + Validity::AllInvalid, |
| 187 | + ) |
| 188 | + .into_array(); |
| 189 | + let second = ConstantArray::new( |
| 190 | + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), |
| 191 | + 2, |
| 192 | + ) |
| 193 | + .into_array(); |
| 194 | + let expected = DecimalArray::new( |
| 195 | + Buffer::<i32>::zeroed(indices.len()), |
| 196 | + dtype, |
| 197 | + Validity::AllInvalid, |
| 198 | + ); |
| 199 | + check(vec![first, second], indices, expected) |
| 200 | +} |
| 201 | + |
| 202 | +#[test] |
| 203 | +fn test_decimal_interleave_empty_sources() -> VortexResult<()> { |
| 204 | + let dtype = DecimalDType::new(9, 0); |
| 205 | + let empty = DecimalArray::new(Buffer::<i8>::empty(), dtype, Validity::NonNullable).into_array(); |
| 206 | + let expected = DecimalArray::new(Buffer::<i32>::empty(), dtype, Validity::NonNullable); |
| 207 | + check(vec![empty.clone(), empty], &[], expected) |
| 208 | +} |
| 209 | + |
| 210 | +#[test] |
| 211 | +fn test_decimal_interleave_executes_lazy_children() -> VortexResult<()> { |
| 212 | + let dtype = DecimalDType::new(2, 0); |
| 213 | + let values = |
| 214 | + DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); |
| 215 | + let lazy = SliceArray::try_new(values.clone(), 0..2)?.into_array(); |
| 216 | + let branches = SliceArray::try_new(PrimitiveArray::from_iter([1u8, 0, 1]).into_array(), 0..3)? |
| 217 | + .into_array(); |
| 218 | + let rows = SliceArray::try_new(PrimitiveArray::from_iter([1u16, 0, 0]).into_array(), 0..3)? |
| 219 | + .into_array(); |
| 220 | + let interleaved = InterleaveArray::try_new(vec![values, lazy], branches, rows)?.into_array(); |
| 221 | + let expected = DecimalArray::new( |
| 222 | + Buffer::copy_from([-2i8, 1, 1]), |
| 223 | + dtype, |
| 224 | + Validity::NonNullable, |
| 225 | + ) |
| 226 | + .into_array(); |
| 227 | + let mut ctx = VortexSession::empty().create_execution_ctx(); |
| 228 | + assert_arrays_eq!(interleaved, expected, &mut ctx); |
| 229 | + |
| 230 | + Ok(()) |
| 231 | +} |
| 232 | + |
| 233 | +#[rstest] |
| 234 | +#[case::array_index(2, 0, "array index out of bounds")] |
| 235 | +#[case::row_index(0, 2, "row index out of bounds")] |
| 236 | +#[case::constant_row_index(1, 2, "row index out of bounds")] |
| 237 | +#[case::large_array_index(u64::MAX, 0, "array index out of bounds")] |
| 238 | +fn test_decimal_interleave_rejects_out_of_bounds_selectors( |
| 239 | + #[case] branch: u64, |
| 240 | + #[case] row: u8, |
| 241 | + #[case] message: &str, |
| 242 | + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, |
| 243 | +) -> VortexResult<()> { |
| 244 | + let dtype = DecimalDType::new(2, 0); |
| 245 | + let values = DecimalArray::new( |
| 246 | + Buffer::copy_from([1i8, -2]), |
| 247 | + dtype, |
| 248 | + Validity::from(nullability), |
| 249 | + ) |
| 250 | + .into_array(); |
| 251 | + let constant = |
| 252 | + ConstantArray::new(Scalar::decimal(1i8.into(), dtype, nullability), 2).into_array(); |
| 253 | + let array = InterleaveArray::try_new( |
| 254 | + vec![values, constant], |
| 255 | + PrimitiveArray::from_iter([branch]).into_array(), |
| 256 | + PrimitiveArray::from_iter([row]).into_array(), |
| 257 | + )? |
| 258 | + .into_array(); |
| 259 | + let error = array |
| 260 | + .execute::<DecimalArray>(&mut VortexSession::empty().create_execution_ctx()) |
| 261 | + .unwrap_err(); |
| 262 | + assert!(error.to_string().contains(message), "{error}"); |
| 263 | + |
| 264 | + Ok(()) |
| 265 | +} |
| 266 | + |
| 267 | +#[test] |
| 268 | +fn test_decimal_interleave_rejects_wide_values_outside_precision() -> VortexResult<()> { |
| 269 | + let dtype = DecimalDType::new(2, 0); |
| 270 | + let source = |
| 271 | + DecimalArray::new(Buffer::copy_from([100i16]), dtype, Validity::NonNullable).into_array(); |
| 272 | + let array = InterleaveArray::try_new( |
| 273 | + vec![source.clone(), source], |
| 274 | + PrimitiveArray::from_iter([0u8]).into_array(), |
| 275 | + PrimitiveArray::from_iter([0u8]).into_array(), |
| 276 | + )? |
| 277 | + .into_array(); |
| 278 | + let error = array |
| 279 | + .execute::<DecimalArray>(&mut VortexSession::empty().create_execution_ctx()) |
| 280 | + .unwrap_err(); |
| 281 | + assert!(error.to_string().contains("precision"), "{error}"); |
| 282 | + |
| 283 | + Ok(()) |
| 284 | +} |
0 commit comments