End DataFusion stream early if file can be pruned away by dynamic expr

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

vortex-datafusion/src/persistent/mod.rs

@@ -9,6 +9,7 @@ pub mod metrics;
 mod opener;
 mod sink;
 mod source;
+mod stream;
 
 pub use access_plan::VortexAccessPlan;
 pub use format::VortexFormat;

vortex-datafusion/src/persistent/opener.rs

@@ -51,6 +51,7 @@ use crate::VortexAccessPlan;
 use crate::convert::exprs::ExpressionConvertor;
 use crate::convert::exprs::can_be_pushed_down;
 use crate::convert::exprs::make_vortex_predicate;
+use crate::persistent::stream::PrunableStream;
 
 #[derive(Clone)]
 pub(crate) struct VortexOpener {
@@ -336,7 +337,11 @@ impl FileOpener for VortexOpener {
                 .map(move |batch| batch.and_then(|b| schema_mapping.map_batch(b)))
                 .boxed();
 
-            Ok(stream)
+            if let Some(file_pruner) = file_pruner {
+                Ok(PrunableStream::new(file_pruner, stream).boxed())
+            } else {
+                Ok(stream)
+            }
         }
         .in_current_span()
         .boxed())

vortex-datafusion/src/persistent/stream.rs

@@ -0,0 +1,39 @@
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use datafusion_common::Result as DFResult;
+use datafusion_common::arrow::array::RecordBatch;
+use datafusion_pruning::FilePruner;
+use futures::Stream;
+use futures::StreamExt;
+use futures::stream::BoxStream;
+
+/// Utility to end a stream early if its backing [`PartitionFile`] can be pruned away by an updated dynamic expression.
+pub(crate) struct PrunableStream {
+    file_pruner: FilePruner,
+    stream: BoxStream<'static, DFResult<RecordBatch>>,
+}
+
+impl PrunableStream {
+    pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult<RecordBatch>>) -> Self {
+        Self {
+            file_pruner,
+            stream,
+        }
+    }
+}
+
+impl Stream for PrunableStream {
+    type Item = DFResult<RecordBatch>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.as_mut().file_pruner.should_prune() {
+            Ok(true) => return Poll::Ready(None),
+            Ok(false) => {}
+            Err(e) => return Poll::Ready(Some(Err(e))),
+        };
+
+        self.stream.poll_next_unpin(cx)
+    }
+}