Skip to content

Commit 73a0415

Browse files
dankingdanking
committed
Stream selected PCO slices and compact range takes without expanding descriptors
Signed-off-by: Daniel King <dan@spiraldb.com>
1 parent c1c91ba commit 73a0415

5 files changed

Lines changed: 280 additions & 6 deletions

File tree

‎encodings/pco/src/array.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,11 @@ impl VTable for Pco {
282282
}
283283

284284
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
285+
if array.slice_start() != 0 || array.slice_stop() != array.unsliced_n_rows() {
286+
return Ok(ExecutionResult::done(
287+
crate::take::slice(array.as_view(), ctx)?.into_array(),
288+
));
289+
}
285290
let unsliced_validity = array.unsliced_validity();
286291
Ok(ExecutionResult::done(
287292
array

‎encodings/pco/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,3 +106,6 @@ mod selection_contract_tests;
106106

107107
#[cfg(test)]
108108
mod selective_take_tests;
109+
110+
#[cfg(test)]
111+
mod range_selection_tests;

‎encodings/pco/src/range_selection_tests.rs‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use std::sync::LazyLock;
5+
6+
use rstest::rstest;
7+
use vortex_array::ArrayRef;
8+
use vortex_array::IntoArray;
9+
use vortex_array::VortexSessionExecute;
10+
use vortex_array::arrays::PiecewiseSequenceArray;
11+
use vortex_array::arrays::PrimitiveArray;
12+
use vortex_array::assert_arrays_eq;
13+
use vortex_array::session::ArraySessionExt;
14+
use vortex_array::validity::Validity;
15+
use vortex_error::VortexExpect;
16+
use vortex_error::VortexResult;
17+
use vortex_session::VortexSession;
18+
19+
use crate::Pco;
20+
21+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
22+
let session = VortexSession::empty();
23+
crate::initialize(&session);
24+
session
25+
});
26+
27+
#[rstest]
28+
#[case(vec![0u64, 40, 0, 70], vec![15u64, 20, 3, 0], vec![1u64, 1, 0, 1])]
29+
#[case(vec![100u64, 0, 20], vec![20u64, 20, 20], vec![1u64, 1, 2])]
30+
#[case(vec![99u64, 0], vec![0u64, 0], vec![u64::MAX, 0])]
31+
fn piecewise_ranges(
32+
#[case] starts: Vec<u64>,
33+
#[case] lengths: Vec<u64>,
34+
#[case] strides: Vec<u64>,
35+
) -> VortexResult<()> {
36+
let mut ctx = SESSION.create_execution_ctx();
37+
let input =
38+
PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i)));
39+
let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array();
40+
let mut expected_indices = Vec::new();
41+
for ((&start, &length), &stride) in starts.iter().zip(&lengths).zip(&strides) {
42+
expected_indices.extend((0..length).map(|i| start + i * stride));
43+
}
44+
let indices = PiecewiseSequenceArray::try_new(
45+
PrimitiveArray::from_iter(starts).into_array(),
46+
PrimitiveArray::from_iter(lengths).into_array(),
47+
PrimitiveArray::from_iter(strides).into_array(),
48+
expected_indices.len(),
49+
)?
50+
.into_array();
51+
assert_arrays_eq!(
52+
encoded.take(indices)?,
53+
input.take(PrimitiveArray::from_iter(expected_indices).into_array())?,
54+
&mut ctx
55+
);
56+
Ok(())
57+
}
58+
59+
#[test]
60+
fn rejects_invalid_indices_and_descriptors() -> VortexResult<()> {
61+
let mut ctx = SESSION.create_execution_ctx();
62+
let input = PrimitiveArray::from_iter(0..64i32);
63+
let encoded = Pco::from_primitive(input.as_view(), 3, 4, &mut ctx)?.into_array();
64+
for indices in [
65+
PrimitiveArray::from_iter([-1i64]).into_array(),
66+
PrimitiveArray::from_iter([64i64]).into_array(),
67+
] {
68+
assert!(
69+
encoded
70+
.take(indices)
71+
.and_then(|array| array.execute::<PrimitiveArray>(&mut ctx))
72+
.is_err()
73+
);
74+
}
75+
for (start, length, stride, declared) in [
76+
(u64::MAX, 2u64, 1u64, 2usize),
77+
(1, 2, u64::MAX, 2),
78+
(1, 2, 1, 3),
79+
] {
80+
let indices = PiecewiseSequenceArray::try_new(
81+
PrimitiveArray::from_iter([start]).into_array(),
82+
PrimitiveArray::from_iter([length]).into_array(),
83+
PrimitiveArray::from_iter([stride]).into_array(),
84+
declared,
85+
)?
86+
.into_array();
87+
assert!(
88+
encoded
89+
.take(indices)
90+
.and_then(|array| array.execute::<PrimitiveArray>(&mut ctx))
91+
.is_err()
92+
);
93+
}
94+
Ok(())
95+
}
96+
97+
#[test]
98+
fn large_public_compressor_selections() -> VortexResult<()> {
99+
let mut ctx = SESSION.create_execution_ctx();
100+
let len = pco::DEFAULT_MAX_PAGE_N + 257;
101+
let input = PrimitiveArray::from_iter(
102+
(0..len).map(|i| u32::try_from(i).vortex_expect("fixture fits u32")),
103+
);
104+
let encoded = Pco::from_primitive(input.as_view(), 0, 4096, &mut ctx)?.into_array();
105+
let indices = PrimitiveArray::from_iter([
106+
0u32,
107+
900,
108+
260,
109+
3000,
110+
u32::try_from(len - 1).vortex_expect("fixture fits u32"),
111+
4095,
112+
900,
113+
])
114+
.into_array();
115+
assert_arrays_eq!(
116+
encoded.take(indices.clone())?,
117+
input.take(indices)?,
118+
&mut ctx
119+
);
120+
assert_arrays_eq!(
121+
encoded.slice(len - 1000..len - 1)?,
122+
input.slice(len - 1000..len - 1)?,
123+
&mut ctx
124+
);
125+
Ok(())
126+
}
127+
128+
#[test]
129+
fn mixed_fragments_narrow_slices_and_compact_repeats() -> VortexResult<()> {
130+
let mut ctx = SESSION.create_execution_ctx();
131+
let input = PrimitiveArray::from_option_iter(
132+
(0..32768i64).map(|row| (row % 5 != 0).then_some(row * 11)),
133+
);
134+
let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array();
135+
let mixed = PrimitiveArray::from_iter(
136+
(0..4096u64).map(|i| if i < 2 { i } else { (i * 104729) % 32768 }),
137+
)
138+
.into_array();
139+
assert_arrays_eq!(encoded.take(mixed.clone())?, input.take(mixed)?, &mut ctx);
140+
let repeated =
141+
PrimitiveArray::from_iter((0..100000u64).map(|i| if i % 2 == 0 { 0 } else { 127 }))
142+
.into_array();
143+
assert_arrays_eq!(
144+
encoded.slice(30000..30128)?.take(repeated.clone())?,
145+
input.slice(30000..30128)?.take(repeated)?,
146+
&mut ctx
147+
);
148+
let compact = PiecewiseSequenceArray::try_new(
149+
PrimitiveArray::from_iter([1u64, 5]).into_array(),
150+
PrimitiveArray::from_iter([500000u64, 500000]).into_array(),
151+
PrimitiveArray::from_iter([0u64, 0]).into_array(),
152+
1000000,
153+
)?
154+
.into_array();
155+
let expected = PrimitiveArray::from_option_iter(
156+
std::iter::repeat_n(Some(11i64), 500000).chain(std::iter::repeat_n(None, 500000)),
157+
);
158+
assert_arrays_eq!(encoded.take(compact)?, expected, &mut ctx);
159+
let nulls = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid);
160+
let encoded_nulls = Pco::from_primitive(nulls.as_view(), 3, 1024, &mut ctx)?.into_array();
161+
let two = PrimitiveArray::from_iter([0u64, 99999]).into_array();
162+
assert_arrays_eq!(encoded_nulls.take(two.clone())?, nulls.take(two)?, &mut ctx);
163+
assert!(
164+
encoded_nulls
165+
.take(PrimitiveArray::from_iter([100000u64]).into_array())?
166+
.execute::<PrimitiveArray>(&mut ctx)
167+
.is_err()
168+
);
169+
Ok(())
170+
}
171+
172+
#[rstest]
173+
fn streamed_slices_preserve_float_bits_and_validity(
174+
#[values(false, true)] dense: bool,
175+
) -> VortexResult<()> {
176+
let mut ctx = SESSION.create_execution_ctx();
177+
let bits = [
178+
0u64,
179+
1u64 << 63,
180+
0x7ff8000000000042,
181+
0x7ff0000000000000,
182+
0xfff0000000000000,
183+
0x3ff4000000000000,
184+
];
185+
let values = (0..2048)
186+
.map(|row| f64::from_bits(bits[row % bits.len()]))
187+
.collect::<Vec<_>>();
188+
let validity = if dense {
189+
Validity::NonNullable
190+
} else {
191+
Validity::from_iter((0..2048).map(|row| row % 7 != 0))
192+
};
193+
let input = PrimitiveArray::new(values.clone(), validity);
194+
let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array();
195+
for range in [1..2047, 89..129, 1000..1001] {
196+
let actual = encoded
197+
.slice(range.clone())?
198+
.execute::<PrimitiveArray>(&mut ctx)?;
199+
assert_eq!(actual.dtype(), input.dtype());
200+
let mask = actual.validity()?.execute_mask(actual.len(), &mut ctx)?;
201+
for (offset, &value) in actual.as_slice::<f64>().iter().enumerate() {
202+
let row = range.start + offset;
203+
assert_eq!(mask.value(offset), dense || row % 7 != 0);
204+
if mask.value(offset) {
205+
assert_eq!(value.to_bits(), values[row].to_bits());
206+
}
207+
}
208+
}
209+
Ok(())
210+
}

‎encodings/pco/src/slice.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,17 @@ use vortex_array::ArrayView;
88
use vortex_array::IntoArray;
99
use vortex_array::arrays::slice::SliceReduce;
1010
use vortex_error::VortexResult;
11+
use vortex_error::vortex_ensure;
1112

1213
use crate::Pco;
1314
use crate::PcoArrayExt;
1415

1516
impl SliceReduce for Pco {
1617
fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
18+
vortex_ensure!(
19+
range.start <= range.end && range.end <= array.len(),
20+
"Pco slice range exceeds array bounds"
21+
);
1722
let unsliced_validity = array.unsliced_validity();
1823
Ok(Some(
1924
Pco::try_new(

‎encodings/pco/src/take.rs‎

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use vortex_array::arrays::ConstantArray;
2121
use vortex_array::arrays::PiecewiseSequence;
2222
use vortex_array::arrays::PrimitiveArray;
2323
use vortex_array::arrays::dict::TakeExecute;
24+
use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt;
2425
use vortex_array::dtype::NativePType;
2526
use vortex_array::dtype::Nullability;
2627
use vortex_array::dtype::half;
@@ -61,12 +62,28 @@ impl TakeExecute for Pco {
6162
.union_nullability(indices.dtype().nullability());
6263
return Ok(Some(Canonical::empty(&dtype).into_array()));
6364
}
64-
// Compact range execution is introduced separately; preserve the normal fallback.
65-
if indices.is::<PiecewiseSequence>() {
66-
return Ok(None);
67-
}
68-
let grouped;
69-
let runs = {
65+
let compact = indices.as_opt::<PiecewiseSequence>().is_some();
66+
let mut grouped = None;
67+
let runs = if let Some(sequence) = indices.as_opt::<PiecewiseSequence>() {
68+
let starts = sequence.starts().clone().execute::<PrimitiveArray>(ctx)?;
69+
let lengths = sequence.lengths().clone().execute::<PrimitiveArray>(ctx)?;
70+
let multipliers = sequence
71+
.multipliers()
72+
.clone()
73+
.execute::<PrimitiveArray>(ctx)?;
74+
let starts = primitive_usizes(&starts)?;
75+
let lengths = primitive_usizes(&lengths)?;
76+
let multipliers = primitive_usizes(&multipliers)?;
77+
let mut runs = Vec::with_capacity(starts.len());
78+
for ((start, length), stride) in starts.into_iter().zip(lengths).zip(multipliers) {
79+
runs.push(Run {
80+
start: Some(start),
81+
length,
82+
stride,
83+
});
84+
}
85+
runs
86+
} else {
7087
let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
7188
let mask = indices.validity()?.execute_mask(indices.len(), ctx)?;
7289
// Reuse native logical-null-aware extrema and their immutable-array stats cache.
@@ -112,6 +129,21 @@ impl TakeExecute for Pco {
112129
})
113130
};
114131
validate_runs(&runs, array.len(), indices.len())?;
132+
if compact {
133+
if let Some(result) = zero_page_take(array, indices.len(), ctx)? {
134+
return Ok(Some(result));
135+
}
136+
// Compact ranges/repeats remain compact. Only genuine strided requests carry
137+
// per-position selection work; output cardinality alone is not setup work.
138+
let random = runs
139+
.iter()
140+
.filter(|run| run.stride > 1)
141+
.map(|run| run.length)
142+
.sum();
143+
if descriptor_bytes(random) >= canonical_bytes(array) && random != 0 {
144+
return Ok(None);
145+
}
146+
}
115147
let nullable = array.dtype().nullability() | indices.dtype().nullability();
116148
let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType<T> => {
117149
gather::<T>(array, &runs, indices.len(), nullable, grouped, ctx)?
@@ -286,6 +318,16 @@ where
286318
Ok(runs)
287319
}
288320

321+
fn primitive_usizes(array: &PrimitiveArray) -> VortexResult<Vec<usize>> {
322+
match_each_integer_ptype!(array.ptype(), |I| {
323+
array
324+
.as_slice::<I>()
325+
.iter()
326+
.map(|&value| checked_index(value))
327+
.collect()
328+
})
329+
}
330+
289331
fn validate_runs(runs: &[Run], source_len: usize, output_len: usize) -> VortexResult<()> {
290332
let mut total = 0usize;
291333
for run in runs {
@@ -697,3 +739,12 @@ fn stream_range<T: Number + NativePType>(
697739
.union_nullability(nullability);
698740
Ok(PrimitiveArray::new(values.freeze(), validity))
699741
}
742+
743+
pub(crate) fn slice(
744+
array: ArrayView<'_, Pco>,
745+
ctx: &mut ExecutionCtx,
746+
) -> VortexResult<PrimitiveArray> {
747+
match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType<T> => {
748+
stream_range::<T>(array, 0..array.len(), array.dtype().nullability(), ctx)
749+
})
750+
}

0 commit comments

Comments
 (0)