Skip to content

Commit dde9039

Browse files
connortsui20connortsui20
committed
Add fixed-size binary support
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 3ba3236 commit dde9039

87 files changed

Lines changed: 2028 additions & 125 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎encodings/sparse/src/canonical.rs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ use vortex_array::ArrayRef;
99
use vortex_array::ExecutionCtx;
1010
use vortex_array::IntoArray;
1111
use vortex_array::arrays::BoolArray;
12+
use vortex_array::arrays::FixedSizeBinary;
13+
use vortex_array::arrays::FixedSizeBinaryArray;
1214
use vortex_array::arrays::FixedSizeList;
1315
use vortex_array::arrays::FixedSizeListArray;
1416
use vortex_array::arrays::ListView;
@@ -20,6 +22,7 @@ use vortex_array::arrays::Struct;
2022
use vortex_array::arrays::StructArray;
2123
use vortex_array::arrays::VarBinView;
2224
use vortex_array::arrays::VarBinViewArray;
25+
use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
2326
use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
2427
use vortex_array::arrays::listview::ListViewArrayExt;
2528
use vortex_array::arrays::primitive::PrimitiveArrayExt;
@@ -54,6 +57,7 @@ use vortex_buffer::BitBuffer;
5457
use vortex_buffer::Buffer;
5558
use vortex_buffer::BufferString;
5659
use vortex_buffer::ByteBuffer;
60+
use vortex_buffer::ByteBufferMut;
5761
use vortex_buffer::buffer;
5862
use vortex_buffer::buffer_mut;
5963
use vortex_error::VortexError;
@@ -140,6 +144,14 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
140144
let fill = fill_value.as_binary().value().cloned();
141145
execute_varbin(&patches, &fill_value, dtype.clone(), fill, len, ctx)?
142146
}
147+
DType::FixedSizeBinary(byte_width, nullability) => execute_sparse_fixed_size_binary(
148+
&patches,
149+
&fill_value,
150+
*byte_width,
151+
*nullability,
152+
len,
153+
ctx,
154+
)?,
143155
DType::List(values_dtype, nullability) => execute_sparse_lists(
144156
&patches,
145157
&fill_value,
@@ -165,6 +177,46 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
165177
})
166178
}
167179

180+
fn execute_sparse_fixed_size_binary(
181+
resolved: &Patches,
182+
fill_scalar: &Scalar,
183+
byte_width: u32,
184+
nullability: Nullability,
185+
len: usize,
186+
ctx: &mut ExecutionCtx,
187+
) -> VortexResult<ArrayRef> {
188+
let byte_width_usize = byte_width as usize;
189+
let fill = fill_scalar
190+
.as_binary()
191+
.value()
192+
.cloned()
193+
.unwrap_or_else(|| ByteBuffer::zeroed(byte_width_usize));
194+
let mut dense = ByteBufferMut::with_capacity(len.saturating_mul(byte_width_usize));
195+
for _ in 0..len {
196+
dense.extend_from_slice(fill.as_slice());
197+
}
198+
199+
let indices = resolved.indices().as_::<Primitive>().into_owned();
200+
let values = resolved.values().as_::<FixedSizeBinary>().into_owned();
201+
let patch_bytes = values.fixed_size_binary_buffer_handle().to_host_sync();
202+
match_each_integer_ptype!(indices.ptype(), |I| {
203+
for (patch_row, patch_index) in indices.as_slice::<I>().iter().enumerate() {
204+
let patch_index = <usize as NumCast>::from(*patch_index)
205+
.vortex_expect("fixed-size binary patch index must fit in usize");
206+
let source = patch_row * byte_width_usize;
207+
let target = patch_index * byte_width_usize;
208+
dense[target..target + byte_width_usize]
209+
.copy_from_slice(&patch_bytes[source..source + byte_width_usize]);
210+
}
211+
});
212+
213+
let validity = sparse_validity(resolved, fill_scalar, nullability, len, ctx)?;
214+
Ok(
215+
FixedSizeBinaryArray::new_fixed_size_binary(dense.freeze(), byte_width, len, validity)
216+
.into_array(),
217+
)
218+
}
219+
168220
fn execute_sparse_lists(
169221
resolved: &Patches,
170222
fill_value: &Scalar,

‎fuzz/src/array/compare.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,10 @@ pub fn compare_canonical_array(
174174
result_nullability,
175175
)
176176
}
177-
DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => {
177+
DType::FixedSizeBinary(..)
178+
| DType::List(..)
179+
| DType::FixedSizeList(..)
180+
| DType::Struct(..) => {
178181
let scalar_vals: Vec<Scalar> = (0..array.len())
179182
.map(|i| array.execute_scalar(i, ctx).vortex_expect("scalar_at"))
180183
.collect();

‎fuzz/src/array/fill_null.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ pub fn fill_null_canonical_array(
4343
fill_primitive_array(array, fill_value, result_nullability, ctx)
4444
}
4545
Canonical::Decimal(array) => fill_decimal_array(array, fill_value, result_nullability, ctx),
46+
Canonical::FixedSizeBinary(array) => array.into_array().fill_null(fill_value.clone())?,
4647
Canonical::VarBinView(array) => {
4748
fill_varbinview_array(array, fill_value, result_nullability, ctx)
4849
}

‎fuzz/src/array/filter.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ pub fn filter_canonical_array(
9797
.collect::<Vec<_>>();
9898
Ok(VarBinViewArray::from_iter(values, array.dtype().clone()).into_array())
9999
}
100-
DType::List(..) | DType::FixedSizeList(..) => {
100+
DType::FixedSizeBinary(..) | DType::List(..) | DType::FixedSizeList(..) => {
101101
let mut indices = Vec::new();
102102
for (idx, bool) in filter.iter().enumerate() {
103103
if *bool {

‎fuzz/src/array/mask.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use vortex_array::arrays::VarBinViewArray;
1818
use vortex_array::arrays::bool::BoolArrayExt;
1919
use vortex_array::arrays::decimal::DecimalArrayExt;
2020
use vortex_array::arrays::extension::ExtensionArrayExt;
21+
use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
2122
use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
2223
use vortex_array::arrays::listview::ListViewArrayExt;
2324
use vortex_array::arrays::struct_::StructArrayExt;
@@ -88,10 +89,20 @@ pub fn mask_canonical_array(
8889
Canonical::Decimal(array) => {
8990
let new_validity = mask_validity(&array.validity()?, mask, ctx);
9091
match_each_decimal_value_type!(array.values_type(), |D| {
91-
DecimalArray::new(array.buffer::<D>(), array.decimal_dtype(), new_validity)
92+
DecimalArray::new_decimal(array.buffer::<D>(), array.decimal_dtype(), new_validity)
9293
.into_array()
9394
})
9495
}
96+
Canonical::FixedSizeBinary(array) => {
97+
let new_validity = mask_validity(&array.validity()?, mask, ctx);
98+
vortex_array::arrays::FixedSizeBinaryArray::new_fixed_size_binary(
99+
array.buffer_handle().to_host_sync(),
100+
array.byte_width(),
101+
array.len(),
102+
new_validity,
103+
)
104+
.into_array()
105+
}
95106
Canonical::VarBinView(array) => {
96107
let new_validity = mask_validity(&array.validity()?, mask, ctx);
97108
VarBinViewArray::new_handle(

‎fuzz/src/array/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -486,7 +486,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
486486
// These support all actions
487487
ActionType::iter().collect()
488488
}
489-
DType::Utf8(_) | DType::Binary(_) => {
489+
DType::Utf8(_) | DType::Binary(_) | DType::FixedSizeBinary(..) => {
490490
// Utf8/Binary supports everything except Sum and FillNull
491491
// Actions: Compress, Slice, Take, SearchSorted, Filter, Compare, Cast, MinMax, Mask, ScalarAt
492492
[

‎fuzz/src/array/scalar_at.rs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use vortex_array::IntoArray;
99
use vortex_array::arrays::bool::BoolArrayExt;
1010
use vortex_array::arrays::decimal::DecimalArrayExt;
1111
use vortex_array::arrays::extension::ExtensionArrayExt;
12+
use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
1213
use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
1314
use vortex_array::arrays::listview::ListViewArrayExt;
1415
use vortex_array::arrays::struct_::StructArrayExt;
@@ -39,19 +40,18 @@ pub fn scalar_at_canonical_array(
3940
array.to_bit_buffer().value(index),
4041
array.dtype().nullability(),
4142
),
42-
Canonical::Primitive(array) => {
43-
match_each_native_ptype!(array.ptype(), |T| {
44-
Scalar::primitive(array.as_slice::<T>()[index], array.dtype().nullability())
45-
})
46-
}
47-
Canonical::Decimal(array) => {
48-
match_each_decimal_value_type!(array.values_type(), |D| {
49-
Scalar::decimal(
50-
DecimalValue::from(array.buffer::<D>()[index]),
51-
array.decimal_dtype(),
52-
array.dtype().nullability(),
53-
)
54-
})
43+
Canonical::Primitive(array) => match_each_native_ptype!(array.ptype(), |T| {
44+
Scalar::primitive(array.as_slice::<T>()[index], array.dtype().nullability())
45+
}),
46+
Canonical::Decimal(array) => match_each_decimal_value_type!(array.values_type(), |D| {
47+
Scalar::decimal(
48+
DecimalValue::from(array.buffer::<D>()[index]),
49+
array.decimal_dtype(),
50+
array.dtype().nullability(),
51+
)
52+
}),
53+
Canonical::FixedSizeBinary(array) => {
54+
Scalar::fixed_size_binary(array.value(index), array.dtype().nullability())
5555
}
5656
Canonical::VarBinView(array) => varbin_scalar(array.bytes_at(index), array.dtype()),
5757
Canonical::List(array) => {

‎fuzz/src/array/search_sorted.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,10 @@ pub fn search_sorted_canonical_array(
143143
};
144144
SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side)
145145
}
146-
DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => {
146+
DType::FixedSizeBinary(..)
147+
| DType::List(..)
148+
| DType::FixedSizeList(..)
149+
| DType::Struct(..) => {
147150
let scalar_vals = (0..array.len())
148151
.map(|i| array.execute_scalar(i, ctx))
149152
.collect::<VortexResult<Vec<_>>>()?;

‎fuzz/src/array/slice.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use vortex_array::arrays::bool::BoolArrayExt;
1515
use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
1616
use vortex_array::arrays::listview::ListViewArrayExt;
1717
use vortex_array::arrays::struct_::StructArrayExt;
18+
use vortex_array::builders::builder_with_capacity;
1819
use vortex_array::dtype::DType;
1920
use vortex_array::match_each_decimal_value_type;
2021
use vortex_array::match_each_native_ptype;
@@ -124,6 +125,13 @@ pub fn slice_canonical_array(
124125
)
125126
.map(|a| a.into_array())
126127
}
128+
DType::FixedSizeBinary(..) => {
129+
let mut builder = builder_with_capacity(array.dtype(), stop - start);
130+
for index in start..stop {
131+
builder.append_scalar(&array.execute_scalar(index, ctx)?)?;
132+
}
133+
Ok(builder.finish())
134+
}
127135
d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => {
128136
unreachable!("DType {d} not supported for fuzzing")
129137
}

‎fuzz/src/array/sort.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,10 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR
9292
opt_values.sort();
9393
Ok(VarBinViewArray::from_iter(opt_values, array.dtype().clone()).into_array())
9494
}
95-
DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => {
95+
DType::FixedSizeBinary(..)
96+
| DType::List(..)
97+
| DType::FixedSizeList(..)
98+
| DType::Struct(..) => {
9699
let mut sort_indices = (0..array.len()).collect::<Vec<_>>();
97100
sort_indices.sort_by(|a, b| {
98101
let lhs = array.execute_scalar(*a, ctx).vortex_expect("scalar_at");

0 commit comments

Comments
 (0)