Skip to content

Commit 815e475

Browse files
dankingdanking
committed
Inline PiecewiseSequence range length checks
Signed-off-by: Daniel King <dan@spiraldb.com>
1 parent 2a33748 commit 815e475

6 files changed

Lines changed: 132 additions & 102 deletions

File tree

‎vortex-array/src/arrays/bool/compute/take.rs‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ use vortex_buffer::BitBufferMut;
88
use vortex_buffer::BitBufferView;
99
use vortex_buffer::get_bit;
1010
use vortex_error::VortexResult;
11+
use vortex_error::vortex_ensure;
12+
use vortex_error::vortex_err;
1113
use vortex_mask::Mask;
1214

1315
use crate::ArrayRef;
@@ -22,8 +24,6 @@ use crate::arrays::bool::BoolArrayExt;
2224
use crate::arrays::dict::TakeExecute;
2325
use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
2426
use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
25-
use crate::arrays::piecewise_sequence::validate_index_ranges;
26-
use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
2727
use crate::builtins::ArrayBuiltins;
2828
use crate::dtype::UnsignedPType;
2929
use crate::executor::ExecutionCtx;
@@ -146,12 +146,22 @@ fn take_bit_slices_constant_length<S>(
146146
where
147147
S: UnsignedPType,
148148
{
149-
validate_index_ranges_constant(source.len(), starts, length, output_len)?;
149+
let computed_len = starts
150+
.len()
151+
.checked_mul(length)
152+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
153+
vortex_ensure!(
154+
computed_len == output_len,
155+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
156+
);
150157

151158
let mut values = BitBufferMut::with_capacity(output_len);
152159
for &start in starts {
153160
let start = start.as_();
154-
values.append_buffer(&source.slice(start..start + length));
161+
let end = start
162+
.checked_add(length)
163+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
164+
values.append_buffer(&source.slice(start..end));
155165
}
156166

157167
Ok(values.freeze())
@@ -167,15 +177,24 @@ where
167177
S: UnsignedPType,
168178
L: UnsignedPType,
169179
{
170-
validate_index_ranges(source.len(), starts, lengths, output_len)?;
171-
172180
let mut values = BitBufferMut::with_capacity(output_len);
181+
let mut computed_len = 0usize;
173182
for (&start, &length) in starts.iter().zip_eq(lengths) {
174183
let start = start.as_();
175184
let length = length.as_();
176-
values.append_buffer(&source.slice(start..start + length));
185+
let end = start
186+
.checked_add(length)
187+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
188+
computed_len = computed_len
189+
.checked_add(length)
190+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
191+
values.append_buffer(&source.slice(start..end));
177192
}
178193

194+
vortex_ensure!(
195+
computed_len == output_len,
196+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
197+
);
179198
Ok(values.freeze())
180199
}
181200

‎vortex-array/src/arrays/decimal/compute/take.rs‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ use itertools::Itertools as _;
55
use vortex_buffer::Buffer;
66
use vortex_buffer::BufferMut;
77
use vortex_error::VortexResult;
8+
use vortex_error::vortex_ensure;
9+
use vortex_error::vortex_err;
810

911
use crate::ArrayRef;
1012
use crate::IntoArray;
@@ -16,8 +18,6 @@ use crate::arrays::PrimitiveArray;
1618
use crate::arrays::dict::TakeExecute;
1719
use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
1820
use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
19-
use crate::arrays::piecewise_sequence::validate_index_ranges;
20-
use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
2121
use crate::dtype::IntegerPType;
2222
use crate::dtype::NativeDecimalType;
2323
use crate::dtype::UnsignedPType;
@@ -186,12 +186,22 @@ where
186186
S: UnsignedPType,
187187
T: NativeDecimalType,
188188
{
189-
validate_index_ranges_constant(values.len(), starts, length, output_len)?;
189+
let computed_len = starts
190+
.len()
191+
.checked_mul(length)
192+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
193+
vortex_ensure!(
194+
computed_len == output_len,
195+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
196+
);
190197

191198
let mut result = BufferMut::<T>::with_capacity(output_len);
192199
for &start in starts {
193200
let start = start.as_();
194-
result.extend_from_slice(&values[start..start + length]);
201+
let end = start
202+
.checked_add(length)
203+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
204+
result.extend_from_slice(&values[start..end]);
195205
}
196206

197207
Ok(result.freeze())
@@ -208,15 +218,24 @@ where
208218
L: UnsignedPType,
209219
T: NativeDecimalType,
210220
{
211-
validate_index_ranges(values.len(), starts, lengths, output_len)?;
212-
213221
let mut result = BufferMut::<T>::with_capacity(output_len);
222+
let mut computed_len = 0usize;
214223
for (&start, &length) in starts.iter().zip_eq(lengths) {
215224
let start = start.as_();
216225
let length = length.as_();
217-
result.extend_from_slice(&values[start..start + length]);
226+
let end = start
227+
.checked_add(length)
228+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
229+
computed_len = computed_len
230+
.checked_add(length)
231+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
232+
result.extend_from_slice(&values[start..end]);
218233
}
219234

235+
vortex_ensure!(
236+
computed_len == output_len,
237+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
238+
);
220239
Ok(result.freeze())
221240
}
222241

‎vortex-array/src/arrays/piecewise_sequence/mod.rs‎

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -134,38 +134,6 @@ fn check_index_array(name: &str, array: &ArrayRef) -> VortexResult<()> {
134134
Ok(())
135135
}
136136

137-
pub(crate) fn validate_index_ranges_constant<S>(
138-
source_len: usize,
139-
starts: &[S],
140-
length: usize,
141-
output_len: usize,
142-
) -> VortexResult<()>
143-
where
144-
S: UnsignedPType,
145-
{
146-
let computed_len = starts
147-
.len()
148-
.checked_mul(length)
149-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
150-
vortex_ensure!(
151-
computed_len == output_len,
152-
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
153-
);
154-
155-
for &start in starts {
156-
let start: usize = start.as_();
157-
let end = start
158-
.checked_add(length)
159-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
160-
vortex_ensure!(
161-
end <= source_len,
162-
"PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}"
163-
);
164-
}
165-
166-
Ok(())
167-
}
168-
169137
pub(crate) fn materialize_ranges<S, L, M>(
170138
starts: &PrimitiveArray,
171139
lengths: &PrimitiveArray,
@@ -210,36 +178,3 @@ where
210178
}
211179
Ok(values)
212180
}
213-
214-
pub(crate) fn validate_index_ranges<S, L>(
215-
source_len: usize,
216-
starts: &[S],
217-
lengths: &[L],
218-
output_len: usize,
219-
) -> VortexResult<()>
220-
where
221-
S: UnsignedPType,
222-
L: UnsignedPType,
223-
{
224-
let mut computed_len = 0usize;
225-
for (&start, &length) in starts.iter().zip_eq(lengths) {
226-
let start: usize = start.as_();
227-
let length: usize = length.as_();
228-
let end = start
229-
.checked_add(length)
230-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
231-
vortex_ensure!(
232-
end <= source_len,
233-
"PiecewiseSequenceArray range {start}..{end} exceeds source length {source_len}"
234-
);
235-
computed_len = computed_len
236-
.checked_add(length)
237-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
238-
}
239-
240-
vortex_ensure!(
241-
computed_len == output_len,
242-
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
243-
);
244-
Ok(())
245-
}

‎vortex-array/src/arrays/primitive/compute/take/mod.rs‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ use vortex_buffer::Buffer;
1111
use vortex_buffer::BufferMut;
1212
use vortex_error::VortexResult;
1313
use vortex_error::vortex_bail;
14+
use vortex_error::vortex_ensure;
15+
use vortex_error::vortex_err;
1416
use vortex_mask::Mask;
1517

1618
use crate::ArrayRef;
@@ -23,8 +25,6 @@ use crate::arrays::PrimitiveArray;
2325
use crate::arrays::dict::TakeExecute;
2426
use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
2527
use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
26-
use crate::arrays::piecewise_sequence::validate_index_ranges;
27-
use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
2828
use crate::builtins::ArrayBuiltins;
2929
use crate::dtype::DType;
3030
use crate::dtype::IntegerPType;
@@ -241,15 +241,24 @@ where
241241
S: UnsignedPType,
242242
L: UnsignedPType,
243243
{
244-
validate_index_ranges(source.len(), starts, lengths, output_len)?;
245-
246244
let mut values = BufferMut::<T>::with_capacity(output_len);
245+
let mut computed_len = 0usize;
247246
for (&start, &length) in starts.iter().zip_eq(lengths) {
248247
let start = start.as_();
249248
let length = length.as_();
250-
values.extend_from_slice(&source[start..start + length]);
249+
let end = start
250+
.checked_add(length)
251+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
252+
computed_len = computed_len
253+
.checked_add(length)
254+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
255+
values.extend_from_slice(&source[start..end]);
251256
}
252257

258+
vortex_ensure!(
259+
computed_len == output_len,
260+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
261+
);
253262
Ok(values.freeze())
254263
}
255264

@@ -296,12 +305,22 @@ where
296305
T: Copy,
297306
S: UnsignedPType,
298307
{
299-
validate_index_ranges_constant(source.len(), starts, length, output_len)?;
308+
let computed_len = starts
309+
.len()
310+
.checked_mul(length)
311+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
312+
vortex_ensure!(
313+
computed_len == output_len,
314+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
315+
);
300316

301317
let mut values = BufferMut::<T>::with_capacity(output_len);
302318
for &start in starts {
303319
let start = start.as_();
304-
values.extend_from_slice(&source[start..start + length]);
320+
let end = start
321+
.checked_add(length)
322+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
323+
values.extend_from_slice(&source[start..end]);
305324
}
306325

307326
Ok(values.freeze())

‎vortex-array/src/arrays/varbin/compute/take.rs‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ use crate::arrays::VarBinArray;
2222
use crate::arrays::dict::TakeExecute;
2323
use crate::arrays::piecewise_sequence::UnitMultiplierLengths;
2424
use crate::arrays::piecewise_sequence::execute_unit_multiplier_index_arrays;
25-
use crate::arrays::piecewise_sequence::validate_index_ranges;
26-
use crate::arrays::piecewise_sequence::validate_index_ranges_constant;
2725
use crate::arrays::primitive::PrimitiveArrayExt;
2826
use crate::arrays::varbin::VarBinArrayExt;
2927
use crate::dtype::DType;
@@ -424,15 +422,24 @@ where
424422
Offset: IntegerPType,
425423
NewOffset: IntegerPType,
426424
{
427-
validate_index_ranges_constant(offsets.len() - 1, starts, length, output_len)?;
425+
let computed_len = starts
426+
.len()
427+
.checked_mul(length)
428+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
429+
vortex_ensure!(
430+
computed_len == output_len,
431+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
432+
);
428433

429434
let mut new_offsets = BufferMut::<NewOffset>::with_capacity(output_len + 1);
430435
new_offsets.push(NewOffset::zero());
431436
let mut output_bytes = 0usize;
432437

433438
for &start in starts {
434439
let start = start.as_();
435-
let end = start + length;
440+
let end = start
441+
.checked_add(length)
442+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
436443
if length == 0 {
437444
continue;
438445
}
@@ -464,7 +471,9 @@ where
464471
let mut new_data = ByteBufferMut::with_capacity(output_bytes);
465472
for &start in starts {
466473
let start = start.as_();
467-
let end = start + length;
474+
let end = start
475+
.checked_add(length)
476+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
468477
if length == 0 {
469478
continue;
470479
}
@@ -497,16 +506,20 @@ where
497506
Offset: IntegerPType,
498507
NewOffset: IntegerPType,
499508
{
500-
validate_index_ranges(offsets.len() - 1, starts, lengths, output_len)?;
501-
502509
let mut new_offsets = BufferMut::<NewOffset>::with_capacity(output_len + 1);
503510
new_offsets.push(NewOffset::zero());
504511
let mut output_bytes = 0usize;
512+
let mut computed_len = 0usize;
505513

506514
for (&start, &length) in starts.iter().zip_eq(lengths) {
507515
let start = start.as_();
508516
let length = length.as_();
509-
let end = start + length;
517+
let end = start
518+
.checked_add(length)
519+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
520+
computed_len = computed_len
521+
.checked_add(length)
522+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
510523
if length == 0 {
511524
continue;
512525
}
@@ -534,12 +547,18 @@ where
534547
.checked_add(byte_end - byte_start)
535548
.ok_or_else(|| vortex_err!("PiecewiseSequence VarBin output byte length overflow"))?;
536549
}
550+
vortex_ensure!(
551+
computed_len == output_len,
552+
"PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
553+
);
537554

538555
let mut new_data = ByteBufferMut::with_capacity(output_bytes);
539556
for (&start, &length) in starts.iter().zip_eq(lengths) {
540557
let start = start.as_();
541558
let length = length.as_();
542-
let end = start + length;
559+
let end = start
560+
.checked_add(length)
561+
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
543562
if length == 0 {
544563
continue;
545564
}

0 commit comments

Comments
 (0)