RowMasks use bitmasks instead of bitmaps (#1346)

spiraldb · Nov 18, 2024 · 7bcda5c · 7bcda5c
1 parent f6980bf
commit 7bcda5c
Show file tree

Hide file tree

Showing 7 changed files with 219 additions and 145 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml
@@ -18,7 +18,6 @@ arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 bytes = { workspace = true }
-croaring = { workspace = true }
 compio = { workspace = true, features = ["bytes", "macros"], optional = true }
 flatbuffers = { workspace = true }
 flume = { workspace = true }

diff --git a/vortex-file/src/read/buffered.rs b/vortex-file/src/read/buffered.rs
@@ -28,9 +28,7 @@ impl BufferedLayoutReader {
     // TODO(robert): Support out of order reads
     fn buffer_read(&mut self, mask: &RowMask) -> VortexResult<Option<Vec<MessageLocator>>> {
         while let Some(((begin, end), layout)) = self.layouts.pop_front() {
-            if mask.begin() <= begin && begin < mask.end()
-                || mask.begin() < end && end <= mask.end()
-            {
+            if mask.end() > begin && mask.begin() <= end {
                 self.layouts.push_front(((begin, end), layout));
                 break;
             }

diff --git a/vortex-file/src/read/layouts/chunked.rs b/vortex-file/src/read/layouts/chunked.rs
@@ -150,13 +150,13 @@ mod tests {
     use std::iter;
     use std::sync::{Arc, RwLock};
 
+    use arrow_buffer::BooleanBufferBuilder;
     use bytes::Bytes;
-    use croaring::Bitmap;
     use flatbuffers::{root_unchecked, FlatBufferBuilder};
     use futures_util::TryStreamExt;
-    use vortex_array::array::{ChunkedArray, PrimitiveArray};
+    use vortex_array::array::{BoolArray, ChunkedArray, PrimitiveArray};
     use vortex_array::{ArrayDType, IntoArrayData, IntoArrayVariant};
-    use vortex_dtype::PType;
+    use vortex_dtype::{Nullability, PType};
     use vortex_expr::{BinaryExpr, Identity, Literal, Operator};
     use vortex_flatbuffers::{footer, WriteFlatBuffer};
     use vortex_ipc::messages::writer::MessageWriter;
@@ -292,7 +292,7 @@ mod tests {
             &mut projection_layout,
             cache,
             &buf,
-            &RowMask::try_new(Bitmap::from_range(0..500), 0, 500).unwrap(),
+            &RowMask::new_valid_between(0, 500),
         );
 
         assert!(arr.is_some());
@@ -309,10 +309,29 @@ mod tests {
         let cache = Arc::new(RwLock::new(LayoutMessageCache::default()));
         let (_, mut projection_layout, buf, _) =
             layout_and_bytes(cache.clone(), Scan::new(None)).await;
+
+        let mut first_range = BooleanBufferBuilder::new(200);
+        first_range.append_n(150, true);
+        first_range.append_n(50, false);
+
+        let mut snd_range = BooleanBufferBuilder::new(200);
+        snd_range.append_n(50, false);
+        snd_range.append_n(100, true);
+        snd_range.append_n(50, false);
         let mut arr = [
-            RowMask::try_new(Bitmap::from_range(0..150), 0, 200).unwrap(),
-            RowMask::try_new(Bitmap::from_range(50..150), 200, 400).unwrap(),
-            RowMask::try_new(Bitmap::from_range(0..100), 400, 500).unwrap(),
+            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::new_valid_between(400, 500),
         ]
         .into_iter()
         .flat_map(|s| read_layout_data(&mut projection_layout, cache.clone(), &buf, &s))

diff --git a/vortex-file/src/read/layouts/test_read.rs b/vortex-file/src/read/layouts/test_read.rs
@@ -2,7 +2,6 @@ use std::collections::{BTreeSet, VecDeque};
 use std::sync::{Arc, RwLock};
 
 use bytes::Bytes;
-use croaring::Bitmap;
 use itertools::Itertools;
 use vortex_array::ArrayData;
 use vortex_error::VortexUnwrap;
@@ -17,9 +16,7 @@ pub fn layout_splits(layout: &mut dyn LayoutReader, length: usize) -> Vec<RowMas
     splits
         .into_iter()
         .tuple_windows::<(usize, usize)>()
-        .map(|(begin, end)| unsafe {
-            RowMask::new_unchecked(Bitmap::from_range(begin as u32..end as u32), 0, end)
-        })
+        .map(|(begin, end)| RowMask::new_valid_between(begin, end))
         .collect::<Vec<_>>()
 }