Filter pushdown over layouts

pyvortex/python/vortex/dataset.py

@@ -96,6 +96,8 @@ def head(
         table : :class:`.pyarrow.Table`
 
         """
+        if batch_size is not None:
+            raise ValueError("batch_size is not supported")
         if batch_readahead is not None:
             raise ValueError("batch_readahead not supported")
         if fragment_readahead is not None:
@@ -109,11 +111,7 @@ def head(
         del memory_pool
         if filter is not None:
             filter = arrow_to_vortex_expr(filter, self.schema)
-        return (
-            self._dataset.to_array(columns=columns, batch_size=batch_size, row_filter=filter)
-            .slice(0, num_rows)
-            .to_arrow_table()
-        )
+        return self._dataset.to_array(columns=columns, row_filter=filter).slice(0, num_rows).to_arrow_table()
 
     def join(
         self,
@@ -232,11 +230,7 @@ def take(
         table : :class:`.pyarrow.Table`
 
         """
-        return (
-            self._dataset.to_array(columns=columns, batch_size=batch_size, row_filter=filter)
-            .take(encoding.array(indices))
-            .to_arrow_table()
-        )
+        return self._dataset.to_array(columns=columns, row_filter=filter).take(encoding.array(indices)).to_arrow_table()
 
     def to_record_batch_reader(
         self,
@@ -276,6 +270,8 @@ def to_record_batch_reader(
         table : :class:`.pyarrow.Table`
 
         """
+        if batch_size is not None:
+            raise ValueError("batch_size is not supported")
         if batch_readahead is not None:
             raise ValueError("batch_readahead not supported")
         if fragment_readahead is not None:
@@ -289,7 +285,7 @@ def to_record_batch_reader(
         del memory_pool
         if filter is not None:
             filter = arrow_to_vortex_expr(filter, self.schema)
-        return self._dataset.to_record_batch_reader(columns=columns, batch_size=batch_size, row_filter=filter)
+        return self._dataset.to_record_batch_reader(columns=columns, row_filter=filter)
 
     def to_batches(
         self,
@@ -383,6 +379,8 @@ def to_table(
         table : :class:`.pyarrow.Table`
 
         """
+        if batch_size is not None:
+            raise ValueError("batch_size is not supported")
         if batch_readahead is not None:
             raise ValueError("batch_readahead not supported")
         if fragment_readahead is not None:
@@ -396,7 +394,7 @@ def to_table(
         del memory_pool
         if filter is not None:
             filter = arrow_to_vortex_expr(filter, self.schema)
-        return self._dataset.to_array(columns=columns, batch_size=batch_size, row_filter=filter).to_arrow_table()
+        return self._dataset.to_array(columns=columns, row_filter=filter).to_arrow_table()
 
 
 def from_path(path: str) -> VortexDataset:

pyvortex/src/dataset.rs

@@ -27,7 +27,6 @@ use crate::{PyArray, TOKIO_RUNTIME};
 pub async fn layout_stream_from_reader<T: VortexReadAt + Unpin>(
     reader: T,
     projection: Projection,
-    batch_size: Option<usize>,
     row_filter: Option<RowFilter>,
 ) -> VortexResult<LayoutBatchStream<T>> {
     let mut builder = LayoutReaderBuilder::new(
@@ -39,10 +38,6 @@ pub async fn layout_stream_from_reader<T: VortexReadAt + Unpin>(
     )
     .with_projection(projection);
 
-    if let Some(batch_size) = batch_size {
-        builder = builder.with_batch_size(batch_size);
-    }
-
     if let Some(row_filter) = row_filter {
         builder = builder.with_row_filter(row_filter);
     }
@@ -53,10 +48,9 @@ pub async fn layout_stream_from_reader<T: VortexReadAt + Unpin>(
 pub async fn read_array_from_reader<T: VortexReadAt + Unpin + 'static>(
     reader: T,
     projection: Projection,
-    batch_size: Option<usize>,
     row_filter: Option<RowFilter>,
 ) -> VortexResult<Array> {
-    layout_stream_from_reader(reader, projection, batch_size, row_filter)
+    layout_stream_from_reader(reader, projection, row_filter)
         .await?
         .read_all()
         .await
@@ -124,13 +118,11 @@ impl TokioFileDataset {
     async fn async_to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyArray> {
         let inner = read_array_from_reader(
             self.file().await?,
             projection_from_python(columns)?,
-            batch_size,
             row_filter_from_python(row_filter),
         )
         .await?;
@@ -140,13 +132,11 @@ impl TokioFileDataset {
     async fn async_to_record_batch_reader(
         self_: PyRef<'_, Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyObject> {
         let layout_reader = layout_stream_from_reader(
             self_.file().await?,
             projection_from_python(columns)?,
-            batch_size,
             row_filter_from_python(row_filter),
         )
         .await?;
@@ -164,25 +154,23 @@ impl TokioFileDataset {
         self_.schema.clone().to_pyarrow(self_.py())
     }
 
-    #[pyo3(signature = (*, columns=None, batch_size=None, row_filter=None))]
+    #[pyo3(signature = (*, columns=None, row_filter=None))]
     pub fn to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyArray> {
-        TOKIO_RUNTIME.block_on(self.async_to_array(columns, batch_size, row_filter))
+        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter))
     }
 
-    #[pyo3(signature = (*, columns=None, batch_size=None, row_filter=None))]
+    #[pyo3(signature = (*, columns=None, row_filter=None))]
     pub fn to_record_batch_reader(
         self_: PyRef<Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyObject> {
         TOKIO_RUNTIME.block_on(Self::async_to_record_batch_reader(
-            self_, columns, batch_size, row_filter,
+            self_, columns, row_filter,
         ))
     }
 }
@@ -208,13 +196,11 @@ impl ObjectStoreUrlDataset {
     async fn async_to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyArray> {
         let inner = read_array_from_reader(
             self.reader().await?,
             projection_from_python(columns)?,
-            batch_size,
             row_filter_from_python(row_filter),
         )
         .await?;
@@ -224,13 +210,11 @@ impl ObjectStoreUrlDataset {
     async fn async_to_record_batch_reader(
         self_: PyRef<'_, Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyObject> {
         let layout_reader = layout_stream_from_reader(
             self_.reader().await?,
             projection_from_python(columns)?,
-            batch_size,
             row_filter_from_python(row_filter),
         )
         .await?;
@@ -248,25 +232,23 @@ impl ObjectStoreUrlDataset {
         self_.schema.clone().to_pyarrow(self_.py())
     }
 
-    #[pyo3(signature = (*, columns=None, batch_size=None, row_filter=None))]
+    #[pyo3(signature = (*, columns=None, row_filter=None))]
     pub fn to_array(
         &self,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyArray> {
-        TOKIO_RUNTIME.block_on(self.async_to_array(columns, batch_size, row_filter))
+        TOKIO_RUNTIME.block_on(self.async_to_array(columns, row_filter))
     }
 
-    #[pyo3(signature = (*, columns=None, batch_size=None, row_filter=None))]
+    #[pyo3(signature = (*, columns=None, row_filter=None))]
     pub fn to_record_batch_reader(
         self_: PyRef<Self>,
         columns: Option<Vec<Bound<'_, PyAny>>>,
-        batch_size: Option<usize>,
         row_filter: Option<&Bound<'_, PyExpr>>,
     ) -> PyResult<PyObject> {
         TOKIO_RUNTIME.block_on(Self::async_to_record_batch_reader(
-            self_, columns, batch_size, row_filter,
+            self_, columns, row_filter,
         ))
     }
 }

pyvortex/src/io.rs

@@ -132,7 +132,7 @@ pub fn read_path(
     row_filter: Option<&Bound<PyExpr>>,
 ) -> PyResult<PyArray> {
     let dataset = TOKIO_RUNTIME.block_on(TokioFileDataset::try_new(path.extract()?))?;
-    dataset.to_array(projection, None, row_filter)
+    dataset.to_array(projection, row_filter)
 }
 
 /// Read a vortex struct array from a URL.
@@ -184,7 +184,7 @@ pub fn read_url(
     row_filter: Option<&Bound<PyExpr>>,
 ) -> PyResult<PyArray> {
     let dataset = TOKIO_RUNTIME.block_on(ObjectStoreUrlDataset::try_new(url.extract()?))?;
-    dataset.to_array(projection, None, row_filter)
+    dataset.to_array(projection, row_filter)
 }
 
 /// Write a vortex struct array to the local filesystem.

pyvortex/test/test_dataset.py

@@ -53,7 +53,7 @@ def test_to_batches(ds):
 
     chunk0 = next(ds.to_batches(columns=["string", "bool"]))
     assert chunk0.to_struct_array() == pa.array(
-        [record(x, columns=["string", "bool"]) for x in range(1 << 16)], type=schema
+        [record(x, columns=["string", "bool"]) for x in range(1_000_000)], type=schema
     )
 
 

vortex-datafusion/src/persistent/execution.rs

@@ -98,7 +98,6 @@ impl ExecutionPlan for VortexExec {
             ctx: self.ctx.clone(),
             object_store,
             projection: self.file_scan_config.projection.clone(),
-            batch_size: None,
             predicate: self.predicate.clone(),
             arrow_schema,
         };

vortex-datafusion/src/persistent/opener.rs

@@ -17,7 +17,6 @@ use vortex_serde::layouts::{
 pub struct VortexFileOpener {
     pub ctx: Arc<Context>,
     pub object_store: Arc<dyn ObjectStore>,
-    pub batch_size: Option<usize>,
     pub projection: Option<Vec<usize>>,
     pub predicate: Option<Arc<dyn PhysicalExpr>>,
     pub arrow_schema: SchemaRef,
@@ -33,10 +32,6 @@ impl FileOpener for VortexFileOpener {
             LayoutDeserializer::new(self.ctx.clone(), Arc::new(LayoutContext::default())),
         );
 
-        if let Some(batch_size) = self.batch_size {
-            builder = builder.with_batch_size(batch_size);
-        }
-
         let row_filter = self
             .predicate
             .clone()

vortex-expr/src/identity.rs

@@ -22,7 +22,7 @@ impl PartialEq<dyn Any> for Identity {
     fn eq(&self, other: &dyn Any) -> bool {
         unbox_any(other)
             .downcast_ref::<Self>()
-            .map(|x| x == other)
+            .map(|x| x == self)
             .unwrap_or(false)
     }
 }

vortex-expr/src/lib.rs

@@ -61,7 +61,7 @@ fn split_inner(expr: &Arc<dyn VortexExpr>, exprs: &mut Vec<Arc<dyn VortexExpr>>)
 }
 
 // Taken from apache-datafusion, necessary since you can't require VortexExpr implement PartialEq<dyn VortexExpr>
-pub(crate) fn unbox_any(any: &dyn Any) -> &dyn Any {
+pub fn unbox_any(any: &dyn Any) -> &dyn Any {
     if any.is::<Arc<dyn VortexExpr>>() {
         any.downcast_ref::<Arc<dyn VortexExpr>>()
             .vortex_expect("any.is::<Arc<dyn VortexExpr>> returned true but downcast_ref failed")

vortex-flatbuffers/flatbuffers/vortex-serde/footer.fbs

@@ -7,6 +7,7 @@ table Layout {
     encoding: uint16;
     buffers: [Buffer];
     children: [Layout];
+    length: uint64;
     metadata: [ubyte];
 }
 

vortex-flatbuffers/src/generated/footer.rs

@@ -152,7 +152,8 @@ impl<'a> Layout<'a> {
   pub const VT_ENCODING: flatbuffers::VOffsetT = 4;
   pub const VT_BUFFERS: flatbuffers::VOffsetT = 6;
   pub const VT_CHILDREN: flatbuffers::VOffsetT = 8;
-  pub const VT_METADATA: flatbuffers::VOffsetT = 10;
+  pub const VT_LENGTH: flatbuffers::VOffsetT = 10;
+  pub const VT_METADATA: flatbuffers::VOffsetT = 12;
 
   #[inline]
   pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
@@ -164,6 +165,7 @@ impl<'a> Layout<'a> {
     args: &'args LayoutArgs<'args>
   ) -> flatbuffers::WIPOffset<Layout<'bldr>> {
     let mut builder = LayoutBuilder::new(_fbb);
+    builder.add_length(args.length);
     if let Some(x) = args.metadata { builder.add_metadata(x); }
     if let Some(x) = args.children { builder.add_children(x); }
     if let Some(x) = args.buffers { builder.add_buffers(x); }
@@ -194,6 +196,13 @@ impl<'a> Layout<'a> {
     unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Layout>>>>(Layout::VT_CHILDREN, None)}
   }
   #[inline]
+  pub fn length(&self) -> u64 {
+    // Safety:
+    // Created from valid Table for this object
+    // which contains a valid value in this slot
+    unsafe { self._tab.get::<u64>(Layout::VT_LENGTH, Some(0)).unwrap()}
+  }
+  #[inline]
   pub fn metadata(&self) -> Option<flatbuffers::Vector<'a, u8>> {
     // Safety:
     // Created from valid Table for this object
@@ -212,6 +221,7 @@ impl flatbuffers::Verifiable for Layout<'_> {
      .visit_field::<u16>("encoding", Self::VT_ENCODING, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, Buffer>>>("buffers", Self::VT_BUFFERS, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, flatbuffers::ForwardsUOffset<Layout>>>>("children", Self::VT_CHILDREN, false)?
+     .visit_field::<u64>("length", Self::VT_LENGTH, false)?
      .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, u8>>>("metadata", Self::VT_METADATA, false)?
      .finish();
     Ok(())
@@ -221,6 +231,7 @@ pub struct LayoutArgs<'a> {
     pub encoding: u16,
     pub buffers: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, Buffer>>>,
     pub children: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, flatbuffers::ForwardsUOffset<Layout<'a>>>>>,
+    pub length: u64,
     pub metadata: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, u8>>>,
 }
 impl<'a> Default for LayoutArgs<'a> {
@@ -230,6 +241,7 @@ impl<'a> Default for LayoutArgs<'a> {
       encoding: 0,
       buffers: None,
       children: None,
+      length: 0,
       metadata: None,
     }
   }
@@ -253,6 +265,10 @@ impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> LayoutBuilder<'a, 'b, A> {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(Layout::VT_CHILDREN, children);
   }
   #[inline]
+  pub fn add_length(&mut self, length: u64) {
+    self.fbb_.push_slot::<u64>(Layout::VT_LENGTH, length, 0);
+  }
+  #[inline]
   pub fn add_metadata(&mut self, metadata: flatbuffers::WIPOffset<flatbuffers::Vector<'b , u8>>) {
     self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(Layout::VT_METADATA, metadata);
   }
@@ -277,6 +293,7 @@ impl core::fmt::Debug for Layout<'_> {
       ds.field("encoding", &self.encoding());
       ds.field("buffers", &self.buffers());
       ds.field("children", &self.children());
+      ds.field("length", &self.length());
       ds.field("metadata", &self.metadata());
       ds.finish()
   }