Compose and drive a V2 scan in DataFusion

benchmarks/compress-bench/src/vortex.rs

@@ -3,22 +3,25 @@
 
 use std::io::Cursor;
 use std::path::Path;
-use std::sync::Arc;
 use std::time::Duration;
 use std::time::Instant;
 
 use anyhow::Result;
 use anyhow::bail;
+use arrow_schema::Field;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures::StreamExt;
 use futures::pin_mut;
 use vortex::array::IntoArray;
+use vortex::array::VortexSessionExecute;
 use vortex::dtype::FieldNames;
 use vortex::expr::root;
 use vortex::expr::select;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
+use vortex::file::v2::ScanDriver;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::scan::v2::ScanBuilder;
 use vortex_arrow::ArrowSessionExt;
 use vortex_bench::Format;
@@ -96,13 +99,17 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> {
             .optimize_recursive()?;
         scan = scan.with_projection(projection);
     }
-    let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?);
+    let schema = SESSION.arrow().to_arrow_schema(scan.dtype())?;
+    let field = Field::new_struct("", schema.fields().clone(), false);
 
-    let stream = scan.into_record_batch_stream(schema)?;
+    let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan);
     pin_mut!(stream);
 
-    while let Some(batch) = stream.next().await {
-        let _batch = batch?;
+    while let Some(chunk) = stream.next().await {
+        let mut ctx = SESSION.create_execution_ctx();
+        let _batch = SESSION
+            .arrow()
+            .execute_arrow(chunk?, Some(&field), &mut ctx)?;
     }
     Ok(())
 }

vortex-bench/src/random_access/take.rs

@@ -35,6 +35,8 @@ use vortex::array::stream::ArrayStreamExt;
 use vortex::buffer::Buffer;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::VortexFile;
+use vortex::file::v2::ScanDriver;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::scan::v2::ScanBuilder;
 use vortex::scan::strict_sorted_buffer::StrictSortedBuffer;
 
@@ -190,9 +192,9 @@ impl RandomAccessor for VortexRandomAccessor {
 
     async fn take(&self, indices: &[u64]) -> anyhow::Result<RandomAccessorRet> {
         let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?;
-        let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file())
-            .with_row_indices(indices)
-            .into_stream()?;
+        let scan =
+            ScanBuilder::new(SESSION.clone(), self.file.scan_file()).with_row_indices(indices);
+        let stream = ScanDriver::new(SESSION.handle(), 4, true).run(scan);
         let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream)
             .read_all()
             .await?;

vortex-datafusion/src/persistent/mod.rs

@@ -28,6 +28,7 @@ mod format;
 pub mod metrics;
 mod opener;
 pub mod reader;
+mod scan_v2;
 mod sink;
 mod source;
 mod stream;

vortex-datafusion/src/persistent/opener.rs

@@ -53,15 +53,20 @@ use vortex::error::VortexExpect;
 use vortex::expr::and_collect;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::io::InstrumentedReadAt;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::LayoutReader;
 use vortex::layout::scan::scan_builder::ScanBuilder;
+use vortex::layout::scan::v2::ChunkPlanner;
+use vortex::layout::scan::v2::PrunePlanner;
 use vortex::layout::scan::v2::ScanBuilder as V2ScanBuilder;
 use vortex::layout::scan::v2::ScanCache;
 use vortex::layout::scan::v2::WeakScanCache;
 use vortex::layout::scan::v2::partition_row_range;
 use vortex::layout::segments::SegmentSource;
+use vortex::mask::Mask;
 use vortex::metrics::Label;
 use vortex::metrics::MetricsRegistry;
+use vortex::scan::planning::planner::Planner;
 use vortex::session::VortexSession;
 use vortex_arrow::ArrowSessionExt;
 use vortex_utils::aliases::dash_map::DashMap;
@@ -77,6 +82,7 @@ use crate::metrics::PARTITION_LABEL;
 use crate::metrics::PATH_LABEL;
 use crate::persistent::cache::CachedVortexMetadata;
 use crate::persistent::reader::VortexReaderFactory;
+use crate::persistent::scan_v2::ScanDriver;
 use crate::persistent::stream::PrunableStream;
 
 #[derive(Clone)]
@@ -116,7 +122,7 @@ pub(crate) struct VortexOpener {
     /// shared layout reader.
     pub segment_sources: Arc<DashMap<Path, Weak<dyn SegmentSource>>>,
     /// The V2 scan cache each open file's partitions share: its lowered layout, filter and
-    /// projection plans, proved zones and dictionary values, built once per file.
+    /// projection plans, pruned zones and dictionary values, built once per file.
     pub scan_caches: Arc<DashMap<Path, WeakScanCache>>,
     /// Shared V1 full-file natural splits keyed by file path.
     pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
@@ -437,8 +443,7 @@ impl FileOpener for VortexOpener {
                 scan_file.cache = Some(shared_scan_cache(&scan_caches, &file.object_meta.location));
                 let mut builder = V2ScanBuilder::new(vxf.session().clone(), scan_file)
                     .with_projection(scan_projection)
-                    .with_some_filter(native_filter)
-                    .with_ordered(has_output_ordering);
+                    .with_some_filter(native_filter);
                 if let Some(bytes) = byte_range {
                     // Every partition maps the same byte boundary to the same row boundary.
                     let rows = partition_row_range(bytes, file.object_meta.size, vxf.row_count());
@@ -462,8 +467,35 @@ impl FileOpener for VortexOpener {
                         .div_ceil(partitions)
                         .max(1)
                 });
-                builder = builder.with_concurrency(concurrency);
-                builder.map(map_batch).into_stream()
+                let handle = vxf.session().handle();
+                // Lowering and planning a wide file is real CPU: it runs where compute runs.
+                let scan = handle
+                    .spawn_cpu(move || builder.prepare())
+                    .await
+                    .map_err(|e| exec_datafusion_err!("Failed to plan Vortex scan: {e}"))?;
+                let scan = Arc::new(scan);
+
+                // The planners, composed: with a filter, its pruning plans first over the whole
+                // partition, continuing into the chunks of the rows they keep; without a
+                // filter, the chunks directly. A chunk plans its conjuncts, then its
+                // projection.
+                let chunks = {
+                    let scan = Arc::clone(&scan);
+                    move |kept: Option<Mask>| -> Box<dyn Planner> {
+                        let rows = scan.rows().clone();
+                        let selection = scan.selection().clone();
+                        Box::new(ChunkPlanner::new(scan, rows, selection, kept))
+                    }
+                };
+                let root: Box<dyn Planner> = if scan.filter().is_some() {
+                    let rows = scan.rows().clone();
+                    Box::new(PrunePlanner::new(Arc::clone(&scan), rows, Box::new(chunks)))
+                } else {
+                    chunks(None)
+                };
+
+                let driver = ScanDriver::new(handle, concurrency, has_output_ordering);
+                Ok(driver.run(scan, root, map_batch))
             } else {
                 // Share V1 layout readers across partitions of the file.
                 let layout_reader = match layout_readers.entry(file.object_meta.location.clone()) {