Polish fixed-size binary implementation

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/src/aggregate_fn/fns/is_constant/fixed_size_binary.rs

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools;
+
+use crate::arrays::FixedSizeBinaryArray;
+use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
+
+pub(super) fn check_fixed_size_binary_constant(array: &FixedSizeBinaryArray) -> bool {
+    let byte_width = array.byte_width() as usize;
+    if byte_width == 0 {
+        return true;
+    }
+
+    array
+        .buffer_handle()
+        .to_host_sync()
+        .as_slice()
+        .chunks_exact(byte_width)
+        .all_equal()
+}

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

@@ -4,13 +4,13 @@
 mod bool;
 mod decimal;
 mod extension;
+mod fixed_size_binary;
 mod fixed_size_list;
 mod list;
 pub mod primitive;
 mod struct_;
 mod varbin;
 
-use itertools::Itertools;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -19,6 +19,7 @@ use vortex_session::registry::CachedId;
 use self::bool::check_bool_constant;
 use self::decimal::check_decimal_constant;
 use self::extension::check_extension_constant;
+use self::fixed_size_binary::check_fixed_size_binary_constant;
 use self::fixed_size_list::check_fixed_size_list_constant;
 use self::list::check_listview_constant;
 use self::primitive::check_primitive_constant;
@@ -36,7 +37,6 @@ 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;
@@ -399,15 +399,7 @@ impl AggregateFnVTable for IsConstant {
                 let batch_is_constant = match c {
                     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::FixedSizeBinary(a) => check_fixed_size_binary_constant(a),
                     Canonical::Bool(b) => check_bool_constant(b),
                     Canonical::VarBinView(v) => check_varbinview_constant(v),
                     Canonical::Struct(s) => check_struct_constant(s, ctx)?,

vortex-array/src/aggregate_fn/fns/is_sorted/fixed_size_binary.rs

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use super::IsSortedIteratorExt;
+use crate::ExecutionCtx;
+use crate::arrays::FixedSizeBinaryArray;
+use crate::arrays::fixed_size_binary::FixedSizeBinaryArrayExt;
+
+pub(super) fn check_fixed_size_binary_sorted(
+    array: &FixedSizeBinaryArray,
+    strict: bool,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<bool> {
+    let byte_width = array.byte_width() as usize;
+    let values = array.buffer_handle().to_host_sync();
+    let validity = array
+        .as_ref()
+        .validity()?
+        .execute_mask(array.len(), ctx)?
+        .to_bit_buffer();
+    let iter = validity.iter().enumerate().map(|(index, is_valid)| {
+        is_valid.then(|| {
+            let start = index * byte_width;
+            &values[start..start + byte_width]
+        })
+    });
+
+    Ok(if strict {
+        iter.is_strict_sorted()
+    } else {
+        iter.is_sorted()
+    })
+}

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

@@ -4,6 +4,7 @@
 mod bool;
 mod decimal;
 mod extension;
+mod fixed_size_binary;
 mod primitive;
 mod varbin;
 
@@ -13,12 +14,12 @@ 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;
 use self::decimal::check_decimal_sorted;
 use self::extension::check_extension_sorted;
+use self::fixed_size_binary::check_fixed_size_binary_sorted;
 use self::primitive::check_primitive_sorted;
 use self::varbin::check_varbinview_sorted;
 use crate::ArrayRef;
@@ -42,44 +43,6 @@ 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<bool> {
-    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<dyn Iterator<Item = bool>>
-        }
-        Mask::AllFalse(len) => Box::new(std::iter::repeat_n(false, *len)),
-        Mask::Values(values) => Box::new(values.bit_buffer().iter()),
-    };
-    let mut previous: Option<Option<&[u8]>> = 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 {

vortex-array/src/arrays/fixed_size_binary/compute.rs

@@ -0,0 +1,223 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Not;
+
+use vortex_buffer::ByteBuffer;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+
+use super::FixedSizeBinary;
+use super::FixedSizeBinaryArray;
+use super::FixedSizeBinaryArrayExt;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::array::OperationsVTable;
+use crate::array::ValidityVTable;
+use crate::arrays::BoolArray;
+use crate::arrays::Masked;
+use crate::arrays::fixed_width::FixedWidthArray;
+use crate::arrays::fixed_width::vtable as fixed_width;
+use crate::arrays::slice::SliceReduce;
+use crate::dtype::DType;
+use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::scalar::Scalar;
+use crate::scalar_fn::fns::cast::CastKernel;
+use crate::scalar_fn::fns::cast::CastReduce;
+use crate::scalar_fn::fns::fill_null::FillNullKernel;
+use crate::scalar_fn::fns::mask::MaskReduce;
+use crate::validity::Validity;
+
+#[derive(Default, Debug)]
+pub(super) struct FixedSizeBinaryMaskedValidityRule;
+
+impl ArrayParentReduceRule<FixedSizeBinary> for FixedSizeBinaryMaskedValidityRule {
+    type Parent = Masked;
+
+    fn reduce_parent(
+        &self,
+        array: ArrayView<'_, FixedSizeBinary>,
+        parent: ArrayView<'_, Masked>,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        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<FixedSizeBinary> for FixedSizeBinary {
+    fn scalar_at(
+        array: ArrayView<'_, FixedSizeBinary>,
+        index: usize,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        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<FixedSizeBinary> for FixedSizeBinary {
+    fn validity(array: ArrayView<'_, FixedSizeBinary>) -> VortexResult<Validity> {
+        fixed_width::validity(array)
+    }
+}
+
+impl CastReduce for FixedSizeBinary {
+    fn cast(
+        array: ArrayView<'_, FixedSizeBinary>,
+        dtype: &DType,
+    ) -> VortexResult<Option<ArrayRef>> {
+        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<Option<ArrayRef>> {
+        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<Option<ArrayRef>> {
+        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::<BoolArray>(ctx)?.into_bit_buffer().not()
+            }
+            _ => unreachable!("checked in entry point"),
+        };
+        let byte_width = array.byte_width() as usize;
+        is_invalid.for_each_set_index(|invalid_index| {
+            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<Option<ArrayRef>> {
+        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<usize>,
+    ) -> VortexResult<Option<ArrayRef>> {
+        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 FixedWidthArray for FixedSizeBinary {
+    fn byte_width(array: ArrayView<'_, Self>) -> usize {
+        array.byte_width() as usize
+    }
+
+    fn values(array: ArrayView<'_, Self>) -> ByteBuffer {
+        array.buffer_handle().to_host_sync()
+    }
+
+    fn with_values(
+        array: ArrayView<'_, Self>,
+        values: ByteBuffer,
+        len: usize,
+        validity: Validity,
+    ) -> VortexResult<FixedSizeBinaryArray> {
+        FixedSizeBinaryArray::try_new(values, array.byte_width(), len, validity)
+    }
+}