feat(scan): select the V2 scan from vortex-file, DataFusion and DuckDB

vortex-datafusion/src/persistent/diagnostics.rs

@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Opt-in cache ablation for benchmarks over immutable files. This retains a cache per object
+//! for the process lifetime and is not a general-purpose, globally bounded application cache.
+
+use std::sync::Arc;
+use std::sync::LazyLock;
+
+use datafusion_common::Result as DFResult;
+use datafusion_common::exec_datafusion_err;
+use datafusion_datasource::PartitionedFile;
+use futures::StreamExt;
+use futures::TryStreamExt;
+use futures::stream;
+use tokio::sync::OnceCell;
+use vortex::file::VortexFile;
+use vortex::io::VortexReadAt;
+use vortex::layout::segments::MokaSegmentCache;
+use vortex::layout::segments::NoOpSegmentCache;
+use vortex::layout::segments::SegmentCache;
+use vortex::layout::segments::SegmentEviction;
+use vortex::layout::segments::SegmentId;
+use vortex_utils::aliases::dash_map::DashMap;
+
+static CAPACITY: LazyLock<Option<String>> =
+    LazyLock::new(|| std::env::var("VORTEX_BENCH_SEGMENT_CACHE_MB").ok());
+static PRELOAD: LazyLock<bool> = LazyLock::new(|| {
+    std::env::var("VORTEX_BENCH_PRELOAD_SEGMENTS").is_ok_and(|value| value == "1")
+});
+static CACHES: LazyLock<DashMap<CacheKey, BenchmarkCache>> = LazyLock::new(DashMap::default);
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+struct CacheKey {
+    uri: Arc<str>,
+    size: u64,
+    modified: String,
+    etag: Option<String>,
+    version: Option<String>,
+}
+
+#[derive(Clone)]
+pub(super) struct BenchmarkCache {
+    pub(super) segments: Arc<dyn SegmentCache>,
+    preloaded: Arc<OnceCell<()>>,
+    capacity: u64,
+}
+
+impl BenchmarkCache {
+    pub(super) async fn preload(&self, file: &VortexFile) -> DFResult<()> {
+        if !*PRELOAD || self.capacity == 0 {
+            return Ok(());
+        }
+        self.preloaded
+            .get_or_try_init(|| async {
+                let source = file.segment_source();
+                let count = u32::try_from(file.footer().segment_map().len())
+                    .map_err(|e| exec_datafusion_err!("Too many segments: {e}"))?;
+                stream::iter(0..count)
+                    .map(|id| source.request(SegmentId::from(id)))
+                    .buffer_unordered(64)
+                    .try_for_each(|_| async { Ok(()) })
+                    .await
+                    .map_err(|e| exec_datafusion_err!("Preloading segments failed: {e}"))
+            })
+            .await?;
+        Ok(())
+    }
+}
+
+pub(super) fn benchmark_segment_cache(
+    reader: &dyn VortexReadAt,
+    file: &PartitionedFile,
+) -> DFResult<Option<BenchmarkCache>> {
+    let Some(value) = CAPACITY.as_ref() else {
+        return Ok(None);
+    };
+    let bytes = value
+        .parse::<u64>()
+        .ok()
+        .and_then(|mb| mb.checked_mul(1024 * 1024))
+        .ok_or_else(|| exec_datafusion_err!("Invalid VORTEX_BENCH_SEGMENT_CACHE_MB: {value}"))?;
+    // Zero measures the segment-source adapter without retaining bytes, controlling for the
+    // different routing used by V2's uncached direct range service.
+    if bytes == 0 {
+        return Ok(Some(BenchmarkCache {
+            segments: Arc::new(NoOpSegmentCache),
+            preloaded: Arc::default(),
+            capacity: 0,
+        }));
+    }
+    if *PRELOAD && bytes < file.object_meta.size {
+        return Err(exec_datafusion_err!(
+            "Preloaded cache must fit the entire file ({} bytes)",
+            file.object_meta.size
+        ));
+    }
+    let uri = reader.uri().ok_or_else(|| {
+        exec_datafusion_err!("Benchmark segment caching requires an identified reader")
+    })?;
+    // Segment ids are file-local. Object metadata also distinguishes file replacements.
+    let key = CacheKey {
+        uri: Arc::clone(uri),
+        size: file.object_meta.size,
+        modified: file.object_meta.last_modified.to_rfc3339(),
+        etag: file.object_meta.e_tag.clone(),
+        version: file.object_meta.version.clone(),
+    };
+    Ok(Some(
+        CACHES
+            .entry(key)
+            .or_insert_with(|| BenchmarkCache {
+                segments: Arc::new(
+                    MokaSegmentCache::new(bytes, SegmentEviction::TinyLfu).for_file(""),
+                ),
+                preloaded: Arc::default(),
+                capacity: bytes,
+            })
+            .value()
+            .clone(),
+    ))
+}

vortex-datafusion/src/persistent/mod.rs

@@ -24,6 +24,7 @@
 
 mod access_plan;
 mod cache;
+mod diagnostics;
 mod format;
 pub mod metrics;
 mod opener;

vortex-datafusion/src/persistent/opener.rs

@@ -49,8 +49,10 @@ use vortex::error::VortexError;
 use vortex::error::VortexExpect;
 use vortex::expr::and_collect;
 use vortex::file::OpenOptionsSessionExt;
+use vortex::file::planning;
 use vortex::io::InstrumentedReadAt;
 use vortex::layout::LayoutReader;
+use vortex::layout::scan;
 use vortex::layout::scan::scan_builder::ScanBuilder;
 use vortex::metrics::Label;
 use vortex::metrics::MetricsRegistry;
@@ -67,6 +69,7 @@ use crate::convert::schema::calculate_physical_schema;
 use crate::metrics::PARTITION_LABEL;
 use crate::metrics::PATH_LABEL;
 use crate::persistent::cache::CachedVortexMetadata;
+use crate::persistent::diagnostics::benchmark_segment_cache;
 use crate::persistent::reader::VortexReaderFactory;
 use crate::persistent::stream::PrunableStream;
 
@@ -204,6 +207,11 @@ impl FileOpener for VortexOpener {
                 .with_metrics_registry(Arc::clone(&metrics_registry))
                 .with_labels(labels);
 
+            let benchmark_cache = benchmark_segment_cache(&reader, &file)?;
+            if let Some(cache) = &benchmark_cache {
+                open_opts = open_opts.with_segment_cache(Arc::clone(&cache.segments));
+            }
+
             let cached_footer = file_metadata_cache
                 .as_ref()
                 .and_then(|cache| cache.get(file.path()))
@@ -226,6 +234,10 @@ impl FileOpener for VortexOpener {
                 .await
                 .map_err(|e| exec_datafusion_err!("Failed to open Vortex file {e}"))?;
 
+            if let Some(cache) = benchmark_cache {
+                cache.preload(&vxf).await?;
+            }
+
             // 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
             // this cache, but only when planning goes through `VortexFormat`.
@@ -452,7 +464,7 @@ impl FileOpener for VortexOpener {
 
             let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false);
             let residual_path = file.path().clone();
-            let stream = scan_builder
+            let scan_builder = scan_builder
                 .with_metrics_registry(metrics_registry)
                 .with_ordered(has_output_ordering)
                 .map(move |chunk| {
@@ -464,8 +476,14 @@ impl FileOpener for VortexOpener {
                         &mut ctx,
                     )?;
                     Ok(RecordBatch::from(arrow.as_struct().clone()))
-                })
-                .into_stream()
+                });
+            let scan_file = planning::scan_file(&vxf);
+            let batches = if scan::v2::enabled() && scan::v2::supports_scan(&scan_builder, &scan_file) {
+                scan::v2::into_stream(scan_builder, scan_file).map(|s| s.boxed())
+            } else {
+                scan_builder.into_stream().map(|s| s.boxed())
+            };
+            let stream = batches
                 .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))?
                 .map_err(move |e: VortexError| {
                     DataFusionError::External(Box::new(e.with_context(format!(

vortex-duckdb/cpp/include/multi_file_reader.hpp

@@ -28,6 +28,14 @@ struct VortexGlobalState final : GlobalTableFunctionState {
     VortexGlobalState() = default;
     ~VortexGlobalState() override = default;
 
+    void FillPrefetch(idx_t file_index);
+
+    mutex prefetch_mutex;
+    idx_t prefetch_window = 0;
+    idx_t next_prefetch_index = 0;
+    MultiFileListScanData prefetch_scan;
+    const MultiFileList *prefetch_files = nullptr;
+    bool prefetch_exhausted = false;
     const void *ffi_bind_data = nullptr; // needed for local state partial accumulation
     unique_ptr<CData> ffi_global_state;
 };
@@ -158,6 +166,7 @@ struct VortexBaseReader final : BaseFileReader {
     }
 
     unique_ptr<CData> ffi_file;
+    bool prefetched = false;
     /*
      * Populated only for first file reader in scan when BindReader() is
      * called on it. Used in GetStatistics() which is called only for first

vortex-duckdb/cpp/multi_file_reader.cpp

@@ -9,6 +9,48 @@
 
 #include "duckdb/execution/operator/scan/physical_table_scan.hpp"
 
+void VortexGlobalState::FillPrefetch(idx_t file_index) {
+    if (!prefetch_window) {
+        return;
+    }
+    lock_guard<mutex> lock(prefetch_mutex);
+    // Out-of-order claims can extend this by at most DuckDB's worker lookahead.
+    while (!prefetch_exhausted && next_prefetch_index < file_index + prefetch_window) {
+        OpenFileInfo file;
+        if (!prefetch_files->Scan(prefetch_scan, file)) {
+            prefetch_exhausted = true;
+            break;
+        }
+        duckdb_vx_error error = nullptr;
+        duckdb_file_prefetch_submit(ffi_global_state->DataPtr(),
+                                    next_prefetch_index,
+                                    file.path.c_str(),
+                                    file.path.size(),
+                                    &error);
+        if (error) {
+            throw IOException(IntoErrString(error));
+        }
+        next_prefetch_index++;
+    }
+}
+
+static void ClaimPrefetched(VortexGlobalState &global, VortexBaseReader &reader, bool skip) {
+    if (!global.prefetch_window || reader.prefetched) {
+        return;
+    }
+    const auto index = reader.file_list_idx.GetIndex();
+    global.FillPrefetch(index);
+    duckdb_vx_error error = nullptr;
+    auto file = duckdb_file_prefetch_take(global.ffi_global_state->DataPtr(), index, skip, &error);
+    if (error) {
+        throw IOException(IntoErrString(error));
+    }
+    if (!skip) {
+        reader.ffi_file = unique_ptr<CData>(reinterpret_cast<CData *>(file));
+    }
+    reader.prefetched = true;
+}
+
 unique_ptr<FunctionData> VortexBindData::Copy() const {
     auto result = make_uniq<VortexBindData>();
     if (ffi_bind_data) {
@@ -49,6 +91,7 @@ VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data,
     D_ASSERT(gstate.global_state != nullptr);
 
     VortexBaseReader &reader = reader_data.reader->Cast<VortexBaseReader>();
+    auto &global = gstate.global_state->Cast<VortexGlobalState>();
     const VortexBindData &bind = bind_data.bind_data->Cast<VortexBindData>();
 
     const void *const ffi_bind = bind.ffi_bind_data->DataPtr();
@@ -80,11 +123,12 @@ VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data,
                                                                                  context,
                                                                                  gstate);
         if (base_skip == ReaderInitializeType::SKIP_READING_FILE) {
+            ClaimPrefetched(global, reader, true);
             return base_skip;
         }
     }
 
-    const VortexGlobalState &global = gstate.global_state->Cast<VortexGlobalState>();
+    ClaimPrefetched(global, reader, false);
 
     duckdb_vx_error error = nullptr;
     const void *const ffi_global = global.ffi_global_state->DataPtr();
@@ -184,6 +228,12 @@ VortexReaderInterface::InitializeGlobalState(ClientContext &context,
     auto result = make_uniq<VortexGlobalState>();
     result->ffi_bind_data = ffi_bind;
     result->ffi_global_state = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_global_state));
+    result->prefetch_window = duckdb_file_prefetch_window(result->ffi_global_state->DataPtr());
+    if (result->prefetch_window) {
+        result->prefetch_files = &input.file_list;
+        input.file_list.InitializeScan(result->prefetch_scan);
+        result->FillPrefetch(0);
+    }
     return result;
 }
 
@@ -214,10 +264,15 @@ static shared_ptr<BaseFileReader> OpenReader(const OpenFileInfo &file) {
 }
 
 shared_ptr<BaseFileReader> VortexReaderInterface::CreateReader(ClientContext &,
-                                                               GlobalTableFunctionState &,
+                                                               GlobalTableFunctionState &gstate,
                                                                const OpenFileInfo &file,
                                                                idx_t,
                                                                const MultiFileBindData &) {
+    if (gstate.Cast<VortexGlobalState>().prefetch_window) {
+        // InitializeReader first applies DuckDB's virtual-column pruning, then claims
+        // the prepared Rust reader. No data is needed for the base column mapping.
+        return make_shared_ptr<VortexBaseReader>(file, nullptr);
+    }
     return OpenReader(file);
 }
 

vortex-duckdb/include/vortex.h

@@ -58,6 +58,21 @@ duckdb_vx_data duckdb_reader_open(const char *file_path,
                                   size_t file_path_len,
                                   duckdb_vx_error *error);
 
+extern size_t duckdb_file_prefetch_window(const void *global);
+
+extern
+void duckdb_file_prefetch_submit(const void *global,
+                                 size_t index,
+                                 const char *path,
+                                 size_t path_len,
+                                 duckdb_vx_error *error);
+
+extern
+duckdb_vx_data duckdb_file_prefetch_take(const void *global,
+                                         size_t index,
+                                         bool skip,
+                                         duckdb_vx_error *error);
+
 extern
 bool duckdb_reader_get_statistics(const void *file,
                                   const void *bind,

vortex-duckdb/src/ffi.rs

@@ -197,6 +197,53 @@ pub unsafe extern "C-unwind" fn duckdb_reader_open(
     })
 }
 
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn duckdb_file_prefetch_window(global: *const c_void) -> usize {
+    let global = unsafe { global.cast::<GlobalState>().as_ref() }.vortex_expect("null pointer");
+    global
+        .file_prefetch
+        .as_ref()
+        .map_or(0, |prefetch| prefetch.window())
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn duckdb_file_prefetch_submit(
+    global: *const c_void,
+    index: usize,
+    path: *const c_char,
+    path_len: usize,
+    error: *mut cpp::duckdb_vx_error,
+) {
+    let global = unsafe { global.cast::<GlobalState>().as_ref() }.vortex_expect("null pointer");
+    let bytes = unsafe { std::slice::from_raw_parts(path.cast::<u8>(), path_len) };
+    try_or(error, || {
+        let path = str::from_utf8(bytes).map_err(|_| vortex_err!("invalid utf-8"))?;
+        global
+            .file_prefetch
+            .as_ref()
+            .vortex_expect("prefetch disabled")
+            .submit(index, path.to_owned())
+    });
+}
+
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn duckdb_file_prefetch_take(
+    global: *const c_void,
+    index: usize,
+    skip: bool,
+    error: *mut cpp::duckdb_vx_error,
+) -> cpp::duckdb_vx_data {
+    let global = unsafe { global.cast::<GlobalState>().as_ref() }.vortex_expect("null pointer");
+    try_or_null(error, || {
+        let file = global
+            .file_prefetch
+            .as_ref()
+            .vortex_expect("prefetch disabled")
+            .take(index, skip)?;
+        Ok(file.map_or(ptr::null_mut(), |file| Data::from(Box::new(file)).as_ptr()))
+    })
+}
+
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn duckdb_reader_get_statistics(
     file: *const c_void,