Skip to content

Commit 59b02f2

Browse files
connortsui20connortsui20
committed
Add fixed-size binary canonical type
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 2e0b1e2 commit 59b02f2

44 files changed

Lines changed: 1605 additions & 31 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.

‎vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ use crate::aggregate_fn::DynAccumulator;
4242
use crate::aggregate_fn::EmptyOptions;
4343
use crate::aggregate_fn::fns::all_non_distinct::variant::check_variant_identical;
4444
use crate::arrays::StructArray;
45+
use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
4546
use crate::arrays::struct_::StructArrayExt;
4647
use crate::dtype::DType;
4748
use crate::dtype::FieldNames;
@@ -257,6 +258,10 @@ fn check_canonical_identical(
257258
check_primitive_identical(lhs, rhs)
258259
}
259260
(Canonical::Decimal(lhs), Canonical::Decimal(rhs)) => check_decimal_identical(lhs, rhs),
261+
(Canonical::FixedSizeBinary(lhs), Canonical::FixedSizeBinary(rhs)) => {
262+
Ok(lhs.buffer_handle().to_host_sync().as_slice()
263+
== rhs.buffer_handle().to_host_sync().as_slice())
264+
}
260265
(Canonical::VarBinView(lhs), Canonical::VarBinView(rhs)) => {
261266
check_varbinview_identical(lhs, rhs)
262267
}

‎vortex-array/src/aggregate_fn/fns/is_constant/mod.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ pub mod primitive;
1010
mod struct_;
1111
mod varbin;
1212

13+
use itertools::Itertools;
1314
use vortex_error::VortexExpect;
1415
use vortex_error::VortexResult;
1516
use vortex_error::vortex_bail;
@@ -35,6 +36,7 @@ use crate::aggregate_fn::DynAccumulator;
3536
use crate::aggregate_fn::EmptyOptions;
3637
use crate::arrays::Constant;
3738
use crate::arrays::Null;
39+
use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
3840
use crate::builtins::ArrayBuiltins;
3941
use crate::dtype::DType;
4042
use crate::dtype::FieldNames;
@@ -395,10 +397,19 @@ impl AggregateFnVTable for IsConstant {
395397
}
396398

397399
let batch_is_constant = match c {
398-
Canonical::Primitive(p) => check_primitive_constant(p),
400+
Canonical::Primitive(a) => check_primitive_constant(a),
401+
Canonical::Decimal(a) => check_decimal_constant(a),
402+
Canonical::FixedSizeBinary(a) => {
403+
let byte_width = a.byte_width() as usize;
404+
let values = a.buffer_handle().to_host_sync();
405+
values
406+
.as_slice()
407+
.chunks_exact(byte_width.max(1))
408+
.map(|value| if byte_width == 0 { &[][..] } else { value })
409+
.all_equal()
410+
}
399411
Canonical::Bool(b) => check_bool_constant(b),
400412
Canonical::VarBinView(v) => check_varbinview_constant(v),
401-
Canonical::Decimal(d) => check_decimal_constant(d),
402413
Canonical::Struct(s) => check_struct_constant(s, ctx)?,
403414
Canonical::Extension(e) => check_extension_constant(e, ctx)?,
404415
Canonical::List(l) => check_listview_constant(l, ctx)?,

‎vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs‎

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use std::fmt::Formatter;
1313

1414
use vortex_error::VortexExpect;
1515
use vortex_error::VortexResult;
16+
use vortex_mask::Mask;
1617
use vortex_session::registry::CachedId;
1718

1819
use self::bool::check_bool_sorted;
@@ -31,6 +32,7 @@ use crate::aggregate_fn::AggregateFnVTable;
3132
use crate::aggregate_fn::DynAccumulator;
3233
use crate::arrays::Constant;
3334
use crate::arrays::Null;
35+
use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
3436
use crate::builtins::ArrayBuiltins;
3537
use crate::dtype::DType;
3638
use crate::dtype::FieldNames;
@@ -41,6 +43,44 @@ use crate::expr::stats::Stat;
4143
use crate::expr::stats::StatsProviderExt;
4244
use crate::scalar::Scalar;
4345

46+
fn check_fixed_size_binary_sorted(
47+
array: &crate::arrays::FixedSizeBinaryArray,
48+
strict: bool,
49+
ctx: &mut ExecutionCtx,
50+
) -> VortexResult<bool> {
51+
let DType::FixedSizeBinary(byte_width, _) = array.dtype() else {
52+
unreachable!()
53+
};
54+
let byte_width = *byte_width as usize;
55+
let values = array.buffer_handle().to_host_sync();
56+
let validity = array.as_ref().validity()?.execute_mask(array.len(), ctx)?;
57+
let mut valid = match &validity {
58+
Mask::AllTrue(len) => {
59+
Box::new(std::iter::repeat_n(true, *len)) as Box<dyn Iterator<Item = bool>>
60+
}
61+
Mask::AllFalse(len) => Box::new(std::iter::repeat_n(false, *len)),
62+
Mask::Values(values) => Box::new(values.bit_buffer().iter()),
63+
};
64+
let mut previous: Option<Option<&[u8]>> = None;
65+
for index in 0..array.len() {
66+
let current = valid.next().unwrap_or(false).then(|| {
67+
let start = index * byte_width;
68+
&values[start..start + byte_width]
69+
});
70+
if let Some(previous) = previous
71+
&& if strict {
72+
previous >= current
73+
} else {
74+
previous > current
75+
}
76+
{
77+
return Ok(false);
78+
}
79+
previous = Some(current);
80+
}
81+
Ok(true)
82+
}
83+
4484
/// Options for the `is_sorted` aggregate function.
4585
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
4686
pub struct IsSortedOptions {
@@ -254,7 +294,8 @@ impl AggregateFnVTable for IsSorted {
254294
| DType::Primitive(..)
255295
| DType::Decimal(..)
256296
| DType::Utf8(_)
257-
| DType::Binary(_) => Some(DType::Bool(Nullability::NonNullable)),
297+
| DType::Binary(_)
298+
| DType::FixedSizeBinary(..) => Some(DType::Bool(Nullability::NonNullable)),
258299
}
259300
}
260301

@@ -271,7 +312,8 @@ impl AggregateFnVTable for IsSorted {
271312
| DType::Primitive(..)
272313
| DType::Decimal(..)
273314
| DType::Utf8(_)
274-
| DType::Binary(_) => Some(make_is_sorted_partial_dtype(input_dtype)),
315+
| DType::Binary(_)
316+
| DType::FixedSizeBinary(..) => Some(make_is_sorted_partial_dtype(input_dtype)),
275317
}
276318
}
277319

@@ -487,10 +529,13 @@ impl AggregateFnVTable for IsSorted {
487529

488530
// Check within-batch sortedness.
489531
let batch_is_sorted = match c {
490-
Canonical::Primitive(p) => check_primitive_sorted(p, partial.strict, ctx)?,
532+
Canonical::Primitive(a) => check_primitive_sorted(a, partial.strict, ctx)?,
533+
Canonical::Decimal(a) => check_decimal_sorted(a, partial.strict, ctx)?,
534+
Canonical::FixedSizeBinary(a) => {
535+
check_fixed_size_binary_sorted(a, partial.strict, ctx)?
536+
}
491537
Canonical::Bool(b) => check_bool_sorted(b, partial.strict, ctx)?,
492538
Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict, ctx)?,
493-
Canonical::Decimal(d) => check_decimal_sorted(d, partial.strict, ctx)?,
494539
Canonical::Extension(e) => check_extension_sorted(e, partial.strict, ctx)?,
495540
Canonical::Null(_) => !partial.strict,
496541
// Struct, List, FixedSizeList should have been filtered out by return_dtype

‎vortex-array/src/aggregate_fn/fns/min_max/mod.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -411,15 +411,16 @@ impl AggregateFnVTable for MinMax {
411411
Ok(())
412412
}
413413
Columnar::Canonical(c) => match c {
414-
Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx),
414+
Canonical::Primitive(a) => accumulate_primitive(partial, a, ctx),
415+
Canonical::Decimal(a) => accumulate_decimal(partial, a, ctx),
415416
Canonical::Bool(b) => accumulate_bool(partial, b, ctx),
416417
Canonical::VarBinView(v) => accumulate_varbinview(partial, v, ctx),
417-
Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx),
418418
Canonical::Extension(e) => accumulate_extension(partial, e, ctx),
419419
Canonical::Null(_) => Ok(()),
420420
Canonical::Struct(_)
421421
| Canonical::List(_)
422422
| Canonical::FixedSizeList(_)
423+
| Canonical::FixedSizeBinary(_)
423424
| Canonical::Variant(_) => {
424425
vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype())
425426
}

‎vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ use crate::aggregate_fn::EmptyOptions;
4343
use crate::array::ArrayView;
4444
use crate::arrays::Constant;
4545
use crate::arrays::ConstantArray;
46+
use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
4647
use crate::arrays::varbinview::BinaryView;
4748
use crate::dtype::DType;
4849
use crate::dtype::DecimalType;
@@ -195,6 +196,16 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
195196
Canonical::Bool(array) => bool_uncompressed_size_in_bytes(array, ctx),
196197
Canonical::Primitive(array) => primitive_uncompressed_size_in_bytes(array, ctx),
197198
Canonical::Decimal(array) => decimal_uncompressed_size_in_bytes(array, ctx),
199+
Canonical::FixedSizeBinary(array) => {
200+
let byte_width = array.byte_width();
201+
let values = checked_len_mul(array.len(), byte_width as usize, "fixed-size binary")?;
202+
let validity = validity_uncompressed_size_in_bytes(
203+
array.as_ref().validity()?.execute_mask(array.len(), ctx)?,
204+
)?;
205+
values
206+
.checked_add(validity)
207+
.ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))
208+
}
198209
Canonical::VarBinView(array) => varbinview_uncompressed_size_in_bytes(array, ctx),
199210
Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
200211
Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
@@ -229,6 +240,9 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
229240
array.len(),
230241
array.scalar().as_binary().value().map(|value| value.len()),
231242
)?,
243+
DType::FixedSizeBinary(byte_width, _) => {
244+
checked_len_mul(array.len(), *byte_width as usize, "fixed-size binary")?
245+
}
232246
DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => {
233247
let canonical = array.array().clone().execute::<Canonical>(ctx)?;
234248
return canonical_uncompressed_size_in_bytes(&canonical, ctx);
@@ -281,6 +295,7 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool {
281295
| DType::Bool(_)
282296
| DType::Primitive(..)
283297
| DType::Decimal(..)
298+
| DType::FixedSizeBinary(..)
284299
| DType::Utf8(_)
285300
| DType::Binary(_) => true,
286301
DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => {

‎vortex-array/src/arrays/arbitrary.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use crate::arrays::VarBinViewArray;
2626
use crate::arrays::primitive::PrimitiveArrayExt;
2727
use crate::builders::ArrayBuilder;
2828
use crate::builders::DecimalBuilder;
29+
use crate::builders::FixedSizeBinaryBuilder;
2930
use crate::builders::FixedSizeListBuilder;
3031
use crate::builders::ListViewBuilder;
3132
use crate::dtype::DType;
@@ -153,6 +154,17 @@ fn random_array_chunk(
153154
}
154155
DType::Utf8(n) => random_string(u, *n, chunk_len),
155156
DType::Binary(n) => random_bytes(u, *n, chunk_len),
157+
d @ DType::FixedSizeBinary(byte_width, n) => {
158+
let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
159+
let mut builder = FixedSizeBinaryBuilder::with_capacity(*byte_width, *n, elem_len);
160+
for _ in 0..elem_len {
161+
let scalar = random_scalar(u, d)?;
162+
builder
163+
.append_scalar(&scalar)
164+
.vortex_expect("random fixed-size binary scalar must match its builder");
165+
}
166+
Ok(builder.finish())
167+
}
156168
DType::List(elem_dtype, null) => random_list(u, elem_dtype, *null, chunk_len),
157169
DType::FixedSizeList(elem_dtype, list_size, null) => {
158170
random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)

‎vortex-array/src/arrays/constant/vtable/canonical.rs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ use std::sync::Arc;
55

66
use vortex_buffer::BitBuffer;
77
use vortex_buffer::Buffer;
8+
use vortex_buffer::ByteBuffer;
9+
use vortex_buffer::ByteBufferMut;
810
use vortex_buffer::buffer;
911
use vortex_error::VortexExpect;
1012
use vortex_error::VortexResult;
@@ -18,6 +20,7 @@ use crate::arrays::Constant;
1820
use crate::arrays::ConstantArray;
1921
use crate::arrays::DecimalArray;
2022
use crate::arrays::ExtensionArray;
23+
use crate::arrays::FixedSizeBinaryArray;
2124
use crate::arrays::FixedSizeListArray;
2225
use crate::arrays::ListViewArray;
2326
use crate::arrays::NullArray;
@@ -125,6 +128,24 @@ pub(crate) fn constant_canonicalize(
125128
array.len(),
126129
))
127130
}
131+
DType::FixedSizeBinary(byte_width, _) => {
132+
let value = scalar
133+
.as_binary()
134+
.value()
135+
.cloned()
136+
.unwrap_or_else(|| ByteBuffer::zeroed(*byte_width as usize));
137+
let mut values =
138+
ByteBufferMut::with_capacity(array.len().saturating_mul(*byte_width as usize));
139+
for _ in 0..array.len() {
140+
values.extend_from_slice(value.as_slice());
141+
}
142+
Canonical::FixedSizeBinary(FixedSizeBinaryArray::new(
143+
values.freeze(),
144+
*byte_width,
145+
array.len(),
146+
validity,
147+
))
148+
}
128149
DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())),
129150
DType::FixedSizeList(element_dtype, list_size, _) => {
130151
let value = scalar.as_list();

‎vortex-array/src/arrays/dict/execute.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ pub(crate) fn take_canonical(
4747
Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?),
4848
Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)),
4949
Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)),
50+
Canonical::FixedSizeBinary(a) => {
51+
Canonical::FixedSizeBinary(take_fixed_size_binary(&a, codes, ctx))
52+
}
5053
Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)),
5154
Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)),
5255
Canonical::FixedSizeList(a) => {
@@ -113,6 +116,19 @@ fn take_decimal(
113116
.into_owned()
114117
}
115118

119+
fn take_fixed_size_binary(
120+
array: &crate::arrays::FixedSizeBinaryArray,
121+
codes: &PrimitiveArray,
122+
ctx: &mut ExecutionCtx,
123+
) -> crate::arrays::FixedSizeBinaryArray {
124+
let codes_ref = codes.clone().into_array();
125+
<crate::arrays::FixedSizeBinary as TakeExecute>::take(array.as_view(), &codes_ref, ctx)
126+
.vortex_expect("take fixed-size binary array")
127+
.vortex_expect("take fixed-size binary should not return None")
128+
.as_::<crate::arrays::FixedSizeBinary>()
129+
.into_owned()
130+
}
131+
116132
fn take_varbinview(
117133
array: &VarBinViewArray,
118134
codes: &PrimitiveArray,

‎vortex-array/src/arrays/filter/execute/fixed_size_binary.rs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use std::sync::Arc;
5+
6+
use vortex_buffer::ByteBufferMut;
7+
use vortex_error::VortexExpect;
8+
use vortex_mask::MaskValues;
9+
10+
use crate::arrays::FixedSizeBinaryArray;
11+
use crate::arrays::filter::execute::filter_validity;
12+
use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
13+
14+
pub fn filter_fixed_size_binary(
15+
array: &FixedSizeBinaryArray,
16+
mask: &Arc<MaskValues>,
17+
) -> FixedSizeBinaryArray {
18+
let logical_byte_width = array.byte_width();
19+
let byte_width = logical_byte_width as usize;
20+
let source = array.buffer_handle().to_host_sync();
21+
let mut values = ByteBufferMut::with_capacity(mask.true_count().saturating_mul(byte_width));
22+
for index in mask.indices() {
23+
let start = *index * byte_width;
24+
values.extend_from_slice(&source[start..start + byte_width]);
25+
}
26+
let validity = filter_validity(
27+
array
28+
.as_ref()
29+
.validity()
30+
.vortex_expect("fixed-size binary validity should be derivable"),
31+
mask,
32+
);
33+
FixedSizeBinaryArray::new(
34+
values.freeze(),
35+
logical_byte_width,
36+
mask.true_count(),
37+
validity,
38+
)
39+
}

‎vortex-array/src/arrays/filter/execute/mod.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ mod bool;
3333
mod buffer;
3434
pub(crate) mod byte_compress;
3535
mod decimal;
36+
mod fixed_size_binary;
3637
mod fixed_size_list;
3738
mod listview;
3839
mod primitive;
@@ -89,6 +90,9 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
8990
Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)),
9091
Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)),
9192
Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)),
93+
Canonical::FixedSizeBinary(a) => {
94+
Canonical::FixedSizeBinary(fixed_size_binary::filter_fixed_size_binary(&a, mask))
95+
}
9296
Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)),
9397
Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)),
9498
Canonical::FixedSizeList(a) => {

0 commit comments

Comments
 (0)