diff --git a/vortex-datafusion/src/persistent/diagnostics.rs b/vortex-datafusion/src/persistent/diagnostics.rs new file mode 100644 index 00000000000..5481edb7280 --- /dev/null +++ b/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> = + LazyLock::new(|| std::env::var("VORTEX_BENCH_SEGMENT_CACHE_MB").ok()); +static PRELOAD: LazyLock = LazyLock::new(|| { + std::env::var("VORTEX_BENCH_PRELOAD_SEGMENTS").is_ok_and(|value| value == "1") +}); +static CACHES: LazyLock> = LazyLock::new(DashMap::default); + +#[derive(Clone, PartialEq, Eq, Hash)] +struct CacheKey { + uri: Arc, + size: u64, + modified: String, + etag: Option, + version: Option, +} + +#[derive(Clone)] +pub(super) struct BenchmarkCache { + pub(super) segments: Arc, + preloaded: Arc>, + 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> { + let Some(value) = CAPACITY.as_ref() else { + return Ok(None); + }; + let bytes = value + .parse::() + .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(), + )) +} diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs index 53fefacafa4..8adfc8f8b02 100644 --- a/vortex-datafusion/src/persistent/mod.rs +++ b/vortex-datafusion/src/persistent/mod.rs @@ -24,6 +24,7 @@ mod access_plan; mod cache; +mod diagnostics; mod format; pub mod metrics; mod opener; diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 62455fd7b6b..2b7653d8059 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/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!( diff --git a/vortex-duckdb/cpp/include/multi_file_reader.hpp b/vortex-duckdb/cpp/include/multi_file_reader.hpp index 6a62bfbf25c..4e2a753de48 100644 --- a/vortex-duckdb/cpp/include/multi_file_reader.hpp +++ b/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 ffi_global_state; }; @@ -158,6 +166,7 @@ struct VortexBaseReader final : BaseFileReader { } unique_ptr 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 diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp index 7de57ff907d..72837f48588 100644 --- a/vortex-duckdb/cpp/multi_file_reader.cpp +++ b/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 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(reinterpret_cast(file)); + } + reader.prefetched = true; +} + unique_ptr VortexBindData::Copy() const { auto result = make_uniq(); 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(); + auto &global = gstate.global_state->Cast(); const VortexBindData &bind = bind_data.bind_data->Cast(); 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(); + 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(); result->ffi_bind_data = ffi_bind; result->ffi_global_state = unique_ptr(reinterpret_cast(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 OpenReader(const OpenFileInfo &file) { } shared_ptr VortexReaderInterface::CreateReader(ClientContext &, - GlobalTableFunctionState &, + GlobalTableFunctionState &gstate, const OpenFileInfo &file, idx_t, const MultiFileBindData &) { + if (gstate.Cast().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(file, nullptr); + } return OpenReader(file); } diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 53b1e0d020c..27d646aa6cf 100644 --- a/vortex-duckdb/include/vortex.h +++ b/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, diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 01e6cf5a7e2..5f42a2df70f 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/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::().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::().as_ref() }.vortex_expect("null pointer"); + let bytes = unsafe { std::slice::from_raw_parts(path.cast::(), 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::().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, diff --git a/vortex-duckdb/src/file_prefetch.rs b/vortex-duckdb/src/file_prefetch.rs new file mode 100644 index 00000000000..1cb6c293716 --- /dev/null +++ b/vortex-duckdb/src/file_prefetch.rs @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Query-scoped, bounded file preparation ahead of DuckDB's scan workers. + +use std::collections::BTreeMap; +use std::thread::JoinHandle; +use std::time::Instant; + +use futures::FutureExt; +use futures::StreamExt; +use futures::channel::mpsc; +use futures::channel::oneshot; +use futures::future::Either; +use futures::future::select; +use parking_lot::Mutex; +use vortex::error::VortexResult; +use vortex::error::vortex_err; +use vortex::expr::BoundExpression; +use vortex::io::runtime::BlockingRuntime; +use vortex::layout::scan; + +use crate::RUNTIME; +use crate::file_reader::OpenFileReader; +use crate::file_reader::prepare_reader; +use crate::projection::Filter; + +type PreparedFile = VortexResult; + +struct Job { + index: usize, + path: String, + result: oneshot::Sender, +} + +pub(crate) struct FilePrefetch { + window: usize, + jobs: mpsc::UnboundedSender, + results: Mutex>>, + cancel: Option>, + driver: Option>, +} + +impl FilePrefetch { + pub(crate) fn new(projection: BoundExpression, filter: Filter) -> VortexResult> { + let window = std::env::var("VORTEX_DUCKDB_FILE_PREFETCH") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(0) + .min(256); + if window == 0 || !scan::v2::enabled() { + return Ok(None); + } + let (jobs, receiver) = mpsc::unbounded::(); + let (cancel, cancelled) = oneshot::channel(); + // The shared current-thread runtime otherwise stops progressing whenever DuckDB + // is processing a result outside reader_scan's block_on. + let driver = std::thread::Builder::new() + .name("vortex-file-prefetch".into()) + .spawn(move || { + RUNTIME.block_on(async move { + let cancelled = cancelled.shared(); + receiver + .take_until(cancelled.clone()) + .for_each_concurrent(window, |job| { + let projection = projection.clone(); + let filter = filter.clone(); + let cancelled = cancelled.clone(); + async move { + // Shared executor tasks can also run on DuckDB workers inside block_on. + // Await each task after cancellation before joining the driver. + RUNTIME + .handle() + .spawn(async move { + let Job { + index, + path, + result, + } = job; + let work = async move { + let started = Instant::now(); + tracing::debug!( + file_index = index, + "file preparation started" + ); + let prepared = async { + let mut file = OpenFileReader::open(path).await?; + prepare_reader(&mut file, &projection, &filter) + .await?; + Ok(file) + } + .await; + tracing::debug!( + file_index = index, + elapsed_us = started.elapsed().as_micros(), + "file preparation finished" + ); + prepared + }; + let work = prepare_while_needed(result, work); + futures::pin_mut!(work); + drop(select(cancelled, work).await); + }) + .await; + } + }) + .await; + }); + })?; + Ok(Some(Self { + window, + jobs, + results: Mutex::new(BTreeMap::new()), + cancel: Some(cancel), + driver: Some(driver), + })) + } + + pub(crate) fn window(&self) -> usize { + self.window + } + + // C++ submits a sliding window using DuckDB's file ordinals. Completed results + // remain in that same window, so a fast producer cannot prepare the whole dataset. + pub(crate) fn submit(&self, index: usize, path: String) -> VortexResult<()> { + let (result, receiver) = oneshot::channel(); + self.results.lock().insert(index, receiver); + self.jobs + .unbounded_send(Job { + index, + path, + result, + }) + .map_err(|_| vortex_err!("file preparation driver stopped")) + } + + pub(crate) fn take(&self, index: usize, skip: bool) -> VortexResult> { + let result = self + .results + .lock() + .remove(&index) + .ok_or_else(|| vortex_err!("missing prepared file {index}"))?; + if skip { + return Ok(None); + } + let started = Instant::now(); + let file = RUNTIME + .block_on(result) + .map_err(|_| vortex_err!("file preparation driver stopped"))??; + tracing::debug!( + file_index = index, + wait_us = started.elapsed().as_micros(), + "prepared file claimed" + ); + Ok(Some(file)) + } +} + +// Dropping a skipped file's receiver cancels both queued and in-flight preparation. +async fn prepare_while_needed(mut result: oneshot::Sender, work: impl Future) { + futures::pin_mut!(work); + let prepared = match select(result.cancellation(), work).await { + Either::Left(_) => None, + Either::Right((prepared, _)) => Some(prepared), + }; + if let Some(prepared) = prepared { + drop(result.send(prepared)); + } +} + +impl Drop for FilePrefetch { + fn drop(&mut self) { + if let Some(cancel) = self.cancel.take() { + let _ = cancel.send(()); + } + // Join before DuckDB destroys the dynamic-filter state referenced by expressions. + if let Some(driver) = self.driver.take() { + drop(driver.join()); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + use std::task::Context; + use std::task::Poll; + use std::task::Waker; + + use futures::FutureExt; + use futures::channel::oneshot; + use futures::executor::block_on; + use futures::future; + use vortex::error::VortexResult; + use vortex::error::vortex_err; + + use super::prepare_while_needed; + + struct Dropped(Arc); + + impl Drop for Dropped { + fn drop(&mut self) { + self.0.store(true, Ordering::Relaxed); + } + } + + #[test] + fn test_skipped_file_does_not_start_preparation() { + let (result, receiver) = oneshot::channel::<()>(); + drop(receiver); + block_on(prepare_while_needed(result, async { + panic!("a skipped file must not be opened"); + })); + } + + #[test] + fn test_skipped_file_cancels_in_flight_preparation() { + let (result, receiver) = oneshot::channel::<()>(); + let dropped = Arc::new(AtomicBool::new(false)); + let guard = Dropped(Arc::clone(&dropped)); + let mut preparation = Box::pin(prepare_while_needed(result, async move { + let _guard = guard; + future::pending::<()>().await; + })); + let mut cx = Context::from_waker(Waker::noop()); + assert_eq!(preparation.poll_unpin(&mut cx), Poll::Pending); + assert!(!dropped.load(Ordering::Relaxed)); + + drop(receiver); + assert_eq!(preparation.poll_unpin(&mut cx), Poll::Ready(())); + assert!(dropped.load(Ordering::Relaxed)); + } + + #[test] + fn test_preparation_error_reaches_the_scan_worker() -> VortexResult<()> { + let (result, receiver) = oneshot::channel::>(); + block_on(prepare_while_needed( + result, + future::ready(Err(vortex_err!("file preparation failed"))), + )); + let error = block_on(receiver) + .map_err(|error| vortex_err!("preparation result channel closed: {error}"))? + .expect_err("preparation should fail"); + assert!(error.to_string().contains("file preparation failed")); + Ok(()) + } +} diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index f90e4d42742..a7b0569d3e5 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -14,10 +14,13 @@ use vortex::dtype::DType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_panic; +use vortex::expr::BoundExpression; use vortex::file::Footer; +use vortex::file::VortexFile; use vortex::file::multi::MultiFileSession; use vortex::file::multi::open_cached; use vortex::file::multi::parse_uri_or_path; +use vortex::file::planning; use vortex::file::v2::FileStatsLayoutReader; use vortex::io::compat::Compat; use vortex::io::filesystem::FileSystemRef; @@ -25,6 +28,7 @@ use vortex::io::object_store::ObjectStoreFileSystem; use vortex::io::runtime::BlockingRuntime as _; use vortex::io::std_file::StdFileSystem; use vortex::layout::LayoutReaderRef; +use vortex::layout::scan; use vortex::layout::scan::scan_builder::ScanBuilder; use vortex::mask::Mask; use vortex::session::SessionExt as _; @@ -98,8 +102,8 @@ fn resolve_filesystem(url: &Url) -> VortexResult<(FileSystemRef, String)> { } pub struct OpenFileReader { - /// The layout reader and conversion cache of the file. `None` after the file is skipped or - /// all its splits are taken. + /// The opened file, its layout reader and conversion cache. `None` after the file is skipped + /// or all its splits are taken. /// /// DuckDB releases its last reference to a file while it holds the global lock of the /// multi-file scan, and dropping this state takes long. So every local state that scans the @@ -109,35 +113,42 @@ pub struct OpenFileReader { /// File splits stored in inverse order pub splits: Vec, total_splits: usize, + prepared_skip: Option, } /// The state of an open file that the threads scanning it share. pub struct FileScanState { + /// The opened file, which the `VORTEX_SCAN_V2` executor reads through. + file: VortexFile, reader: LayoutReaderRef, cache: ConversionCache, } impl OpenFileReader { - async fn open(path: String) -> VortexResult { + pub(crate) async fn open(path: String) -> VortexResult { let (fs, fs_path) = resolve_filesystem(&parse_uri_or_path(&path)?)?; let source = fs.open_read(&fs_path).await?; let file = open_cached(&SESSION, Some(&path), source, None, &|options| options).await?; Ok(OpenFileReader { scan_state: Some(Arc::new(FileScanState { reader: file.layout_reader()?, + file, cache: ConversionCache::default(), })), splits: vec![], total_splits: 0, + prepared_skip: None, }) } - fn reader(&self) -> &LayoutReaderRef { - &self - .scan_state + fn scan_state(&self) -> &FileScanState { + self.scan_state .as_ref() .vortex_expect("file used after its scan state was released") - .reader + } + + fn reader(&self) -> &LayoutReaderRef { + &self.scan_state().reader } fn can_skip(&self, filter: &Filter) -> VortexResult { @@ -185,31 +196,49 @@ pub fn reader_bind(file: &OpenFileReader, result: &mut BindResultRef) -> VortexR } /// Called once per file by one thread under file-local lock. Determines -/// whether the opened file should be skipped. If this function returns false, +/// whether the opened file should be skipped. If this function returns true, /// duckdb closes the file and doesn't call reader_try_initialize_scan on it. pub fn reader_initialize(file: &mut OpenFileReader, global: &GlobalState) -> VortexResult { - if file.can_skip(&global.filter)? { - // DuckDB releases a skipped file under the global lock. Here only the - // file-local lock is held, so drop the scan state now. + if let Some(skip) = file.prepared_skip { + return Ok(skip); + } + RUNTIME.block_on(prepare_reader(file, &global.projection, &global.filter)) +} + +pub(crate) async fn prepare_reader( + file: &mut OpenFileReader, + projection: &BoundExpression, + filter: &Filter, +) -> VortexResult { + if file.can_skip(filter)? { + // DuckDB releases a skipped file under the global lock. Here the global + // lock is not held, so drop the scan state now. file.scan_state = None; + file.prepared_skip = Some(true); return Ok(true); } // Getting splits is non-trivial work so we prefer doing it here under file // lock and not in reader_try_initialize_scan under global lock. let reader = Arc::clone(file.reader()); - let filter = &global.filter; let builder = ScanBuilder::new(SESSION.clone(), reader) - .with_projection(global.projection.clone()) + .with_projection(projection.clone()) .with_some_filter(filter.filter.clone()) .with_selection(filter.row_selection.clone()); - let scan = builder.prepare()?; - let mut splits = scan.execute(filter.row_range.clone())?; + let scan_file = planning::scan_file(&file.scan_state().file); + let mut splits = if scan::v2::enabled() && scan::v2::supports_scan(&builder, &scan_file) { + scan::v2::prepare(builder, scan_file)? + .execute_pruned(filter.row_range.clone()) + .await? + } else { + builder.prepare()?.execute(filter.row_range.clone())? + }; // threads take last element of file.splits so we need to reverse splits.reverse(); file.total_splits = splits.len(); file.splits = splits; + file.prepared_skip = Some(false); Ok(false) } diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index 1209213135f..e1dbe11bd9b 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -28,6 +28,7 @@ mod convert; pub mod duckdb; mod exporter; mod ffi; +mod file_prefetch; mod file_reader; mod projection; mod table_function; diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index adecd2dbbdf..71aefb64ceb 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -187,6 +187,7 @@ impl Projection { } } +#[derive(Clone)] pub struct Filter { pub filter: Option, pub row_selection: Selection, diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 197057b80f3..cb86b0d5097 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -52,6 +52,7 @@ use crate::duckdb::TableFilterSetRef; use crate::duckdb::TableInitInput; use crate::duckdb::Value; use crate::exporter::ArrayExporter; +use crate::file_prefetch::FilePrefetch; use crate::file_reader::FileScanState; use crate::projection::DuckdbField; use crate::projection::Filter; @@ -139,6 +140,7 @@ impl<'a> TableInitInput<'a> { } pub struct GlobalState { + pub file_prefetch: Option, pub projection: BoundExpression, pub filter: Filter, pub file_row_number_column_pos: Option, @@ -299,6 +301,7 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult { let projection = optimize_and_bind(projection, &bind_data.dtype)?; Ok(GlobalState { + file_prefetch: FilePrefetch::new(projection.clone(), filter.clone())?, projection, filter, aggregates: bind_data.aggregates.clone(), diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index fefe824ecf4..a7b3e6892a0 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -22,6 +22,7 @@ use vortex_layout::LayoutReader; use vortex_layout::scan::layout::LayoutReaderDataSource; use vortex_layout::scan::scan_builder::ScanBuilder; use vortex_layout::scan::split_by::SplitBy; +use vortex_layout::scan::v2::ScanFile; use vortex_layout::segments::SegmentSource; use vortex_scan::DataSourceRef; use vortex_session::VortexSession; @@ -49,6 +50,8 @@ pub struct VortexFile { metadata: Arc>, /// None id LayoutReader caching is turned off layout_reader_cache: Option>>, + /// The V2 equivalent of the cached reader tree, enabled by the same option. + pub(crate) scan_file_cache: Option>, /// Serves the planning protocol's reads, when the file was opened over a reader. scan_io: Option>, } @@ -87,12 +90,16 @@ impl VortexFile { session, metadata: Arc::new(HashMap::new()), layout_reader_cache: None, + scan_file_cache: None, scan_io: None, } } pub(crate) fn with_scan_io(mut self, scan_io: Arc) -> Self { self.scan_io = Some(scan_io); + if let Some(cache) = &mut self.scan_file_cache { + cache.take(); + } self } @@ -111,6 +118,7 @@ impl VortexFile { /// /// Repeated calls to [`layout_reader`](Self::layout_reader), [`scan`](Self::scan), and /// [`data_source`](Self::data_source) will share the same reader tree. + /// V2 scans retain the equivalent lowered layout plans. pub fn with_caching(self) -> Self { Self { footer: self.footer, @@ -118,6 +126,7 @@ impl VortexFile { session: self.session, metadata: self.metadata, layout_reader_cache: Some(OnceLock::new()), + scan_file_cache: Some(OnceLock::new()), scan_io: self.scan_io, } } @@ -170,7 +179,7 @@ impl VortexFile { /// Replace the segment source used by this file. /// - /// Any cached layout reader is cleared so that subsequent scans construct readers over the + /// Any cached layout reader and V2 plans are cleared so that subsequent scans read from the /// replacement source. pub fn with_segment_source(mut self, segment_source: Arc) -> Self { self.segment_source = segment_source; @@ -179,6 +188,9 @@ impl VortexFile { if self.layout_reader_cache.is_some() { self.layout_reader_cache = Some(OnceLock::new()); } + if let Some(cache) = &mut self.scan_file_cache { + cache.take(); + } self } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 707de8bb47e..25926b78e85 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -109,6 +109,7 @@ pub mod flatbuffers; mod footer; pub mod multi; mod open; +pub mod planning; mod pruning; mod read; /// Segment sources, caches, and sinks used by file readers and writers. diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index aa1f0fff3c5..a0a0c1dc0ce 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -541,9 +541,14 @@ mod tests { use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; + use futures::TryStreamExt; use futures::future::BoxFuture; use parking_lot::Mutex; + use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ChunkedArray; + use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; use vortex_array::memory::BufferAllocatorRef; use vortex_array::memory::MemorySessionExt; @@ -553,6 +558,9 @@ mod tests { use vortex_buffer::ByteBufferMut; use vortex_error::vortex_bail; use vortex_io::session::RuntimeSession; + use vortex_layout::scan::v2; + use vortex_layout::segments::MokaSegmentCache; + use vortex_layout::segments::SegmentEviction; use vortex_layout::session::LayoutSession; use vortex_session::registry::Id; use vortex_session::registry::ReadContext; @@ -560,6 +568,61 @@ mod tests { use super::*; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; + use crate::planning::scan_file; + + #[tokio::test] + async fn v2_reuses_configured_segment_cache_across_opens() -> VortexResult<()> { + let session = test_session().with_tokio(); + let expected = Buffer::from_iter(0..65_536_i32).into_array(); + let mut bytes = ByteBufferMut::empty(); + session + .write_options() + .write(&mut bytes, expected.clone().to_array_stream()) + .await?; + let bytes = ByteBuffer::from(bytes); + let footer = session + .open_options() + .open_buffer(bytes.clone())? + .footer() + .clone(); + let total_read = Arc::new(AtomicUsize::new(0)); + let reader = CountingRead { + inner: bytes.clone(), + total_read: Arc::clone(&total_read), + first_read_len: Arc::default(), + reads: Arc::default(), + }; + let cache: Arc = Arc::new( + MokaSegmentCache::new(bytes.len() as u64, SegmentEviction::TinyLfu) + .for_file("test-file"), + ); + let mut after_warmup = 0; + for iteration in 0..2 { + let file = session + .open_options() + .with_footer(footer.clone()) + .with_segment_cache(Arc::clone(&cache)) + .open_read(reader.clone()) + .await?; + assert!(file.scan_io().is_none()); + let chunks: Vec = v2::into_stream(file.scan()?, scan_file(&file))? + .try_collect() + .await?; + assert_arrays_eq!( + ChunkedArray::try_new(chunks, expected.dtype().clone())?, + expected, + &mut session.create_execution_ctx() + ); + let read = total_read.load(Ordering::Relaxed); + if iteration == 0 { + assert!(read > 0); + after_warmup = read; + } else { + assert_eq!(read, after_warmup, "warm scan must not read segment bytes"); + } + } + Ok(()) + } fn test_session() -> VortexSession { let session = vortex_array::array_session() diff --git a/vortex-file/src/planning/footer_open/mod.rs b/vortex-file/src/planning/footer_open/mod.rs new file mode 100644 index 00000000000..378aa6d8b61 --- /dev/null +++ b/vortex-file/src/planning/footer_open/mod.rs @@ -0,0 +1,220 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Footer open: the only stage in this slice whose IO goes through the protocol. + +use std::sync::Arc; + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::VortexReadAt; +use vortex_io::request::IoBatch; +use vortex_io::request::IoConsumer; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoSlot; +use vortex_io::request::IoTarget; +use vortex_scan::planning::next::Next; +use vortex_scan::planning::planner::Planner; +use vortex_scan::planning::planner::PlannerOutput; +use vortex_scan::planning::planner::State; +use vortex_scan::planning::planner::WorkScope; +use vortex_session::VortexSession; + +use crate::DeserializeStep; +use crate::EOF_SIZE; +use crate::Footer; +use crate::FooterDeserializer; +use crate::MAX_POSTSCRIPT_SIZE; +use crate::planning::FileSource; +use crate::planning::OpenedFile; + +/// The opener's default tail read: enough to hold the largest postscript and the EOF marker. +pub const DEFAULT_INITIAL_READ_SIZE: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE; + +/// Opens a file: discovers its size if unknown, reads and parses the footer unless one was +/// supplied, validates the footer against the size, and hands an [`OpenedFile`] to `next`. +/// +/// One request is outstanding at a time, published by the compute that discovers it. A delivery +/// the stage did not ask for is a driver bug and panics. +pub struct FooterOpen { + read: Arc, + initial_read_size: usize, + session: VortexSession, + next: Next, + phase: Phase, + io: IoSlot, +} + +enum Phase { + /// Nothing is requested yet: the next compute emits the child when the footer and size are + /// both known, and asks for whichever is missing otherwise. + Start { + size: Option, + footer: Option