Reduce V2 duplicate scan requests

gatesn · gatesn · commit 72f39c74147f · 2026-06-21T19:23:53.000-04:00
Signed-off-by: "Nicholas Gates" &lt;nick@nickgates.com&gt;
diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs
@@ -103,6 +103,8 @@ pub(crate) struct VortexOpener {
     pub layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Shared full-file natural split ranges keyed by file path.
     pub natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
+    /// Shared V2 file handles keyed by file path.
+    pub vortex_files: Arc<DashMap<Path, Arc<VortexFile>>>,
     /// Whether the query has output ordering specified
     pub has_output_ordering: bool,
 
@@ -139,6 +141,7 @@ impl FileOpener for VortexOpener {
         let limit = self.limit;
         let layout_readers = Arc::clone(&self.layout_readers);
         let natural_split_ranges = Arc::clone(&self.natural_split_ranges);
+        let vortex_files = Arc::clone(&self.vortex_files);
         let has_output_ordering = self.has_output_ordering;
         let scan_concurrency = self.scan_concurrency;
 
@@ -216,10 +219,24 @@ impl FileOpener for VortexOpener {
                 open_opts = open_opts.with_footer(footer);
             }
 
-            let vxf = open_opts
-                .open_read(reader)
-                .await
-                .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?;
+            let vxf = if let Some(hit) = vortex_files.get(&file.object_meta.location) {
+                Arc::clone(hit.value())
+            } else {
+                let opened = Arc::new(
+                    open_opts
+                        .open_read(reader)
+                        .await
+                        .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?,
+                );
+
+                match vortex_files.entry(file.object_meta.location.clone()) {
+                    Entry::Occupied(entry) => Arc::clone(entry.get()),
+                    Entry::Vacant(entry) => {
+                        entry.insert(Arc::clone(&opened));
+                        opened
+                    }
+                }
+            };
 
             // On a miss, cache the parsed footer so other partitions and later executions
             // skip the footer fetch and parse. `infer_schema`/`infer_stats` also populate
@@ -915,6 +932,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1111,6 +1129,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1199,6 +1218,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1357,6 +1377,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1418,6 +1439,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1628,6 +1650,7 @@ mod tests {
             metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
             layout_readers: Default::default(),
             natural_split_ranges: Default::default(),
+            vortex_files: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs
@@ -32,6 +32,7 @@ use object_store::ObjectStore;
 use object_store::path::Path;
 use vortex::error::VortexExpect;
 use vortex::file::VORTEX_FILE_EXTENSION;
+use vortex::file::VortexFile;
 use vortex::layout::LayoutReader;
 use vortex::layout::scan::v2::scan2_enabled;
 use vortex::metrics::DefaultMetricsRegistry;
@@ -200,6 +201,8 @@ pub struct VortexSource {
     layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Shared full-file natural split ranges keyed by path.
     natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
+    /// Shared V2 file handles keyed by path.
+    vortex_files: Arc<DashMap<Path, Arc<VortexFile>>>,
     expression_convertor: Arc<dyn ExpressionConvertor>,
     pub(crate) vortex_reader_factory: Option<Arc<dyn VortexReaderFactory>>,
     pub(crate) ordered: bool,
@@ -233,6 +236,7 @@ impl VortexSource {
             _unused_df_metrics: Default::default(),
             layout_readers: Arc::new(DashMap::default()),
             natural_split_ranges: Arc::new(DashMap::default()),
+            vortex_files: Arc::new(DashMap::default()),
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             vortex_reader_factory: None,
             vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
@@ -367,6 +371,7 @@ impl VortexSource {
             metrics_registry: Arc::clone(&self.vx_metrics_registry),
             layout_readers: Arc::clone(&self.layout_readers),
             natural_split_ranges: Arc::clone(&self.natural_split_ranges),
+            vortex_files: Arc::clone(&self.vortex_files),
             has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered,
             expression_convertor: Arc::clone(&self.expression_convertor),
             file_metadata_cache: self.file_metadata_cache.clone(),
diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs
@@ -30,7 +30,10 @@ use vortex_layout::segments::SegmentInfo;
 use vortex_layout::segments::SegmentSource;
 use vortex_scan::DataSourceRef;
 use vortex_scan::ScanRequest;
+use vortex_scan::plan::PreparedStateCache;
+use vortex_scan::plan::PreparedStateCacheRef;
 use vortex_scan::plan::ScanPlanRef;
+use vortex_scan::segments::SegmentFutureCache;
 use vortex_session::VortexSession;
 
 use crate::FileStatistics;
@@ -58,6 +61,10 @@ pub struct VortexFile {
     layout_reader_cache: Option<OnceLock<Arc<dyn LayoutReader>>>,
     /// Shared cache for the v2 physical scan plan root.
     scan_plan_root_cache: Arc<OnceLock<ScanPlanRef>>,
+    /// Shared cache for v2 prepared state across row-range scans of this file.
+    scan_plan_state_cache: PreparedStateCacheRef,
+    /// Shared cache for v2 in-flight segment futures across row-range scans of this file.
+    scan_plan_segment_future_cache: Arc<SegmentFutureCache>,
 }
 
 fn layout_reader(
@@ -104,6 +111,8 @@ impl VortexFile {
             session,
             layout_reader_cache: None,
             scan_plan_root_cache: Arc::new(OnceLock::new()),
+            scan_plan_state_cache: Arc::new(PreparedStateCache::default()),
+            scan_plan_segment_future_cache: Arc::new(SegmentFutureCache::new()),
         }
     }
 
@@ -116,6 +125,8 @@ impl VortexFile {
             session: self.session,
             layout_reader_cache: Some(OnceLock::new()),
             scan_plan_root_cache: self.scan_plan_root_cache,
+            scan_plan_state_cache: self.scan_plan_state_cache,
+            scan_plan_segment_future_cache: self.scan_plan_segment_future_cache,
         }
     }
 
@@ -203,6 +214,14 @@ impl VortexFile {
         Ok(root)
     }
 
+    pub(crate) fn scan_plan_state_cache(&self) -> PreparedStateCacheRef {
+        Arc::clone(&self.scan_plan_state_cache)
+    }
+
+    pub(crate) fn scan_plan_segment_future_cache(&self) -> Arc<SegmentFutureCache> {
+        Arc::clone(&self.scan_plan_segment_future_cache)
+    }
+
     /// Create a [`DataSource`](vortex_scan::DataSource) from this file for scanning.
     ///
     /// Wraps the file's layout reader with [`FileStatsLayoutReader`] (when file-level
diff --git a/vortex-file/src/multi/scan_v2.rs b/vortex-file/src/multi/scan_v2.rs
@@ -1089,8 +1089,9 @@ impl<T: Send + 'static> Work<T> {
     ) -> Self {
         let known_bytes = registered.bytes();
         let future = async move {
-            let _registered = registered;
-            future.await
+            let result = future.await;
+            drop(registered);
+            result
         }
         .boxed();
         Self {
@@ -1862,7 +1863,7 @@ impl PreparedScanPlanFile {
             },
         );
         let scheduled_segment_source = Arc::clone(&registered_source.source);
-        let segment_future_cache = Arc::new(SegmentFutureCache::new());
+        let segment_future_cache = file.scan_plan_segment_future_cache();
         let reader = FileReader::new(
             Arc::new(ScheduledSegmentSourceReader::new(
                 segment_source_id,
@@ -1872,7 +1873,8 @@ impl PreparedScanPlanFile {
             session.clone(),
         );
 
-        let mut prepare_ctx = PrepareCtx::new(session.clone());
+        let mut prepare_ctx =
+            PrepareCtx::with_state_cache(session.clone(), file.scan_plan_state_cache());
         let projection_pushed = push_expr(&root, &projection, file.dtype(), &session)?;
         let mut split_hints = Vec::new();
         extend_split_hints(&projection_pushed, &mut split_hints);
diff --git a/vortex-layout/src/scan/v2/layouts/flat.rs b/vortex-layout/src/scan/v2/layouts/flat.rs
@@ -13,6 +13,7 @@ use std::fmt;
 use std::ops::Range;
 use std::sync::Arc;
 
+use futures::FutureExt;
 use futures::future::BoxFuture;
 use parking_lot::Mutex;
 use vortex_array::ArrayRef;
@@ -21,6 +22,7 @@ use vortex_array::IntoArray;
 use vortex_array::arrays::SliceArray;
 use vortex_array::expr::Expression;
 use vortex_array::serde::SerializedArray;
+use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
@@ -44,6 +46,7 @@ use vortex_session::VortexSession;
 use crate::layout_v2::Flat;
 use crate::layout_v2::Layout;
 use crate::layout_v2::LayoutRef;
+use crate::layouts::SharedArrayFuture;
 use crate::segments::SegmentPlanCtx;
 use crate::segments::SegmentRequests;
 
@@ -63,19 +66,38 @@ pub struct FlatScanPlan {
     layout: LayoutRef,
 }
 
-/// Per-query cache of the parsed (still lazy) array. Concurrent decodes
-/// are benign: the segment fetch is deduplicated by the shared segment
-/// source, and last-write-wins on the parsed array.
+/// Per-query cache of the parsed (still lazy) array.
 #[derive(Default)]
 pub struct FlatScanState {
-    array: Mutex<Option<ArrayRef>>,
+    array: Mutex<Option<SharedArrayFuture>>,
 }
 
 struct FlatPreparedRead {
     node: Arc<FlatScanPlan>,
     state: Arc<FlatScanState>,
 }
 
+impl FlatScanPlan {
+    fn array(&self, io: &FileReader, state: &FlatScanState) -> SharedArrayFuture {
+        if let Some(hit) = state.array.lock().clone() {
+            return hit;
+        }
+
+        let mut guard = state.array.lock();
+        if let Some(hit) = guard.clone() {
+            return hit;
+        }
+
+        let layout = self.layout.clone();
+        let io = io.clone();
+        let future = async move { decode_flat(&layout, &io).await.map_err(Arc::new) }
+            .boxed()
+            .shared();
+        *guard = Some(future.clone());
+        future
+    }
+}
+
 impl ScanPlan for FlatScanPlan {
     fn init_state(&self, _cx: &mut StateCtx<'_>) -> VortexResult<ScanStateRef> {
         Ok(Arc::new(FlatScanState::default()))
@@ -90,7 +112,10 @@ impl ScanPlan for FlatScanPlan {
     }
 
     fn prepare_read(self: Arc<Self>, cx: &mut PrepareCtx) -> VortexResult<Option<PreparedReadRef>> {
-        let key = PreparedStateKey::new::<FlatScanState>(Arc::as_ptr(&self) as *const () as usize);
+        let flat = self.layout.as_opt::<Flat>().ok_or_else(|| {
+            vortex_err!("expected flat layout, got {}", self.layout.encoding_id())
+        })?;
+        let key = PreparedStateKey::new::<FlatScanState>(*flat.data().segment_id() as usize);
         let state = cx.shared_state(key, || Ok(FlatScanState::default()))?;
         Ok(Some(Arc::new(FlatPreparedRead { node: self, state })))
     }
@@ -120,13 +145,11 @@ impl PreparedRead for FlatPreparedRead {
         _local: &'a mut ExecutionCtx,
     ) -> BoxFuture<'a, VortexResult<ArrayRef>> {
         Box::pin(async move {
-            let array = if let Some(hit) = self.state.array.lock().clone() {
-                hit
-            } else {
-                let decoded = decode_flat(&self.node.layout, io).await?;
-                *self.state.array.lock() = Some(decoded.clone());
-                decoded
-            };
+            let array = self
+                .node
+                .array(io, &self.state)
+                .await
+                .map_err(VortexError::from)?;
             let dense = slice_to_range(array, &range)?;
             if rows.selection.len() != dense.len() {
                 vortex_bail!(
diff --git a/vortex-scan/src/plan/mod.rs b/vortex-scan/src/plan/mod.rs
@@ -1242,6 +1242,20 @@ impl PreparedRead for MaskPreparedRead {
         })
     }
 
+    fn segment_requests(
+        &self,
+        range: Range<u64>,
+        rows: RowScope<'_>,
+        cx: &mut SegmentPlanCtx,
+    ) -> VortexResult<SegmentRequests> {
+        let mut requests = self.input.segment_requests(range.clone(), rows, cx)?;
+        if requests.is_unknown() {
+            return Ok(requests);
+        }
+        requests.extend(self.validity.segment_requests(range, rows, cx)?);
+        Ok(requests)
+    }
+
     fn release(&self, frontier: u64) -> VortexResult<()> {
         self.input.release(frontier)?;
         self.validity.release(frontier)
