Skip to content

Commit 66a5896

Browse files
dankingdanking
committed
Simplify PiecewiseSequence take range handling
Signed-off-by: Daniel King <dan@spiraldb.com>
1 parent 815e475 commit 66a5896

6 files changed

Lines changed: 29 additions & 73 deletions

File tree

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -158,10 +158,7 @@ where
158158
let mut values = BitBufferMut::with_capacity(output_len);
159159
for &start in starts {
160160
let start = start.as_();
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));
161+
values.append_buffer(&source.slice(start..).slice(..length));
165162
}
166163

167164
Ok(values.freeze())
@@ -182,13 +179,10 @@ where
182179
for (&start, &length) in starts.iter().zip_eq(lengths) {
183180
let start = start.as_();
184181
let length = length.as_();
185-
let end = start
186-
.checked_add(length)
187-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
188182
computed_len = computed_len
189183
.checked_add(length)
190184
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
191-
values.append_buffer(&source.slice(start..end));
185+
values.append_buffer(&source.slice(start..).slice(..length));
192186
}
193187

194188
vortex_ensure!(

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -198,10 +198,7 @@ where
198198
let mut result = BufferMut::<T>::with_capacity(output_len);
199199
for &start in starts {
200200
let start = start.as_();
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]);
201+
result.extend_from_slice(&values[start..][..length]);
205202
}
206203

207204
Ok(result.freeze())
@@ -223,13 +220,10 @@ where
223220
for (&start, &length) in starts.iter().zip_eq(lengths) {
224221
let start = start.as_();
225222
let length = length.as_();
226-
let end = start
227-
.checked_add(length)
228-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
229223
computed_len = computed_len
230224
.checked_add(length)
231225
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
232-
result.extend_from_slice(&values[start..end]);
226+
result.extend_from_slice(&values[start..][..length]);
233227
}
234228

235229
vortex_ensure!(

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

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ use crate::arrays::primitive::PrimitiveArrayExt;
2424
use crate::builders::builder_with_capacity;
2525
use crate::dtype::DType;
2626
use crate::dtype::IntegerPType;
27-
use crate::dtype::Nullability;
2827
use crate::executor::ExecutionCtx;
2928
use crate::match_each_integer_ptype;
3029
use crate::optimizer::ArrayOptimizer;
@@ -94,7 +93,7 @@ fn take_non_empty_fsl(
9493
) -> VortexResult<ArrayRef> {
9594
debug_assert!(!array.is_empty());
9695

97-
let DType::Primitive(ptype, nullability) = indices.dtype() else {
96+
let DType::Primitive(ptype, _) = indices.dtype() else {
9897
vortex_bail!("Invalid indices dtype: {}", indices.dtype())
9998
};
10099
if !ptype.is_int() {
@@ -107,13 +106,7 @@ fn take_non_empty_fsl(
107106

108107
let indices_array = indices.clone().execute::<PrimitiveArray>(ctx)?;
109108
match_each_integer_ptype!(indices_array.ptype(), |I| {
110-
take_non_empty_non_degenerate_fsl::<I>(
111-
array,
112-
indices,
113-
indices_array.as_view(),
114-
*nullability,
115-
ctx,
116-
)
109+
take_non_empty_non_degenerate_fsl::<I>(array, indices, indices_array.as_view(), ctx)
117110
})
118111
}
119112

@@ -154,22 +147,17 @@ fn take_non_empty_non_degenerate_fsl<I: IntegerPType>(
154147
array: ArrayView<'_, FixedSizeList>,
155148
indices: &ArrayRef,
156149
indices_array: ArrayView<'_, Primitive>,
157-
indices_nullability: Nullability,
158150
ctx: &mut ExecutionCtx,
159151
) -> VortexResult<ArrayRef> {
160152
debug_assert!(!array.is_empty());
161153
debug_assert_ne!(array.list_size(), 0);
162154

163155
let (new_elements, new_len) =
164156
take_non_empty_non_degenerate_elements::<I>(array, indices_array, ctx)?;
165-
let new_validity = if array.dtype().is_nullable() || indices_nullability.is_nullable() {
166-
array.validity()?.take(indices)?
167-
} else {
168-
Validity::NonNullable
169-
};
157+
let new_validity = array.validity()?.take(indices)?;
170158

171-
// SAFETY: `new_elements` has `new_len * list_size` elements. `new_validity` is either
172-
// non-nullable or was produced by `Validity::take` for `new_len`.
159+
// SAFETY: `new_elements` has `new_len * list_size` elements, and `Validity::take` produces
160+
// validity for `new_len`.
173161
unsafe {
174162
FixedSizeListArray::new_unchecked(new_elements, array.list_size(), new_validity, new_len)
175163
}

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -246,13 +246,10 @@ where
246246
for (&start, &length) in starts.iter().zip_eq(lengths) {
247247
let start = start.as_();
248248
let length = length.as_();
249-
let end = start
250-
.checked_add(length)
251-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
252249
computed_len = computed_len
253250
.checked_add(length)
254251
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
255-
values.extend_from_slice(&source[start..end]);
252+
values.extend_from_slice(&source[start..][..length]);
256253
}
257254

258255
vortex_ensure!(
@@ -317,10 +314,7 @@ where
317314
let mut values = BufferMut::<T>::with_capacity(output_len);
318315
for &start in starts {
319316
let start = start.as_();
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]);
317+
values.extend_from_slice(&source[start..][..length]);
324318
}
325319

326320
Ok(values.freeze())

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

Lines changed: 16 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -437,22 +437,20 @@ where
437437

438438
for &start in starts {
439439
let start = start.as_();
440-
let end = start
441-
.checked_add(length)
442-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
443440
if length == 0 {
444441
continue;
445442
}
446443

447-
let byte_start = offsets[start].as_();
448-
let byte_end = offsets[end].as_();
444+
let offset_range = &offsets[start..][..=length];
445+
let byte_start = offset_range[0].as_();
446+
let byte_end = offset_range[length].as_();
449447
vortex_ensure!(
450448
byte_start <= byte_end && byte_end <= data.len(),
451449
"VarBin offsets range {byte_start}..{byte_end} exceeds data length {}",
452450
data.len()
453451
);
454452

455-
for &offset in &offsets[start + 1..=end] {
453+
for &offset in &offset_range[1..] {
456454
let offset = offset.as_();
457455
let relative = offset.checked_sub(byte_start).ok_or_else(|| {
458456
vortex_err!("VarBin offsets are not monotonic at offset {offset}")
@@ -471,16 +469,14 @@ where
471469
let mut new_data = ByteBufferMut::with_capacity(output_bytes);
472470
for &start in starts {
473471
let start = start.as_();
474-
let end = start
475-
.checked_add(length)
476-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
477472
if length == 0 {
478473
continue;
479474
}
480475

481-
let byte_start = offsets[start].as_();
482-
let byte_end = offsets[end].as_();
483-
new_data.extend_from_slice(&data[byte_start..byte_end]);
476+
let offset_range = &offsets[start..][..=length];
477+
let byte_start = offset_range[0].as_();
478+
let byte_end = offset_range[length].as_();
479+
new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]);
484480
}
485481

486482
let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
@@ -514,25 +510,23 @@ where
514510
for (&start, &length) in starts.iter().zip_eq(lengths) {
515511
let start = start.as_();
516512
let length = length.as_();
517-
let end = start
518-
.checked_add(length)
519-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
520513
computed_len = computed_len
521514
.checked_add(length)
522515
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
523516
if length == 0 {
524517
continue;
525518
}
526519

527-
let byte_start = offsets[start].as_();
528-
let byte_end = offsets[end].as_();
520+
let offset_range = &offsets[start..][..=length];
521+
let byte_start = offset_range[0].as_();
522+
let byte_end = offset_range[length].as_();
529523
vortex_ensure!(
530524
byte_start <= byte_end && byte_end <= data.len(),
531525
"VarBin offsets range {byte_start}..{byte_end} exceeds data length {}",
532526
data.len()
533527
);
534528

535-
for &offset in &offsets[start + 1..=end] {
529+
for &offset in &offset_range[1..] {
536530
let offset = offset.as_();
537531
let relative = offset.checked_sub(byte_start).ok_or_else(|| {
538532
vortex_err!("VarBin offsets are not monotonic at offset {offset}")
@@ -556,16 +550,14 @@ where
556550
for (&start, &length) in starts.iter().zip_eq(lengths) {
557551
let start = start.as_();
558552
let length = length.as_();
559-
let end = start
560-
.checked_add(length)
561-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
562553
if length == 0 {
563554
continue;
564555
}
565556

566-
let byte_start = offsets[start].as_();
567-
let byte_end = offsets[end].as_();
568-
new_data.extend_from_slice(&data[byte_start..byte_end]);
557+
let offset_range = &offsets[start..][..=length];
558+
let byte_start = offset_range[0].as_();
559+
let byte_end = offset_range[length].as_();
560+
new_data.extend_from_slice(&data[byte_start..][..byte_end - byte_start]);
569561
}
570562

571563
let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -173,10 +173,7 @@ where
173173
let mut views = BufferMut::<BinaryView>::with_capacity(output_len);
174174
for &start in starts {
175175
let start = start.as_();
176-
let end = start
177-
.checked_add(length)
178-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
179-
views.extend_from_slice(&source[start..end]);
176+
views.extend_from_slice(&source[start..][..length]);
180177
}
181178

182179
Ok(views.freeze())
@@ -197,13 +194,10 @@ where
197194
for (&start, &length) in starts.iter().zip_eq(lengths) {
198195
let start = start.as_();
199196
let length = length.as_();
200-
let end = start
201-
.checked_add(length)
202-
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
203197
computed_len = computed_len
204198
.checked_add(length)
205199
.ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
206-
views.extend_from_slice(&source[start..end]);
200+
views.extend_from_slice(&source[start..][..length]);
207201
}
208202

209203
vortex_ensure!(

0 commit comments

Comments
 (0)