End DataFusion stream early if file can be pruned away by dynamic expr (#5874)

This PR adds an additional pruning check on every poll of the File
stream, so if the underlying dynamic expression is updated, the stream
will stop.

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,40 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+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>> {
+        if self.as_mut().file_pruner.should_prune()? {
+            Poll::Ready(None)
+        } else {
+            self.stream.poll_next_unpin(cx)
+        }
+    }
+}