enable FilterFn on VarBin -> fix issue with FSST filter

spiraldb · Oct 18, 2024 · 3c1af26 · 3c1af26
1 parent 548f63c
commit 3c1af26
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs
@@ -3,13 +3,13 @@ use std::sync::Arc;
 
 use fsst::{Decompressor, Symbol};
 use serde::{Deserialize, Serialize};
-use vortex::array::VarBinArray;
+use vortex::array::{VarBin, VarBinArray};
 use vortex::encoding::ids;
 use vortex::stats::{ArrayStatisticsCompute, StatsSet};
 use vortex::validity::{ArrayValidity, LogicalValidity, Validity};
 use vortex::variants::{ArrayVariants, BinaryArrayTrait, Utf8ArrayTrait};
 use vortex::visitor::AcceptArrayVisitor;
-use vortex::{impl_encoding, Array, ArrayDType, ArrayTrait, IntoCanonical};
+use vortex::{impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, IntoCanonical};
 use vortex_dtype::{DType, Nullability, PType};
 use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 
@@ -73,6 +73,13 @@ impl FSSTArray {
             vortex_bail!(InvalidArgument: "uncompressed_lengths must have integer type and cannot be nullable");
         }
 
+        if codes.encoding().id() != VarBin::ID {
+            vortex_bail!(
+                InvalidArgument: "codes must have varbin encoding, was {}",
+                codes.encoding().id()
+            );
+        }
+
         // Check: strings must be a Binary array.
         if !matches!(codes.dtype(), DType::Binary(_)) {
             vortex_bail!(InvalidArgument: "codes array must be DType::Binary type");

diff --git a/vortex-array/src/array/varbin/compute/filter.rs b/vortex-array/src/array/varbin/compute/filter.rs
@@ -38,7 +38,7 @@ fn filter_select_var_bin_by_slice(
     predicate: &dyn BoolArrayTrait,
     selection_count: usize,
 ) -> VortexResult<VarBinArray> {
-    let offsets = values.offsets().as_primitive();
+    let offsets = values.offsets().into_primitive()?;
     match_each_integer_ptype!(offsets.ptype(), |$O| {
         filter_select_var_bin_by_slice_primitive_offset(
             values.dtype().clone(),
@@ -137,7 +137,7 @@ fn filter_select_var_bin_by_index(
     predicate: &dyn BoolArrayTrait,
     selection_count: usize,
 ) -> VortexResult<VarBinArray> {
-    let offsets = values.offsets().as_primitive();
+    let offsets = values.offsets().into_primitive()?;
     match_each_integer_ptype!(offsets.ptype(), |$O| {
         filter_select_var_bin_by_index_primitive_offset(
             values.dtype().clone(),

diff --git a/vortex-array/src/array/varbin/compute/mod.rs b/vortex-array/src/array/varbin/compute/mod.rs
@@ -3,14 +3,18 @@ use vortex_scalar::Scalar;
 
 use crate::array::varbin::{varbin_scalar, VarBinArray};
 use crate::compute::unary::ScalarAtFn;
-use crate::compute::{ArrayCompute, SliceFn, TakeFn};
+use crate::compute::{ArrayCompute, FilterFn, SliceFn, TakeFn};
 use crate::ArrayDType;
 
 mod filter;
 mod slice;
 mod take;
 
 impl ArrayCompute for VarBinArray {
+    fn filter(&self) -> Option<&dyn FilterFn> {
+        Some(self)
+    }
+
     fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
         Some(self)
     }