Skip to content

Commit b2e089c

Browse files
joseph-isaacsjoseph-isaacs
authored
perf[array]: with_children vtable function (#5694)
Signed-off-by: Joe Isaacs <[email protected]>
1 parent f95caac commit b2e089c

File tree

51 files changed

+880
-87
lines changed

Some content is hidden

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

51 files changed

+880
-87
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/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 {

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
use fastlanes::FastLanes;
55
use prost::Message;
6+
use vortex_array::ArrayRef;
67
use vortex_array::ProstMetadata;
78
use vortex_array::serde::ArrayChildren;
89
use vortex_array::vtable;
@@ -17,6 +18,7 @@ use vortex_dtype::DType;
1718
use vortex_dtype::PType;
1819
use vortex_dtype::match_each_unsigned_integer_ptype;
1920
use vortex_error::VortexResult;
21+
use vortex_error::vortex_ensure;
2022
use vortex_error::vortex_err;
2123

2224
use crate::DeltaArray;
@@ -59,6 +61,23 @@ impl VTable for DeltaVTable {
5961
DeltaVTable.as_vtable()
6062
}
6163

64+
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
65+
// DeltaArray children order (from visit_children):
66+
// 1. bases
67+
// 2. deltas
68+
69+
vortex_ensure!(
70+
children.len() == 2,
71+
"Expected 2 children for Delta encoding, got {}",
72+
children.len()
73+
);
74+
75+
array.bases = children[0].clone();
76+
array.deltas = children[1].clone();
77+
78+
Ok(())
79+
}
80+
6281
fn metadata(array: &DeltaArray) -> VortexResult<Self::Metadata> {
6382
Ok(ProstMetadata(DeltaMetadata {
6483
deltas_len: array.deltas().len() as u64,

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@ pub mod for_decompress;
1717
/// storage requirements when values are clustered around a specific point.
1818
#[derive(Clone, Debug)]
1919
pub struct FoRArray {
20-
encoded: ArrayRef,
21-
reference: Scalar,
22-
stats_set: ArrayStats,
20+
pub(super) encoded: ArrayRef,
21+
pub(super) reference: Scalar,
22+
pub(super) stats_set: ArrayStats,
2323
}
2424

2525
impl FoRArray {

0 commit comments

Comments
 (0)