Skip to content

Commit 3b7bea5

Browse files
committed
feat: execute interleaved decimal arrays
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent d8f5edf commit 3b7bea5

5 files changed

Lines changed: 496 additions & 80 deletions

File tree

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Execution for decimal [`Interleave`](super::super::Interleave) values.
5+
//!
6+
//! Sources can use different stored widths for the same logical dtype. Decimal execution
7+
//! normalizes them to the logical canonical width before gathering rows.
8+
9+
use vortex_buffer::Buffer;
10+
use vortex_error::VortexResult;
11+
use vortex_error::vortex_err;
12+
13+
use super::super::Interleave;
14+
use super::super::InterleaveArrayExt;
15+
use super::super::InterleaveSlots;
16+
use super::fixed_width::Source;
17+
use super::fixed_width::gather;
18+
use crate::AnyColumnar;
19+
use crate::array::Array;
20+
use crate::arrays::Constant;
21+
use crate::arrays::Decimal;
22+
use crate::arrays::DecimalArray;
23+
use crate::arrays::Primitive;
24+
use crate::arrays::decimal::DecimalArrayExt;
25+
use crate::arrays::decimal::copy_to_canonical_width;
26+
use crate::dtype::DecimalDType;
27+
use crate::dtype::NativeDecimalType;
28+
use crate::executor::ExecutionCtx;
29+
use crate::executor::ExecutionResult;
30+
use crate::match_each_decimal_value_type;
31+
use crate::match_each_unsigned_integer_ptype;
32+
use crate::require_child;
33+
34+
pub(super) fn execute(
35+
mut array: Array<Interleave>,
36+
ctx: &mut ExecutionCtx,
37+
) -> VortexResult<ExecutionResult> {
38+
let num_values = array.num_values();
39+
array =
40+
require_child!(array, array.array_indices(), InterleaveSlots::ARRAY_INDICES => Primitive);
41+
array = require_child!(array, array.row_indices(), InterleaveSlots::ROW_INDICES => Primitive);
42+
for i in 0..num_values {
43+
array = require_child!(array, array.value(i), InterleaveSlots::VALUES_OFFSET + i => AnyColumnar);
44+
}
45+
46+
let validity = array.as_ref().validity()?;
47+
let dtype = DecimalDType::try_from(array.dtype())?;
48+
let output = match_each_decimal_value_type!(dtype.canonical_values_type(), |T| {
49+
DecimalArray::new(gather_values::<T>(&array, dtype, ctx)?, dtype, validity)
50+
});
51+
52+
Ok(ExecutionResult::done(output))
53+
}
54+
55+
fn gather_values<T: NativeDecimalType>(
56+
array: &Array<Interleave>,
57+
dtype: DecimalDType,
58+
ctx: &mut ExecutionCtx,
59+
) -> VortexResult<Buffer<T>> {
60+
let allocator = ctx.allocator().clone();
61+
let values = (0..array.num_values())
62+
.map(|i| {
63+
let value = array.value(i);
64+
let len = value.len();
65+
let (data, row_mask) = if let Some(constant) = value.as_opt::<Constant>() {
66+
let payload = match constant.scalar().as_decimal().decimal_value() {
67+
Some(value) => value.cast::<T>().ok_or_else(|| {
68+
vortex_err!("interleave decimal constant must fit {dtype}, got {value}")
69+
})?,
70+
None => T::default(),
71+
};
72+
(Buffer::full_in(payload, 1, allocator.clone()), 0)
73+
} else {
74+
let decimal = value.as_::<Decimal>();
75+
let data = if decimal.values_type() == T::DECIMAL_TYPE {
76+
decimal.buffer::<T>()
77+
} else {
78+
// A decimal source can have noncanonical storage. Normalize before indexing
79+
// rows so that every source has the same physical width. A narrowing copy
80+
// checks the precision and zeros arbitrary null payloads.
81+
let mask = decimal.validity()?.execute_mask(len, ctx)?;
82+
let copy = match_each_decimal_value_type!(decimal.values_type(), |S| {
83+
copy_to_canonical_width(decimal.buffer::<S>().as_slice(), &mask, dtype)?
84+
});
85+
Buffer::<T>::from_byte_buffer(copy.unwrap_host())
86+
};
87+
(data, usize::MAX)
88+
};
89+
90+
Ok(Source {
91+
data,
92+
len,
93+
row_mask,
94+
})
95+
})
96+
.collect::<VortexResult<Vec<_>>>()?;
97+
let branches = array.array_indices().as_::<Primitive>();
98+
let rows = array.row_indices().as_::<Primitive>();
99+
100+
match_each_unsigned_integer_ptype!(branches.ptype(), |A| {
101+
match_each_unsigned_integer_ptype!(rows.ptype(), |R| {
102+
gather(
103+
&values,
104+
branches.as_slice::<A>(),
105+
rows.as_slice::<R>(),
106+
&allocator,
107+
)
108+
})
109+
})
110+
}
111+
112+
#[cfg(test)]
113+
mod tests;
Lines changed: 284 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,284 @@
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

Comments
 (0)