Skip to content

Commit 5112fd9

Browse files
connortsui20connortsui20
committed
Add fixed-size binary canonical type
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 2880b5f commit 5112fd9

48 files changed

Lines changed: 1552 additions & 72 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: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,10 @@ fn check_canonical_identical(
257257
check_primitive_identical(lhs, rhs)
258258
}
259259
(Canonical::Decimal(lhs), Canonical::Decimal(rhs)) => check_decimal_identical(lhs, rhs),
260+
(Canonical::FixedSizeBinary(lhs), Canonical::FixedSizeBinary(rhs)) => {
261+
Ok(lhs.buffer_handle().to_host_sync().as_slice()
262+
== rhs.buffer_handle().to_host_sync().as_slice())
263+
}
260264
(Canonical::VarBinView(lhs), Canonical::VarBinView(rhs)) => {
261265
check_varbinview_identical(lhs, rhs)
262266
}

‎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: 48 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;
@@ -41,6 +42,44 @@ use crate::expr::stats::Stat;
4142
use crate::expr::stats::StatsProviderExt;
4243
use crate::scalar::Scalar;
4344

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

@@ -271,7 +311,8 @@ impl AggregateFnVTable for IsSorted {
271311
| DType::Primitive(..)
272312
| DType::Decimal(..)
273313
| DType::Utf8(_)
274-
| DType::Binary(_) => Some(make_is_sorted_partial_dtype(input_dtype)),
314+
| DType::Binary(_)
315+
| DType::FixedSizeBinary(..) => Some(make_is_sorted_partial_dtype(input_dtype)),
275316
}
276317
}
277318

@@ -487,10 +528,13 @@ impl AggregateFnVTable for IsSorted {
487528

488529
// Check within-batch sortedness.
489530
let batch_is_sorted = match c {
490-
Canonical::Primitive(p) => check_primitive_sorted(p, partial.strict, ctx)?,
531+
Canonical::Primitive(a) => check_primitive_sorted(a, partial.strict, ctx)?,
532+
Canonical::Decimal(a) => check_decimal_sorted(a, partial.strict, ctx)?,
533+
Canonical::FixedSizeBinary(a) => {
534+
check_fixed_size_binary_sorted(a, partial.strict, ctx)?
535+
}
491536
Canonical::Bool(b) => check_bool_sorted(b, partial.strict, ctx)?,
492537
Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict, ctx)?,
493-
Canonical::Decimal(d) => check_decimal_sorted(d, partial.strict, ctx)?,
494538
Canonical::Extension(e) => check_extension_sorted(e, partial.strict, ctx)?,
495539
Canonical::Null(_) => !partial.strict,
496540
// 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/array/erased.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use crate::arrays::Constant;
3939
use crate::arrays::Decimal;
4040
use crate::arrays::DictArray;
4141
use crate::arrays::FilterArray;
42+
use crate::arrays::FixedSizeBinary;
4243
use crate::arrays::Null;
4344
use crate::arrays::Primitive;
4445
use crate::arrays::SliceArray;
@@ -450,6 +451,7 @@ impl ArrayRef {
450451
|| self.is::<Bool>()
451452
|| self.is::<Primitive>()
452453
|| self.is::<Decimal>()
454+
|| self.is::<FixedSizeBinary>()
453455
|| self.is::<VarBin>()
454456
|| self.is::<VarBinView>()
455457
}
@@ -810,6 +812,7 @@ mod tests {
810812

811813
use super::*;
812814
use crate::arrays::DecimalArray;
815+
use crate::arrays::FixedSizeBinaryArray;
813816
use crate::dtype::DecimalDType;
814817

815818
#[test]
@@ -821,12 +824,23 @@ mod tests {
821824
Validity::NonNullable,
822825
)
823826
.into_array();
827+
let fixed_size_binary = FixedSizeBinaryArray::new(
828+
buffer![1u8, 2].into_byte_buffer(),
829+
2,
830+
1,
831+
Validity::NonNullable,
832+
)
833+
.into_array();
824834

825835
assert!(primitive.is::<Primitive>());
826836
assert!(!primitive.is::<Decimal>());
827837
assert!(decimal.is::<Decimal>());
828838
assert!(!decimal.is::<Primitive>());
839+
assert!(fixed_size_binary.is::<FixedSizeBinary>());
840+
assert!(!fixed_size_binary.is::<Primitive>());
841+
assert!(!fixed_size_binary.is::<Decimal>());
829842
assert!(primitive.is_arrow());
830843
assert!(decimal.is_arrow());
844+
assert!(fixed_size_binary.is_arrow());
831845
}
832846
}

‎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/decimal/array.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,9 @@ pub trait DecimalArrayExt: TypedArrayRef<Decimal> {
143143
FixedWidthType::Primitive(_) => {
144144
unreachable!("DecimalArrayExt requires decimal fixed-width storage")
145145
}
146+
FixedWidthType::FixedSizeBinary(_) => {
147+
unreachable!("DecimalArrayExt requires decimal fixed-width storage")
148+
}
146149
}
147150
}
148151

@@ -338,6 +341,9 @@ impl FixedWidthData {
338341
FixedWidthType::Primitive(_) => {
339342
vortex_panic!("primitive fixed-width data does not contain decimal values")
340343
}
344+
FixedWidthType::FixedSizeBinary(_) => {
345+
vortex_panic!("fixed-size binary data does not contain decimal values")
346+
}
341347
};
342348
if values_type != T::DECIMAL_TYPE {
343349
vortex_panic!(
@@ -356,6 +362,9 @@ impl FixedWidthData {
356362
FixedWidthType::Primitive(_) => {
357363
vortex_panic!("primitive fixed-width data does not contain decimal values")
358364
}
365+
FixedWidthType::FixedSizeBinary(_) => {
366+
vortex_panic!("fixed-size binary data does not contain decimal values")
367+
}
359368
}
360369
}
361370
}

‎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,

0 commit comments

Comments
 (0)