feat: teach LayoutBatchStream to filter by indices (#1242)

danking · web-flow · commit 8f5b74911737 · 2024-11-07T13:48:30.000-05:00
I have some thoughts on how to make the stream.rs logic easier to
follow; however, this, I think, is the fewest changes to achieve indices
filtering.

This change evaluates the filter on any row in a chunk containing at
least one kept index. I have some thoughts on how to avoid evaluating
the filter at a sub-chunk level but they require more substantial
changes.
diff --git a/pyvortex/python/vortex/dataset.py b/pyvortex/python/vortex/dataset.py
@@ -230,7 +230,9 @@ def take(
         table : :class:`.pyarrow.Table`
 
         """
-        return self._dataset.to_array(columns=columns, row_filter=filter).take(encoding.array(indices)).to_arrow_table()
+        return self._dataset.to_array(
+            columns=columns, row_filter=filter, indices=encoding.array(indices)
+        ).to_arrow_table()
 
     def to_record_batch_reader(
         self,
diff --git a/pyvortex/src/dataset.rs b/pyvortex/src/dataset.rs
@@ -28,6 +28,7 @@ pub async fn layout_stream_from_reader<T: VortexReadAt + Unpin>(
     reader: T,
     projection: Projection,
     row_filter: Option<RowFilter>,
+    indices: Option<Array>,
 ) -> VortexResult<LayoutBatchStream<T>> {
     let mut builder = LayoutBatchStreamBuilder::new(
         reader,
@@ -42,15 +43,20 @@ pub async fn layout_stream_from_reader<T: VortexReadAt + Unpin>(
         builder = builder.with_row_filter(row_filter);
     }
 
+    if let Some(indices) = indices {
+        builder = builder.with_indices(indices);
+    }
+
     builder.build().await
 }
 
 pub async fn read_array_from_reader<T: VortexReadAt + Unpin + 'static>(
     reader: T,
     projection: Projection,
     row_filter: Option<RowFilter>,
+    indices: Option<Array>,
 ) -> VortexResult<Array> {
-    layout_stream_from_reader(reader, projection, row_filter)
+    layout_stream_from_reader(reader, projection, row_filter, indices)
         .await?
         .read_all()
         .await
@@ -119,11 +125,13 @@ impl TokioFileDataset {
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyArray> {
         let inner = read_array_from_reader(
             self.file().await?,
             projection_from_python(columns)?,
             row_filter_from_python(row_filter),
+            indices.map(PyArray::unwrap).cloned(),
         )
         .await?;
         Ok(PyArray::new(inner))
@@ -133,11 +141,13 @@ impl TokioFileDataset {
         self_: PyRef<'_, Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyObject> {
         let layout_reader = layout_stream_from_reader(
             self_.file().await?,
             projection_from_python(columns)?,
             row_filter_from_python(row_filter),
+            indices.map(PyArray::unwrap).cloned(),
         )
         .await?;
 
@@ -154,23 +164,25 @@ impl TokioFileDataset {
         self_.schema.clone().to_pyarrow(self_.py())
     }
 
-    #[pyo3(signature = (*, columns=None, row_filter=None))]
+    #[pyo3(signature = (*, columns = None, row_filter = None, indices = None))]
     pub fn to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyArray> {
-        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter))
+        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter, indices))
     }
 
-    #[pyo3(signature = (*, columns=None, row_filter=None))]
+    #[pyo3(signature = (*, columns = None, row_filter = None, indices = None))]
     pub fn to_record_batch_reader(
         self_: PyRef<Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyObject> {
         TOKIO_RUNTIME.block_on(Self::async_to_record_batch_reader(
-            self_, columns, row_filter,
+            self_, columns, row_filter, indices,
         ))
     }
 }
@@ -197,11 +209,13 @@ impl ObjectStoreUrlDataset {
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyArray> {
         let inner = read_array_from_reader(
             self.reader().await?,
             projection_from_python(columns)?,
             row_filter_from_python(row_filter),
+            indices.map(PyArray::unwrap).cloned(),
         )
         .await?;
         Ok(PyArray::new(inner))
@@ -211,11 +225,13 @@ impl ObjectStoreUrlDataset {
         self_: PyRef<'_, Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyObject> {
         let layout_reader = layout_stream_from_reader(
             self_.reader().await?,
             projection_from_python(columns)?,
             row_filter_from_python(row_filter),
+            indices.map(PyArray::unwrap).cloned(),
         )
         .await?;
 
@@ -232,23 +248,25 @@ impl ObjectStoreUrlDataset {
         self_.schema.clone().to_pyarrow(self_.py())
     }
 
-    #[pyo3(signature = (*, columns=None, row_filter=None))]
+    #[pyo3(signature = (*, columns = None, row_filter = None, indices = None))]
     pub fn to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyArray> {
-        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter))
+        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter, indices))
     }
 
-    #[pyo3(signature = (*, columns=None, row_filter=None))]
+    #[pyo3(signature = (*, columns = None, row_filter = None, indices = None))]
     pub fn to_record_batch_reader(
         self_: PyRef<Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
         row_filter: Option<&Bound<'_, PyExpr>>,
+        indices: Option<&PyArray>,
     ) -> PyResult<PyObject> {
         TOKIO_RUNTIME.block_on(Self::async_to_record_batch_reader(
-            self_, columns, row_filter,
+            self_, columns, row_filter, indices,
         ))
     }
 }
diff --git a/pyvortex/src/io.rs b/pyvortex/src/io.rs
@@ -125,14 +125,15 @@ use crate::{PyArray, TOKIO_RUNTIME};
 /// >>> # b.to_arrow_array()
 ///
 #[pyfunction]
-#[pyo3(signature = (path, *, projection = None, row_filter = None))]
+#[pyo3(signature = (path, *, projection = None, row_filter = None, indices = None))]
 pub fn read_path(
     path: Bound<PyString>,
     projection: Option<Vec<Bound<PyAny>>>,
     row_filter: Option<&Bound<PyExpr>>,
+    indices: Option<&PyArray>,
 ) -> PyResult<PyArray> {
     let dataset = TOKIO_RUNTIME.block_on(TokioFileDataset::try_new(path.extract()?))?;
-    dataset.to_array(projection, row_filter)
+    dataset.to_array(projection, row_filter, indices)
 }
 
 /// Read a vortex struct array from a URL.
@@ -177,14 +178,15 @@ pub fn read_path(
 /// >>> a = vortex.io.read_url("file:/path/to/dataset.vortex")  # doctest: +SKIP
 ///
 #[pyfunction]
-#[pyo3(signature = (url, *, projection = None, row_filter = None))]
+#[pyo3(signature = (url, *, projection = None, row_filter = None, indices = None))]
 pub fn read_url(
     url: Bound<PyString>,
     projection: Option<Vec<Bound<PyAny>>>,
     row_filter: Option<&Bound<PyExpr>>,
+    indices: Option<&PyArray>,
 ) -> PyResult<PyArray> {
     let dataset = TOKIO_RUNTIME.block_on(ObjectStoreUrlDataset::try_new(url.extract()?))?;
-    dataset.to_array(projection, row_filter)
+    dataset.to_array(projection, row_filter, indices)
 }
 
 /// Write a vortex struct array to the local filesystem.
diff --git a/pyvortex/test/test_dataset.py b/pyvortex/test/test_dataset.py
@@ -24,8 +24,8 @@ def ds(tmpdir_factory) -> vortex.dataset.VortexDataset:
     if not os.path.exists(fname):
         a = pa.array([record(x) for x in range(1_000_000)])
         arr = vortex.encoding.compress(vortex.array(a))
-        vortex.io.write_path(arr, "/tmp/foo.vortex")
-    return vortex.dataset.VortexDataset.from_path("/tmp/foo.vortex")
+        vortex.io.write_path(arr, str(fname))
+    return vortex.dataset.VortexDataset.from_path(str(fname))
 
 
 def test_schema(ds):
diff --git a/vortex-serde/src/layouts/read/buffered.rs b/vortex-serde/src/layouts/read/buffered.rs
@@ -28,6 +28,15 @@ impl BufferedLayoutReader {
 
     // TODO(robert): Support out of order reads
     fn buffer_read(&mut self, mask: &RowMask) -> VortexResult<Option<Vec<Message>>> {
+        while let Some(((begin, end), layout)) = self.layouts.pop_front() {
+            if mask.begin() <= begin && begin < mask.end()
+                || mask.begin() < end && end <= mask.end()
+            {
+                self.layouts.push_front(((begin, end), layout));
+                break;
+            }
+        }
+
         while let Some(((begin, end), mut layout)) = self.layouts.pop_front() {
             // This selection doesn't know about rows in this chunk, we should put it back and wait for another request with different range
             if mask.end() <= begin || mask.begin() > end {
diff --git a/vortex-serde/src/layouts/read/builder.rs b/vortex-serde/src/layouts/read/builder.rs
@@ -5,6 +5,7 @@ use vortex_error::VortexResult;
 use vortex_expr::Select;
 use vortex_schema::projection::Projection;
 
+use super::RowMask;
 use crate::io::VortexReadAt;
 use crate::layouts::read::cache::{LayoutMessageCache, LazyDeserializedDType, RelativeLayoutCache};
 use crate::layouts::read::context::LayoutDeserializer;
@@ -18,7 +19,7 @@ pub struct LayoutBatchStreamBuilder<R> {
     layout_serde: LayoutDeserializer,
     projection: Option<Projection>,
     size: Option<u64>,
-    indices: Option<Array>,
+    row_mask: Option<Array>,
     row_filter: Option<RowFilter>,
 }
 
@@ -28,9 +29,9 @@ impl<R: VortexReadAt> LayoutBatchStreamBuilder<R> {
             reader,
             layout_serde,
             projection: None,
-            row_filter: None,
             size: None,
-            indices: None,
+            row_mask: None,
+            row_filter: None,
         }
     }
 
@@ -45,12 +46,12 @@ impl<R: VortexReadAt> LayoutBatchStreamBuilder<R> {
     }
 
     pub fn with_indices(mut self, array: Array) -> Self {
-        // TODO(#441): Allow providing boolean masks
         assert!(
-            array.dtype().is_int(),
-            "Mask arrays have to be integer arrays"
+            !array.dtype().is_nullable() && (array.dtype().is_int() || array.dtype().is_boolean()),
+            "Mask arrays have to be non-nullable integer or boolean arrays"
         );
-        self.indices = Some(array);
+
+        self.row_mask = Some(array);
         self
     }
 
@@ -87,22 +88,34 @@ impl<R: VortexReadAt> LayoutBatchStreamBuilder<R> {
 
         let filter_reader = self
             .row_filter
-            .as_ref()
-            .map(|filter| {
+            .map(|row_filter| {
                 footer.layout(
-                    Scan::new(Some(Arc::new(filter.clone()))),
+                    Scan::new(Some(Arc::new(row_filter))),
                     RelativeLayoutCache::new(message_cache.clone(), footer_dtype),
                 )
             })
             .transpose()?;
 
+        let row_mask = self
+            .row_mask
+            .as_ref()
+            .map(|row_mask| {
+                if row_mask.dtype().is_int() {
+                    RowMask::from_index_array(row_mask, 0, row_count as usize)
+                } else {
+                    RowMask::from_mask_array(row_mask, 0, row_count as usize)
+                }
+            })
+            .transpose()?;
+
         Ok(LayoutBatchStream::new(
             self.reader,
             data_reader,
             filter_reader,
             message_cache,
             projected_dtype,
             row_count,
+            row_mask,
         ))
     }
 
diff --git a/vortex-serde/src/layouts/read/mask.rs b/vortex-serde/src/layouts/read/mask.rs
@@ -121,6 +121,21 @@ impl RowMask {
         }
     }
 
+    /// Unset, in place, any bits that are unset in `other`.
+    pub fn and_inplace(&mut self, other: &RowMask) -> VortexResult<()> {
+        if self.begin != other.begin || self.end != other.end {
+            vortex_bail!(
+                "begin and ends must match: {}-{} {}-{}",
+                self.begin,
+                self.end,
+                other.begin,
+                other.end
+            );
+        }
+        self.values.and_inplace(&other.values);
+        Ok(())
+    }
+
     /// Filter array with this `RowMask`.
     ///
     /// This function assumes that Array is no longer than the mask length and that the mask starts on same offset as the array,
diff --git a/vortex-serde/src/layouts/read/stream.rs b/vortex-serde/src/layouts/read/stream.rs
diff --git a/vortex-serde/src/layouts/tests.rs b/vortex-serde/src/layouts/tests.rs
