• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

vortex-data / vortex / 16001746832

01 Jul 2025 02:09PM UTC coverage: 77.952% (+0.03%) from 77.923%
16001746832

push

github

web-flow
fix: SparseArray handles take with null indices if the fill value is not null (#3677)

Signed-off-by: Robert Kruszewski <github@robertk.io>

117 of 123 new or added lines in 2 files covered. (95.12%)

5 existing lines in 1 file now uncovered.

42908 of 55044 relevant lines covered (77.95%)

56213.14 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

82.15
/vortex-array/src/patches.rs
1
use std::cmp::Ordering;
2
use std::fmt::Debug;
3
use std::hash::Hash;
4

5
use itertools::Itertools as _;
6
use num_traits::{NumCast, ToPrimitive};
7
use serde::{Deserialize, Serialize};
8
use vortex_buffer::BufferMut;
9
use vortex_dtype::Nullability::NonNullable;
10
use vortex_dtype::{
11
    DType, NativePType, PType, match_each_integer_ptype, match_each_unsigned_integer_ptype,
12
};
13
use vortex_error::{
14
    VortexError, VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_err,
15
};
16
use vortex_mask::{AllOr, Mask};
17
use vortex_scalar::{PValue, Scalar};
18
use vortex_utils::aliases::hash_map::HashMap;
19

20
use crate::arrays::PrimitiveArray;
21
use crate::compute::{cast, filter, take};
22
use crate::search_sorted::{SearchResult, SearchSorted, SearchSortedSide};
23
use crate::vtable::ValidityHelper;
24
use crate::{Array, ArrayRef, IntoArray, ToCanonical};
25

26
#[derive(Copy, Clone, Serialize, Deserialize, prost::Message)]
148✔
27
pub struct PatchesMetadata {
28
    #[prost(uint64, tag = "1")]
29
    len: u64,
30
    #[prost(uint64, tag = "2")]
31
    offset: u64,
32
    #[prost(enumeration = "PType", tag = "3")]
33
    indices_ptype: i32,
34
}
35

36
impl PatchesMetadata {
37
    pub fn new(len: usize, offset: usize, indices_ptype: PType) -> Self {
74✔
38
        Self {
74✔
39
            len: len as u64,
74✔
40
            offset: offset as u64,
74✔
41
            indices_ptype: indices_ptype as i32,
74✔
42
        }
74✔
43
    }
74✔
44

45
    #[inline]
46
    pub fn len(&self) -> usize {
×
47
        usize::try_from(self.len).vortex_expect("len is a valid usize")
×
48
    }
×
49

50
    #[inline]
51
    pub fn is_empty(&self) -> bool {
×
52
        self.len == 0
×
53
    }
×
54

55
    #[inline]
56
    pub fn offset(&self) -> usize {
×
57
        usize::try_from(self.offset).vortex_expect("offset is a valid usize")
×
58
    }
×
59

60
    #[inline]
61
    pub fn indices_dtype(&self) -> DType {
×
62
        assert!(
×
63
            self.indices_ptype().is_unsigned_int(),
×
64
            "Patch indices must be unsigned integers"
×
65
        );
66
        DType::Primitive(self.indices_ptype(), NonNullable)
×
67
    }
×
68
}
69

70
/// A helper for working with patched arrays.
71
#[derive(Debug, Clone)]
72
pub struct Patches {
73
    array_len: usize,
74
    offset: usize,
75
    indices: ArrayRef,
76
    values: ArrayRef,
77
}
78

79
impl Patches {
80
    pub fn new(array_len: usize, offset: usize, indices: ArrayRef, values: ArrayRef) -> Self {
6,370✔
81
        assert_eq!(
6,370✔
82
            indices.len(),
6,370✔
83
            values.len(),
6,370✔
84
            "Patch indices and values must have the same length"
×
85
        );
86
        assert!(
6,370✔
87
            indices.dtype().is_unsigned_int(),
6,370✔
88
            "Patch indices must be unsigned integers"
×
89
        );
90
        assert!(
6,370✔
91
            indices.len() <= array_len,
6,370✔
92
            "Patch indices must be shorter than the array length"
×
93
        );
94
        assert!(!indices.is_empty(), "Patch indices must not be empty");
6,370✔
95
        let max = usize::try_from(
6,370✔
96
            &indices
6,370✔
97
                .scalar_at(indices.len() - 1)
6,370✔
98
                .vortex_expect("indices are not empty"),
6,370✔
99
        )
6,370✔
100
        .vortex_expect("indices must be a number");
6,370✔
101
        assert!(
6,370✔
102
            max - offset < array_len,
6,370✔
103
            "Patch indices {max:?}, offset {offset} are longer than the array length {array_len}"
×
104
        );
105
        Self::new_unchecked(array_len, offset, indices, values)
6,370✔
106
    }
6,370✔
107

108
    /// Construct new patches without validating any of the arguments
109
    ///
110
    /// # Safety
111
    ///
112
    /// Users have to assert that
113
    /// * Indices and values have the same length
114
    /// * Indices is an unsigned integer type
115
    /// * Indices must be sorted
116
    /// * Last value in indices is smaller than array_len
117
    pub fn new_unchecked(
7,998✔
118
        array_len: usize,
7,998✔
119
        offset: usize,
7,998✔
120
        indices: ArrayRef,
7,998✔
121
        values: ArrayRef,
7,998✔
122
    ) -> Self {
7,998✔
123
        Self {
7,998✔
124
            array_len,
7,998✔
125
            offset,
7,998✔
126
            indices,
7,998✔
127
            values,
7,998✔
128
        }
7,998✔
129
    }
7,998✔
130

131
    pub fn array_len(&self) -> usize {
32,304✔
132
        self.array_len
32,304✔
133
    }
32,304✔
134

135
    pub fn num_patches(&self) -> usize {
1,296✔
136
        self.indices.len()
1,296✔
137
    }
1,296✔
138

139
    pub fn dtype(&self) -> &DType {
1,739✔
140
        self.values.dtype()
1,739✔
141
    }
1,739✔
142

143
    pub fn indices(&self) -> &ArrayRef {
8,126✔
144
        &self.indices
8,126✔
145
    }
8,126✔
146

147
    pub fn into_indices(self) -> ArrayRef {
×
148
        self.indices
×
149
    }
×
150

151
    pub fn indices_mut(&mut self) -> &mut ArrayRef {
×
152
        &mut self.indices
×
153
    }
×
154

155
    pub fn values(&self) -> &ArrayRef {
18,004✔
156
        &self.values
18,004✔
157
    }
18,004✔
158

159
    pub fn into_values(self) -> ArrayRef {
37✔
160
        self.values
37✔
161
    }
37✔
162

163
    pub fn values_mut(&mut self) -> &mut ArrayRef {
×
164
        &mut self.values
×
165
    }
×
166

167
    pub fn offset(&self) -> usize {
83,143✔
168
        self.offset
83,143✔
169
    }
83,143✔
170

UNCOV
171
    pub fn indices_ptype(&self) -> PType {
×
UNCOV
172
        PType::try_from(self.indices.dtype()).vortex_expect("primitive indices")
×
UNCOV
173
    }
×
174

175
    pub fn to_metadata(&self, len: usize, dtype: &DType) -> VortexResult<PatchesMetadata> {
74✔
176
        if self.indices.len() > len {
74✔
177
            vortex_bail!(
×
178
                "Patch indices {} are longer than the array length {}",
×
179
                self.indices.len(),
×
180
                len
×
181
            );
×
182
        }
74✔
183
        if self.values.dtype() != dtype {
74✔
184
            vortex_bail!(
×
185
                "Patch values dtype {} does not match array dtype {}",
×
186
                self.values.dtype(),
×
187
                dtype
×
188
            );
×
189
        }
74✔
190
        Ok(PatchesMetadata {
74✔
191
            len: self.indices.len() as u64,
74✔
192
            offset: self.offset as u64,
74✔
193
            indices_ptype: PType::try_from(self.indices.dtype()).vortex_expect("primitive indices")
74✔
194
                as i32,
74✔
195
        })
74✔
196
    }
74✔
197

198
    pub fn cast_values(self, values_dtype: &DType) -> VortexResult<Self> {
1,628✔
199
        Ok(Self::new_unchecked(
1,628✔
200
            self.array_len,
1,628✔
201
            self.offset,
1,628✔
202
            self.indices,
1,628✔
203
            cast(&self.values, values_dtype)?,
1,628✔
204
        ))
205
    }
1,628✔
206

207
    /// Get the patched value at a given index if it exists.
208
    pub fn get_patched(&self, index: usize) -> VortexResult<Option<Scalar>> {
394,198✔
209
        if let Some(patch_idx) = self.search_index(index)?.to_found() {
394,198✔
210
            self.values().scalar_at(patch_idx).map(Some)
6,253✔
211
        } else {
212
            Ok(None)
387,945✔
213
        }
214
    }
394,198✔
215

216
    /// Return the insertion point of `index` in the [Self::indices].
217
    pub fn search_index(&self, index: usize) -> VortexResult<SearchResult> {
395,094✔
218
        Ok(self.indices.as_primitive_typed().search_sorted(
395,094✔
219
            &PValue::U64((index + self.offset) as u64),
395,094✔
220
            SearchSortedSide::Left,
395,094✔
221
        ))
395,094✔
222
    }
395,094✔
223

224
    /// Return the search_sorted result for the given target re-mapped into the original indices.
225
    pub fn search_sorted<T: Into<Scalar>>(
11✔
226
        &self,
11✔
227
        target: T,
11✔
228
        side: SearchSortedSide,
11✔
229
    ) -> VortexResult<SearchResult> {
11✔
230
        let target = target.into();
11✔
231

232
        let sr = if self.values().dtype().is_primitive() {
11✔
233
            self.values()
11✔
234
                .as_primitive_typed()
11✔
235
                .search_sorted(&target.as_primitive().pvalue(), side)
11✔
236
        } else {
237
            self.values().search_sorted(&target, side)
×
238
        };
239

240
        let index_idx = sr.to_offsets_index(self.indices().len(), side);
11✔
241
        let index = usize::try_from(&self.indices().scalar_at(index_idx)?)? - self.offset;
11✔
242
        Ok(match sr {
11✔
243
            // If we reached the end of patched values when searching then the result is one after the last patch index
244
            SearchResult::Found(i) => SearchResult::Found(
3✔
245
                if i == self.indices().len() || side == SearchSortedSide::Right {
3✔
246
                    index + 1
2✔
247
                } else {
248
                    index
1✔
249
                },
250
            ),
251
            // If the result is NotFound we should return index that's one after the nearest not found index for the corresponding value
252
            SearchResult::NotFound(i) => {
8✔
253
                SearchResult::NotFound(if i == 0 { index } else { index + 1 })
8✔
254
            }
255
        })
256
    }
11✔
257

258
    /// Returns the minimum patch index
259
    pub fn min_index(&self) -> VortexResult<usize> {
408✔
260
        Ok(usize::try_from(&self.indices().scalar_at(0)?)? - self.offset)
408✔
261
    }
408✔
262

263
    /// Returns the maximum patch index
264
    pub fn max_index(&self) -> VortexResult<usize> {
408✔
265
        Ok(usize::try_from(&self.indices().scalar_at(self.indices().len() - 1)?)? - self.offset)
408✔
266
    }
408✔
267

268
    /// Filter the patches by a mask, resulting in new patches for the filtered array.
269
    pub fn filter(&self, mask: &Mask) -> VortexResult<Option<Self>> {
149✔
270
        match mask.indices() {
149✔
271
            AllOr::All => Ok(Some(self.clone())),
×
272
            AllOr::None => Ok(None),
×
273
            AllOr::Some(mask_indices) => {
149✔
274
                let flat_indices = self.indices().to_primitive()?;
149✔
275
                match_each_unsigned_integer_ptype!(flat_indices.ptype(), |I| {
149✔
276
                    filter_patches_with_mask(
×
277
                        flat_indices.as_slice::<I>(),
×
278
                        self.offset(),
×
279
                        self.values(),
×
280
                        mask_indices,
×
281
                    )
×
282
                })
283
            }
284
        }
285
    }
149✔
286

287
    /// Slice the patches by a range of the patched array.
288
    pub fn slice(&self, start: usize, stop: usize) -> VortexResult<Option<Self>> {
448✔
289
        let patch_start = self.search_index(start)?.to_index();
448✔
290
        let patch_stop = self.search_index(stop)?.to_index();
448✔
291

448✔
292
        if patch_start == patch_stop {
448✔
293
            return Ok(None);
37✔
294
        }
411✔
295

296
        // Slice out the values and indices
297
        let values = self.values().slice(patch_start, patch_stop)?;
411✔
298
        let indices = self.indices().slice(patch_start, patch_stop)?;
411✔
299

300
        Ok(Some(Self::new(
411✔
301
            stop - start,
411✔
302
            start + self.offset(),
411✔
303
            indices,
411✔
304
            values,
411✔
305
        )))
411✔
306
    }
448✔
307

308
    // https://docs.google.com/spreadsheets/d/1D9vBZ1QJ6mwcIvV5wIL0hjGgVchcEnAyhvitqWu2ugU
309
    const PREFER_MAP_WHEN_PATCHES_OVER_INDICES_LESS_THAN: f64 = 5.0;
310

311
    fn is_map_faster_than_search(&self, take_indices: &PrimitiveArray) -> bool {
482✔
312
        (self.num_patches() as f64 / take_indices.len() as f64)
482✔
313
            < Self::PREFER_MAP_WHEN_PATCHES_OVER_INDICES_LESS_THAN
482✔
314
    }
482✔
315

316
    /// Take the indicies from the patches
317
    ///
318
    /// Any nulls in take_indices are added to the resulting patches.
319
    pub fn take_with_nulls(&self, take_indices: &dyn Array) -> VortexResult<Option<Self>> {
74✔
320
        if take_indices.is_empty() {
74✔
NEW
321
            return Ok(None);
×
322
        }
74✔
323

324
        let take_indices = take_indices.to_primitive()?;
74✔
325
        if self.is_map_faster_than_search(&take_indices) {
74✔
326
            self.take_map(take_indices, true)
74✔
327
        } else {
NEW
328
            self.take_search(take_indices, true)
×
329
        }
330
    }
74✔
331

332
    /// Take the indices from the patches.
333
    ///
334
    /// Any nulls in take_indices are ignored.
335
    pub fn take(&self, take_indices: &dyn Array) -> VortexResult<Option<Self>> {
408✔
336
        if take_indices.is_empty() {
408✔
337
            return Ok(None);
×
338
        }
408✔
339

340
        let take_indices = take_indices.to_primitive()?;
408✔
341
        if self.is_map_faster_than_search(&take_indices) {
408✔
342
            self.take_map(take_indices, false)
334✔
343
        } else {
344
            self.take_search(take_indices, false)
74✔
345
        }
346
    }
408✔
347

348
    pub fn take_search(
74✔
349
        &self,
74✔
350
        take_indices: PrimitiveArray,
74✔
351
        include_nulls: bool,
74✔
352
    ) -> VortexResult<Option<Self>> {
74✔
353
        let indices = self.indices.to_primitive()?;
74✔
354
        let new_length = take_indices.len();
74✔
355

356
        let Some((new_indices, values_indices)) =
74✔
357
            match_each_unsigned_integer_ptype!(indices.ptype(), |Indices| {
74✔
358
                match_each_integer_ptype!(take_indices.ptype(), |TakeIndices| {
×
359
                    take_search::<_, TakeIndices>(
×
360
                        indices.as_slice::<Indices>(),
×
361
                        take_indices,
×
362
                        self.offset(),
×
NEW
363
                        include_nulls,
×
UNCOV
364
                    )?
×
365
                })
366
            })
367
        else {
368
            return Ok(None);
×
369
        };
370

371
        Ok(Some(Self::new(
372
            new_length,
74✔
373
            0,
74✔
374
            new_indices,
74✔
375
            take(self.values(), &values_indices)?,
74✔
376
        )))
377
    }
74✔
378

379
    pub fn take_map(
408✔
380
        &self,
408✔
381
        take_indices: PrimitiveArray,
408✔
382
        include_nulls: bool,
408✔
383
    ) -> VortexResult<Option<Self>> {
408✔
384
        let indices = self.indices.to_primitive()?;
408✔
385
        let new_length = take_indices.len();
408✔
386

387
        let Some((new_sparse_indices, value_indices)) =
371✔
388
            match_each_unsigned_integer_ptype!(indices.ptype(), |Indices| {
408✔
389
                match_each_integer_ptype!(take_indices.ptype(), |TakeIndices| {
37✔
390
                    take_map::<_, TakeIndices>(
391
                        indices.as_slice::<Indices>(),
×
392
                        take_indices,
×
393
                        self.offset(),
×
394
                        self.min_index()?,
×
395
                        self.max_index()?,
×
NEW
396
                        include_nulls,
×
UNCOV
397
                    )?
×
398
                })
399
            })
400
        else {
401
            return Ok(None);
37✔
402
        };
403

404
        Ok(Some(Patches::new(
405
            new_length,
371✔
406
            0,
371✔
407
            new_sparse_indices,
371✔
408
            take(self.values(), &value_indices)?,
371✔
409
        )))
410
    }
408✔
411

412
    pub fn map_values<F>(self, f: F) -> VortexResult<Self>
16✔
413
    where
16✔
414
        F: FnOnce(ArrayRef) -> VortexResult<ArrayRef>,
16✔
415
    {
16✔
416
        let values = f(self.values)?;
16✔
417
        if self.indices.len() != values.len() {
16✔
418
            vortex_bail!(
×
419
                "map_values must preserve length: expected {} received {}",
×
420
                self.indices.len(),
×
421
                values.len()
×
422
            )
×
423
        }
16✔
424
        Ok(Self::new(self.array_len, self.offset, self.indices, values))
16✔
425
    }
16✔
426
}
427

428
fn take_search<I: NativePType + NumCast + PartialOrd, T: NativePType + NumCast>(
74✔
429
    indices: &[I],
74✔
430
    take_indices: PrimitiveArray,
74✔
431
    indices_offset: usize,
74✔
432
    include_nulls: bool,
74✔
433
) -> VortexResult<Option<(ArrayRef, ArrayRef)>>
74✔
434
where
74✔
435
    usize: TryFrom<T>,
74✔
436
    VortexError: From<<usize as TryFrom<T>>::Error>,
74✔
437
{
74✔
438
    let take_indices_validity = take_indices.validity();
74✔
439
    let indices_offset = I::from(indices_offset).vortex_expect("indices_offset out of range");
74✔
440

74✔
441
    let (values_indices, new_indices): (BufferMut<u64>, BufferMut<u64>) = take_indices
74✔
442
        .as_slice::<T>()
74✔
443
        .iter()
74✔
444
        .enumerate()
74✔
445
        .filter_map(|(i, &v)| {
259✔
446
            I::from(v)
259✔
447
                .and_then(|v| {
259✔
448
                    // If we have to take nulls the take index doesn't matter, make it 0 for consistency
259✔
449
                    if include_nulls && take_indices_validity.is_null(i).vortex_unwrap() {
259✔
NEW
450
                        Some(0)
×
451
                    } else {
452
                        indices
259✔
453
                            .search_sorted(&(v + indices_offset), SearchSortedSide::Left)
259✔
454
                            .to_found()
259✔
455
                            .map(|patch_idx| patch_idx as u64)
259✔
456
                    }
457
                })
259✔
458
                .map(|patch_idx| (patch_idx, i as u64))
259✔
459
        })
259✔
460
        .unzip();
74✔
461

74✔
462
    if new_indices.is_empty() {
74✔
463
        return Ok(None);
×
464
    }
74✔
465

74✔
466
    let new_indices = new_indices.into_array();
74✔
467
    let values_validity = take_indices_validity.take(&new_indices)?;
74✔
468
    Ok(Some((
74✔
469
        new_indices,
74✔
470
        PrimitiveArray::new(values_indices, values_validity).into_array(),
74✔
471
    )))
74✔
472
}
74✔
473

474
fn take_map<I: NativePType + Hash + Eq + TryFrom<usize>, T: NativePType>(
408✔
475
    indices: &[I],
408✔
476
    take_indices: PrimitiveArray,
408✔
477
    indices_offset: usize,
408✔
478
    min_index: usize,
408✔
479
    max_index: usize,
408✔
480
    include_nulls: bool,
408✔
481
) -> VortexResult<Option<(ArrayRef, ArrayRef)>>
408✔
482
where
408✔
483
    usize: TryFrom<T>,
408✔
484
    VortexError: From<<I as TryFrom<usize>>::Error>,
408✔
485
{
408✔
486
    let take_indices_validity = take_indices.validity();
408✔
487
    let take_indices = take_indices.as_slice::<T>();
408✔
488
    let offset_i = I::try_from(indices_offset)?;
408✔
489

490
    let sparse_index_to_value_index: HashMap<I, usize> = indices
408✔
491
        .iter()
408✔
492
        .copied()
408✔
493
        .map(|idx| idx - offset_i)
965✔
494
        .enumerate()
408✔
495
        .map(|(value_index, sparse_index)| (sparse_index, value_index))
965✔
496
        .collect();
408✔
497

498
    let (new_sparse_indices, value_indices): (BufferMut<u64>, BufferMut<u64>) = take_indices
408✔
499
        .iter()
408✔
500
        .copied()
408✔
501
        .map(usize::try_from)
408✔
502
        .process_results(|iter| {
408✔
503
            iter.enumerate()
408✔
504
                .filter_map(|(idx_in_take, ti)| {
1,149✔
505
                    // If we have to take nulls the take index doesn't matter, make it 0 for consistency
1,149✔
506
                    if include_nulls && take_indices_validity.is_null(idx_in_take).vortex_unwrap() {
1,149✔
507
                        Some((idx_in_take as u64, 0))
74✔
508
                    } else if ti < min_index || ti > max_index {
1,075✔
509
                        None
371✔
510
                    } else {
511
                        sparse_index_to_value_index
704✔
512
                            .get(
704✔
513
                                &I::try_from(ti)
704✔
514
                                    .vortex_expect("take index is between min and max index"),
704✔
515
                            )
704✔
516
                            .map(|value_index| (idx_in_take as u64, *value_index as u64))
704✔
517
                    }
518
                })
1,149✔
519
                .unzip()
408✔
520
        })
408✔
521
        .map_err(|_| vortex_err!("Failed to convert index to usize"))?;
408✔
522

523
    if new_sparse_indices.is_empty() {
408✔
524
        return Ok(None);
37✔
525
    }
371✔
526

371✔
527
    let new_sparse_indices = new_sparse_indices.into_array();
371✔
528
    let values_validity = take_indices_validity.take(&new_sparse_indices)?;
371✔
529
    Ok(Some((
371✔
530
        new_sparse_indices,
371✔
531
        PrimitiveArray::new(value_indices, values_validity).into_array(),
371✔
532
    )))
371✔
533
}
408✔
534

535
/// Filter patches with the provided mask (in flattened space).
536
///
537
/// The filter mask may contain indices that are non-patched. The return value of this function
538
/// is a new set of `Patches` with the indices relative to the provided `mask` rank, and the
539
/// patch values.
540
fn filter_patches_with_mask<T: ToPrimitive + Copy + Ord>(
149✔
541
    patch_indices: &[T],
149✔
542
    offset: usize,
149✔
543
    patch_values: &dyn Array,
149✔
544
    mask_indices: &[usize],
149✔
545
) -> VortexResult<Option<Patches>> {
149✔
546
    let true_count = mask_indices.len();
149✔
547
    let mut new_patch_indices = BufferMut::<u64>::with_capacity(true_count);
149✔
548
    let mut new_mask_indices = Vec::with_capacity(true_count);
149✔
549

550
    // Attempt to move the window by `STRIDE` elements on each iteration. This assumes that
551
    // the patches are relatively sparse compared to the overall mask, and so many indices in the
552
    // mask will end up being skipped.
553
    const STRIDE: usize = 4;
554

555
    let mut mask_idx = 0usize;
149✔
556
    let mut true_idx = 0usize;
149✔
557

558
    while mask_idx < patch_indices.len() && true_idx < true_count {
522✔
559
        // NOTE: we are searching for overlaps between sorted, unaligned indices in `patch_indices`
560
        //  and `mask_indices`. We assume that Patches are sparse relative to the global space of
561
        //  the mask (which covers both patch and non-patch values of the parent array), and so to
562
        //  quickly jump through regions with no overlap, we attempt to move our pointers by STRIDE
563
        //  elements on each iteration. If we cannot rule out overlap due to min/max values, we
564
        //  fallback to performing a two-way iterator merge.
565
        if (mask_idx + STRIDE) < patch_indices.len() && (true_idx + STRIDE) < mask_indices.len() {
373✔
566
            // Load a vector of each into our registers.
567
            let left_min = patch_indices[mask_idx].to_usize().vortex_expect("left_min") - offset;
×
568
            let left_max = patch_indices[mask_idx + STRIDE]
×
569
                .to_usize()
×
570
                .vortex_expect("left_max")
×
571
                - offset;
×
572
            let right_min = mask_indices[true_idx];
×
573
            let right_max = mask_indices[true_idx + STRIDE];
×
574

×
575
            if left_min > right_max {
×
576
                // Advance right side
577
                true_idx += STRIDE;
×
578
                continue;
×
579
            } else if right_min > left_max {
×
580
                mask_idx += STRIDE;
×
581
                continue;
×
582
            } else {
×
583
                // Fallthrough to direct comparison path.
×
584
            }
×
585
        }
373✔
586

587
        // Two-way sorted iterator merge:
588

589
        let left = patch_indices[mask_idx].to_usize().vortex_expect("left") - offset;
373✔
590
        let right = mask_indices[true_idx];
373✔
591

373✔
592
        match left.cmp(&right) {
373✔
593
            Ordering::Less => {
38✔
594
                mask_idx += 1;
38✔
595
            }
38✔
596
            Ordering::Greater => {
148✔
597
                true_idx += 1;
148✔
598
            }
148✔
599
            Ordering::Equal => {
187✔
600
                // Save the mask index as well as the positional index.
187✔
601
                new_mask_indices.push(mask_idx);
187✔
602
                new_patch_indices.push(true_idx as u64);
187✔
603

187✔
604
                mask_idx += 1;
187✔
605
                true_idx += 1;
187✔
606
            }
187✔
607
        }
608
    }
609

610
    if new_mask_indices.is_empty() {
149✔
611
        return Ok(None);
×
612
    }
149✔
613

149✔
614
    let new_patch_indices = new_patch_indices.into_array();
149✔
615
    let new_patch_values = filter(
149✔
616
        patch_values,
149✔
617
        &Mask::from_indices(patch_values.len(), new_mask_indices),
149✔
618
    )?;
149✔
619

620
    Ok(Some(Patches::new(
149✔
621
        true_count,
149✔
622
        0,
149✔
623
        new_patch_indices,
149✔
624
        new_patch_values,
149✔
625
    )))
149✔
626
}
149✔
627

628
#[cfg(test)]
629
mod test {
630
    use rstest::{fixture, rstest};
631
    use vortex_buffer::buffer;
632
    use vortex_mask::Mask;
633

634
    use crate::arrays::PrimitiveArray;
635
    use crate::patches::Patches;
636
    use crate::search_sorted::{SearchResult, SearchSortedSide};
637
    use crate::validity::Validity;
638
    use crate::{IntoArray, ToCanonical};
639

640
    #[test]
641
    fn test_filter() {
1✔
642
        let patches = Patches::new(
1✔
643
            100,
1✔
644
            0,
1✔
645
            buffer![10u32, 11, 20].into_array(),
1✔
646
            buffer![100, 110, 200].into_array(),
1✔
647
        );
1✔
648

1✔
649
        let filtered = patches
1✔
650
            .filter(&Mask::from_indices(100, vec![10, 20, 30]))
1✔
651
            .unwrap()
1✔
652
            .unwrap();
1✔
653

1✔
654
        let indices = filtered.indices().to_primitive().unwrap();
1✔
655
        let values = filtered.values().to_primitive().unwrap();
1✔
656
        assert_eq!(indices.as_slice::<u64>(), &[0, 1]);
1✔
657
        assert_eq!(values.as_slice::<i32>(), &[100, 200]);
1✔
658
    }
1✔
659

660
    #[fixture]
18✔
661
    fn patches() -> Patches {
662
        Patches::new(
663
            20,
664
            0,
665
            buffer![2u64, 9, 15].into_array(),
666
            PrimitiveArray::new(buffer![33_i32, 44, 55], Validity::AllValid).into_array(),
667
        )
668
    }
669

670
    #[rstest]
671
    fn search_larger_than(patches: Patches) {
672
        let res = patches.search_sorted(66, SearchSortedSide::Left).unwrap();
673
        assert_eq!(res, SearchResult::NotFound(16));
674
    }
675

676
    #[rstest]
677
    fn search_less_than(patches: Patches) {
678
        let res = patches.search_sorted(22, SearchSortedSide::Left).unwrap();
679
        assert_eq!(res, SearchResult::NotFound(2));
680
    }
681

682
    #[rstest]
683
    fn search_found(patches: Patches) {
684
        let res = patches.search_sorted(44, SearchSortedSide::Left).unwrap();
685
        assert_eq!(res, SearchResult::Found(9));
686
    }
687

688
    #[rstest]
689
    fn search_not_found_right(patches: Patches) {
690
        let res = patches.search_sorted(56, SearchSortedSide::Right).unwrap();
691
        assert_eq!(res, SearchResult::NotFound(16));
692
    }
693

694
    #[rstest]
695
    fn search_sliced(patches: Patches) {
696
        let sliced = patches.slice(7, 20).unwrap().unwrap();
697
        assert_eq!(
698
            sliced.search_sorted(22, SearchSortedSide::Left).unwrap(),
699
            SearchResult::NotFound(2)
700
        );
701
    }
702

703
    #[test]
704
    fn search_right() {
1✔
705
        let patches = Patches::new(
1✔
706
            6,
1✔
707
            0,
1✔
708
            buffer![0u8, 1, 4, 5].into_array(),
1✔
709
            buffer![-128i8, -98, 8, 50].into_array(),
1✔
710
        );
1✔
711

1✔
712
        assert_eq!(
1✔
713
            patches.search_sorted(-98, SearchSortedSide::Right).unwrap(),
1✔
714
            SearchResult::Found(2)
1✔
715
        );
1✔
716
        assert_eq!(
1✔
717
            patches.search_sorted(50, SearchSortedSide::Right).unwrap(),
1✔
718
            SearchResult::Found(6),
1✔
719
        );
1✔
720
        assert_eq!(
1✔
721
            patches.search_sorted(7, SearchSortedSide::Right).unwrap(),
1✔
722
            SearchResult::NotFound(2),
1✔
723
        );
1✔
724
        assert_eq!(
1✔
725
            patches.search_sorted(51, SearchSortedSide::Right).unwrap(),
1✔
726
            SearchResult::NotFound(6)
1✔
727
        );
1✔
728
    }
1✔
729

730
    #[test]
731
    fn search_left() {
1✔
732
        let patches = Patches::new(
1✔
733
            20,
1✔
734
            0,
1✔
735
            buffer![0u64, 1, 17, 18, 19].into_array(),
1✔
736
            buffer![11i32, 22, 33, 44, 55].into_array(),
1✔
737
        );
1✔
738
        assert_eq!(
1✔
739
            patches.search_sorted(30, SearchSortedSide::Left).unwrap(),
1✔
740
            SearchResult::NotFound(2)
1✔
741
        );
1✔
742
        assert_eq!(
1✔
743
            patches.search_sorted(54, SearchSortedSide::Left).unwrap(),
1✔
744
            SearchResult::NotFound(19)
1✔
745
        );
1✔
746
    }
1✔
747

748
    #[rstest]
749
    fn take_wit_nulls(patches: Patches) {
750
        let taken = patches
751
            .take(
752
                &PrimitiveArray::new(buffer![9, 0], Validity::from_iter(vec![true, false]))
753
                    .into_array(),
754
            )
755
            .unwrap()
756
            .unwrap();
757
        let primitive_values = taken.values().to_primitive().unwrap();
758
        assert_eq!(taken.array_len(), 2);
759
        assert_eq!(primitive_values.as_slice::<i32>(), [44]);
760
        assert_eq!(
761
            primitive_values.validity_mask().unwrap(),
762
            Mask::from_iter(vec![true])
763
        );
764
    }
765

766
    #[test]
767
    fn test_slice() {
1✔
768
        let values = buffer![15_u32, 135, 13531, 42].into_array();
1✔
769
        let indices = buffer![10_u64, 11, 50, 100].into_array();
1✔
770

1✔
771
        let patches = Patches::new(101, 0, indices, values);
1✔
772

1✔
773
        let sliced = patches.slice(15, 100).unwrap().unwrap();
1✔
774
        assert_eq!(sliced.array_len(), 100 - 15);
1✔
775
        let primitive = sliced.values().to_primitive().unwrap();
1✔
776

1✔
777
        assert_eq!(primitive.as_slice::<u32>(), &[13531]);
1✔
778
    }
1✔
779

780
    #[test]
781
    fn doubly_sliced() {
1✔
782
        let values = buffer![15_u32, 135, 13531, 42].into_array();
1✔
783
        let indices = buffer![10_u64, 11, 50, 100].into_array();
1✔
784

1✔
785
        let patches = Patches::new(101, 0, indices, values);
1✔
786

1✔
787
        let sliced = patches.slice(15, 100).unwrap().unwrap();
1✔
788
        assert_eq!(sliced.array_len(), 100 - 15);
1✔
789
        let primitive = sliced.values().to_primitive().unwrap();
1✔
790

1✔
791
        assert_eq!(primitive.as_slice::<u32>(), &[13531]);
1✔
792

793
        let doubly_sliced = sliced.slice(35, 36).unwrap().unwrap();
1✔
794
        let primitive_doubly_sliced = doubly_sliced.values().to_primitive().unwrap();
1✔
795

1✔
796
        assert_eq!(primitive_doubly_sliced.as_slice::<u32>(), &[13531]);
1✔
797
    }
1✔
798
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE TRIAL · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc