Skip to content

Commit fd440d7

Browse files
gatesngatesn
authored
SearchSorted non-null (#3249)
1 parent 7a0b912 commit fd440d7

File tree

7 files changed

+85
-37
lines changed

7 files changed

+85
-37
lines changed

‎encodings/runend/src/array.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ impl RunEndArray {
8989
.as_primitive_typed()
9090
.vortex_expect("ends array must be primitive")
9191
.search_sorted(
92-
&Some(PValue::from(index + self.offset())),
92+
&PValue::from(index + self.offset()),
9393
SearchSortedSide::Right,
9494
)
9595
.to_ends_index(self.ends().len()))
@@ -107,7 +107,7 @@ impl RunEndArray {
107107
.as_primitive_typed()
108108
.vortex_expect("ends array must be primitive")
109109
.search_sorted_many(
110-
indices.into_iter().map(|i| Some(PValue::from(i))),
110+
indices.into_iter().map(PValue::from),
111111
SearchSortedSide::Right,
112112
)
113113
.map(|result| result.to_ends_index(self.ends().len()))

‎vortex-array/src/arrays/chunked/variants.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ impl NullArrayTrait for ChunkedArray {}
5555
impl BoolArrayTrait for ChunkedArray {}
5656

5757
impl PrimitiveArrayTrait for ChunkedArray {
58-
fn value(&self, idx: usize) -> Option<PValue> {
58+
fn value_unchecked(&self, idx: usize) -> PValue {
5959
let (chunk, offset_in_chunk) = self.find_chunk_idx(idx);
6060
let chunk = self
6161
.chunks()
@@ -64,7 +64,7 @@ impl PrimitiveArrayTrait for ChunkedArray {
6464
chunk
6565
.as_primitive_typed()
6666
.vortex_expect("Chunk was not a PrimitiveArray")
67-
.value(offset_in_chunk)
67+
.value_unchecked(offset_in_chunk)
6868
}
6969
}
7070

‎vortex-array/src/arrays/constant/variants.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,13 @@ impl PrimitiveArrayTrait for ConstantArray {
5656
fn value(&self, _idx: usize) -> Option<PValue> {
5757
self.scalar().as_primitive().pvalue()
5858
}
59+
60+
fn value_unchecked(&self, _idx: usize) -> PValue {
61+
self.scalar()
62+
.as_primitive()
63+
.pvalue()
64+
.unwrap_or_else(|| PValue::zero(self.ptype()))
65+
}
5966
}
6067

6168
impl Utf8ArrayTrait for ConstantArray {}

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

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ mod accessor;
66
use arrow_buffer::BooleanBufferBuilder;
77
use vortex_buffer::{Buffer, BufferMut, ByteBuffer};
88
use vortex_dtype::{DType, NativePType, Nullability, PType, match_each_native_ptype};
9-
use vortex_error::{VortexExpect, VortexResult, vortex_panic};
9+
use vortex_error::{VortexResult, vortex_panic};
1010
use vortex_mask::Mask;
1111

1212
use crate::array::{ArrayCanonicalImpl, ArrayValidityImpl};
@@ -289,15 +289,10 @@ impl ArrayVariantsImpl for PrimitiveArray {
289289
}
290290

291291
impl PrimitiveArrayTrait for PrimitiveArray {
292-
fn value(&self, idx: usize) -> Option<PValue> {
293-
self.validity
294-
.is_valid(idx)
295-
.vortex_expect("validity check")
296-
.then(|| {
297-
match_each_native_ptype!(self.ptype(), |$T| {
298-
PValue::from(self.as_slice::<$T>()[idx])
299-
})
300-
})
292+
fn value_unchecked(&self, idx: usize) -> PValue {
293+
match_each_native_ptype!(self.ptype(), |$T| {
294+
PValue::from(self.as_slice::<$T>()[idx])
295+
})
301296
}
302297
}
303298

‎vortex-array/src/patches.rs‎

Lines changed: 34 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,13 @@ use vortex_dtype::Nullability::NonNullable;
1010
use vortex_dtype::{DType, NativePType, PType, match_each_integer_ptype};
1111
use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
1212
use vortex_mask::{AllOr, Mask};
13-
use vortex_scalar::Scalar;
13+
use vortex_scalar::{PValue, Scalar};
1414

1515
use crate::aliases::hash_map::HashMap;
1616
use crate::arrays::PrimitiveArray;
17-
use crate::compute::{
18-
SearchResult, SearchSorted, SearchSortedSide, cast, filter, search_sorted, take,
19-
};
17+
use crate::compute::{SearchResult, SearchSorted, SearchSortedSide, cast, filter, take};
2018
use crate::variants::PrimitiveArrayTrait;
21-
use crate::{Array, ArrayRef, IntoArray, ToCanonical};
19+
use crate::{Array, ArrayRef, ArrayVariants, IntoArray, ToCanonical};
2220

2321
#[derive(Copy, Clone, Serialize, Deserialize, prost::Message)]
2422
pub struct PatchesMetadata {
@@ -220,7 +218,14 @@ impl Patches {
220218

221219
/// Return the insertion point of `index` in the [Self::indices].
222220
pub fn search_index(&self, index: usize) -> VortexResult<SearchResult> {
223-
search_sorted(&self.indices, index + self.offset, SearchSortedSide::Left)
221+
Ok(self
222+
.indices
223+
.as_primitive_typed()
224+
.vortex_expect("must be primitive")
225+
.search_sorted(
226+
&PValue::U64((index + self.offset) as u64),
227+
SearchSortedSide::Left,
228+
))
224229
}
225230

226231
/// Return the search_sorted result for the given target re-mapped into the original indices.
@@ -229,23 +234,29 @@ impl Patches {
229234
target: T,
230235
side: SearchSortedSide,
231236
) -> VortexResult<SearchResult> {
232-
search_sorted(self.values(), target.into(), side).and_then(|sr| {
233-
let index_idx = sr.to_offsets_index(self.indices().len(), side);
234-
let index = usize::try_from(&self.indices().scalar_at(index_idx)?)? - self.offset;
235-
Ok(match sr {
236-
// If we reached the end of patched values when searching then the result is one after the last patch index
237-
SearchResult::Found(i) => SearchResult::Found(
238-
if i == self.indices().len() || side == SearchSortedSide::Right {
239-
index + 1
240-
} else {
241-
index
242-
},
243-
),
244-
// If the result is NotFound we should return index that's one after the nearest not found index for the corresponding value
245-
SearchResult::NotFound(i) => {
246-
SearchResult::NotFound(if i == 0 { index } else { index + 1 })
247-
}
248-
})
237+
let target = target.into();
238+
239+
let sr = if let Some(parray) = self.values().as_primitive_typed() {
240+
parray.search_sorted(&target.as_primitive().pvalue(), side)
241+
} else {
242+
self.values().search_sorted(&target, side)
243+
};
244+
245+
let index_idx = sr.to_offsets_index(self.indices().len(), side);
246+
let index = usize::try_from(&self.indices().scalar_at(index_idx)?)? - self.offset;
247+
Ok(match sr {
248+
// If we reached the end of patched values when searching then the result is one after the last patch index
249+
SearchResult::Found(i) => SearchResult::Found(
250+
if i == self.indices().len() || side == SearchSortedSide::Right {
251+
index + 1
252+
} else {
253+
index
254+
},
255+
),
256+
// If the result is NotFound we should return index that's one after the nearest not found index for the corresponding value
257+
SearchResult::NotFound(i) => {
258+
SearchResult::NotFound(if i == 0 { index } else { index + 1 })
259+
}
249260
})
250261
}
251262

‎vortex-array/src/variants.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,18 @@ pub trait PrimitiveArrayTrait: Array {
3838

3939
/// Return the primitive value at the given index.
4040
fn value(&self, idx: usize) -> Option<PValue> {
41+
self.is_valid(idx)
42+
.vortex_expect("is valid")
43+
.then(|| self.value_unchecked(idx))
44+
}
45+
46+
/// Return the primitive value at the given index, ignoring nullability.
47+
fn value_unchecked(&self, idx: usize) -> PValue {
4148
self.scalar_at(idx)
4249
.vortex_expect("scalar at index")
4350
.as_primitive()
4451
.pvalue()
52+
.unwrap_or_else(|| PValue::zero(self.ptype()))
4553
}
4654
}
4755

@@ -55,6 +63,17 @@ impl IndexOrd<Option<PValue>> for dyn PrimitiveArrayTrait + '_ {
5563
}
5664
}
5765

66+
impl IndexOrd<PValue> for dyn PrimitiveArrayTrait + '_ {
67+
fn index_cmp(&self, idx: usize, elem: &PValue) -> Option<Ordering> {
68+
assert!(self.all_valid().vortex_expect("all valid"));
69+
self.value_unchecked(idx).partial_cmp(elem)
70+
}
71+
72+
fn index_len(&self) -> usize {
73+
Array::len(self)
74+
}
75+
}
76+
5877
pub trait Utf8ArrayTrait: Array {}
5978

6079
pub trait BinaryArrayTrait: Array {}

‎vortex-scalar/src/pvalue.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,22 @@ macro_rules! as_primitive {
112112
}
113113

114114
impl PValue {
115+
pub fn zero(ptype: PType) -> PValue {
116+
match ptype {
117+
PType::U8 => PValue::U8(0),
118+
PType::U16 => PValue::U16(0),
119+
PType::U32 => PValue::U32(0),
120+
PType::U64 => PValue::U64(0),
121+
PType::I8 => PValue::I8(0),
122+
PType::I16 => PValue::I16(0),
123+
PType::I32 => PValue::I32(0),
124+
PType::I64 => PValue::I64(0),
125+
PType::F16 => PValue::F16(f16::from_f32(0.0)),
126+
PType::F32 => PValue::F32(0.0),
127+
PType::F64 => PValue::F64(0.0),
128+
}
129+
}
130+
115131
pub fn ptype(&self) -> PType {
116132
match self {
117133
Self::U8(_) => PType::U8,

0 commit comments

Comments
 (0)