Use FilterMask in RowMask (#1616)

spiraldb · Dec 9, 2024 · b49fc1b · b49fc1b
1 parent 0fbc411
commit b49fc1b
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 155 deletions.
diff --git a/vortex-array/src/compute/filter.rs b/vortex-array/src/compute/filter.rs
@@ -1,5 +1,5 @@
 use std::iter::TrustedLen;
-use std::sync::OnceLock;
+use std::sync::{Arc, OnceLock};
 
 use arrow_array::BooleanArray;
 use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer};
@@ -118,9 +118,9 @@ pub struct FilterMask {
     array: ArrayData,
     true_count: usize,
     range_selectivity: f64,
-    indices: OnceLock<Vec<usize>>,
-    slices: OnceLock<Vec<(usize, usize)>>,
-    buffer: OnceLock<BooleanBuffer>,
+    indices: Arc<OnceLock<Vec<usize>>>,
+    slices: Arc<OnceLock<Vec<(usize, usize)>>>,
+    buffer: Arc<OnceLock<BooleanBuffer>>,
 }
 
 /// We implement Clone manually to trigger population of our cached indices or slices.
@@ -328,9 +328,9 @@ impl TryFrom<ArrayData> for FilterMask {
             array,
             true_count,
             range_selectivity: selectivity,
-            indices: OnceLock::new(),
-            slices: OnceLock::new(),
-            buffer: OnceLock::new(),
+            indices: Arc::new(OnceLock::new()),
+            slices: Arc::new(OnceLock::new()),
+            buffer: Arc::new(OnceLock::new()),
         })
     }
 }

diff --git a/vortex-file/src/read/layouts/chunked.rs b/vortex-file/src/read/layouts/chunked.rs
@@ -409,9 +409,10 @@ mod tests {
     use bytes::Bytes;
     use flatbuffers::{root, FlatBufferBuilder};
     use futures_util::TryStreamExt;
-    use vortex_array::array::{BoolArray, ChunkedArray, PrimitiveArray};
+    use vortex_array::array::{ChunkedArray, PrimitiveArray};
+    use vortex_array::compute::FilterMask;
     use vortex_array::{ArrayDType, ArrayLen, IntoArrayData, IntoArrayVariant};
-    use vortex_dtype::{Nullability, PType};
+    use vortex_dtype::PType;
     use vortex_expr::{BinaryExpr, Identity, Literal, Operator};
     use vortex_flatbuffers::{footer, WriteFlatBuffer};
     use vortex_ipc::messages::writer::MessageWriter;
@@ -582,18 +583,8 @@ mod tests {
         snd_range.append_n(100, true);
         snd_range.append_n(50, false);
         let mut arr = [
-            RowMask::try_new(
-                BoolArray::new(first_range.finish(), Nullability::NonNullable).into_array(),
-                0,
-                200,
-            )
-            .unwrap(),
-            RowMask::try_new(
-                BoolArray::new(snd_range.finish(), Nullability::NonNullable).into_array(),
-                200,
-                400,
-            )
-            .unwrap(),
+            RowMask::try_new(FilterMask::from(first_range.finish()), 0, 200).unwrap(),
+            RowMask::try_new(FilterMask::from(snd_range.finish()), 200, 400).unwrap(),
             RowMask::new_valid_between(400, 500),
         ]
         .into_iter()