diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 0be6d0a4162..454f396bc9e 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -9,6 +9,8 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::FixedSizeBinary; +use vortex_array::arrays::FixedSizeBinaryArray; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListView; @@ -20,6 +22,7 @@ use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -54,6 +57,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_buffer::buffer_mut; use vortex_error::VortexError; @@ -140,6 +144,14 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort let fill = fill_value.as_binary().value().cloned(); execute_varbin(&patches, &fill_value, dtype.clone(), fill, len, ctx)? } + DType::FixedSizeBinary(byte_width, nullability) => execute_sparse_fixed_size_binary( + &patches, + &fill_value, + *byte_width, + *nullability, + len, + ctx, + )?, DType::List(values_dtype, nullability) => execute_sparse_lists( &patches, &fill_value, @@ -165,6 +177,43 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort }) } +fn execute_sparse_fixed_size_binary( + resolved: &Patches, + fill_scalar: &Scalar, + byte_width: u32, + nullability: Nullability, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let byte_width_usize = byte_width as usize; + let fill = fill_scalar + .as_binary() + .value() + .cloned() + .unwrap_or_else(|| ByteBuffer::zeroed(byte_width_usize)); + let mut dense = ByteBufferMut::with_capacity(len.saturating_mul(byte_width_usize)); + for _ in 0..len { + dense.extend_from_slice(fill.as_slice()); + } + + let indices = resolved.indices().as_::().into_owned(); + let values = resolved.values().as_::().into_owned(); + let patch_bytes = values.buffer_handle().to_host_sync(); + match_each_integer_ptype!(indices.ptype(), |I| { + for (patch_row, patch_index) in indices.as_slice::().iter().enumerate() { + let patch_index = ::from(*patch_index) + .vortex_expect("fixed-size binary patch index must fit in usize"); + let source = patch_row * byte_width_usize; + let target = patch_index * byte_width_usize; + dense[target..target + byte_width_usize] + .copy_from_slice(&patch_bytes[source..source + byte_width_usize]); + } + }); + + let validity = sparse_validity(resolved, fill_scalar, nullability, len, ctx)?; + Ok(FixedSizeBinaryArray::new(dense.freeze(), byte_width, len, validity).into_array()) +} + fn execute_sparse_lists( resolved: &Patches, fill_value: &Scalar, diff --git a/fuzz/src/array/compare.rs b/fuzz/src/array/compare.rs index f594a9b4102..6e9a9b46031 100644 --- a/fuzz/src/array/compare.rs +++ b/fuzz/src/array/compare.rs @@ -174,7 +174,10 @@ pub fn compare_canonical_array( result_nullability, ) } - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => { + DType::FixedSizeBinary(..) + | DType::List(..) + | DType::FixedSizeList(..) + | DType::Struct(..) => { let scalar_vals: Vec = (0..array.len()) .map(|i| array.execute_scalar(i, ctx).vortex_expect("scalar_at")) .collect(); diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index 95c85f95d1e..debb52429fa 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -43,6 +43,7 @@ pub fn fill_null_canonical_array( fill_primitive_array(array, fill_value, result_nullability, ctx) } Canonical::Decimal(array) => fill_decimal_array(array, fill_value, result_nullability, ctx), + Canonical::FixedSizeBinary(array) => array.into_array().fill_null(fill_value.clone())?, Canonical::VarBinView(array) => { fill_varbinview_array(array, fill_value, result_nullability, ctx) } diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index aa54b0d9488..4608ab3b7d8 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -97,7 +97,7 @@ pub fn filter_canonical_array( .collect::>(); Ok(VarBinViewArray::from_iter(values, array.dtype().clone()).into_array()) } - DType::List(..) | DType::FixedSizeList(..) => { + DType::FixedSizeBinary(..) | DType::List(..) | DType::FixedSizeList(..) => { let mut indices = Vec::new(); for (idx, bool) in filter.iter().enumerate() { if *bool { diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 49471dffada..4f9a1692dc5 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -18,6 +18,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::decimal::DecimalArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::struct_::StructArrayExt; @@ -92,6 +93,16 @@ pub fn mask_canonical_array( .into_array() }) } + Canonical::FixedSizeBinary(array) => { + let new_validity = mask_validity(&array.validity()?, mask, ctx); + vortex_array::arrays::FixedSizeBinaryArray::new( + array.buffer_handle().to_host_sync(), + array.byte_width(), + array.len(), + new_validity, + ) + .into_array() + } Canonical::VarBinView(array) => { let new_validity = mask_validity(&array.validity()?, mask, ctx); VarBinViewArray::new_handle( diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index b6402cb1e1d..aa075eec8a0 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -486,7 +486,7 @@ fn actions_for_dtype(dtype: &DType) -> HashSet { // These support all actions ActionType::iter().collect() } - DType::Utf8(_) | DType::Binary(_) => { + DType::Utf8(_) | DType::Binary(_) | DType::FixedSizeBinary(..) => { // Utf8/Binary supports everything except Sum and FillNull // Actions: Compress, Slice, Take, SearchSorted, Filter, Compare, Cast, MinMax, Mask, ScalarAt [ diff --git a/fuzz/src/array/scalar_at.rs b/fuzz/src/array/scalar_at.rs index f7a9e26c9af..7e77d31b558 100644 --- a/fuzz/src/array/scalar_at.rs +++ b/fuzz/src/array/scalar_at.rs @@ -9,6 +9,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::decimal::DecimalArrayExt; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::struct_::StructArrayExt; @@ -39,19 +40,18 @@ pub fn scalar_at_canonical_array( array.to_bit_buffer().value(index), array.dtype().nullability(), ), - Canonical::Primitive(array) => { - match_each_native_ptype!(array.ptype(), |T| { - Scalar::primitive(array.as_slice::()[index], array.dtype().nullability()) - }) - } - Canonical::Decimal(array) => { - match_each_decimal_value_type!(array.values_type(), |D| { - Scalar::decimal( - DecimalValue::from(array.buffer::()[index]), - array.decimal_dtype(), - array.dtype().nullability(), - ) - }) + Canonical::Primitive(array) => match_each_native_ptype!(array.ptype(), |T| { + Scalar::primitive(array.as_slice::()[index], array.dtype().nullability()) + }), + Canonical::Decimal(array) => match_each_decimal_value_type!(array.values_type(), |D| { + Scalar::decimal( + DecimalValue::from(array.buffer::()[index]), + array.decimal_dtype(), + array.dtype().nullability(), + ) + }), + Canonical::FixedSizeBinary(array) => { + Scalar::fixed_size_binary(array.value(index), array.dtype().nullability()) } Canonical::VarBinView(array) => varbin_scalar(array.bytes_at(index), array.dtype()), Canonical::List(array) => { diff --git a/fuzz/src/array/search_sorted.rs b/fuzz/src/array/search_sorted.rs index 762182b35fb..2f9bc07256f 100644 --- a/fuzz/src/array/search_sorted.rs +++ b/fuzz/src/array/search_sorted.rs @@ -143,7 +143,10 @@ pub fn search_sorted_canonical_array( }; SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side) } - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => { + DType::FixedSizeBinary(..) + | DType::List(..) + | DType::FixedSizeList(..) + | DType::Struct(..) => { let scalar_vals = (0..array.len()) .map(|i| array.execute_scalar(i, ctx)) .collect::>>()?; diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 2a151002805..e194cf61c89 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -15,6 +15,7 @@ use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; use vortex_array::match_each_decimal_value_type; use vortex_array::match_each_native_ptype; @@ -124,6 +125,13 @@ pub fn slice_canonical_array( ) .map(|a| a.into_array()) } + DType::FixedSizeBinary(..) => { + let mut builder = builder_with_capacity(array.dtype(), stop - start); + for index in start..stop { + builder.append_scalar(&array.execute_scalar(index, ctx)?)?; + } + Ok(builder.finish()) + } d @ (DType::Null | DType::Union(..) | DType::Variant(_) | DType::Extension(_)) => { unreachable!("DType {d} not supported for fuzzing") } diff --git a/fuzz/src/array/sort.rs b/fuzz/src/array/sort.rs index 3f00fccd06e..81696abfc5e 100644 --- a/fuzz/src/array/sort.rs +++ b/fuzz/src/array/sort.rs @@ -92,7 +92,10 @@ pub fn sort_canonical_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexR opt_values.sort(); Ok(VarBinViewArray::from_iter(opt_values, array.dtype().clone()).into_array()) } - DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) => { + DType::FixedSizeBinary(..) + | DType::List(..) + | DType::FixedSizeList(..) + | DType::Struct(..) => { let mut sort_indices = (0..array.len()).collect::>(); sort_indices.sort_by(|a, b| { let lhs = array.execute_scalar(*a, ctx).vortex_expect("scalar_at"); diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index c67ae6184d8..08c77de7b48 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -114,7 +114,7 @@ pub fn take_canonical_array( ) .into_array()) } - DType::List(..) | DType::FixedSizeList(..) => { + DType::FixedSizeBinary(..) | DType::List(..) | DType::FixedSizeList(..) => { let mut builder = builder_with_capacity( &array.dtype().union_nullability(nullable), indices_slice_non_opt.len(), diff --git a/lang/cpp/include/vortex/dtype.hpp b/lang/cpp/include/vortex/dtype.hpp index 1be5007c6d6..cee3fd7a409 100644 --- a/lang/cpp/include/vortex/dtype.hpp +++ b/lang/cpp/include/vortex/dtype.hpp @@ -37,6 +37,8 @@ enum class DataTypeVariant { Decimal = DTYPE_DECIMAL, // Nested fixed-size list FixedSizeList = DTYPE_FIXED_SIZE_LIST, + // Fixed-size binary data + FixedSizeBinary = DTYPE_FIXED_SIZE_BINARY, }; // Primitive type @@ -95,6 +97,7 @@ class DataType { DataType list_element() const; DataType fixed_size_list_element() const; uint32_t fixed_size_list_size() const; + uint32_t fixed_size_binary_size() const; private: friend struct detail::Access; @@ -135,6 +138,7 @@ DataType float32(bool nullable = false); DataType float64(bool nullable = false); DataType utf8(bool nullable = false); DataType binary(bool nullable = false); +DataType fixed_size_binary(uint32_t byte_width, bool nullable = false); DataType decimal(uint8_t precision, int8_t scale, bool nullable = false); DataType list(DataType element, bool nullable = false); DataType fixed_size_list(DataType element, uint32_t size, bool nullable = false); diff --git a/lang/cpp/src/dtype.cpp b/lang/cpp/src/dtype.cpp index 808ead6fd50..d25b2d1fe60 100644 --- a/lang/cpp/src/dtype.cpp +++ b/lang/cpp/src/dtype.cpp @@ -124,6 +124,10 @@ uint32_t DataType::fixed_size_list_size() const { return vx_dtype_fixed_size_list_size(handle_.get()); } +uint32_t DataType::fixed_size_binary_size() const { + return vx_dtype_fixed_size_binary_size(handle_.get()); +} + namespace dtype { DataType null() { @@ -174,6 +178,9 @@ DataType utf8(bool nullable) { DataType binary(bool nullable) { return Access::adopt(vx_dtype_new_binary(nullable)); } +DataType fixed_size_binary(uint32_t byte_width, bool nullable) { + return Access::adopt(vx_dtype_new_fixed_size_binary(byte_width, nullable)); +} DataType decimal(uint8_t precision, int8_t scale, bool nullable) { return Access::adopt(vx_dtype_new_decimal(precision, scale, nullable)); } diff --git a/lang/cpp/tests/dtype.cpp b/lang/cpp/tests/dtype.cpp index 1b5e3573092..df07e72500e 100644 --- a/lang/cpp/tests/dtype.cpp +++ b/lang/cpp/tests/dtype.cpp @@ -26,6 +26,13 @@ TEST_CASE("Decimal dtype", "[dtype]") { REQUIRE_THROWS_AS(d.list_element(), VortexException); } +TEST_CASE("Fixed-size binary dtype", "[dtype]") { + auto d = dtype::fixed_size_binary(16, dtype::Nullable); + REQUIRE(d.variant() == DataTypeVariant::FixedSizeBinary); + REQUIRE(d.fixed_size_binary_size() == 16); + REQUIRE(d.nullable()); +} + TEST_CASE("copy dtype", "[dtype]") { auto d = dtype::int32(true); DataType d2 = d; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index f030cc810e2..f276b6aff12 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -257,6 +257,10 @@ fn check_canonical_identical( check_primitive_identical(lhs, rhs) } (Canonical::Decimal(lhs), Canonical::Decimal(rhs)) => check_decimal_identical(lhs, rhs), + (Canonical::FixedSizeBinary(lhs), Canonical::FixedSizeBinary(rhs)) => { + Ok(lhs.buffer_handle().to_host_sync().as_slice() + == rhs.buffer_handle().to_host_sync().as_slice()) + } (Canonical::VarBinView(lhs), Canonical::VarBinView(rhs)) => { check_varbinview_identical(lhs, rhs) } diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 490ede7f640..c6ab2edd03d 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -10,6 +10,7 @@ pub mod primitive; mod struct_; mod varbin; +use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -35,6 +36,7 @@ use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::arrays::Constant; use crate::arrays::Null; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -395,10 +397,19 @@ impl AggregateFnVTable for IsConstant { } let batch_is_constant = match c { - Canonical::Primitive(p) => check_primitive_constant(p), + Canonical::Primitive(a) => check_primitive_constant(a), + Canonical::Decimal(a) => check_decimal_constant(a), + Canonical::FixedSizeBinary(a) => { + let byte_width = a.byte_width() as usize; + let values = a.buffer_handle().to_host_sync(); + values + .as_slice() + .chunks_exact(byte_width.max(1)) + .map(|value| if byte_width == 0 { &[][..] } else { value }) + .all_equal() + } Canonical::Bool(b) => check_bool_constant(b), Canonical::VarBinView(v) => check_varbinview_constant(v), - Canonical::Decimal(d) => check_decimal_constant(d), Canonical::Struct(s) => check_struct_constant(s, ctx)?, Canonical::Extension(e) => check_extension_constant(e, ctx)?, Canonical::List(l) => check_listview_constant(l, ctx)?, diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index ada9e6eb2e5..6acf6340f00 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -13,6 +13,7 @@ use std::fmt::Formatter; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use self::bool::check_bool_sorted; @@ -41,6 +42,44 @@ use crate::expr::stats::Stat; use crate::expr::stats::StatsProviderExt; use crate::scalar::Scalar; +fn check_fixed_size_binary_sorted( + array: &crate::arrays::FixedSizeBinaryArray, + strict: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let DType::FixedSizeBinary(byte_width, _) = array.dtype() else { + unreachable!() + }; + let byte_width = *byte_width as usize; + let values = array.buffer_handle().to_host_sync(); + let validity = array.as_ref().validity()?.execute_mask(array.len(), ctx)?; + let mut valid = match &validity { + Mask::AllTrue(len) => { + Box::new(std::iter::repeat_n(true, *len)) as Box> + } + Mask::AllFalse(len) => Box::new(std::iter::repeat_n(false, *len)), + Mask::Values(values) => Box::new(values.bit_buffer().iter()), + }; + let mut previous: Option> = None; + for index in 0..array.len() { + let current = valid.next().unwrap_or(false).then(|| { + let start = index * byte_width; + &values[start..start + byte_width] + }); + if let Some(previous) = previous + && if strict { + previous >= current + } else { + previous > current + } + { + return Ok(false); + } + previous = Some(current); + } + Ok(true) +} + /// Options for the `is_sorted` aggregate function. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct IsSortedOptions { @@ -254,7 +293,8 @@ impl AggregateFnVTable for IsSorted { | DType::Primitive(..) | DType::Decimal(..) | DType::Utf8(_) - | DType::Binary(_) => Some(DType::Bool(Nullability::NonNullable)), + | DType::Binary(_) + | DType::FixedSizeBinary(..) => Some(DType::Bool(Nullability::NonNullable)), } } @@ -271,7 +311,8 @@ impl AggregateFnVTable for IsSorted { | DType::Primitive(..) | DType::Decimal(..) | DType::Utf8(_) - | DType::Binary(_) => Some(make_is_sorted_partial_dtype(input_dtype)), + | DType::Binary(_) + | DType::FixedSizeBinary(..) => Some(make_is_sorted_partial_dtype(input_dtype)), } } @@ -487,10 +528,13 @@ impl AggregateFnVTable for IsSorted { // Check within-batch sortedness. let batch_is_sorted = match c { - Canonical::Primitive(p) => check_primitive_sorted(p, partial.strict, ctx)?, + Canonical::Primitive(a) => check_primitive_sorted(a, partial.strict, ctx)?, + Canonical::Decimal(a) => check_decimal_sorted(a, partial.strict, ctx)?, + Canonical::FixedSizeBinary(a) => { + check_fixed_size_binary_sorted(a, partial.strict, ctx)? + } Canonical::Bool(b) => check_bool_sorted(b, partial.strict, ctx)?, Canonical::VarBinView(v) => check_varbinview_sorted(v, partial.strict, ctx)?, - Canonical::Decimal(d) => check_decimal_sorted(d, partial.strict, ctx)?, Canonical::Extension(e) => check_extension_sorted(e, partial.strict, ctx)?, Canonical::Null(_) => !partial.strict, // Struct, List, FixedSizeList should have been filtered out by return_dtype diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 51601ce76fd..dd504e48360 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -411,15 +411,16 @@ impl AggregateFnVTable for MinMax { Ok(()) } Columnar::Canonical(c) => match c { - Canonical::Primitive(p) => accumulate_primitive(partial, p, ctx), + Canonical::Primitive(a) => accumulate_primitive(partial, a, ctx), + Canonical::Decimal(a) => accumulate_decimal(partial, a, ctx), Canonical::Bool(b) => accumulate_bool(partial, b, ctx), Canonical::VarBinView(v) => accumulate_varbinview(partial, v, ctx), - Canonical::Decimal(d) => accumulate_decimal(partial, d, ctx), Canonical::Extension(e) => accumulate_extension(partial, e, ctx), Canonical::Null(_) => Ok(()), Canonical::Struct(_) | Canonical::List(_) | Canonical::FixedSizeList(_) + | Canonical::FixedSizeBinary(_) | Canonical::Variant(_) => { vortex_bail!("Unsupported canonical type for min_max: {}", batch.dtype()) } diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 1d04a074d6f..2ff34d1322e 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -43,6 +43,7 @@ use crate::aggregate_fn::EmptyOptions; use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use crate::arrays::varbinview::BinaryView; use crate::dtype::DType; use crate::dtype::DecimalType; @@ -195,6 +196,16 @@ pub(crate) fn canonical_uncompressed_size_in_bytes( Canonical::Bool(array) => bool_uncompressed_size_in_bytes(array, ctx), Canonical::Primitive(array) => primitive_uncompressed_size_in_bytes(array, ctx), Canonical::Decimal(array) => decimal_uncompressed_size_in_bytes(array, ctx), + Canonical::FixedSizeBinary(array) => { + let byte_width = array.byte_width(); + let values = checked_len_mul(array.len(), byte_width as usize, "fixed-size binary")?; + let validity = validity_uncompressed_size_in_bytes( + array.as_ref().validity()?.execute_mask(array.len(), ctx)?, + )?; + values + .checked_add(validity) + .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64")) + } Canonical::VarBinView(array) => varbinview_uncompressed_size_in_bytes(array, ctx), Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx), Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx), @@ -229,6 +240,9 @@ pub(crate) fn constant_uncompressed_size_in_bytes( array.len(), array.scalar().as_binary().value().map(|value| value.len()), )?, + DType::FixedSizeBinary(byte_width, _) => { + checked_len_mul(array.len(), *byte_width as usize, "fixed-size binary")? + } DType::List(..) | DType::FixedSizeList(..) | DType::Struct(..) | DType::Extension(_) => { let canonical = array.array().clone().execute::(ctx)?; return canonical_uncompressed_size_in_bytes(&canonical, ctx); @@ -281,6 +295,7 @@ fn supports_uncompressed_size_in_bytes(dtype: &DType) -> bool { | DType::Bool(_) | DType::Primitive(..) | DType::Decimal(..) + | DType::FixedSizeBinary(..) | DType::Utf8(_) | DType::Binary(_) => true, DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => { diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 58056f6f577..05ad766f176 100644 --- a/vortex-array/src/array/erased.rs +++ b/vortex-array/src/array/erased.rs @@ -812,6 +812,8 @@ mod tests { use super::*; use crate::arrays::Decimal; use crate::arrays::DecimalArray; + use crate::arrays::FixedSizeBinary; + use crate::arrays::FixedSizeBinaryArray; use crate::arrays::Primitive; use crate::dtype::DecimalDType; @@ -824,10 +826,20 @@ mod tests { Validity::NonNullable, ) .into_array(); + let fixed_size_binary = FixedSizeBinaryArray::new( + buffer![1u8, 2].into_byte_buffer(), + 2, + 1, + Validity::NonNullable, + ) + .into_array(); assert!(primitive.is::()); assert!(!primitive.is::()); assert!(decimal.is::()); assert!(!decimal.is::()); + assert!(fixed_size_binary.is::()); + assert!(!fixed_size_binary.is::()); + assert!(!fixed_size_binary.is::()); } } diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index 5d5d9f60bf9..7e29d3f3c27 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -26,6 +26,7 @@ use crate::arrays::VarBinViewArray; use crate::arrays::primitive::PrimitiveArrayExt; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; +use crate::builders::FixedSizeBinaryBuilder; use crate::builders::FixedSizeListBuilder; use crate::builders::ListViewBuilder; use crate::dtype::DType; @@ -153,6 +154,17 @@ fn random_array_chunk( } DType::Utf8(n) => random_string(u, *n, chunk_len), DType::Binary(n) => random_bytes(u, *n, chunk_len), + d @ DType::FixedSizeBinary(byte_width, n) => { + let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?); + let mut builder = FixedSizeBinaryBuilder::with_capacity(*byte_width, *n, elem_len); + for _ in 0..elem_len { + let scalar = random_scalar(u, d)?; + builder + .append_scalar(&scalar) + .vortex_expect("random fixed-size binary scalar must match its builder"); + } + Ok(builder.finish()) + } DType::List(elem_dtype, null) => random_list(u, elem_dtype, *null, chunk_len), DType::FixedSizeList(elem_dtype, list_size, null) => { random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..ad1bdc5378a 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -5,6 +5,8 @@ use std::sync::Arc; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -18,6 +20,7 @@ use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeBinaryArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::NullArray; @@ -125,6 +128,24 @@ pub(crate) fn constant_canonicalize( array.len(), )) } + DType::FixedSizeBinary(byte_width, _) => { + let value = scalar + .as_binary() + .value() + .cloned() + .unwrap_or_else(|| ByteBuffer::zeroed(*byte_width as usize)); + let mut values = + ByteBufferMut::with_capacity(array.len().saturating_mul(*byte_width as usize)); + for _ in 0..array.len() { + values.extend_from_slice(value.as_slice()); + } + Canonical::FixedSizeBinary(FixedSizeBinaryArray::new( + values.freeze(), + *byte_width, + array.len(), + validity, + )) + } DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())), DType::FixedSizeList(element_dtype, list_size, _) => { let value = scalar.as_list(); diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs index 94607cc4eb3..e01f5d9118c 100644 --- a/vortex-array/src/arrays/decimal/array.rs +++ b/vortex-array/src/arrays/decimal/array.rs @@ -143,6 +143,9 @@ pub trait DecimalArrayExt: TypedArrayRef { FixedWidthType::Primitive(_) => { unreachable!("DecimalArrayExt requires decimal fixed-width storage") } + FixedWidthType::FixedSizeBinary(_) => { + unreachable!("DecimalArrayExt requires decimal fixed-width storage") + } } } @@ -338,6 +341,9 @@ impl FixedWidthData { FixedWidthType::Primitive(_) => { vortex_panic!("primitive fixed-width data does not contain decimal values") } + FixedWidthType::FixedSizeBinary(_) => { + vortex_panic!("fixed-size binary data does not contain decimal values") + } }; if values_type != T::DECIMAL_TYPE { vortex_panic!( @@ -356,6 +362,9 @@ impl FixedWidthData { FixedWidthType::Primitive(_) => { vortex_panic!("primitive fixed-width data does not contain decimal values") } + FixedWidthType::FixedSizeBinary(_) => { + vortex_panic!("fixed-size binary data does not contain decimal values") + } } } } diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 47a9a051666..b55daf89784 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/vortex-array/src/arrays/dict/execute.rs @@ -47,6 +47,9 @@ pub(crate) fn take_canonical( Canonical::Bool(a) => Canonical::Bool(take_bool(&a, codes, ctx)?), Canonical::Primitive(a) => Canonical::Primitive(take_primitive(&a, codes, ctx)), Canonical::Decimal(a) => Canonical::Decimal(take_decimal(&a, codes, ctx)), + Canonical::FixedSizeBinary(a) => { + Canonical::FixedSizeBinary(take_fixed_size_binary(&a, codes, ctx)) + } Canonical::VarBinView(a) => Canonical::VarBinView(take_varbinview(&a, codes, ctx)), Canonical::List(a) => Canonical::List(take_listview(&a, codes, ctx)), Canonical::FixedSizeList(a) => { @@ -113,6 +116,19 @@ fn take_decimal( .into_owned() } +fn take_fixed_size_binary( + array: &crate::arrays::FixedSizeBinaryArray, + codes: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> crate::arrays::FixedSizeBinaryArray { + let codes_ref = codes.clone().into_array(); + ::take(array.as_view(), &codes_ref, ctx) + .vortex_expect("take fixed-size binary array") + .vortex_expect("take fixed-size binary should not return None") + .as_::() + .into_owned() +} + fn take_varbinview( array: &VarBinViewArray, codes: &PrimitiveArray, diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_binary.rs b/vortex-array/src/arrays/filter/execute/fixed_size_binary.rs new file mode 100644 index 00000000000..f520bafd3b7 --- /dev/null +++ b/vortex-array/src/arrays/filter/execute/fixed_size_binary.rs @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_mask::MaskValues; + +use crate::arrays::FixedSizeBinaryArray; +use crate::arrays::filter::execute::filter_validity; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; + +pub fn filter_fixed_size_binary( + array: &FixedSizeBinaryArray, + mask: &Arc, +) -> FixedSizeBinaryArray { + let logical_byte_width = array.byte_width(); + let byte_width = logical_byte_width as usize; + let source = array.buffer_handle().to_host_sync(); + let mut values = ByteBufferMut::with_capacity(mask.true_count().saturating_mul(byte_width)); + for index in mask.indices() { + let start = *index * byte_width; + values.extend_from_slice(&source[start..start + byte_width]); + } + let validity = filter_validity( + array + .as_ref() + .validity() + .vortex_expect("fixed-size binary validity should be derivable"), + mask, + ); + FixedSizeBinaryArray::new( + values.freeze(), + logical_byte_width, + mask.true_count(), + validity, + ) +} diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 52b0373ac49..86b4e2b6452 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -33,6 +33,7 @@ mod bool; mod buffer; pub(crate) mod byte_compress; mod decimal; +mod fixed_size_binary; mod fixed_size_list; mod listview; mod primitive; @@ -89,6 +90,9 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)), Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)), Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)), + Canonical::FixedSizeBinary(a) => { + Canonical::FixedSizeBinary(fixed_size_binary::filter_fixed_size_binary(&a, mask)) + } Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)), Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)), Canonical::FixedSizeList(a) => { diff --git a/vortex-array/src/arrays/fixed_size_binary/mod.rs b/vortex-array/src/arrays/fixed_size_binary/mod.rs new file mode 100644 index 00000000000..ebc980ec96f --- /dev/null +++ b/vortex-array/src/arrays/fixed_size_binary/mod.rs @@ -0,0 +1,702 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonical contiguous storage for fixed-size binary values. + +use std::ops::Not; + +use num_traits::AsPrimitive; +use smallvec::smallvec; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::ArrayRef; +use crate::ArraySlots; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::TypedArrayRef; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::validity_to_child; +use crate::arrays::ConstantArray; +use crate::arrays::Dict; +use crate::arrays::Masked; +use crate::arrays::PrimitiveArray; +use crate::arrays::dict::TakeExecute; +use crate::arrays::dict::TakeExecuteAdaptor; +use crate::arrays::fixed_width; +use crate::arrays::fixed_width::FixedWidthData; +use crate::arrays::fixed_width::FixedWidthType; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::arrays::slice::SliceReduce; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::builders::FixedSizeBinaryBuilder; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::match_each_integer_ptype; +use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::rules::ArrayParentReduceRule; +use crate::optimizer::rules::ParentRuleSet; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::cast::CastExecuteAdaptor; +use crate::scalar_fn::fns::cast::CastKernel; +use crate::scalar_fn::fns::cast::CastReduce; +use crate::scalar_fn::fns::cast::CastReduceAdaptor; +use crate::scalar_fn::fns::fill_null::FillNull; +use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::fill_null::FillNullKernel; +use crate::scalar_fn::fns::mask::MaskReduce; +use crate::scalar_fn::fns::mask::MaskReduceAdaptor; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Cast.id(), + FixedSizeBinary, + CastExecuteAdaptor(FixedSizeBinary), + ); + kernels.register_execute_parent_kernel( + FillNull.id(), + FixedSizeBinary, + FillNullExecuteAdaptor(FixedSizeBinary), + ); + kernels.register_execute_parent_kernel( + Dict.id(), + FixedSizeBinary, + TakeExecuteAdaptor(FixedSizeBinary), + ); +} + +static RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&FixedSizeBinaryMaskedValidityRule), + ParentRuleSet::lift(&CastReduceAdaptor(FixedSizeBinary)), + ParentRuleSet::lift(&MaskReduceAdaptor(FixedSizeBinary)), + ParentRuleSet::lift(&SliceReduceAdaptor(FixedSizeBinary)), +]); + +/// A canonical array of fixed-size byte strings. +pub type FixedSizeBinaryArray = Array; + +/// The canonical VTable for fixed-size binary arrays. +#[derive(Clone, Debug)] +pub struct FixedSizeBinary; + +impl VTable for FixedSizeBinary { + type TypedArrayData = FixedWidthData; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.fixed_size_binary"); + *ID + } + + fn nbuffers(array: ArrayView<'_, Self>) -> usize { + fixed_width::nbuffers(array) + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + fixed_width::buffer(array, idx) + } + + fn buffer_name(array: ArrayView<'_, Self>, idx: usize) -> Option { + fixed_width::buffer_name(array, idx) + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + fixed_width::with_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn validate( + &self, + data: &FixedWidthData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let DType::FixedSizeBinary(byte_width, nullability) = dtype else { + vortex_error::vortex_bail!("Expected fixed-size binary dtype, got {dtype:?}"); + }; + vortex_ensure!( + data.physical_type() == FixedWidthType::FixedSizeBinary(*byte_width), + "Fixed-size binary dtype width {byte_width} does not match physical type {:?}", + data.physical_type(), + ); + fixed_width::validate_layout(data, len, slots, *nullability, "FixedSizeBinaryArray") + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!( + metadata.is_empty(), + "FixedSizeBinaryArray expects empty metadata" + ); + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let DType::FixedSizeBinary(byte_width, _) = dtype else { + vortex_error::vortex_bail!("Expected fixed-size binary dtype, got {dtype:?}"); + }; + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; + let slots = FixedWidthData::make_fixed_size_binary_slots(&validity, len); + let data = FixedWidthData::try_new_fixed_size_binary(buffers[0].clone(), *byte_width, len)?; + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + let Some(builder) = builder + .as_any_mut() + .downcast_mut::() + else { + vortex_error::vortex_bail!( + "append_to_builder for FixedSizeBinary requires a FixedSizeBinaryBuilder" + ); + }; + builder.append_fixed_size_binary_array(&array.into_owned(), ctx) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + fixed_width::slot_name(idx) + } + + fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done(array)) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +#[derive(Default, Debug)] +struct FixedSizeBinaryMaskedValidityRule; + +impl ArrayParentReduceRule for FixedSizeBinaryMaskedValidityRule { + type Parent = Masked; + + fn reduce_parent( + &self, + array: ArrayView<'_, FixedSizeBinary>, + parent: ArrayView<'_, Masked>, + _child_idx: usize, + ) -> VortexResult> { + let validity = array.validity()?.and(parent.validity()?)?; + Ok(Some( + FixedSizeBinaryArray::try_new_handle( + array.buffer_handle().clone(), + array.byte_width(), + array.len(), + validity, + )? + .into_array(), + )) + } +} + +impl OperationsVTable for FixedSizeBinary { + fn scalar_at( + array: ArrayView<'_, FixedSizeBinary>, + index: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + let byte_width = array.byte_width() as usize; + let values = array.buffer_handle().to_host_sync(); + let start = index * byte_width; + Ok(Scalar::fixed_size_binary( + values.slice(start..start + byte_width), + array.dtype().nullability(), + )) + } +} + +impl ValidityVTable for FixedSizeBinary { + fn validity(array: ArrayView<'_, FixedSizeBinary>) -> VortexResult { + fixed_width::validity(array) + } +} + +impl CastReduce for FixedSizeBinary { + fn cast( + array: ArrayView<'_, FixedSizeBinary>, + dtype: &DType, + ) -> VortexResult> { + let DType::FixedSizeBinary(byte_width, nullability) = dtype else { + return Ok(None); + }; + if *byte_width != array.byte_width() { + return Ok(None); + } + let Some(validity) = array + .validity()? + .trivially_cast_nullability(*nullability, array.len())? + else { + return Ok(None); + }; + Ok(Some( + FixedSizeBinaryArray::try_new_handle( + array.buffer_handle().clone(), + *byte_width, + array.len(), + validity, + )? + .into_array(), + )) + } +} + +impl CastKernel for FixedSizeBinary { + fn cast( + array: ArrayView<'_, FixedSizeBinary>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let DType::FixedSizeBinary(byte_width, nullability) = dtype else { + return Ok(None); + }; + if *byte_width != array.byte_width() { + return Ok(None); + } + let validity = array + .validity()? + .cast_nullability(*nullability, array.len(), ctx)?; + Ok(Some( + FixedSizeBinaryArray::try_new_handle( + array.buffer_handle().clone(), + *byte_width, + array.len(), + validity, + )? + .into_array(), + )) + } +} + +impl FillNullKernel for FixedSizeBinary { + fn fill_null( + array: ArrayView<'_, FixedSizeBinary>, + fill_value: &Scalar, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let result_validity = Validity::from(fill_value.dtype().nullability()); + let mut values = array.buffer_handle().to_host_sync().into_mut(); + let fill = fill_value + .as_binary() + .value() + .vortex_expect("top-level fill_null ensure non-null fill value"); + let is_invalid = match array.validity()? { + Validity::Array(is_valid) => is_valid + .execute::(ctx)? + .into_bit_buffer() + .not(), + _ => unreachable!("checked in entry point"), + }; + let byte_width = array.byte_width() as usize; + for invalid_index in is_invalid.set_indices() { + let start = invalid_index * byte_width; + values[start..start + byte_width].copy_from_slice(fill.as_slice()); + } + Ok(Some( + FixedSizeBinaryArray::new( + values.freeze(), + array.byte_width(), + array.len(), + result_validity, + ) + .into_array(), + )) + } +} + +impl MaskReduce for FixedSizeBinary { + fn mask( + array: ArrayView<'_, FixedSizeBinary>, + mask: &ArrayRef, + ) -> VortexResult> { + Ok(Some( + FixedSizeBinaryArray::try_new_handle( + array.buffer_handle().clone(), + array.byte_width(), + array.len(), + array.validity()?.and(Validity::Array(mask.clone()))?, + )? + .into_array(), + )) + } +} + +impl SliceReduce for FixedSizeBinary { + fn slice( + array: ArrayView<'_, FixedSizeBinary>, + range: std::ops::Range, + ) -> VortexResult> { + let byte_width = array.byte_width(); + let width = byte_width as usize; + let values = array + .buffer_handle() + .slice(range.start * width..range.end * width); + let len = range.len(); + let validity = array.validity()?.slice(range)?; + Ok(Some( + FixedSizeBinaryArray::try_new_handle(values, byte_width, len, validity)?.into_array(), + )) + } +} + +impl TakeExecute for FixedSizeBinary { + fn take( + array: ArrayView<'_, FixedSizeBinary>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_error::vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + let indices_validity = indices.validity()?; + let indices_nulls_zeroed = match indices_validity.execute_mask(indices.len(), ctx)? { + Mask::AllTrue(_) => indices.clone(), + Mask::AllFalse(_) => { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + } + Mask::Values(_) => indices + .clone() + .fill_null(Scalar::from(0).cast(indices.dtype())?)?, + }; + let indices = if ptype.is_unsigned_int() { + indices_nulls_zeroed.execute::(ctx)? + } else { + indices_nulls_zeroed + .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? + .execute::(ctx)? + }; + let validity = array + .validity()? + .take(&indices.clone().into_array())? + .and(indices_validity)?; + let byte_width = array.byte_width(); + let width = byte_width as usize; + let source = array.buffer_handle().to_host_sync(); + let values = match_each_integer_ptype!(indices.ptype(), |I| { + let mut values = BufferMut::::with_capacity(indices.len().saturating_mul(width)); + for index in indices.as_slice::() { + let index: usize = index.as_(); + let start = index * width; + values.extend_from_slice(&source[start..start + width]); + } + values.freeze().into_byte_buffer() + }); + Ok(Some( + FixedSizeBinaryArray::new(values, byte_width, indices.len(), validity).into_array(), + )) + } +} + +/// Typed accessors for a [`FixedSizeBinaryArray`]. +pub trait FixedSizeBinaryArrayExt: TypedArrayRef { + /// The number of bytes in every value. + fn byte_width(&self) -> u32 { + match self.as_ref().dtype() { + DType::FixedSizeBinary(byte_width, _) => *byte_width, + _ => unreachable!("FixedSizeBinaryArrayExt requires a fixed-size binary dtype"), + } + } + + /// The handle for the contiguous values buffer. + fn buffer_handle(&self) -> &BufferHandle { + &self.buffer + } + + /// Copies the value at `index` into a standalone byte buffer. + fn value(&self, index: usize) -> ByteBuffer { + assert!(index < self.len(), "fixed-size binary index out of bounds"); + let byte_width = self.byte_width() as usize; + let start = index * byte_width; + self.buffer_handle() + .to_host_sync() + .slice(start..start + byte_width) + } +} + +impl> FixedSizeBinaryArrayExt for T {} + +impl FixedWidthData { + fn make_fixed_size_binary_slots(validity: &Validity, len: usize) -> ArraySlots { + smallvec![validity_to_child(validity, len)] + } + + fn try_new_fixed_size_binary( + buffer: BufferHandle, + byte_width: u32, + len: usize, + ) -> VortexResult { + let expected_len = len + .checked_mul(byte_width as usize) + .ok_or_else(|| vortex_error::vortex_err!("Fixed-size binary buffer length overflow"))?; + vortex_ensure!( + buffer.len() == expected_len, + InvalidArgument: "Fixed-size binary buffer length {} does not match {len} values of width {byte_width}", + buffer.len(), + ); + Ok(Self { + physical_type: FixedWidthType::FixedSizeBinary(byte_width), + buffer, + len, + }) + } +} + +impl Array { + /// Creates a fixed-size binary array from one contiguous values buffer. + /// + /// `len` is explicit so that arrays with a zero-byte value width can retain their row count. + pub fn new( + buffer: impl Into, + byte_width: u32, + len: usize, + validity: Validity, + ) -> Self { + Self::try_new(buffer.into(), byte_width, len, validity) + .vortex_expect("FixedSizeBinaryArray construction failed") + } + + /// Tries to create a fixed-size binary array from one contiguous values buffer. + pub fn try_new( + buffer: ByteBuffer, + byte_width: u32, + len: usize, + validity: Validity, + ) -> VortexResult { + if let Some(validity_len) = validity.maybe_len() { + vortex_ensure!( + validity_len == len, + InvalidArgument: "Fixed-size binary validity length {validity_len} does not match array length {len}", + ); + } + Self::try_new_handle(BufferHandle::new_host(buffer), byte_width, len, validity) + } + + pub(crate) fn try_new_handle( + buffer: BufferHandle, + byte_width: u32, + len: usize, + validity: Validity, + ) -> VortexResult { + if let Some(validity_len) = validity.maybe_len() { + vortex_ensure!( + validity_len == len, + InvalidArgument: "Fixed-size binary validity length {validity_len} does not match array length {len}", + ); + } + let dtype = DType::FixedSizeBinary(byte_width, validity.nullability()); + let slots = FixedWidthData::make_fixed_size_binary_slots(&validity, len); + let data = FixedWidthData::try_new_fixed_size_binary(buffer, byte_width, len)?; + Array::try_from_parts(ArrayParts::new(FixedSizeBinary, dtype, len, data).with_slots(slots)) + } + + /// Creates an empty fixed-size binary array. + pub fn empty(byte_width: u32, nullability: Nullability) -> Self { + Self::new( + ByteBuffer::empty(), + byte_width, + 0, + Validity::from(nullability), + ) + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::ByteBuffer; + use vortex_buffer::ByteBufferMut; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_mask::Mask; + use vortex_session::registry::ReadContext; + + use crate::ArrayContext; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::FixedSizeBinaryArray; + use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::scalar::Scalar; + use crate::serde::SerializeOptions; + use crate::serde::SerializedArray; + use crate::validity::Validity; + + #[test] + fn scalar_at_and_zero_width() -> VortexResult<()> { + let values = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4, 5, 6].into_byte_buffer(), + 2, + 3, + Validity::NonNullable, + ); + assert_eq!(values.value(1).as_slice(), &[3, 4]); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + values.execute_scalar(2, &mut ctx)?, + Scalar::fixed_size_binary(vec![5u8, 6], Nullability::NonNullable) + ); + + let empty_values = + FixedSizeBinaryArray::new(ByteBuffer::empty(), 0, 4, Validity::NonNullable); + assert_eq!(empty_values.len(), 4); + assert_eq!( + empty_values.dtype(), + &DType::FixedSizeBinary(0, Nullability::NonNullable) + ); + Ok(()) + } + + #[test] + fn slice_filter_and_take() -> VortexResult<()> { + let values = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4, 5, 6, 7, 8].into_byte_buffer(), + 2, + 4, + Validity::from_iter([true, false, true, true]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let sliced = values.slice(1..4)?; + assert_eq!( + sliced.execute_scalar(0, &mut ctx)?, + Scalar::null(DType::FixedSizeBinary(2, Nullability::Nullable)) + ); + assert_eq!( + sliced.execute_scalar(2, &mut ctx)?, + Scalar::fixed_size_binary(vec![7u8, 8], Nullability::Nullable) + ); + + let filtered = values.filter(Mask::from_iter([true, true, false, true]))?; + assert_eq!(filtered.len(), 3); + assert_eq!( + filtered.execute_scalar(2, &mut ctx)?, + Scalar::fixed_size_binary(vec![7u8, 8], Nullability::Nullable) + ); + + let taken = values.take(buffer![3u32, 0].into_array())?; + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::fixed_size_binary(vec![7u8, 8], Nullability::Nullable) + ); + assert_eq!( + taken.execute_scalar(1, &mut ctx)?, + Scalar::fixed_size_binary(vec![1u8, 2], Nullability::Nullable) + ); + Ok(()) + } + + #[test] + fn take_null_index_ignores_out_of_bounds_physical_value() -> VortexResult<()> { + let values = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4].into_byte_buffer(), + 2, + 2, + Validity::NonNullable, + ); + let indices = crate::arrays::PrimitiveArray::new( + buffer![1u64, 2], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + let taken = values.take(indices.into_array())?; + let mut ctx = array_session().create_execution_ctx(); + + assert_eq!( + taken.execute_scalar(0, &mut ctx)?, + Scalar::fixed_size_binary(vec![3u8, 4], Nullability::Nullable) + ); + assert_eq!( + taken.execute_scalar(1, &mut ctx)?, + Scalar::null(DType::FixedSizeBinary(2, Nullability::Nullable)) + ); + Ok(()) + } + + #[test] + fn nullable_serde_roundtrip() -> VortexResult<()> { + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let array = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4, 5, 6].into_byte_buffer(), + 2, + 3, + Validity::from_iter([true, false, true]), + ); + let dtype = array.dtype().clone(); + let len = array.len(); + + let array_ctx = ArrayContext::empty(); + let serialized = array.clone().into_array().serialize( + &array_ctx, + &session, + &SerializeOptions::default(), + )?; + let mut concat = ByteBufferMut::empty(); + for buffer in serialized { + concat.extend_from_slice(buffer.as_ref()); + } + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + + assert!(decoded.is::()); + assert_arrays_eq!(decoded, array, &mut ctx); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/fixed_width.rs b/vortex-array/src/arrays/fixed_width.rs index 1c133ab8d5f..17d95fe9de0 100644 --- a/vortex-array/src/arrays/fixed_width.rs +++ b/vortex-array/src/arrays/fixed_width.rs @@ -3,8 +3,8 @@ //! Shared physical storage helpers for canonical fixed-width arrays. //! -//! Primitive and decimal arrays have distinct VTables and logical behavior, but use the same -//! one-buffer-plus-validity physical representation. +//! Primitive, decimal, and fixed-size binary arrays have distinct VTables and logical behavior, +//! but use the same one-buffer-plus-validity physical representation. use std::fmt::Display; use std::fmt::Formatter; @@ -37,6 +37,8 @@ pub enum FixedWidthType { Primitive(PType), /// An integer storage type used for logical decimal values. Decimal(DecimalType), + /// An opaque byte string with the given width. + FixedSizeBinary(u32), } impl FixedWidthType { @@ -45,11 +47,12 @@ impl FixedWidthType { match self { Self::Primitive(ptype) => ptype.byte_width(), Self::Decimal(decimal_type) => decimal_type.byte_width(), + Self::FixedSizeBinary(byte_width) => byte_width as usize, } } } -/// The shared physical data stored by the primitive and decimal VTables. +/// The shared physical data stored by the primitive, decimal, and fixed-size-binary VTables. #[derive(Clone, Debug)] pub struct FixedWidthData { pub(crate) physical_type: FixedWidthType, @@ -62,6 +65,7 @@ impl Display for FixedWidthData { match self.physical_type { FixedWidthType::Primitive(ptype) => write!(f, "ptype: {ptype}"), FixedWidthType::Decimal(values_type) => write!(f, "values_type: {values_type}"), + FixedWidthType::FixedSizeBinary(byte_width) => write!(f, "byte_width: {byte_width}"), } } } diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index 7b3fcdf3118..30eba8646d2 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -22,6 +22,7 @@ use crate::arrays::VariantArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::struct_::StructArrayExt; @@ -45,6 +46,18 @@ pub fn mask_validity_canonical( Canonical::Bool(a) => Canonical::Bool(mask_validity_bool(a, validity)?), Canonical::Primitive(a) => Canonical::Primitive(mask_validity_primitive(a, validity)?), Canonical::Decimal(a) => Canonical::Decimal(mask_validity_decimal(a, validity)?), + Canonical::FixedSizeBinary(a) => { + let byte_width = a.byte_width(); + let len = a.len(); + let buffer = a.buffer_handle().clone(); + let new_validity = Validity::and(a.validity()?, validity)?; + Canonical::FixedSizeBinary(crate::arrays::FixedSizeBinaryArray::try_new_handle( + buffer, + byte_width, + len, + new_validity, + )?) + } Canonical::VarBinView(a) => Canonical::VarBinView(mask_validity_varbinview(a, validity)?), Canonical::List(a) => Canonical::List(mask_validity_listview(a, validity)?), Canonical::FixedSizeList(a) => { diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index d8fb102aeef..deaba1bbfb7 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -4,9 +4,9 @@ //! Built-in array encodings. //! //! Canonical arrays are the default uncompressed representation for a logical dtype: -//! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], [`VarBinViewArray`], -//! [`ListViewArray`], [`FixedSizeListArray`], [`StructArray`], [`ExtensionArray`], and -//! [`VariantArray`]. +//! [`NullArray`], [`BoolArray`], [`PrimitiveArray`], [`DecimalArray`], +//! [`FixedSizeBinaryArray`], [`VarBinViewArray`], [`ListViewArray`], [`FixedSizeListArray`], +//! [`StructArray`], [`ExtensionArray`], and [`VariantArray`]. //! //! Utility and lazy arrays represent common transformations without immediately materializing //! their result. Examples include [`ChunkedArray`] for concatenation, [`ConstantArray`] for repeated @@ -62,6 +62,11 @@ pub use filter::FilterArray; pub mod fixed_width; +pub mod fixed_size_binary; +pub use fixed_size_binary::FixedSizeBinary; +pub use fixed_size_binary::FixedSizeBinaryArray; +pub use fixed_size_binary::FixedSizeBinaryArrayExt; + pub mod fixed_size_list; pub use fixed_size_list::FixedSizeList; pub use fixed_size_list::FixedSizeListArray; @@ -130,10 +135,11 @@ use vortex_session::VortexSession; pub(crate) fn initialize(session: &VortexSession) { bool::initialize(session); chunked::initialize(session); - decimal::initialize(session); dict::initialize(session); + decimal::initialize(session); extension::initialize(session); filter::initialize(session); + fixed_size_binary::initialize(session); fixed_size_list::initialize(session); list::initialize(session); listview::initialize(session); diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index bd448994f72..dd7287e501a 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -512,6 +512,9 @@ impl FixedWidthData { FixedWidthType::Decimal(_) => { vortex_panic!("decimal fixed-width data does not have a primitive type") } + FixedWidthType::FixedSizeBinary(_) => { + vortex_panic!("fixed-size binary data does not have a primitive type") + } } } diff --git a/vortex-array/src/arrays/validation_tests.rs b/vortex-array/src/arrays/validation_tests.rs index 30758a5c29e..e66188d2b57 100644 --- a/vortex-array/src/arrays/validation_tests.rs +++ b/vortex-array/src/arrays/validation_tests.rs @@ -17,9 +17,13 @@ mod tests { use crate::IntoArray; use crate::arrays::ChunkedArray; + use crate::arrays::Decimal; use crate::arrays::DecimalArray; + use crate::arrays::FixedSizeBinary; + use crate::arrays::FixedSizeBinaryArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; + use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; @@ -90,6 +94,40 @@ mod tests { assert!(result.is_err()); } + #[test] + fn fixed_width_logical_types_have_distinct_vtables() { + let primitive = buffer![1i32, 2].into_array(); + let decimal = DecimalArray::new( + Buffer::from_iter([1i128, 2]), + crate::dtype::DecimalDType::new(10, 2), + Validity::NonNullable, + ) + .into_array(); + let fixed_size_binary = FixedSizeBinaryArray::new( + ByteBuffer::from(vec![1u8, 2, 3, 4]), + 2, + 2, + Validity::NonNullable, + ) + .into_array(); + + assert!(primitive.is::()); + assert!(!primitive.is::()); + assert!(!primitive.is::()); + + assert!(decimal.is::()); + assert!(!decimal.is::()); + assert!(!decimal.is::()); + + assert!(fixed_size_binary.is::()); + assert!(!fixed_size_binary.is::()); + assert!(!fixed_size_binary.is::()); + + assert_ne!(primitive.encoding_id(), decimal.encoding_id()); + assert_ne!(primitive.encoding_id(), fixed_size_binary.encoding_id()); + assert_ne!(decimal.encoding_id(), fixed_size_binary.encoding_id()); + } + #[test] fn test_varbin_array_validation_success() { // Valid case: offsets are monotonically increasing and within bounds. diff --git a/vortex-array/src/builders/fixed_size_binary.rs b/vortex-array/src/builders/fixed_size_binary.rs new file mode 100644 index 00000000000..67664cf3cb3 --- /dev/null +++ b/vortex-array/src/builders/fixed_size_binary.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::FixedSizeBinaryArray; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; +use crate::builders::ArrayBuilder; +use crate::builders::LazyBitBufferBuilder; +use crate::canonical::Canonical; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; + +/// A builder for canonical fixed-size binary arrays. +pub struct FixedSizeBinaryBuilder { + dtype: DType, + values: BufferMut, + nulls: LazyBitBufferBuilder, +} + +impl FixedSizeBinaryBuilder { + /// Creates a builder with capacity for `capacity` values. + pub fn with_capacity(byte_width: u32, nullability: Nullability, capacity: usize) -> Self { + Self { + dtype: DType::FixedSizeBinary(byte_width, nullability), + values: BufferMut::with_capacity(capacity.saturating_mul(byte_width as usize)), + nulls: LazyBitBufferBuilder::new(capacity), + } + } + + /// Appends one non-null fixed-size value. + pub fn append_value(&mut self, value: impl AsRef<[u8]>) -> VortexResult<()> { + let value = value.as_ref(); + vortex_ensure!( + value.len() == self.byte_width(), + "FixedSizeBinaryBuilder expected {} bytes, got {}", + self.byte_width(), + value.len(), + ); + self.values.extend_from_slice(value); + self.nulls.append_non_null(); + Ok(()) + } + + fn byte_width(&self) -> usize { + self.byte_width_u32() as usize + } + + fn byte_width_u32(&self) -> u32 { + let DType::FixedSizeBinary(byte_width, _) = self.dtype else { + unreachable!() + }; + byte_width + } + + pub(crate) fn append_fixed_size_binary_array( + &mut self, + array: &FixedSizeBinaryArray, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.dtype(), + "FixedSizeBinaryBuilder expected array with dtype {}, got {}", + self.dtype(), + array.dtype(), + ); + self.values + .extend_from_slice(array.buffer_handle().to_host_sync().as_slice()); + self.nulls + .append_validity_mask(&array.as_ref().validity()?.execute_mask(array.len(), ctx)?); + Ok(()) + } + + /// Finishes this builder directly into a fixed-size binary array. + pub fn finish_into_fixed_size_binary(&mut self) -> FixedSizeBinaryArray { + let len = self.nulls.len(); + let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + FixedSizeBinaryArray::new( + std::mem::take(&mut self.values).freeze().into_byte_buffer(), + self.byte_width_u32(), + len, + validity, + ) + } +} + +impl ArrayBuilder for FixedSizeBinaryBuilder { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn len(&self) -> usize { + self.nulls.len() + } + + fn append_zeros(&mut self, n: usize) { + self.values.push_n(0, n.saturating_mul(self.byte_width())); + self.nulls.append_n_non_nulls(n); + } + + unsafe fn append_nulls_unchecked(&mut self, n: usize) { + self.values.push_n(0, n.saturating_mul(self.byte_width())); + self.nulls.append_n_nulls(n); + } + + fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + vortex_ensure!( + scalar.dtype() == self.dtype(), + "FixedSizeBinaryBuilder expected scalar with dtype {}, got {}", + self.dtype(), + scalar.dtype(), + ); + match scalar.as_binary().value() { + Some(value) => self.append_value(value), + None => { + self.append_null(); + Ok(()) + } + } + } + + fn reserve_exact(&mut self, additional: usize) { + self.values + .reserve(additional.saturating_mul(self.byte_width())); + self.nulls.reserve_exact(additional); + } + + unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); + } + + fn finish(&mut self) -> ArrayRef { + self.finish_into_fixed_size_binary().into_array() + } + + fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical { + Canonical::FixedSizeBinary(self.finish_into_fixed_size_binary()) + } +} diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 561efc2711e..41dfc7c97c5 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -56,6 +56,7 @@ mod bool; mod decimal; pub mod dict; mod extension; +mod fixed_size_binary; mod fixed_size_list; mod list; mod listview; @@ -67,6 +68,7 @@ mod varbinview; pub use bool::*; pub use decimal::*; pub use extension::*; +pub use fixed_size_binary::*; pub use fixed_size_list::*; pub use list::*; pub use listview::*; @@ -282,6 +284,11 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box Box::new(FixedSizeBinaryBuilder::with_capacity( + *byte_width, + *n, + capacity, + )), DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity( Arc::clone(dtype), *n, diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 138e6941def..caf8169b9fd 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -607,6 +607,9 @@ fn create_test_scalars_for_dtype(dtype: &DType, count: usize) -> Vec { } DType::Utf8(n) => Scalar::utf8(format!("test_string_{}", i), *n), DType::Binary(n) => Scalar::binary(format!("bytes_{}", i).into_bytes(), *n), + DType::FixedSizeBinary(size, n) => { + Scalar::fixed_size_binary(vec![i as u8; *size as usize], *n) + } DType::List(element_dtype, n) => { // Create list scalars with a few elements. let elements: Vec = (0..=i) diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index e4e385746ce..35b41b2328f 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -25,6 +25,8 @@ use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::Extension; use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeBinary; +use crate::arrays::FixedSizeBinaryArray; use crate::arrays::FixedSizeList; use crate::arrays::FixedSizeListArray; use crate::arrays::ListView; @@ -40,12 +42,11 @@ use crate::arrays::VarBinViewArray; use crate::arrays::Variant; use crate::arrays::VariantArray; use crate::arrays::bool::BoolDataParts; -use crate::arrays::decimal::DecimalDataParts; use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewDataParts; use crate::arrays::listview::ListViewRebuildMode; -use crate::arrays::primitive::PrimitiveDataParts; use crate::arrays::struct_::StructDataParts; use crate::arrays::varbinview::VarBinViewDataParts; use crate::arrays::variant::VariantArrayExt; @@ -124,6 +125,7 @@ pub enum Canonical { Bool(BoolArray), Primitive(PrimitiveArray), Decimal(DecimalArray), + FixedSizeBinary(FixedSizeBinaryArray), VarBinView(VarBinViewArray), List(ListViewArray), FixedSizeList(FixedSizeListArray), @@ -142,6 +144,7 @@ macro_rules! match_each_canonical { Canonical::Bool($ident) => $eval, Canonical::Primitive($ident) => $eval, Canonical::Decimal($ident) => $eval, + Canonical::FixedSizeBinary($ident) => $eval, Canonical::VarBinView($ident) => $eval, Canonical::List($ident) => $eval, Canonical::FixedSizeList($ident) => $eval, @@ -197,6 +200,9 @@ impl Canonical { Validity::from(n), ) }), + DType::FixedSizeBinary(byte_width, n) => { + Canonical::FixedSizeBinary(FixedSizeBinaryArray::empty(*byte_width, *n)) + } DType::List(dtype, n) => Canonical::List(unsafe { ListViewArray::new_unchecked( Canonical::empty(dtype).into_array(), @@ -337,6 +343,22 @@ impl Canonical { } } + pub fn as_fixed_size_binary(&self) -> &FixedSizeBinaryArray { + if let Canonical::FixedSizeBinary(a) = self { + a + } else { + vortex_panic!("Cannot get FixedSizeBinaryArray from {:?}", &self) + } + } + + pub fn into_fixed_size_binary(self) -> FixedSizeBinaryArray { + if let Canonical::FixedSizeBinary(a) = self { + a + } else { + vortex_panic!("Cannot unwrap FixedSizeBinaryArray from {:?}", &self) + } + } + pub fn as_varbinview(&self) -> &VarBinViewArray { if let Canonical::VarBinView(a) = self { a @@ -575,32 +597,36 @@ impl Executable for CanonicalValidity { )?, ))) } - Canonical::Primitive(p) => { - let PrimitiveDataParts { - ptype, - buffer, - validity, - } = p.into_data_parts(); + Canonical::Primitive(array) => { + let validity = + child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()) + .execute(ctx)?; + let ptype = array.ptype(); + let data = array.into_data(); Ok(CanonicalValidity(Canonical::Primitive(unsafe { - PrimitiveArray::new_unchecked_from_handle(buffer, ptype, validity.execute(ctx)?) + PrimitiveArray::new_unchecked_from_handle(data.buffer, ptype, validity) }))) } - Canonical::Decimal(d) => { - let DecimalDataParts { - decimal_dtype, - values, - values_type, - validity, - } = d.into_data_parts(); + Canonical::Decimal(array) => { + let parts = array.into_data_parts(); Ok(CanonicalValidity(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( - values, - values_type, - decimal_dtype, - validity.execute(ctx)?, + parts.values, + parts.values_type, + parts.decimal_dtype, + parts.validity.execute(ctx)?, ) }))) } + Canonical::FixedSizeBinary(array) => { + let byte_width = array.byte_width(); + let len = array.len(); + let validity = array.validity()?.execute(ctx)?; + let data = array.into_data(); + Ok(CanonicalValidity(Canonical::FixedSizeBinary( + FixedSizeBinaryArray::try_new_handle(data.buffer, byte_width, len, validity)?, + ))) + } Canonical::VarBinView(vbv) => { let VarBinViewDataParts { dtype, @@ -732,32 +758,36 @@ impl Executable for RecursiveCanonical { )?, ))) } - Canonical::Primitive(p) => { - let PrimitiveDataParts { - ptype, - buffer, - validity, - } = p.into_data_parts(); + Canonical::Primitive(array) => { + let validity = + child_to_validity(array.slots()[0].as_ref(), array.dtype().nullability()) + .execute(ctx)?; + let ptype = array.ptype(); + let data = array.into_data(); Ok(RecursiveCanonical(Canonical::Primitive(unsafe { - PrimitiveArray::new_unchecked_from_handle(buffer, ptype, validity.execute(ctx)?) + PrimitiveArray::new_unchecked_from_handle(data.buffer, ptype, validity) }))) } - Canonical::Decimal(d) => { - let DecimalDataParts { - decimal_dtype, - values, - values_type, - validity, - } = d.into_data_parts(); + Canonical::Decimal(array) => { + let parts = array.into_data_parts(); Ok(RecursiveCanonical(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( - values, - values_type, - decimal_dtype, - validity.execute(ctx)?, + parts.values, + parts.values_type, + parts.decimal_dtype, + parts.validity.execute(ctx)?, ) }))) } + Canonical::FixedSizeBinary(array) => { + let byte_width = array.byte_width(); + let len = array.len(); + let validity = array.validity()?.execute(ctx)?; + let data = array.into_data(); + Ok(RecursiveCanonical(Canonical::FixedSizeBinary( + FixedSizeBinaryArray::try_new_handle(data.buffer, byte_width, len, validity)?, + ))) + } Canonical::VarBinView(vbv) => { let VarBinViewDataParts { dtype, @@ -903,6 +933,16 @@ impl Executable for DecimalArray { } } +/// Execute the array to canonical form and unwrap as a [`FixedSizeBinaryArray`]. +impl Executable for FixedSizeBinaryArray { + fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + match array.try_downcast::() { + Ok(fixed_size_binary) => Ok(fixed_size_binary), + Err(array) => Ok(Canonical::execute(array, ctx)?.into_fixed_size_binary()), + } + } +} + /// Execute the array to canonical form and unwrap as a [`BoolArray`]. /// /// This will panic if the array's dtype is not bool. @@ -1026,6 +1066,7 @@ pub enum CanonicalView<'a> { Bool(ArrayView<'a, Bool>), Primitive(ArrayView<'a, Primitive>), Decimal(ArrayView<'a, Decimal>), + FixedSizeBinary(ArrayView<'a, FixedSizeBinary>), VarBinView(ArrayView<'a, VarBinView>), List(ArrayView<'a, ListView>), FixedSizeList(ArrayView<'a, FixedSizeList>), @@ -1041,6 +1082,7 @@ impl From> for Canonical { CanonicalView::Bool(a) => Canonical::Bool(a.into_owned()), CanonicalView::Primitive(a) => Canonical::Primitive(a.into_owned()), CanonicalView::Decimal(a) => Canonical::Decimal(a.into_owned()), + CanonicalView::FixedSizeBinary(a) => Canonical::FixedSizeBinary(a.into_owned()), CanonicalView::VarBinView(a) => Canonical::VarBinView(a.into_owned()), CanonicalView::List(a) => Canonical::List(a.into_owned()), CanonicalView::FixedSizeList(a) => Canonical::FixedSizeList(a.into_owned()), @@ -1059,6 +1101,7 @@ impl CanonicalView<'_> { CanonicalView::Bool(a) => a.array().clone(), CanonicalView::Primitive(a) => a.array().clone(), CanonicalView::Decimal(a) => a.array().clone(), + CanonicalView::FixedSizeBinary(a) => a.array().clone(), CanonicalView::VarBinView(a) => a.array().clone(), CanonicalView::List(a) => a.array().clone(), CanonicalView::FixedSizeList(a) => a.array().clone(), @@ -1080,6 +1123,7 @@ impl Matcher for AnyCanonical { || array.is::() || array.is::() || array.is::() + || array.is::() || array.is::() || array.is::() || array.is::() @@ -1098,6 +1142,8 @@ impl Matcher for AnyCanonical { Some(CanonicalView::Primitive(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::Decimal(a)) + } else if let Some(a) = array.as_opt::() { + Some(CanonicalView::FixedSizeBinary(a)) } else if let Some(a) = array.as_opt::() { Some(CanonicalView::Struct(a)) } else if let Some(a) = array.as_opt::() { diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index 73f98cce4a7..fd214968a0c 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1236,7 +1236,7 @@ fn test_cast_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { } DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => unimplemented!(), - DType::Extension(_) => vec![], // Extension types typically only cast to themselves + DType::FixedSizeBinary(..) | DType::Extension(_) => vec![], /* These types typically only cast to themselves */ }; // Test each target dtype diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index 47bc8f1bd82..340751445f0 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -78,6 +78,18 @@ fn test_take_all(array: &ArrayRef, ctx: &mut ExecutionCtx) { result_prim.buffer_handle().to_host_sync() ); } + (Canonical::Decimal(orig_prim), Canonical::Decimal(result_prim)) => { + assert_eq!( + orig_prim.buffer_handle().to_host_sync(), + result_prim.buffer_handle().to_host_sync() + ); + } + (Canonical::FixedSizeBinary(orig_prim), Canonical::FixedSizeBinary(result_prim)) => { + assert_eq!( + orig_prim.buffer_handle().to_host_sync(), + result_prim.buffer_handle().to_host_sync() + ); + } _ => { // For non-primitive types, check scalar values for i in 0..len { diff --git a/vortex-array/src/dtype/arbitrary/mod.rs b/vortex-array/src/dtype/arbitrary/mod.rs index c8ccbbb25d4..2e4e15a7d60 100644 --- a/vortex-array/src/dtype/arbitrary/mod.rs +++ b/vortex-array/src/dtype/arbitrary/mod.rs @@ -49,11 +49,12 @@ fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result { 3 => DType::Decimal(u.arbitrary()?, u.arbitrary()?), 4 => DType::Utf8(u.arbitrary()?), 5 => DType::Binary(u.arbitrary()?), + 6 => DType::FixedSizeBinary(u.int_in_range(0..=32)?, u.arbitrary()?), // container types - 6 => DType::Struct(random_struct_dtype(u, depth - 1)?, u.arbitrary()?), - 7 => DType::List(Arc::new(random_dtype(u, depth - 1)?), u.arbitrary()?), - 8 => DType::FixedSizeList( + 7 => DType::Struct(random_struct_dtype(u, depth - 1)?, u.arbitrary()?), + 8 => DType::List(Arc::new(random_dtype(u, depth - 1)?), u.arbitrary()?), + 9 => DType::FixedSizeList( Arc::new(random_dtype(u, depth - 1)?), // We limit the list size to 3 rather (following random struct fields). u.choose_index(3)?.try_into().vortex_expect("impossible"), diff --git a/vortex-array/src/dtype/dtype_impl.rs b/vortex-array/src/dtype/dtype_impl.rs index 7de16913c6b..391a6a8e022 100644 --- a/vortex-array/src/dtype/dtype_impl.rs +++ b/vortex-array/src/dtype/dtype_impl.rs @@ -63,6 +63,7 @@ impl DType { | Decimal(_, null) | Utf8(null) | Binary(null) + | FixedSizeBinary(_, null) | List(_, null) | FixedSizeList(_, _, null) | Struct(_, null) @@ -93,6 +94,7 @@ impl DType { Decimal(ddt, _) => Decimal(*ddt, nullability), Utf8(_) => Utf8(nullability), Binary(_) => Binary(nullability), + FixedSizeBinary(size, _) => FixedSizeBinary(*size, nullability), List(edt, _) => List(Arc::clone(edt), nullability), FixedSizeList(edt, size, _) => FixedSizeList(Arc::clone(edt), *size, nullability), Struct(sf, _) => Struct(sf.clone(), nullability), @@ -117,6 +119,7 @@ impl DType { (Decimal(lhs, _), Decimal(rhs, _)) => lhs == rhs, (Utf8(_), Utf8(_)) => true, (Binary(_), Binary(_)) => true, + (FixedSizeBinary(lhs, _), FixedSizeBinary(rhs, _)) => lhs == rhs, (List(lhs_dtype, _), List(rhs_dtype, _)) => lhs_dtype.eq_ignore_nullability(rhs_dtype), (FixedSizeList(lhs_dtype, lhs_size, _), FixedSizeList(rhs_dtype, rhs_size, _)) => { lhs_size == rhs_size && lhs_dtype.eq_ignore_nullability(rhs_dtype) @@ -139,6 +142,7 @@ impl DType { match self { Null | Bool(_) | Utf8(_) | Binary(_) | Variant(_) => {} + FixedSizeBinary(size, _) => size.hash(state), Primitive(ptype, _) => ptype.hash(state), Decimal(decimal, _) => decimal.hash(state), List(element, _) => element.hash_ignore_nullability(state), @@ -252,6 +256,11 @@ impl DType { matches!(self, Binary(_)) } + /// Check if `self` is a [`DType::FixedSizeBinary`]. + pub fn is_fixed_size_binary(&self) -> bool { + matches!(self, FixedSizeBinary(..)) + } + /// Check if `self` is a [`DType::List`]. pub fn is_list(&self) -> bool { matches!(self, List(_, _)) @@ -306,6 +315,7 @@ impl DType { Some(DecimalType::smallest_decimal_value_type(decimal).byte_width()) } Utf8(_) | Binary(_) | List(..) => None, + FixedSizeBinary(size, _) => Some(*size as usize), FixedSizeList(elem_dtype, list_size, _) => { elem_dtype.element_size().map(|s| s * *list_size as usize) } @@ -506,6 +516,7 @@ impl Display for DType { Decimal(ddt, null) => write!(f, "{ddt}{null}"), Utf8(null) => write!(f, "utf8{null}"), Binary(null) => write!(f, "binary{null}"), + FixedSizeBinary(size, null) => write!(f, "fixed_size_binary[{size}]{null}"), List(edt, null) => write!(f, "list({edt}){null}"), FixedSizeList(edt, size, null) => write!(f, "fixed_size_list({edt})[{size}]{null}"), Struct(sf, null) => write!( diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index 8b579c70d39..26b02f20ee0 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -90,6 +90,9 @@ pub enum DType { /// Logical binary data. Binary(Nullability), + /// Logical binary data whose values all have the same byte width. + FixedSizeBinary(u32, Nullability), + /// A logical variable-length list type. /// /// This is parameterized by a single `DType` that represents the element type of the inner @@ -141,6 +144,7 @@ impl PartialEq for DType { (Self::Decimal(da, na), Self::Decimal(db, nb)) => da == db && na == nb, (Self::Utf8(a), Self::Utf8(b)) => a == b, (Self::Binary(a), Self::Binary(b)) => a == b, + (Self::FixedSizeBinary(sa, na), Self::FixedSizeBinary(sb, nb)) => sa == sb && na == nb, (Self::List(da, na), Self::List(db, nb)) => { na == nb && (Arc::ptr_eq(da, db) || da == db) } @@ -161,6 +165,7 @@ impl PartialEq for DType { | (Self::Decimal(..), _) | (Self::Utf8(_), _) | (Self::Binary(_), _) + | (Self::FixedSizeBinary(..), _) | (Self::List(..), _) | (Self::FixedSizeList(..), _) | (Self::Struct(..), _) diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index a796490f58a..4266006cdb8 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -183,6 +183,12 @@ impl TryFrom for DType { .nullable() .into(), )), + fb::Type::FixedSizeBinary => { + let fixed = fb.type__as_fixed_size_binary().ok_or_else(|| { + vortex_err!("failed to parse fixed-size binary from flatbuffer") + })?; + Ok(Self::FixedSizeBinary(fixed.size(), fixed.nullable().into())) + } fb::Type::List => { let fb_list = fb .type__as_list() @@ -334,6 +340,14 @@ impl WriteFlatBuffer for DType { }, ) .as_union_value(), + Self::FixedSizeBinary(size, n) => fb::FixedSizeBinary::create( + fbb, + &fb::FixedSizeBinaryArgs { + size: *size, + nullable: (*n).into(), + }, + ) + .as_union_value(), Self::List(edt, n) => { let element_type = Some(edt.as_ref().write_flatbuffer(fbb)?); fb::List::create( @@ -446,6 +460,7 @@ impl WriteFlatBuffer for DType { Self::Decimal(..) => fb::Type::Decimal, Self::Utf8(_) => fb::Type::Utf8, Self::Binary(_) => fb::Type::Binary, + Self::FixedSizeBinary(..) => fb::Type::FixedSizeBinary, Self::List(..) => fb::Type::List, Self::FixedSizeList(..) => fb::Type::FixedSizeList, Self::Struct(..) => fb::Type::Struct_, diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 8389a90200a..e19420306f6 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -46,6 +46,7 @@ impl DType { )), DtypeType::Utf8(u) => Ok(Self::Utf8(u.nullable.into())), DtypeType::Binary(b) => Ok(Self::Binary(b.nullable.into())), + DtypeType::FixedSizeBinary(b) => Ok(Self::FixedSizeBinary(b.size, b.nullable.into())), DtypeType::List(l) => { let nullable = l.nullable.into(); Ok(Self::List( @@ -154,6 +155,12 @@ impl TryFrom<&DType> for pb::DType { DType::Binary(null) => DtypeType::Binary(pb::Binary { nullable: (*null).into(), }), + DType::FixedSizeBinary(size, null) => { + DtypeType::FixedSizeBinary(pb::FixedSizeBinary { + size: *size, + nullable: (*null).into(), + }) + } DType::List(edt, null) => DtypeType::List(Box::new(pb::List { element_type: Some(Box::new(edt.as_ref().try_into()?)), nullable: (*null).into(), @@ -287,6 +294,8 @@ mod tests { DType::Primitive(PType::F64, Nullability::NonNullable), DType::Utf8(Nullability::Nullable), DType::Binary(Nullability::NonNullable), + DType::FixedSizeBinary(0, Nullability::NonNullable), + DType::FixedSizeBinary(16, Nullability::Nullable), ]; for dtype in test_cases { diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index b3870b8c03e..bc88b1d2c61 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -96,6 +96,13 @@ impl Serialize for DType { } DType::Utf8(n) => serializer.serialize_newtype_variant("DType", 4, "Utf8", n), DType::Binary(n) => serializer.serialize_newtype_variant("DType", 5, "Binary", n), + DType::FixedSizeBinary(size, n) => { + let mut state = + serializer.serialize_tuple_variant("DType", 12, "FixedSizeBinary", 2)?; + state.serialize_field(size)?; + state.serialize_field(n)?; + state.end() + } DType::List(element_dtype, n) => { let mut state = serializer.serialize_tuple_variant("DType", 6, "List", 2)?; state.serialize_field(element_dtype.as_ref())?; @@ -175,6 +182,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { "Decimal", "Utf8", "Binary", + "FixedSizeBinary", "List", "FixedSizeList", "Struct", @@ -228,6 +236,12 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, DType> { let n = access.newtype_variant()?; Ok(DType::Binary(n)) } + "FixedSizeBinary" => { + #[derive(Deserialize)] + struct Fields(u32, Nullability); + let Fields(size, n) = access.newtype_variant()?; + Ok(DType::FixedSizeBinary(size, n)) + } "List" => access.newtype_variant_seed(ListFieldsSeed { session: self.session, }), diff --git a/vortex-array/src/scalar/arbitrary.rs b/vortex-array/src/scalar/arbitrary.rs index 49528f05a68..cbfd50f61ac 100644 --- a/vortex-array/src/scalar/arbitrary.rs +++ b/vortex-array/src/scalar/arbitrary.rs @@ -64,6 +64,15 @@ pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result { ))), ) .vortex_expect("unable to construct random `Scalar`_"), + DType::FixedSizeBinary(size, _) => Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Binary(ByteBuffer::from( + (0..*size) + .map(|_| u.arbitrary::()) + .collect::>>()?, + ))), + ) + .vortex_expect("unable to construct random `Scalar`_"), DType::List(edt, _) => Scalar::try_new( dtype.clone(), Some(ScalarValue::Tuple( diff --git a/vortex-array/src/scalar/cast.rs b/vortex-array/src/scalar/cast.rs index a4bb04e05e0..f79baaa4cef 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -55,7 +55,7 @@ impl Scalar { DType::Primitive(..) => self.as_primitive().cast(target_dtype), DType::Decimal(..) => self.as_decimal().cast(target_dtype), DType::Utf8(_) => self.as_utf8().cast(target_dtype), - DType::Binary(_) => self.as_binary().cast(target_dtype), + DType::Binary(_) | DType::FixedSizeBinary(..) => self.as_binary().cast(target_dtype), DType::List(..) | DType::FixedSizeList(..) => self.as_list().cast(target_dtype), DType::Struct(..) => self.as_struct().cast(target_dtype), DType::Union(..) => vortex_bail!( diff --git a/vortex-array/src/scalar/constructor.rs b/vortex-array/src/scalar/constructor.rs index 234e95441a7..5d80da4881c 100644 --- a/vortex-array/src/scalar/constructor.rs +++ b/vortex-array/src/scalar/constructor.rs @@ -106,6 +106,18 @@ impl Scalar { .vortex_expect("unable to construct a binary `Scalar`") } + /// Creates a new fixed-size binary scalar from a byte buffer. + pub fn fixed_size_binary(buffer: impl Into, nullability: Nullability) -> Self { + let buffer = buffer.into(); + let size = u32::try_from(buffer.len()) + .vortex_expect("fixed-size binary values must fit in a u32 byte width"); + Self::try_new( + DType::FixedSizeBinary(size, nullability), + Some(ScalarValue::Binary(buffer)), + ) + .vortex_expect("unable to construct a fixed-size binary `Scalar`") + } + /// Creates a new list scalar with the given element type and children. /// /// # Panics diff --git a/vortex-array/src/scalar/display.rs b/vortex-array/src/scalar/display.rs index 2947afe5a51..0d9b849f2b6 100644 --- a/vortex-array/src/scalar/display.rs +++ b/vortex-array/src/scalar/display.rs @@ -17,7 +17,9 @@ impl Display for Scalar { DType::Primitive(..) => write!(f, "{}", self.as_primitive()), DType::Decimal(..) => write!(f, "{}", self.as_decimal()), DType::Utf8(_) => write!(f, "{}", self.as_utf8()), - DType::Binary(_) => write!(f, "{}", self.as_binary()), + DType::Binary(_) | DType::FixedSizeBinary(..) => { + write!(f, "{}", self.as_binary()) + } DType::List(..) | DType::FixedSizeList(..) => write!(f, "{}", self.as_list()), DType::Struct(..) => write!(f, "{}", self.as_struct()), DType::Union(..) => write!(f, "{}", self.as_union()), diff --git a/vortex-array/src/scalar/proto.rs b/vortex-array/src/scalar/proto.rs index d3f7894f86e..4450525377c 100644 --- a/vortex-array/src/scalar/proto.rs +++ b/vortex-array/src/scalar/proto.rs @@ -403,7 +403,9 @@ fn string_from_proto(s: &str, dtype: &DType) -> VortexResult { fn bytes_from_proto(bytes: &[u8], dtype: &DType) -> VortexResult { match dtype { DType::Utf8(_) => Ok(ScalarValue::Utf8(BufferString::try_from(bytes)?)), - DType::Binary(_) => Ok(ScalarValue::Binary(ByteBuffer::copy_from(bytes))), + DType::Binary(_) | DType::FixedSizeBinary(..) => { + Ok(ScalarValue::Binary(ByteBuffer::copy_from(bytes))) + } // TODO(connor): This is incorrect, we need to verify this matches the inner decimal_dtype. DType::Decimal(..) => Ok(ScalarValue::Decimal(match bytes.len() { 1 => DecimalValue::I8(bytes[0] as i8), @@ -668,6 +670,18 @@ mod tests { )); } + #[test] + fn test_scalar_value_serde_roundtrip_fixed_size_binary() { + round_trip(Scalar::fixed_size_binary( + ByteBuffer::copy_from(b"0123456789abcdef"), + Nullability::Nullable, + )); + round_trip(Scalar::null(DType::FixedSizeBinary( + 16, + Nullability::Nullable, + ))); + } + #[test] fn test_scalar_value_serde_roundtrip_utf8() { round_trip(Scalar::utf8("hello", Nullability::NonNullable)); diff --git a/vortex-array/src/scalar/scalar_impl.rs b/vortex-array/src/scalar/scalar_impl.rs index a4351932152..3c863062f9b 100644 --- a/vortex-array/src/scalar/scalar_impl.rs +++ b/vortex-array/src/scalar/scalar_impl.rs @@ -211,6 +211,7 @@ impl Scalar { DType::Decimal(..) => value.as_decimal().is_zero(), DType::Utf8(_) => value.as_utf8().is_empty(), DType::Binary(_) => value.as_binary().is_empty(), + DType::FixedSizeBinary(..) => value.as_binary().iter().all(|byte| *byte == 0), DType::List(..) => value.as_list().is_empty(), // A fixed-size list is zero only if it has the expected number of elements and every // element is itself a non-null zero value.1 @@ -290,7 +291,7 @@ impl Scalar { DType::Utf8(_) => self .value() .map_or_else(|| 0, |value| value.as_utf8().len()), - DType::Binary(_) => self + DType::Binary(_) | DType::FixedSizeBinary(..) => self .value() .map_or_else(|| 0, |value| value.as_binary().len()), DType::List(..) | DType::FixedSizeList(..) => self diff --git a/vortex-array/src/scalar/scalar_value.rs b/vortex-array/src/scalar/scalar_value.rs index c898832d96c..22a438bdba7 100644 --- a/vortex-array/src/scalar/scalar_value.rs +++ b/vortex-array/src/scalar/scalar_value.rs @@ -60,6 +60,7 @@ impl ScalarValue { DType::Decimal(dt, ..) => Self::Decimal(DecimalValue::zero(dt)), DType::Utf8(_) => Self::Utf8(BufferString::empty()), DType::Binary(_) => Self::Binary(ByteBuffer::empty()), + DType::FixedSizeBinary(size, _) => Self::Binary(ByteBuffer::zeroed(*size as usize)), DType::List(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { let elements = (0..*size) @@ -111,6 +112,7 @@ impl ScalarValue { DType::Decimal(dt, ..) => Self::Decimal(DecimalValue::zero(dt)), DType::Utf8(_) => Self::Utf8(BufferString::empty()), DType::Binary(_) => Self::Binary(ByteBuffer::empty()), + DType::FixedSizeBinary(size, _) => Self::Binary(ByteBuffer::zeroed(*size as usize)), DType::List(..) => Self::Tuple(vec![]), DType::FixedSizeList(edt, size, _) => { let elements = (0..*size) diff --git a/vortex-array/src/scalar/typed_view/binary.rs b/vortex-array/src/scalar/typed_view/binary.rs index cb107e3613c..d62115da41c 100644 --- a/vortex-array/src/scalar/typed_view/binary.rs +++ b/vortex-array/src/scalar/typed_view/binary.rs @@ -68,8 +68,10 @@ impl<'a> BinaryScalar<'a> { /// /// Returns an error if the data type is not a binary type. pub fn try_new(dtype: &'a DType, value: Option<&'a ScalarValue>) -> VortexResult { - if !matches!(dtype, DType::Binary(..)) { - vortex_bail!("Can only construct binary scalar from binary dtype, found {dtype}") + if !matches!(dtype, DType::Binary(..) | DType::FixedSizeBinary(..)) { + vortex_bail!( + "Can only construct binary scalar from binary or fixed-size binary dtype, found {dtype}" + ) } Ok(Self { @@ -92,9 +94,9 @@ impl<'a> BinaryScalar<'a> { /// Casts this scalar to the given `dtype`. pub(crate) fn cast(&self, dtype: &DType) -> VortexResult { - if !matches!(dtype, DType::Binary(..)) { + if !matches!(dtype, DType::Binary(..) | DType::FixedSizeBinary(..)) { vortex_bail!( - "Cannot cast binary to {dtype}: binary scalars can only be cast to binary types with different nullability" + "Cannot cast binary to {dtype}: binary scalars can only be cast to binary or fixed-size binary types" ) } Scalar::try_new( diff --git a/vortex-array/src/scalar/validate.rs b/vortex-array/src/scalar/validate.rs index f81147376b8..fb92a16efe2 100644 --- a/vortex-array/src/scalar/validate.rs +++ b/vortex-array/src/scalar/validate.rs @@ -74,6 +74,17 @@ impl Scalar { "binary dtype expected Binary value, got {value}", ); } + DType::FixedSizeBinary(size, _) => { + let ScalarValue::Binary(buffer) = value else { + vortex_bail!("fixed-size binary dtype expected Binary value, got {value}"); + }; + vortex_ensure_eq!( + buffer.len(), + *size as usize, + "fixed-size binary dtype expected {size} bytes, got {}", + buffer.len(), + ); + } DType::List(elem_dtype, _) => { let ScalarValue::Tuple(elements) = value else { vortex_bail!("list dtype expected Tuple value, got {value}"); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 70e98639b54..2b4a9213c71 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -214,9 +214,10 @@ fn compare_arrays( DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), - DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { - nested::compare_nested(lhs, rhs, op, nullability, ctx) - } + DType::FixedSizeBinary(..) + | DType::Struct(..) + | DType::List(..) + | DType::FixedSizeList(..) => nested::compare_nested(lhs, rhs, op, nullability, ctx), DType::Union(..) | DType::Variant(_) | DType::Extension(_) => { vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs index 56d8cbf0496..058177a54a5 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -161,6 +161,24 @@ fn build_values_comparator( let rhs = rhs.clone().execute::(ctx)?; Box::new(move |i, j| view_bytes(&lhs, i).cmp(view_bytes(&rhs, j))) } + DType::FixedSizeBinary(byte_width, _) => { + let byte_width = *byte_width as usize; + let lhs = lhs + .clone() + .execute::(ctx)? + .buffer_handle() + .to_host_sync(); + let rhs = rhs + .clone() + .execute::(ctx)? + .buffer_handle() + .to_host_sync(); + Box::new(move |i, j| { + let lhs_start = i * byte_width; + let rhs_start = j * byte_width; + lhs[lhs_start..lhs_start + byte_width].cmp(&rhs[rhs_start..rhs_start + byte_width]) + }) + } DType::Struct(..) => { let lhs = lhs.clone().execute::(ctx)?; let rhs = rhs.clone().execute::(ctx)?; diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 20852779d42..d20b3437c2a 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -24,6 +24,7 @@ use crate::arrays::Bool; use crate::arrays::Constant; use crate::arrays::Decimal; use crate::arrays::Extension; +use crate::arrays::FixedSizeBinary; use crate::arrays::FixedSizeList; use crate::arrays::ListView; use crate::arrays::Null; @@ -181,6 +182,7 @@ fn cast_canonical( CanonicalView::Bool(a) => ::cast(a, dtype, ctx), CanonicalView::Primitive(a) => ::cast(a, dtype, ctx), CanonicalView::Decimal(a) => ::cast(a, dtype, ctx), + CanonicalView::FixedSizeBinary(a) => ::cast(a, dtype, ctx), CanonicalView::VarBinView(a) => ::cast(a, dtype, ctx), CanonicalView::List(a) => ::cast(a, dtype, ctx), CanonicalView::FixedSizeList(a) => ::cast(a, dtype, ctx), diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index c9eae72ad08..819a7fb2102 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -18,6 +18,7 @@ use crate::ColumnarView; use crate::ExecutionCtx; use crate::arrays::Bool; use crate::arrays::Decimal; +use crate::arrays::FixedSizeBinary; use crate::arrays::Primitive; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -166,6 +167,10 @@ fn fill_null_canonical( } CanonicalView::Decimal(a) => ::fill_null(a, fill_value, ctx)? .ok_or_else(|| vortex_err!("FillNullKernel for DecimalArray returned None")), + CanonicalView::FixedSizeBinary(a) => { + ::fill_null(a, fill_value, ctx)? + .ok_or_else(|| vortex_err!("FillNullKernel for FixedSizeBinaryArray returned None")) + } other => vortex_bail!( "No FillNullKernel for canonical array {}", other.to_array_ref().encoding_id() diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 98a769ceea9..6e16797529c 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -20,6 +20,7 @@ use crate::arrays::Constant; use crate::arrays::Decimal; use crate::arrays::Dict; use crate::arrays::Extension; +use crate::arrays::FixedSizeBinary; use crate::arrays::FixedSizeList; use crate::arrays::List; use crate::arrays::ListView; @@ -69,6 +70,7 @@ impl Default for ArraySession { this.register(Bool); this.register(Primitive); this.register(Decimal); + this.register(FixedSizeBinary); this.register(VarBinView); this.register(ListView); this.register(FixedSizeList); diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index e507140ef05..737a8f38262 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -8,6 +8,7 @@ use arrow_array::Array as ArrowArray; use arrow_array::ArrowPrimitiveType; use arrow_array::BooleanArray as ArrowBooleanArray; use arrow_array::DictionaryArray; +use arrow_array::FixedSizeBinaryArray as ArrowFixedSizeBinaryArray; use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray; use arrow_array::GenericByteArray; use arrow_array::GenericByteViewArray; @@ -62,6 +63,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::DictArray; +use vortex_array::arrays::FixedSizeBinaryArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; @@ -223,6 +225,16 @@ impl FromArrowArray<&ArrowPrimitiveArray> for ArrayRef { } } +impl FromArrowArray<&ArrowFixedSizeBinaryArray> for ArrayRef { + fn from_arrow(value: &ArrowFixedSizeBinaryArray, nullable: bool) -> VortexResult { + let byte_width = u32::try_from(value.value_length()) + .map_err(|_| vortex_err!("Arrow fixed-size binary width cannot be negative"))?; + let values = ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::()); + let validity = nulls(value.nulls(), nullable)?; + Ok(FixedSizeBinaryArray::new(values, byte_width, value.len(), validity).into_array()) + } +} + macro_rules! impl_from_arrow_temporal { ($T:path) => { impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef { @@ -546,6 +558,9 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), + DataType::FixedSizeBinary(..) => { + Self::from_arrow(array.as_fixed_size_binary(), nullable) + } DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { diff --git a/vortex-arrow/src/dtype.rs b/vortex-arrow/src/dtype.rs index a87c931f3df..8aa1fb00dea 100644 --- a/vortex-arrow/src/dtype.rs +++ b/vortex-arrow/src/dtype.rs @@ -198,6 +198,12 @@ impl TryFromArrowType<(&DataType, Nullability)> for DType { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { DType::Binary(nullability) } + DataType::FixedSizeBinary(byte_width) => DType::FixedSizeBinary( + u32::try_from(*byte_width).map_err(|_| { + vortex_err!("Arrow fixed-size binary width cannot be negative: {byte_width}") + })?, + nullability, + ), DataType::Date32 => DType::Extension(Date::new(TimeUnit::Days, nullability).erased()), DataType::Date64 => { DType::Extension(Date::new(TimeUnit::Milliseconds, nullability).erased()) @@ -354,6 +360,10 @@ pub(crate) fn to_data_type_naive(dtype: &DType) -> VortexResult { } DType::Utf8(_) => DataType::Utf8View, DType::Binary(_) => DataType::BinaryView, + DType::FixedSizeBinary(byte_width, _) => DataType::FixedSizeBinary( + i32::try_from(*byte_width) + .map_err(|_| vortex_err!("Fixed-size binary width exceeds Arrow i32 range"))?, + ), // There are four kinds of lists: List (32-bit offsets), Large List (64-bit), List View // (32-bit), Large List View (64-bit). We cannot both guarantee zero-copy and commit to an // Arrow dtype because we do not how large our offsets are. @@ -643,7 +653,6 @@ mod test { #[rstest] #[case::duration(DataType::Duration(ArrowTimeUnit::Microsecond))] #[case::interval(DataType::Interval(arrow_schema::IntervalUnit::DayTime))] - #[case::fixed_size_binary(DataType::FixedSizeBinary(3))] fn test_try_from_arrow_unsupported_type_errors(#[case] data_type: DataType) { let err = DType::try_from_arrow((&data_type, Nullability::NonNullable)) .expect_err("unsupported Arrow type should not convert") diff --git a/vortex-arrow/src/executor/fixed_size_binary.rs b/vortex-arrow/src/executor/fixed_size_binary.rs new file mode 100644 index 00000000000..f1051e4d03d --- /dev/null +++ b/vortex-arrow/src/executor/fixed_size_binary.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::FixedSizeBinaryArray as ArrowFixedSizeBinaryArray; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::FixedSizeBinaryArray; +use vortex_array::arrays::FixedSizeBinaryArrayExt; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::null_buffer::to_null_buffer; + +pub(super) fn to_arrow_fixed_size_binary( + array: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let array = array.execute::(ctx)?; + let byte_width = i32::try_from(array.byte_width()) + .map_err(|_| vortex_err!("Fixed-size binary width exceeds Arrow i32 range"))?; + let nulls = to_null_buffer(array.as_ref().validity()?.execute_mask(array.len(), ctx)?); + Ok(Arc::new(ArrowFixedSizeBinaryArray::new( + byte_width, + array.buffer_handle().to_host_sync().into_arrow_buffer(), + nulls, + ))) +} + +#[cfg(test)] +mod tests { + use arrow_array::Array as ArrowArray; + use arrow_array::cast::AsArray; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::FixedSizeBinaryArray; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::to_arrow_fixed_size_binary; + + #[test] + fn exports_values_and_validity() -> VortexResult<()> { + let array = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4, 5, 6].into_byte_buffer(), + 2, + 3, + Validity::from_iter([true, false, true]), + ); + let mut ctx = array_session().create_execution_ctx(); + let arrow = to_arrow_fixed_size_binary(array.into_array(), &mut ctx)?; + let arrow = arrow.as_fixed_size_binary(); + + assert_eq!(arrow.value_length(), 2); + assert_eq!(arrow.value(0), &[1, 2]); + assert!(arrow.is_null(1)); + assert_eq!(arrow.value(2), &[5, 6]); + Ok(()) + } +} diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index 875577788b4..b62eda81587 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -11,6 +11,7 @@ mod byte; pub mod byte_view; mod decimal; mod dictionary; +mod fixed_size_binary; mod fixed_size_list; mod list; mod list_view; @@ -48,6 +49,7 @@ use crate::executor::byte::to_arrow_byte_array; use crate::executor::byte_view::to_arrow_byte_view; use crate::executor::decimal::to_arrow_decimal; use crate::executor::dictionary::to_arrow_dictionary; +use crate::executor::fixed_size_binary::to_arrow_fixed_size_binary; use crate::executor::fixed_size_list::to_arrow_fixed_list; use crate::executor::list::to_arrow_list; use crate::executor::list_view::to_arrow_list_view; @@ -162,6 +164,7 @@ pub(crate) fn execute_arrow_naive( DataType::FixedSizeList(elements_field, list_size) => { to_arrow_fixed_list(array, *list_size, elements_field, ctx) } + DataType::FixedSizeBinary(_) => to_arrow_fixed_size_binary(array, ctx), // TODO(joe): pass down preferred DataType::ListView(elements_field) => to_arrow_list_view::(array, elements_field, ctx), // TODO(joe): pass down preferred @@ -191,11 +194,7 @@ pub(crate) fn execute_arrow_naive( dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), - DataType::FixedSizeBinary(_) - | DataType::Map(..) - | DataType::Duration(_) - | DataType::Interval(_) - | DataType::Union(..) => { + DataType::Map(..) | DataType::Duration(_) | DataType::Interval(_) | DataType::Union(..) => { vortex_bail!("Conversion to Arrow type {resolved_type} is not supported"); } }?; diff --git a/vortex-arrow/src/scalar.rs b/vortex-arrow/src/scalar.rs index aeac56f43d3..80a08e0ee54 100644 --- a/vortex-arrow/src/scalar.rs +++ b/vortex-arrow/src/scalar.rs @@ -7,6 +7,8 @@ use std::sync::Arc; use arrow_array::Scalar as ArrowScalar; use arrow_array::*; +use arrow_buffer::Buffer; +use arrow_buffer::NullBuffer; use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::extension::datetime::AnyTemporal; @@ -68,6 +70,18 @@ impl ToArrowDatum for Scalar { DType::Decimal(..) => decimal_to_arrow(value.as_decimal()), DType::Utf8(_) => utf8_to_arrow(value.as_utf8()), DType::Binary(_) => binary_to_arrow(value.as_binary()), + DType::FixedSizeBinary(byte_width, _) => { + let byte_width = i32::try_from(*byte_width) + .map_err(|_| vortex_err!("Fixed-size binary width exceeds Arrow i32 range"))?; + match value.as_binary().value() { + Some(bytes) => Ok(Arc::new(FixedSizeBinaryArray::new_scalar(bytes))), + None => Ok(Arc::new(ArrowScalar::new(FixedSizeBinaryArray::new( + byte_width, + Buffer::from(vec![0; byte_width as usize]), + Some(NullBuffer::new_null(SCALAR_ARRAY_LEN)), + )))), + } + } DType::List(..) => unimplemented!("list scalar conversion"), DType::FixedSizeList(..) => unimplemented!("fixed-size list scalar conversion"), DType::Struct(..) => unimplemented!("struct scalar conversion"), diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..ff745a1dd96 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -110,11 +110,14 @@ impl CascadingCompressor { Canonical::Bool(bool_array) => { self.choose_and_compress(Canonical::Bool(bool_array), compress_ctx, exec_ctx) } - Canonical::Primitive(primitive) => { - self.choose_and_compress(Canonical::Primitive(primitive), compress_ctx, exec_ctx) + Canonical::Primitive(array) => { + self.choose_and_compress(Canonical::Primitive(array), compress_ctx, exec_ctx) } - Canonical::Decimal(decimal) => { - self.choose_and_compress(Canonical::Decimal(decimal), compress_ctx, exec_ctx) + Canonical::Decimal(array) => { + self.choose_and_compress(Canonical::Decimal(array), compress_ctx, exec_ctx) + } + Canonical::FixedSizeBinary(array) => { + self.choose_and_compress(Canonical::FixedSizeBinary(array), compress_ctx, exec_ctx) } Canonical::Struct(struct_array) => { let fields = struct_array diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index c677d96ed0a..6e6d6f5ea4d 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -260,6 +260,15 @@ fn export_canonical( export_fixed_size(buffer, len, 0, validity_buffer, null_count, ctx) } + Canonical::FixedSizeBinary(array) => { + let len = array.len(); + let validity = array.validity()?; + let buffer = array.buffer_handle().clone(); + let (validity_buffer, null_count) = + export_arrow_validity_buffer(validity, len, 0, ctx).await?; + let buffer = ctx.ensure_on_device(buffer).await?; + export_fixed_size(buffer, len, 0, validity_buffer, null_count, ctx) + } Canonical::Null(null_array) => { let len = null_array.len(); diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 555bc17d037..45424c7a593 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -11,12 +11,14 @@ use vortex::array::VortexSessionExecute; use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::FixedSizeBinaryArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::decimal::DecimalDataParts; use vortex::array::arrays::extension::ExtensionArrayExt; +use vortex::array::arrays::fixed_size_binary::FixedSizeBinaryArrayExt; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; use vortex::array::arrays::varbinview::BinaryView; @@ -115,6 +117,15 @@ impl CanonicalCudaExt for Canonical { ) })) } + Canonical::FixedSizeBinary(array) => { + let byte_width = array.byte_width(); + let len = array.len(); + let validity = array.validity()?; + let values = array.buffer_handle().clone().try_into_host()?.await?; + Ok(Canonical::FixedSizeBinary(FixedSizeBinaryArray::new( + values, byte_width, len, validity, + ))) + } Canonical::VarBinView(varbinview) => { let VarBinViewDataParts { views, diff --git a/vortex-cuda/src/kernel/filter/mod.rs b/vortex-cuda/src/kernel/filter/mod.rs index 3c504113b11..d206510af0e 100644 --- a/vortex-cuda/src/kernel/filter/mod.rs +++ b/vortex-cuda/src/kernel/filter/mod.rs @@ -69,16 +69,19 @@ impl CudaExecute for FilterExecutor { m @ Mask::Values(_) => { let canonical = child.execute_cuda(ctx).await?; match canonical { - Canonical::Primitive(prim) => { - match_each_native_simd_ptype!(prim.ptype(), |T| { - filter_primitive::(prim, m, ctx).await + Canonical::Primitive(array) => { + match_each_native_simd_ptype!(array.ptype(), |T| { + filter_primitive::(array, m, ctx).await }) } - Canonical::Decimal(decimal) => { - match_each_decimal_value_type!(decimal.values_type(), |D| { - filter_decimal::(decimal, m, ctx).await + Canonical::Decimal(array) => { + match_each_decimal_value_type!(array.values_type(), |D| { + filter_decimal::(array, m, ctx).await }) } + Canonical::FixedSizeBinary(array) => { + unimplemented!("CUDA filter for {}", array.dtype()) + } Canonical::VarBinView(varbinview) => { filter_varbinview(varbinview, m, ctx).await } diff --git a/vortex-cuda/src/kernel/slice/mod.rs b/vortex-cuda/src/kernel/slice/mod.rs index 30b1fdaf640..98cbd488332 100644 --- a/vortex-cuda/src/kernel/slice/mod.rs +++ b/vortex-cuda/src/kernel/slice/mod.rs @@ -45,11 +45,15 @@ impl CudaExecute for SliceExecutor { .into_array() .slice(range)? .execute::(ctx.execution_ctx()), - Canonical::Primitive(prim_array) => prim_array + Canonical::Primitive(array) => array .into_array() .slice(range)? .execute::(ctx.execution_ctx()), - Canonical::Decimal(decimal_array) => decimal_array + Canonical::Decimal(array) => array + .into_array() + .slice(range)? + .execute::(ctx.execution_ctx()), + Canonical::FixedSizeBinary(array) => array .into_array() .slice(range)? .execute::(ctx.execution_ctx()), diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index ec262b479d9..1901a90bde7 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -119,6 +119,15 @@ impl TryToDataFusion for Scalar { .cloned() .map(|b| Vec::::from(b.into_inner())), ), + DType::FixedSizeBinary(byte_width, _) => ScalarValue::FixedSizeBinary( + i32::try_from(*byte_width).map_err(|_| { + vortex_err!("Fixed-size binary width exceeds DataFusion i32 range") + })?, + self.as_binary() + .value() + .cloned() + .map(|b| Vec::::from(b.into_inner())), + ), dtype @ DType::List(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" ), @@ -224,13 +233,22 @@ impl FromDataFusion for Scalar { .as_ref() .map(|s| Scalar::from(s.as_str())) .unwrap_or_else(|| Scalar::null(DType::Utf8(Nullability::Nullable))), - ScalarValue::Binary(b) - | ScalarValue::BinaryView(b) - | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => b + ScalarValue::Binary(b) | ScalarValue::BinaryView(b) | ScalarValue::LargeBinary(b) => b .as_ref() .map(|b| Scalar::binary(ByteBuffer::from(b.clone()), Nullability::Nullable)) .unwrap_or_else(|| Scalar::null(DType::Binary(Nullability::Nullable))), + ScalarValue::FixedSizeBinary(byte_width, b) => b + .as_ref() + .map(|b| { + Scalar::fixed_size_binary(ByteBuffer::from(b.clone()), Nullability::Nullable) + }) + .unwrap_or_else(|| { + Scalar::null(DType::FixedSizeBinary( + u32::try_from(*byte_width) + .vortex_expect("DataFusion fixed-size binary width is non-negative"), + Nullability::Nullable, + )) + }), ScalarValue::Date32(v) | ScalarValue::Time32Second(v) | ScalarValue::Time32Millisecond(v) => { diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index 138a6f1bc69..43dfc3fa053 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -232,6 +232,7 @@ impl TryFrom<&DType> for LogicalType { } DType::Utf8(_) => DUCKDB_TYPE::DUCKDB_TYPE_VARCHAR, DType::Binary(_) => DUCKDB_TYPE::DUCKDB_TYPE_BLOB, + DType::FixedSizeBinary(..) => DUCKDB_TYPE::DUCKDB_TYPE_BLOB, DType::List(element_dtype, _) => { let element_logical_type = LogicalType::try_from(element_dtype.as_ref())?; return LogicalType::list_type(element_logical_type); diff --git a/vortex-duckdb/src/convert/scalar.rs b/vortex-duckdb/src/convert/scalar.rs index 7ad12e89ca3..69a3d1f4341 100644 --- a/vortex-duckdb/src/convert/scalar.rs +++ b/vortex-duckdb/src/convert/scalar.rs @@ -78,6 +78,7 @@ impl ToDuckDBScalar for Scalar { DType::Decimal(..) => self.as_decimal().try_to_duckdb_scalar(), DType::Utf8(_) => self.as_utf8().try_to_duckdb_scalar(), DType::Binary(_) => self.as_binary().try_to_duckdb_scalar(), + DType::FixedSizeBinary(..) => self.as_binary().try_to_duckdb_scalar(), DType::List(..) => vortex_bail!("Vortex List scalars aren't supported"), DType::FixedSizeList(..) => { vortex_bail!("Vortex FixedSizeList scalars aren't supported") diff --git a/vortex-duckdb/src/exporter/canonical.rs b/vortex-duckdb/src/exporter/canonical.rs index 7ee47a40642..e3c20151ba8 100644 --- a/vortex-duckdb/src/exporter/canonical.rs +++ b/vortex-duckdb/src/exporter/canonical.rs @@ -28,6 +28,9 @@ pub(crate) fn new_exporter( Canonical::Bool(array) => bool::new_exporter(array, ctx), Canonical::Primitive(array) => primitive::new_exporter(array, ctx), Canonical::Decimal(array) => decimal::new_exporter(array, ctx), + Canonical::FixedSizeBinary(_) => { + vortex_bail!("DuckDB export does not support fixed-size binary arrays") + } Canonical::VarBinView(array) => varbinview::new_exporter(array, ctx), Canonical::List(array) => list_view::new_exporter(array, cache, ctx), Canonical::FixedSizeList(array) => fixed_size_list::new_exporter(array, cache, ctx), diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index b4c0039657e..757480294f1 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -111,6 +111,10 @@ typedef enum { * Nested fixed-size list type. */ DTYPE_FIXED_SIZE_LIST = 9, + /** + * Fixed-size binary data. + */ + DTYPE_FIXED_SIZE_BINARY = 10, } vx_dtype_variant; /** @@ -1021,6 +1025,11 @@ const vx_dtype *vx_dtype_new_utf8(bool is_nullable); */ const vx_dtype *vx_dtype_new_binary(bool is_nullable); +/** + * Create a new fixed-size binary data type. + */ +const vx_dtype *vx_dtype_new_fixed_size_binary(uint32_t byte_width, bool is_nullable); + /** * Create a new list data type. * @@ -1096,6 +1105,11 @@ const vx_dtype *vx_dtype_fixed_size_list_element(const vx_dtype *dtype); */ uint32_t vx_dtype_fixed_size_list_size(const vx_dtype *dtype); +/** + * Returns the byte width of a fixed-size binary type. + */ +uint32_t vx_dtype_fixed_size_binary_size(const vx_dtype *dtype); + /** * Checks if the type is time. */ diff --git a/vortex-ffi/examples/dtype.c b/vortex-ffi/examples/dtype.c index eb60cb72c4d..54a6364f9e1 100644 --- a/vortex-ffi/examples/dtype.c +++ b/vortex-ffi/examples/dtype.c @@ -86,6 +86,10 @@ void print_decimal_dtype(const vx_dtype *dtype) { printf("decimal(precision=%u, scale=%d)", precision, scale); } +void print_fixed_size_binary_dtype(const vx_dtype *dtype) { + printf("fixed binary(size=%u)", vx_dtype_fixed_size_binary_size(dtype)); +} + void print_dtype(const vx_dtype *dtype) { switch (vx_dtype_get_variant(dtype)) { case DTYPE_NULL: @@ -118,6 +122,9 @@ void print_dtype(const vx_dtype *dtype) { case DTYPE_DECIMAL: print_decimal_dtype(dtype); break; + case DTYPE_FIXED_SIZE_BINARY: + print_fixed_size_binary_dtype(dtype); + break; } printf("%c\n", vx_dtype_is_nullable(dtype) ? '?' : ' '); } diff --git a/vortex-ffi/src/dtype.rs b/vortex-ffi/src/dtype.rs index d38a8c62ec8..e90500d7c25 100644 --- a/vortex-ffi/src/dtype.rs +++ b/vortex-ffi/src/dtype.rs @@ -62,6 +62,8 @@ pub enum vx_dtype_variant { DTYPE_DECIMAL = 8, /// Nested fixed-size list type. DTYPE_FIXED_SIZE_LIST = 9, + /// Fixed-size binary data. + DTYPE_FIXED_SIZE_BINARY = 10, } // TODO(connor)[Union]: Do we need to add union and variant here? @@ -74,6 +76,7 @@ impl From<&DType> for vx_dtype_variant { DType::Decimal(..) => vx_dtype_variant::DTYPE_DECIMAL, DType::Utf8(_) => vx_dtype_variant::DTYPE_UTF8, DType::Binary(_) => vx_dtype_variant::DTYPE_BINARY, + DType::FixedSizeBinary(..) => vx_dtype_variant::DTYPE_FIXED_SIZE_BINARY, DType::List(..) => vx_dtype_variant::DTYPE_LIST, DType::FixedSizeList(..) => vx_dtype_variant::DTYPE_FIXED_SIZE_LIST, DType::Struct(..) => vx_dtype_variant::DTYPE_STRUCT, @@ -117,6 +120,18 @@ pub unsafe extern "C-unwind" fn vx_dtype_new_binary(is_nullable: bool) -> *const vx_dtype::new(Arc::new(DType::Binary(is_nullable.into()))) } +/// Create a new fixed-size binary data type. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_dtype_new_fixed_size_binary( + byte_width: u32, + is_nullable: bool, +) -> *const vx_dtype { + vx_dtype::new(Arc::new(DType::FixedSizeBinary( + byte_width, + is_nullable.into(), + ))) +} + /// Create a new list data type. /// /// Takes ownership of the `element` pointer. @@ -254,6 +269,15 @@ pub unsafe extern "C-unwind" fn vx_dtype_fixed_size_list_size(dtype: *const vx_d } } +/// Returns the byte width of a fixed-size binary type. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_dtype_fixed_size_binary_size(dtype: *const vx_dtype) -> u32 { + match vx_dtype::as_ref(dtype) { + DType::FixedSizeBinary(size, _) => *size, + _ => vortex_panic!("not a fixed-size binary dtype"), + } +} + /// Checks if the type is time. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_dtype_is_time(dtype: *const DType) -> bool { @@ -631,6 +655,20 @@ mod tests { } } + #[test] + fn test_dtype_fixed_size_binary() { + unsafe { + let dtype = vx_dtype_new_fixed_size_binary(16, true); + assert_eq!( + vx_dtype_get_variant(dtype), + vx_dtype_variant::DTYPE_FIXED_SIZE_BINARY + ); + assert!(vx_dtype_is_nullable(dtype)); + assert_eq!(vx_dtype_fixed_size_binary_size(dtype), 16); + vx_dtype_free(dtype); + } + } + #[test] fn test_dtype_primitive_ptype() { unsafe { @@ -656,6 +694,7 @@ mod tests { DType::Decimal(DecimalDType::new(10, 2), true.into()), DType::Utf8(false.into()), DType::Binary(true.into()), + DType::FixedSizeBinary(16, true.into()), DType::FixedSizeList( Arc::new(DType::Primitive(vortex::dtype::PType::U8, false.into())), 4, @@ -672,6 +711,9 @@ mod tests { DType::Decimal(..) => assert_eq!(variant, vx_dtype_variant::DTYPE_DECIMAL), DType::Utf8(_) => assert_eq!(variant, vx_dtype_variant::DTYPE_UTF8), DType::Binary(_) => assert_eq!(variant, vx_dtype_variant::DTYPE_BINARY), + DType::FixedSizeBinary(..) => { + assert_eq!(variant, vx_dtype_variant::DTYPE_FIXED_SIZE_BINARY) + } DType::FixedSizeList(..) => { assert_eq!(variant, vx_dtype_variant::DTYPE_FIXED_SIZE_LIST) } diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 3cedbc1d68c..845a6353132 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -17,6 +17,7 @@ use vortex_array::arrays::Constant; use vortex_array::arrays::Decimal; use vortex_array::arrays::Dict; use vortex_array::arrays::Extension; +use vortex_array::arrays::FixedSizeBinary; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::List; use vortex_array::arrays::ListView; @@ -85,6 +86,7 @@ pub static ALLOWED_ENCODINGS: LazyLock> = LazyLock::new(|| { allowed.insert(Bool.id()); allowed.insert(Primitive.id()); allowed.insert(Decimal.id()); + allowed.insert(FixedSizeBinary.id()); allowed.insert(VarBin.id()); allowed.insert(VarBinView.id()); allowed.insert(List.id()); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..665ce21aead 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -20,6 +20,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Dict; +use vortex_array::arrays::FixedSizeBinaryArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -150,6 +151,14 @@ async fn test_round_trip_many_types() { let numbers = buffer![1u32, 2, 3].into_array(); + let fixed_binary = FixedSizeBinaryArray::new( + buffer![1u8, 2, 3, 4, 5, 6].into_byte_buffer(), + 2, + 3, + Validity::from_iter([true, false, true]), + ) + .into_array(); + let decimal_2 = DecimalArray::new( buffer![100i8, 10i8, 2i8], DecimalDType::new(2, 1), @@ -188,6 +197,7 @@ async fn test_round_trip_many_types() { let st = StructArray::from_fields(&[ ("strings", strings), ("numbers", numbers), + ("fixed_binary", fixed_binary), ("decimal_2", decimal_2), ("decimal_4", decimal_4), ("decimal_9", decimal_9), diff --git a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs index fb3b19513de..44f41d05646 100644 --- a/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs +++ b/vortex-flatbuffers/flatbuffers/vortex-dtype/dtype.fbs @@ -40,6 +40,11 @@ table Binary { nullable: bool; } +table FixedSizeBinary { + size: uint32; + nullable: bool; +} + table Struct_ { names: [string]; dtypes: [DType]; @@ -87,6 +92,7 @@ union Type { FixedSizeList = 10, // This is after `Extension` for backwards compatibility. Variant = 11, Union = 12, + FixedSizeBinary = 13, } table DType { diff --git a/vortex-flatbuffers/src/generated/dtype.rs b/vortex-flatbuffers/src/generated/dtype.rs index 56536ad0b94..538ab6a83f4 100644 --- a/vortex-flatbuffers/src/generated/dtype.rs +++ b/vortex-flatbuffers/src/generated/dtype.rs @@ -126,10 +126,10 @@ impl ::flatbuffers::SimpleToVerifyInSlice for PType {} #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] pub const ENUM_MIN_TYPE: u8 = 0; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] -pub const ENUM_MAX_TYPE: u8 = 12; +pub const ENUM_MAX_TYPE: u8 = 13; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_TYPE: [Type; 13] = [ +pub const ENUM_VALUES_TYPE: [Type; 14] = [ Type::NONE, Type::Null, Type::Bool, @@ -143,6 +143,7 @@ pub const ENUM_VALUES_TYPE: [Type; 13] = [ Type::FixedSizeList, Type::Variant, Type::Union, + Type::FixedSizeBinary, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -163,9 +164,10 @@ impl Type { pub const FixedSizeList: Self = Self(10); pub const Variant: Self = Self(11); pub const Union: Self = Self(12); + pub const FixedSizeBinary: Self = Self(13); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 12; + pub const ENUM_MAX: u8 = 13; pub const ENUM_VALUES: &'static [Self] = &[ Self::NONE, Self::Null, @@ -180,6 +182,7 @@ impl Type { Self::FixedSizeList, Self::Variant, Self::Union, + Self::FixedSizeBinary, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -197,6 +200,7 @@ impl Type { Self::FixedSizeList => Some("FixedSizeList"), Self::Variant => Some("Variant"), Self::Union => Some("Union"), + Self::FixedSizeBinary => Some("FixedSizeBinary"), _ => None, } } @@ -862,6 +866,119 @@ impl ::core::fmt::Debug for Binary<'_> { ds.finish() } } +pub enum FixedSizeBinaryOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct FixedSizeBinary<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for FixedSizeBinary<'a> { + type Inner = FixedSizeBinary<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> FixedSizeBinary<'a> { + pub const VT_SIZE: ::flatbuffers::VOffsetT = 4; + pub const VT_NULLABLE: ::flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + FixedSizeBinary { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args FixedSizeBinaryArgs + ) -> ::flatbuffers::WIPOffset> { + let mut builder = FixedSizeBinaryBuilder::new(_fbb); + builder.add_size(args.size); + builder.add_nullable(args.nullable); + builder.finish() + } + + + #[inline] + pub fn size(&self) -> u32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(FixedSizeBinary::VT_SIZE, Some(0)).unwrap()} + } + #[inline] + pub fn nullable(&self) -> bool { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(FixedSizeBinary::VT_NULLABLE, Some(false)).unwrap()} + } +} + +impl ::flatbuffers::Verifiable for FixedSizeBinary<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::("size", Self::VT_SIZE, false)? + .visit_field::("nullable", Self::VT_NULLABLE, false)? + .finish(); + Ok(()) + } +} +pub struct FixedSizeBinaryArgs { + pub size: u32, + pub nullable: bool, +} +impl<'a> Default for FixedSizeBinaryArgs { + #[inline] + fn default() -> Self { + FixedSizeBinaryArgs { + size: 0, + nullable: false, + } + } +} + +pub struct FixedSizeBinaryBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FixedSizeBinaryBuilder<'a, 'b, A> { + #[inline] + pub fn add_size(&mut self, size: u32) { + self.fbb_.push_slot::(FixedSizeBinary::VT_SIZE, size, 0); + } + #[inline] + pub fn add_nullable(&mut self, nullable: bool) { + self.fbb_.push_slot::(FixedSizeBinary::VT_NULLABLE, nullable, false); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FixedSizeBinaryBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + FixedSizeBinaryBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for FixedSizeBinary<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("FixedSizeBinary"); + ds.field("size", &self.size()); + ds.field("nullable", &self.nullable()); + ds.finish() + } +} pub enum Struct_Offset {} #[derive(Copy, Clone, PartialEq)] @@ -1837,6 +1954,21 @@ impl<'a> DType<'a> { } } + #[inline] + #[allow(non_snake_case)] + pub fn type__as_fixed_size_binary(&self) -> Option> { + if self.type_type() == Type::FixedSizeBinary { + self.type_().map(|t| { + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + unsafe { FixedSizeBinary::init_from_table(t) } + }) + } else { + None + } + } + } impl ::flatbuffers::Verifiable for DType<'_> { @@ -1859,6 +1991,7 @@ impl ::flatbuffers::Verifiable for DType<'_> { Type::FixedSizeList => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::FixedSizeList", pos), Type::Variant => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::Variant", pos), Type::Union => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::Union", pos), + Type::FixedSizeBinary => v.verify_union_variant::<::flatbuffers::ForwardsUOffset>("Type::FixedSizeBinary", pos), _ => Ok(()), } })? @@ -1997,6 +2130,13 @@ impl ::core::fmt::Debug for DType<'_> { ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") } }, + Type::FixedSizeBinary => { + if let Some(x) = self.type__as_fixed_size_binary() { + ds.field("type_", &x) + } else { + ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") + } + }, _ => { let x: Option<()> = None; ds.field("type_", &x) diff --git a/vortex-proto/proto/dtype.proto b/vortex-proto/proto/dtype.proto index 11feef63a8c..c1dcec25ba1 100644 --- a/vortex-proto/proto/dtype.proto +++ b/vortex-proto/proto/dtype.proto @@ -47,6 +47,11 @@ message Binary { bool nullable = 1; } +message FixedSizeBinary { + uint32 size = 1; + bool nullable = 2; +} + message Struct { repeated string names = 1; repeated DType dtypes = 2; @@ -95,6 +100,7 @@ message DType { FixedSizeList fixed_size_list = 10; // This is after `Extension` for backwards compatibility. Variant variant = 11; Union union = 12; + FixedSizeBinary fixed_size_binary = 13; } } diff --git a/vortex-proto/src/generated/vortex.dtype.rs b/vortex-proto/src/generated/vortex.dtype.rs index 16188687ee5..ba468f4682e 100644 --- a/vortex-proto/src/generated/vortex.dtype.rs +++ b/vortex-proto/src/generated/vortex.dtype.rs @@ -32,6 +32,13 @@ pub struct Binary { #[prost(bool, tag = "1")] pub nullable: bool, } +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct FixedSizeBinary { + #[prost(uint32, tag = "1")] + pub size: u32, + #[prost(bool, tag = "2")] + pub nullable: bool, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct Struct { #[prost(string, repeated, tag = "1")] @@ -85,7 +92,10 @@ pub struct Union { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct DType { - #[prost(oneof = "d_type::DtypeType", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12")] + #[prost( + oneof = "d_type::DtypeType", + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13" + )] pub dtype_type: ::core::option::Option, } /// Nested message and enum types in `DType`. @@ -117,6 +127,8 @@ pub mod d_type { Variant(super::Variant), #[prost(message, tag = "12")] Union(super::Union), + #[prost(message, tag = "13")] + FixedSizeBinary(super::FixedSizeBinary), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/vortex-python/src/arrays/builtins/fixed_size_binary.rs b/vortex-python/src/arrays/builtins/fixed_size_binary.rs new file mode 100644 index 00000000000..4c1d2fe2b7f --- /dev/null +++ b/vortex-python/src/arrays/builtins/fixed_size_binary.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use pyo3::pyclass; +use vortex::array::arrays::FixedSizeBinary; + +use crate::arrays::native::EncodingSubclass; +use crate::arrays::native::PyNativeArray; + +/// Concrete class for fixed-size binary arrays with the `vortex.fixed_size_binary` encoding. +#[pyclass(name = "FixedSizeBinaryArray", module = "vortex", extends=PyNativeArray, frozen)] +pub(crate) struct PyFixedSizeBinaryArray; + +impl EncodingSubclass for PyFixedSizeBinaryArray { + type VTable = FixedSizeBinary; +} diff --git a/vortex-python/src/arrays/builtins/mod.rs b/vortex-python/src/arrays/builtins/mod.rs index 38f2e3886ac..cbb477abb03 100644 --- a/vortex-python/src/arrays/builtins/mod.rs +++ b/vortex-python/src/arrays/builtins/mod.rs @@ -4,12 +4,14 @@ mod chunked; mod constant; mod decimal; +mod fixed_size_binary; mod primitive; mod struct_; pub(crate) use chunked::*; pub(crate) use constant::*; pub(crate) use decimal::*; +pub(crate) use fixed_size_binary::*; pub(crate) use primitive::*; use pyo3::prelude::*; pub(crate) use struct_::*; diff --git a/vortex-python/src/arrays/mod.rs b/vortex-python/src/arrays/mod.rs index 11b28f7e704..46f886bf19e 100644 --- a/vortex-python/src/arrays/mod.rs +++ b/vortex-python/src/arrays/mod.rs @@ -193,6 +193,8 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/vortex-python/src/arrays/native.rs b/vortex-python/src/arrays/native.rs index 95310224ec3..793df78c4d0 100644 --- a/vortex-python/src/arrays/native.rs +++ b/vortex-python/src/arrays/native.rs @@ -13,6 +13,7 @@ use vortex::array::arrays::Constant; use vortex::array::arrays::Decimal; use vortex::array::arrays::Dict; use vortex::array::arrays::Extension; +use vortex::array::arrays::FixedSizeBinary; use vortex::array::arrays::FixedSizeList; use vortex::array::arrays::List; use vortex::array::arrays::Null; @@ -41,6 +42,7 @@ use crate::arrays::builtins::PyChunkedArray; use crate::arrays::builtins::PyConstantArray; use crate::arrays::builtins::PyDecimalArray; use crate::arrays::builtins::PyExtensionArray; +use crate::arrays::builtins::PyFixedSizeBinaryArray; use crate::arrays::builtins::PyFixedSizeListArray; use crate::arrays::builtins::PyListArray; use crate::arrays::builtins::PyNullArray; @@ -95,6 +97,10 @@ impl PyNativeArray { return Self::with_subclass(py, array, PyDecimalArray); } + if array.is::() { + return Self::with_subclass(py, array, PyFixedSizeBinaryArray); + } + if array.is::() { return Self::with_subclass(py, array, PyVarBinArray); } diff --git a/vortex-python/src/dtype/factory.rs b/vortex-python/src/dtype/factory.rs index e337aaf5aec..7607d001939 100644 --- a/vortex-python/src/dtype/factory.rs +++ b/vortex-python/src/dtype/factory.rs @@ -322,6 +322,17 @@ pub(super) fn dtype_binary(py: Python<'_>, nullable: bool) -> PyResult, + byte_width: u32, + nullable: bool, +) -> PyResult> { + PyDType::init(py, DType::FixedSizeBinary(byte_width, nullable.into())) +} + /// Construct a struct data type. /// /// Parameters diff --git a/vortex-python/src/dtype/mod.rs b/vortex-python/src/dtype/mod.rs index ce18ceb6d83..432687c0e20 100644 --- a/vortex-python/src/dtype/mod.rs +++ b/vortex-python/src/dtype/mod.rs @@ -84,6 +84,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { m.add_function(wrap_pyfunction!(factory::dtype_float, &m)?)?; m.add_function(wrap_pyfunction!(factory::dtype_utf8, &m)?)?; m.add_function(wrap_pyfunction!(factory::dtype_binary, &m)?)?; + m.add_function(wrap_pyfunction!(factory::dtype_fixed_size_binary, &m)?)?; m.add_function(wrap_pyfunction!(factory::dtype_struct, &m)?)?; m.add_function(wrap_pyfunction!(factory::dtype_list, &m)?)?; m.add_function(wrap_pyfunction!(factory::dtype_fixed_size_list, &m)?)?; @@ -132,6 +133,7 @@ impl PyDType { DType::Decimal(..) => Self::with_subclass(py, dtype, PyDecimalDType), DType::Utf8(..) => Self::with_subclass(py, dtype, PyUtf8DType), DType::Binary(..) => Self::with_subclass(py, dtype, PyBinaryDType), + DType::FixedSizeBinary(..) => Self::with_subclass(py, dtype, PyBinaryDType), DType::List(..) => Self::with_subclass(py, dtype, PyListDType), DType::FixedSizeList(..) => Self::with_subclass(py, dtype, PyFixedSizeListDType), DType::Struct(..) => Self::with_subclass(py, dtype, PyStructDType), diff --git a/vortex-python/src/python_repr.rs b/vortex-python/src/python_repr.rs index 56819184dca..8604b409e89 100644 --- a/vortex-python/src/python_repr.rs +++ b/vortex-python/src/python_repr.rs @@ -67,6 +67,12 @@ impl Display for DTypePythonRepr<'_> { } DType::Utf8(n) => write!(f, "utf8(nullable={})", n.python_repr()), DType::Binary(n) => write!(f, "binary(nullable={})", n.python_repr()), + DType::FixedSizeBinary(byte_width, n) => write!( + f, + "fixed_size_binary({}, nullable={})", + byte_width, + n.python_repr() + ), DType::List(edt, n) => write!( f, "list({}, nullable={})", diff --git a/vortex-python/src/scalar/into_py.rs b/vortex-python/src/scalar/into_py.rs index a3b81ab6b2e..3b2a880df26 100644 --- a/vortex-python/src/scalar/into_py.rs +++ b/vortex-python/src/scalar/into_py.rs @@ -73,7 +73,7 @@ impl<'py> IntoPyObject<'py> for PyVortex<&'_ Scalar> { .cloned() .map(PyVortex) .into_pyobject(py), - DType::Binary(_) => self + DType::Binary(_) | DType::FixedSizeBinary(..) => self .0 .as_binary() .value() diff --git a/vortex-python/src/scalar/mod.rs b/vortex-python/src/scalar/mod.rs index eea66dccc52..9b510695a15 100644 --- a/vortex-python/src/scalar/mod.rs +++ b/vortex-python/src/scalar/mod.rs @@ -111,6 +111,7 @@ impl PyScalar { DType::Decimal(..) => Self::with_subclass(py, scalar, PyDecimalScalar), DType::Utf8(..) => Self::with_subclass(py, scalar, PyUtf8Scalar), DType::Binary(..) => Self::with_subclass(py, scalar, PyBinaryScalar), + DType::FixedSizeBinary(..) => Self::with_subclass(py, scalar, PyBinaryScalar), DType::List(..) | DType::FixedSizeList(..) => { // We represent both lists and fixed-size lists with `PyListScalar` since the notion // of "fixed-size" only applies to full arrays, not scalars. diff --git a/vortex-python/test/test_array.py b/vortex-python/test/test_array.py index 7458f738e05..e73d39506de 100644 --- a/vortex-python/test/test_array.py +++ b/vortex-python/test/test_array.py @@ -25,6 +25,12 @@ def test_varbin_array_round_trip(): assert arr.to_arrow_array() == a +def test_fixed_size_binary_array_round_trip(): + a = pa.array([b"abc", None, b"def"], type=pa.binary(3)) + arr = vortex.array(a) + assert arr.to_arrow_array() == a + + def test_varbin_array_take(): a = vortex.array(pa.array(["a", "b", "c", "d"], type=pa.string_view())) assert a.take(vortex.array(pa.array([0, 2]))).to_arrow_array() == pa.array( @@ -50,7 +56,6 @@ def test_scalar_at_out_of_bounds(): [ pa.duration("us"), pa.month_day_nano_interval(), - pa.binary(3), ], ) def test_unsupported_arrow_type_raises_value_error(arrow_type: pa.DataType): diff --git a/vortex-row/src/codec.rs b/vortex-row/src/codec.rs index fd14bb254b8..d99d158c5e7 100644 --- a/vortex-row/src/codec.rs +++ b/vortex-row/src/codec.rs @@ -190,6 +190,9 @@ pub(crate) fn row_width_for_dtype(dtype: &DType) -> VortexResult { vt.byte_width(), )))) } + DType::FixedSizeBinary(..) => { + vortex_bail!("row encoding does not support FixedSizeBinary arrays") + } DType::Utf8(_) | DType::Binary(_) => Ok(RowWidth::Variable), DType::FixedSizeList(elem, n, _) => match row_width_for_dtype(elem)? { // FSL is fixed iff its element type is fixed. Add a sentinel byte for the FSL