Add Not vortex expr and use it in pruning expression conversion (#1213)

robert3005 · web-flow · commit 635ea6254fba · 2024-11-07T13:57:22.000-05:00
diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs
@@ -9,7 +9,9 @@ use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use vortex_array::variants::{ArrayVariants, BoolArrayTrait};
-use vortex_array::{impl_encoding, ArrayTrait, Canonical, IntoCanonical, TypedArray};
+use vortex_array::{
+    impl_encoding, Array, ArrayTrait, Canonical, IntoArray, IntoCanonical, TypedArray,
+};
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect as _, VortexResult};
@@ -90,6 +92,14 @@ impl ArrayVariants for ByteBoolArray {
 }
 
 impl BoolArrayTrait for ByteBoolArray {
+    fn invert(&self) -> VortexResult<Array> {
+        ByteBoolArray::try_from_vec(
+            self.maybe_null_slice().iter().map(|v| !v).collect(),
+            self.validity(),
+        )
+        .map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a> {
         todo!()
     }
diff --git a/encodings/roaring/src/boolean/mod.rs b/encodings/roaring/src/boolean/mod.rs
@@ -97,6 +97,11 @@ impl ArrayVariants for RoaringBoolArray {
 }
 
 impl BoolArrayTrait for RoaringBoolArray {
+    fn invert(&self) -> VortexResult<Array> {
+        RoaringBoolArray::try_new(self.bitmap().flip(0..(self.len() as u32)), self.len())
+            .map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a> {
         todo!()
     }
diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs
@@ -9,7 +9,8 @@ use vortex_array::stats::{ArrayStatistics, ArrayStatisticsCompute, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use vortex_array::variants::{ArrayVariants, BoolArrayTrait};
 use vortex_array::{
-    impl_encoding, Array, ArrayDType, ArrayTrait, Canonical, IntoArrayVariant, IntoCanonical,
+    impl_encoding, Array, ArrayDType, ArrayTrait, Canonical, IntoArray, IntoArrayVariant,
+    IntoCanonical,
 };
 use vortex_dtype::{DType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
@@ -115,6 +116,11 @@ impl RunEndBoolArray {
 }
 
 impl BoolArrayTrait for RunEndBoolArray {
+    fn invert(&self) -> VortexResult<Array> {
+        RunEndBoolArray::try_new(self.ends(), !self.start(), self.validity())
+            .map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a> {
         todo!()
     }
diff --git a/vortex-array/src/array/bool/mod.rs b/vortex-array/src/array/bool/mod.rs
@@ -13,7 +13,7 @@ use crate::encoding::ids;
 use crate::stats::StatsSet;
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use crate::variants::{ArrayVariants, BoolArrayTrait};
-use crate::{impl_encoding, ArrayTrait, Canonical, IntoCanonical, TypedArray};
+use crate::{impl_encoding, Array, ArrayTrait, Canonical, IntoArray, IntoCanonical, TypedArray};
 
 mod accessors;
 mod compute;
@@ -96,6 +96,10 @@ impl ArrayVariants for BoolArray {
 }
 
 impl BoolArrayTrait for BoolArray {
+    fn invert(&self) -> VortexResult<Array> {
+        BoolArray::try_new(!&self.boolean_buffer(), self.validity()).map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a> {
         Box::new(BitIndexIterator::new(self.buffer(), 0, self.len()))
     }
diff --git a/vortex-array/src/array/chunked/variants.rs b/vortex-array/src/array/chunked/variants.rs
@@ -47,6 +47,20 @@ impl ArrayVariants for ChunkedArray {
 impl NullArrayTrait for ChunkedArray {}
 
 impl BoolArrayTrait for ChunkedArray {
+    fn invert(&self) -> VortexResult<Array> {
+        let chunks = self
+            .chunks()
+            .map(|c| {
+                c.with_dyn(|a| {
+                    a.as_bool_array()
+                        .ok_or_else(|| vortex_err!("Child was not a bool array"))
+                        .and_then(|b| b.invert())
+                })
+            })
+            .collect::<VortexResult<Vec<_>>>()?;
+        ChunkedArray::try_new(chunks, self.dtype().clone()).map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter(&self) -> Box<dyn Iterator<Item = usize>> {
         todo!()
     }
diff --git a/vortex-array/src/array/sparse/variants.rs b/vortex-array/src/array/sparse/variants.rs
@@ -48,6 +48,21 @@ impl ArrayVariants for SparseArray {
 impl NullArrayTrait for SparseArray {}
 
 impl BoolArrayTrait for SparseArray {
+    fn invert(&self) -> VortexResult<Array> {
+        let inverted_fill = self.fill_value().as_bool()?.map(|v| !v);
+        SparseArray::try_new(
+            self.indices(),
+            self.values().with_dyn(|a| {
+                a.as_bool_array()
+                    .ok_or_else(|| vortex_err!("Not a bool array"))
+                    .and_then(|b| b.invert())
+            })?,
+            self.len(),
+            inverted_fill.into(),
+        )
+        .map(|a| a.into_array())
+    }
+
     fn maybe_null_indices_iter(&self) -> Box<dyn Iterator<Item = usize>> {
         Box::new(self.resolved_indices().into_iter())
     }
@@ -120,3 +135,33 @@ impl ExtensionArrayTrait for SparseArray {
         .into_array()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::array::{BoolArray, PrimitiveArray, SparseArray};
+    use crate::{IntoArray, IntoArrayVariant};
+
+    #[test]
+    fn invert_bools_non_null_fill() {
+        let sparse_bools = SparseArray::try_new(
+            PrimitiveArray::from(vec![0u64]).into_array(),
+            BoolArray::from(vec![false]).into_array(),
+            2,
+            true.into(),
+        )
+        .unwrap()
+        .into_array();
+        let inverted = sparse_bools
+            .with_dyn(|a| a.as_bool_array_unchecked().invert())
+            .unwrap();
+        assert_eq!(
+            inverted
+                .into_bool()
+                .unwrap()
+                .boolean_buffer()
+                .iter()
+                .collect::<Vec<_>>(),
+            vec![true, false]
+        );
+    }
+}
diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs
@@ -3,16 +3,14 @@
 //! When callers only want to make assumptions about the DType, and not about any specific
 //! encoding, they can use these traits to write encoding-agnostic code.
 
-use std::ops::Not;
 use std::sync::Arc;
 
 use vortex_dtype::field::Field;
 use vortex_dtype::{DType, ExtDType, FieldNames, PType};
 use vortex_error::{vortex_panic, VortexExpect as _, VortexResult};
 
-use crate::array::BoolArray;
 use crate::iter::{AccessorRef, VectorizedArrayIter};
-use crate::{Array, ArrayTrait, IntoArray, IntoArrayVariant};
+use crate::{Array, ArrayTrait};
 
 pub trait ArrayVariants {
     fn as_null_array(&self) -> Option<&dyn NullArrayTrait> {
@@ -90,15 +88,7 @@ pub trait BoolArrayTrait: ArrayTrait {
     /// True -> False
     /// False -> True
     /// Null -> Null
-    fn invert(&self) -> VortexResult<Array>
-    where
-        Self: Clone,
-    {
-        let bool_array = self.clone().into_bool()?;
-        let validity = bool_array.validity();
-
-        BoolArray::try_new(bool_array.boolean_buffer().not(), validity).map(|a| a.into_array())
-    }
+    fn invert(&self) -> VortexResult<Array>;
 
     fn true_count(&self) -> usize {
         self.statistics()
diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs
@@ -9,13 +9,15 @@ mod column;
 pub mod datafusion;
 mod identity;
 mod literal;
+mod not;
 mod operators;
 mod select;
 
 pub use binary::*;
 pub use column::*;
 pub use identity::*;
 pub use literal::*;
+pub use not::*;
 pub use operators::*;
 pub use select::*;
 use vortex_array::Array;
diff --git a/vortex-expr/src/literal.rs b/vortex-expr/src/literal.rs
@@ -16,6 +16,10 @@ impl Literal {
     pub fn new(value: Scalar) -> Self {
         Self { value }
     }
+
+    pub fn value(&self) -> &Scalar {
+        &self.value
+    }
 }
 
 impl VortexExpr for Literal {
diff --git a/vortex-expr/src/not.rs b/vortex-expr/src/not.rs
@@ -0,0 +1,79 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use vortex_array::aliases::hash_set::HashSet;
+use vortex_array::Array;
+use vortex_dtype::field::Field;
+use vortex_error::{vortex_err, VortexResult};
+
+use crate::{unbox_any, VortexExpr};
+
+#[derive(Debug)]
+pub struct Not {
+    child: Arc<dyn VortexExpr>,
+}
+
+impl Not {
+    pub fn new(child: Arc<dyn VortexExpr>) -> Self {
+        Self { child }
+    }
+
+    pub fn child(&self) -> &Arc<dyn VortexExpr> {
+        &self.child
+    }
+}
+
+impl VortexExpr for Not {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn evaluate(&self, batch: &Array) -> VortexResult<Array> {
+        let child_result = self.child.evaluate(batch)?;
+        child_result.with_dyn(|a| {
+            a.as_bool_array()
+                .ok_or_else(|| vortex_err!("Child was not a bool array"))
+                .and_then(|b| b.invert())
+        })
+    }
+
+    fn collect_references<'a>(&'a self, references: &mut HashSet<&'a Field>) {
+        self.child.collect_references(references)
+    }
+}
+
+impl PartialEq<dyn Any> for Not {
+    fn eq(&self, other: &dyn Any) -> bool {
+        unbox_any(other)
+            .downcast_ref::<Self>()
+            .map(|x| x.child.eq(&self.child))
+            .unwrap_or(false)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use vortex_array::array::BoolArray;
+    use vortex_array::IntoArrayVariant;
+
+    use crate::{Identity, Not, VortexExpr};
+
+    #[test]
+    fn invert_booleans() {
+        let not_expr = Not::new(Arc::new(Identity));
+        let bools = BoolArray::from(vec![false, true, false, false, true, true]);
+        assert_eq!(
+            not_expr
+                .evaluate(bools.as_ref())
+                .unwrap()
+                .into_bool()
+                .unwrap()
+                .boolean_buffer()
+                .iter()
+                .collect::<Vec<_>>(),
+            vec![true, false, true, true, false, false]
+        );
+    }
+}
diff --git a/vortex-serde/src/layouts/pruning.rs b/vortex-serde/src/layouts/pruning.rs
diff --git a/vortex-serde/src/layouts/read/expr_project.rs b/vortex-serde/src/layouts/read/expr_project.rs

Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,11 @@ impl ArrayVariants for RoaringBoolArray {`
`97`	`97`	`}`
`98`	`98`
`99`	`99`	`impl BoolArrayTrait for RoaringBoolArray {`
	`100`	`+ fn invert(&self) -> VortexResult<Array> {`
	`101`	`+ RoaringBoolArray::try_new(self.bitmap().flip(0..(self.len() as u32)), self.len())`
	`102`	`+ .map(\|a\| a.into_array())`
	`103`	`+ }`
	`104`	`+`
`100`	`105`	`fn maybe_null_indices_iter<'a>(&'a self) -> Box<dyn Iterator<Item = usize> + 'a> {`
`101`	`106`	`todo!()`
`102`	`107`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,10 @@ impl Literal {`
`16`	`16`	`pub fn new(value: Scalar) -> Self {`
`17`	`17`	`Self { value }`
`18`	`18`	`}`
	`19`	`+`
	`20`	`+ pub fn value(&self) -> &Scalar {`
	`21`	`+ &self.value`
	`22`	`+ }`
`19`	`23`	`}`
`20`	`24`
`21`	`25`	`impl VortexExpr for Literal {`