Skip to content

Commit 0b80a82

Browse files
gatesngatesn
authored
Merge branch 'develop' into operators-bench
2 parents b753bfc + af0af75 commit 0b80a82

File tree

60 files changed

+1227
-224
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

60 files changed

+1227
-224
lines changed

‎encodings/alp/src/alp/array.rs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ use vortex_error::VortexExpect;
4343
use vortex_error::VortexResult;
4444
use vortex_error::vortex_bail;
4545
use vortex_error::vortex_ensure;
46+
use vortex_error::vortex_err;
4647

4748
use crate::ALPFloat;
4849
use crate::alp::Exponents;
@@ -141,6 +142,51 @@ impl VTable for ALPVTable {
141142
)
142143
}
143144

145+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
146+
// Children: encoded, patches (if present): indices, values, chunk_offsets (optional)
147+
let patches_info = array
148+
.patches
149+
.as_ref()
150+
.map(|p| (p.array_len(), p.offset(), p.chunk_offsets().is_some()));
151+
152+
let expected_children = match &patches_info {
153+
Some((_, _, has_chunk_offsets)) => 1 + 2 + if *has_chunk_offsets { 1 } else { 0 },
154+
None => 1,
155+
};
156+
157+
vortex_ensure!(
158+
children.len() == expected_children,
159+
"ALPArray expects {} children, got {}",
160+
expected_children,
161+
children.len()
162+
);
163+
164+
let mut children_iter = children.into_iter();
165+
array.encoded = children_iter
166+
.next()
167+
.ok_or_else(|| vortex_err!("Expected encoded child"))?;
168+
169+
if let Some((array_len, offset, _has_chunk_offsets)) = patches_info {
170+
let indices = children_iter
171+
.next()
172+
.ok_or_else(|| vortex_err!("Expected patch indices child"))?;
173+
let values = children_iter
174+
.next()
175+
.ok_or_else(|| vortex_err!("Expected patch values child"))?;
176+
let chunk_offsets = children_iter.next();
177+
178+
array.patches = Some(Patches::new(
179+
array_len,
180+
offset,
181+
indices,
182+
values,
183+
chunk_offsets,
184+
));
185+
}
186+
187+
Ok(())
188+
}
189+
144190
fn bind_kernel(array: &ALPArray, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
145191
let encoded = array.encoded().bind_kernel(ctx)?;
146192
let patches_kernels = if let Some(patches) = array.patches() {

‎encodings/alp/src/alp_rd/array.rs‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ use vortex_error::VortexError;
4545
use vortex_error::VortexExpect;
4646
use vortex_error::VortexResult;
4747
use vortex_error::vortex_bail;
48+
use vortex_error::vortex_ensure;
4849
use vortex_error::vortex_err;
4950

5051
use crate::alp_rd::alp_rd_decode;
@@ -185,6 +186,47 @@ impl VTable for ALPRDVTable {
185186
left_parts_patches,
186187
)
187188
}
189+
190+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
191+
// Children: left_parts, right_parts, patches (if present): indices, values
192+
let patches_info = array
193+
.left_parts_patches
194+
.as_ref()
195+
.map(|p| (p.array_len(), p.offset()));
196+
197+
let expected_children = if patches_info.is_some() { 4 } else { 2 };
198+
199+
vortex_ensure!(
200+
children.len() == expected_children,
201+
"ALPRDArray expects {} children, got {}",
202+
expected_children,
203+
children.len()
204+
);
205+
206+
let mut children_iter = children.into_iter();
207+
array.left_parts = children_iter
208+
.next()
209+
.ok_or_else(|| vortex_err!("Expected left_parts child"))?;
210+
array.right_parts = children_iter
211+
.next()
212+
.ok_or_else(|| vortex_err!("Expected right_parts child"))?;
213+
214+
if let Some((array_len, offset)) = patches_info {
215+
let indices = children_iter
216+
.next()
217+
.ok_or_else(|| vortex_err!("Expected patch indices child"))?;
218+
let values = children_iter
219+
.next()
220+
.ok_or_else(|| vortex_err!("Expected patch values child"))?;
221+
222+
array.left_parts_patches = Some(Patches::new(
223+
array_len, offset, indices, values,
224+
None, // chunk_offsets not currently supported for ALPRD
225+
));
226+
}
227+
228+
Ok(())
229+
}
188230
}
189231

190232
#[derive(Clone, Debug)]

‎encodings/bytebool/src/array.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,10 @@ use vortex_buffer::BitBuffer;
3535
use vortex_buffer::BufferHandle;
3636
use vortex_buffer::ByteBuffer;
3737
use vortex_dtype::DType;
38+
use vortex_error::VortexExpect;
3839
use vortex_error::VortexResult;
3940
use vortex_error::vortex_bail;
41+
use vortex_error::vortex_ensure;
4042
use vortex_error::vortex_panic;
4143
use vortex_scalar::Scalar;
4244

@@ -99,6 +101,22 @@ impl VTable for ByteBoolVTable {
99101

100102
Ok(ByteBoolArray::new(buffer, validity))
101103
}
104+
105+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
106+
vortex_ensure!(
107+
children.len() <= 1,
108+
"ByteBoolArray expects at most 1 child (validity), got {}",
109+
children.len()
110+
);
111+
112+
array.validity = if children.is_empty() {
113+
Validity::from(array.dtype.nullability())
114+
} else {
115+
Validity::Array(children.into_iter().next().vortex_expect("checked"))
116+
};
117+
118+
Ok(())
119+
}
102120
}
103121

104122
#[derive(Clone, Debug)]

‎encodings/datetime-parts/src/array.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,10 @@ use vortex_buffer::BufferHandle;
3434
use vortex_dtype::DType;
3535
use vortex_dtype::Nullability;
3636
use vortex_dtype::PType;
37+
use vortex_error::VortexExpect;
3738
use vortex_error::VortexResult;
3839
use vortex_error::vortex_bail;
40+
use vortex_error::vortex_ensure;
3941
use vortex_error::vortex_err;
4042

4143
vtable!(DateTimeParts);
@@ -142,6 +144,21 @@ impl VTable for DateTimePartsVTable {
142144

143145
DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)
144146
}
147+
148+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
149+
vortex_ensure!(
150+
children.len() == 3,
151+
"DateTimePartsArray expects exactly 3 children (days, seconds, subseconds), got {}",
152+
children.len()
153+
);
154+
155+
let mut children_iter = children.into_iter();
156+
array.days = children_iter.next().vortex_expect("checked");
157+
array.seconds = children_iter.next().vortex_expect("checked");
158+
array.subseconds = children_iter.next().vortex_expect("checked");
159+
160+
Ok(())
161+
}
145162
}
146163

147164
#[derive(Clone, Debug)]

‎encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ use vortex_dtype::match_each_signed_integer_ptype;
4444
use vortex_error::VortexExpect;
4545
use vortex_error::VortexResult;
4646
use vortex_error::vortex_bail;
47+
use vortex_error::vortex_ensure;
4748
use vortex_scalar::DecimalValue;
4849
use vortex_scalar::Scalar;
4950

@@ -116,6 +117,16 @@ impl VTable for DecimalBytePartsVTable {
116117

117118
DecimalBytePartsArray::try_new(msp, *decimal_dtype)
118119
}
120+
121+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
122+
vortex_ensure!(
123+
children.len() == 1,
124+
"DecimalBytePartsArray expects exactly 1 child (msp), got {}",
125+
children.len()
126+
);
127+
array.msp = children.into_iter().next().vortex_expect("checked");
128+
Ok(())
129+
}
119130
}
120131

121132
/// This array encodes decimals as between 1-4 columns of primitive typed children.

‎encodings/fastlanes/src/bitpacking/vtable/mod.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_array::ArrayRef;
45
use vortex_array::DeserializeMetadata;
56
use vortex_array::ProstMetadata;
67
use vortex_array::SerializeMetadata;
@@ -24,6 +25,7 @@ use vortex_dtype::PType;
2425
use vortex_error::VortexError;
2526
use vortex_error::VortexResult;
2627
use vortex_error::vortex_bail;
28+
use vortex_error::vortex_ensure;
2729
use vortex_error::vortex_err;
2830
use vortex_vector::VectorMutOps;
2931

@@ -70,6 +72,75 @@ impl VTable for BitPackedVTable {
7072
BitPackedVTable.as_vtable()
7173
}
7274

75+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
76+
// Children: patches (if present): indices, values, chunk_offsets; then validity (if present)
77+
let patches_info = array
78+
.patches()
79+
.map(|p| (p.offset(), p.chunk_offsets().is_some()));
80+
81+
let mut child_idx = 0;
82+
let patches = if let Some((patch_offset, has_chunk_offsets)) = patches_info {
83+
let patch_indices = children
84+
.get(child_idx)
85+
.ok_or_else(|| vortex_err!("Expected patch_indices child at index {}", child_idx))?
86+
.clone();
87+
child_idx += 1;
88+
89+
let patch_values = children
90+
.get(child_idx)
91+
.ok_or_else(|| vortex_err!("Expected patch_values child at index {}", child_idx))?
92+
.clone();
93+
child_idx += 1;
94+
95+
let patch_chunk_offsets = if has_chunk_offsets {
96+
let offsets = children
97+
.get(child_idx)
98+
.ok_or_else(|| {
99+
vortex_err!("Expected patch_chunk_offsets child at index {}", child_idx)
100+
})?
101+
.clone();
102+
child_idx += 1;
103+
Some(offsets)
104+
} else {
105+
None
106+
};
107+
108+
Some(Patches::new(
109+
array.len(),
110+
patch_offset,
111+
patch_indices,
112+
patch_values,
113+
patch_chunk_offsets,
114+
))
115+
} else {
116+
None
117+
};
118+
119+
let validity = if child_idx < children.len() {
120+
Validity::Array(children[child_idx].clone())
121+
} else {
122+
Validity::from(array.dtype().nullability())
123+
};
124+
125+
let expected_children = child_idx
126+
+ if matches!(validity, Validity::Array(_)) {
127+
1
128+
} else {
129+
0
130+
};
131+
vortex_ensure!(
132+
children.len() == expected_children,
133+
"Expected {} children, got {}",
134+
expected_children,
135+
children.len()
136+
);
137+
138+
array.patches = patches;
139+
array.validity = validity;
140+
141+
Ok(())
142+
}
143+
73144
fn metadata(array: &BitPackedArray) -> VortexResult<Self::Metadata> {
74145
Ok(ProstMetadata(BitPackedMetadata {
75146
bit_width: array.bit_width() as u32,

‎encodings/fastlanes/src/delta/array/delta_decompress.rs‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ use fastlanes::Delta;
77
use fastlanes::FastLanes;
88
use fastlanes::Transpose;
99
use num_traits::WrappingAdd;
10-
use vortex_array::Array;
1110
use vortex_array::ToCanonical;
1211
use vortex_array::arrays::PrimitiveArray;
1312
use vortex_array::validity::Validity;
@@ -21,26 +20,32 @@ use crate::DeltaArray;
2120
pub fn delta_decompress(array: &DeltaArray) -> PrimitiveArray {
2221
let bases = array.bases().to_primitive();
2322
let deltas = array.deltas().to_primitive();
24-
let decoded = match_each_unsigned_integer_ptype!(deltas.ptype(), |T| {
23+
24+
let start = array.offset();
25+
let end = start + array.len();
26+
27+
// TODO(connor): This is incorrect, we need to untranspose the validity!!!
28+
29+
let validity = Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability());
30+
let validity = validity.slice(start..end);
31+
32+
match_each_unsigned_integer_ptype!(deltas.ptype(), |T| {
2533
const LANES: usize = T::LANES;
2634

27-
PrimitiveArray::new(
28-
decompress_primitive::<T, LANES>(bases.as_slice(), deltas.as_slice()),
29-
Validity::from_mask(array.deltas().validity_mask(), array.dtype().nullability()),
30-
)
31-
});
35+
let buffer = decompress_primitive::<T, LANES>(bases.as_slice(), deltas.as_slice());
36+
let buffer = buffer.slice(start..end);
3237

33-
decoded
34-
.slice(array.offset()..array.offset() + array.len())
35-
.to_primitive()
38+
PrimitiveArray::new(buffer, validity)
39+
})
3640
}
3741

3842
// TODO(ngates): can we re-use the deltas buffer for the result? Might be tricky given the
3943
// traversal ordering, but possibly doable.
40-
fn decompress_primitive<T: NativePType + Delta + Transpose + WrappingAdd, const LANES: usize>(
41-
bases: &[T],
42-
deltas: &[T],
43-
) -> Buffer<T> {
44+
/// Performs the low-level delta decompression on primitive values.
45+
pub(crate) fn decompress_primitive<T, const LANES: usize>(bases: &[T], deltas: &[T]) -> Buffer<T>
46+
where
47+
T: NativePType + Delta + Transpose + WrappingAdd,
48+
{
4449
// How many fastlanes vectors we will process.
4550
let num_chunks = deltas.len() / 1024;
4651

‎encodings/fastlanes/src/delta/array/mod.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,12 +55,12 @@ pub mod delta_decompress;
5555
/// Note the validity is stored in the deltas array.
5656
#[derive(Clone, Debug)]
5757
pub struct DeltaArray {
58-
offset: usize,
59-
len: usize,
60-
dtype: DType,
61-
bases: ArrayRef,
62-
deltas: ArrayRef,
63-
stats_set: ArrayStats,
58+
pub(super) offset: usize,
59+
pub(super) len: usize,
60+
pub(super) dtype: DType,
61+
pub(super) bases: ArrayRef,
62+
pub(super) deltas: ArrayRef,
63+
pub(super) stats_set: ArrayStats,
6464
}
6565

6666
impl DeltaArray {

0 commit comments

Comments
 (0)