Added a bool true count function and updated error handling

spiraldb · Jul 24, 2024 · 25fe1e5 · 25fe1e5
1 parent 20d5b91
commit 25fe1e5
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 51 deletions.
diff --git a/vortex-array/src/array/bool/compute/filter.rs b/vortex-array/src/array/bool/compute/filter.rs
@@ -1,9 +1,8 @@
 use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{vortex_err, VortexResult};
 
 use crate::array::bool::BoolArray;
 use crate::compute::FilterFn;
-use crate::stats::ArrayStatistics;
 use crate::validity::filter_validity;
 use crate::variants::BoolArrayTrait;
 use crate::{Array, IntoArray};
@@ -15,27 +14,19 @@ impl FilterFn for BoolArray {
 }
 
 fn filter_select_bool(arr: &BoolArray, predicate: &Array) -> VortexResult<BoolArray> {
-    let Some(selection_count) = predicate.statistics().compute_true_count() else {
-        vortex_bail!(
-            NotImplemented: "compute_true_count",
-            predicate.encoding().id()
-        )
-    };
     predicate.with_dyn(|b| {
         let validity = filter_validity(arr.validity(), predicate)?;
-        if let Some(predicate) = b.as_bool_array() {
-            let out = if selection_count * 2 > arr.len() {
-                filter_select_bool_by_slice(&arr.boolean_buffer(), predicate, selection_count)
-            } else {
-                filter_select_bool_by_index(&arr.boolean_buffer(), predicate, selection_count)
-            };
-            BoolArray::try_new(out, validity)
+        let predicate = b.as_bool_array().ok_or(vortex_err!(
+            NotImplemented: "as_bool_array",
+            predicate.encoding().id()
+        ))?;
+        let selection_count = predicate.true_count();
+        let out = if selection_count * 2 > arr.len() {
+            filter_select_bool_by_slice(&arr.boolean_buffer(), predicate, selection_count)
         } else {
-            vortex_bail!(
-                NotImplemented: "as_bool_array",
-                predicate.encoding().id()
-            )
-        }
+            filter_select_bool_by_index(&arr.boolean_buffer(), predicate, selection_count)
+        };
+        BoolArray::try_new(out, validity)
     })
 }
 

diff --git a/vortex-array/src/array/constant/compute.rs b/vortex-array/src/array/constant/compute.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 use arrow_array::Datum;
 use vortex_dtype::Nullability;
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use vortex_expr::Operator;
 use vortex_scalar::Scalar;
 
@@ -69,7 +69,14 @@ impl FilterFn for ConstantArray {
     fn filter(&self, predicate: &Array) -> VortexResult<Array> {
         Ok(Self::new(
             self.scalar().clone(),
-            predicate.statistics().compute_true_count().unwrap(),
+            predicate.with_dyn(|p| {
+                p.as_bool_array()
+                    .ok_or(vortex_err!(
+                        NotImplemented: "as_bool_array",
+                        predicate.encoding().id()
+                    ))
+                    .map(|x| x.true_count())
+            })?,
         )
         .into_array())
     }

diff --git a/vortex-array/src/array/primitive/compute/filter.rs b/vortex-array/src/array/primitive/compute/filter.rs
@@ -1,9 +1,8 @@
 use vortex_dtype::{match_each_native_ptype, NativePType};
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{vortex_err, VortexResult};
 
 use crate::array::primitive::PrimitiveArray;
 use crate::compute::FilterFn;
-use crate::stats::ArrayStatistics;
 use crate::validity::filter_validity;
 use crate::variants::BoolArrayTrait;
 use crate::{Array, IntoArray};
@@ -18,25 +17,17 @@ fn filter_select_primitive(
     arr: &PrimitiveArray,
     predicate: &Array,
 ) -> VortexResult<PrimitiveArray> {
-    let Some(selection_count) = predicate.statistics().compute_true_count() else {
-        vortex_bail!(
-            NotImplemented: "compute_true_count",
-            predicate.encoding().id()
-        )
-    };
     predicate.with_dyn(|b| {
         let validity = filter_validity(arr.validity(), predicate)?;
-        if let Some(bb) = b.as_bool_array() {
-            match_each_native_ptype!(arr.ptype(), |$T| {
-                let slice = arr.maybe_null_slice::<$T>();
-                Ok(PrimitiveArray::from_vec(filter_primitive_slice(slice, bb, selection_count), validity))
-            })
-        } else {
-            vortex_bail!(
+        let predicate = b.as_bool_array().ok_or_else(||vortex_err!(
                 NotImplemented: "as_bool_array",
                 predicate.encoding().id()
-            )
-        }
+            ))?;
+        let selection_count = predicate.true_count();
+        match_each_native_ptype!(arr.ptype(), |$T| {
+            let slice = arr.maybe_null_slice::<$T>();
+            Ok(PrimitiveArray::from_vec(filter_primitive_slice(slice, predicate, selection_count), validity))
+        })
     })
 }
 

diff --git a/vortex-array/src/array/varbin/compute/filter.rs b/vortex-array/src/array/varbin/compute/filter.rs
@@ -1,12 +1,11 @@
 use itertools::Itertools;
 use num_traits::{AsPrimitive, Zero};
 use vortex_dtype::{match_each_integer_ptype, DType, NativePType};
-use vortex_error::VortexResult;
+use vortex_error::{vortex_err, VortexResult};
 
 use crate::array::varbin::builder::VarBinBuilder;
 use crate::array::varbin::VarBinArray;
 use crate::compute::FilterFn;
-use crate::stats::ArrayStatistics;
 use crate::validity::Validity;
 use crate::variants::BoolArrayTrait;
 use crate::{Array, ArrayDType, IntoArray, IntoArrayVariant};
@@ -18,10 +17,14 @@ impl FilterFn for VarBinArray {
 }
 
 fn filter_select_var_bin(arr: &VarBinArray, predicate: &Array) -> VortexResult<VarBinArray> {
-    let selection_count = predicate.statistics().compute_true_count().unwrap();
     predicate.with_dyn(|p| {
-        let predicate = p.as_bool_array_unchecked();
-
+        let predicate = p.as_bool_array().ok_or_else(|| {
+            vortex_err!(
+                NotImplemented: "as_bool_array",
+                predicate.encoding().id()
+            )
+        })?;
+        let selection_count = predicate.true_count();
         if selection_count * 2 > predicate.len() {
             filter_select_var_bin_by_slice(arr, predicate, selection_count)
         } else {

diff --git a/vortex-array/src/compute/unary/mod.rs b/vortex-array/src/compute/unary/mod.rs
@@ -1,9 +1,9 @@
-mod cast;
-mod fill_forward;
-mod scalar_at;
-mod scalar_subtract;
-
 pub use cast::{try_cast, CastFn};
 pub use fill_forward::{fill_forward, FillForwardFn};
 pub use scalar_at::{scalar_at, ScalarAtFn};
 pub use scalar_subtract::{subtract_scalar, SubtractScalarFn};
+
+mod cast;
+mod fill_forward;
+mod scalar_at;
+mod scalar_subtract;
diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs
@@ -13,12 +13,11 @@ use std::fmt::{Debug, Display, Formatter};
 use std::future::ready;
 
 pub use ::paste;
-use itertools::Itertools;
-
 pub use canonical::*;
 pub use context::*;
 pub use data::*;
 pub use implementation::*;
+use itertools::Itertools;
 pub use metadata::*;
 pub use typed::*;
 pub use view::*;

diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs
@@ -75,6 +75,12 @@ pub trait ArrayVariants {
 pub trait NullArrayTrait: ArrayTrait {}
 
 pub trait BoolArrayTrait: ArrayTrait {
+    fn true_count(&self) -> usize {
+        self.statistics()
+            .compute_true_count()
+            .unwrap_or_else(|| self.maybe_null_indices_iter().count())
+    }
+
     // An iterator over the sorted indices of set values in the underlying boolean array
     // good to array with low number of set values.
     fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a>;