diff --git a/Cargo.lock b/Cargo.lock index 9cd5d12d5bca4..17fcfa9495a36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -165,8 +165,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b952ca5a8046ad741b60f142d6eca4aeebcad615694202bc64c5341f23e32c5b" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-arith", "arrow-array", @@ -188,8 +187,7 @@ dependencies = [ [[package]] name = "arrow-arith" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64a13b8d3008c4e9063c597a08f46446fe3fd5789277127672d6c0bdbb43b1ff" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -202,8 +200,7 @@ dependencies = [ [[package]] name = "arrow-array" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9486151b2f0785bafc6fa04fc5c99fcb4495455662e58787ea32eaaed33c4192" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "ahash", "arrow-buffer", @@ -221,8 +218,7 @@ dependencies = [ [[package]] name = "arrow-avro" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e4f9b23a0d7b613acb59fa20bdbe0f80ffdae6411498378340b3915e45f5b84" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -245,8 +241,7 @@ dependencies = [ [[package]] name = "arrow-buffer" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4776577a87794bfdf0b4e90e2ea12454fa7738ea2823c4be5b9d1851da7b434" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "bytes", "half", @@ -257,8 +252,7 @@ dependencies = [ [[package]] name = "arrow-cast" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9ad451ce4f98710828a455b96991b8f031deb2e67f5fcad6773f017e4a69c3a" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -279,8 +273,7 @@ dependencies = [ [[package]] name = "arrow-csv" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8aa7bf96d6141a7bcca2eed57c7c9767d2a2175281857b8a7b68308992864784" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-cast", @@ -294,8 +287,7 @@ dependencies = [ [[package]] name = "arrow-data" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b38fe43e2e8704360f1464e6e8cc4fc381ef02cc4fb0192afa8df1aaa0115c66" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-buffer", "arrow-schema", @@ -307,8 +299,7 @@ dependencies = [ [[package]] name = "arrow-flight" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42115e09dbb694b5955da998912121451c6910b338228cb80a5701370dba43ff" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-arith", "arrow-array", @@ -335,8 +326,7 @@ dependencies = [ [[package]] name = "arrow-ipc" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29dac499fcbc6ba74ee0324057821d381929a48526a3966bd9dffb44aa06d98c" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -351,8 +341,7 @@ dependencies = [ [[package]] name = "arrow-json" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fe05e916ddc50f4c7a363cd69c0ef5894fcee063517e9a0b8582f0c56746af6" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -376,8 +365,7 @@ dependencies = [ [[package]] name = "arrow-ord" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e13dbdc2a9c053c10c7baa6e30faee04a180aa7ce88e471835850ce37abd20b" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -389,8 +377,7 @@ dependencies = [ [[package]] name = "arrow-row" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d5a1f8c733d15260b305683472ee8ad89c62cbd706703ca873b90d051b41592" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -402,8 +389,7 @@ dependencies = [ [[package]] name = "arrow-schema" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9e4969dc350d571766247143ab36a5187d095d3d3690970408bc630d47c69e5" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "bitflags", "serde", @@ -414,8 +400,7 @@ dependencies = [ [[package]] name = "arrow-select" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402770dba90865359d98d1ef92ef16e23d75c0cca9c2c880c8a05468b7743bf9" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "ahash", "arrow-array", @@ -428,8 +413,7 @@ dependencies = [ [[package]] name = "arrow-string" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2b0afbb8b9016700938291123df30838b89decc3213dba00852021988b170d3" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "arrow-array", "arrow-buffer", @@ -4064,9 +4048,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.0" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" dependencies = [ "twox-hash", ] @@ -4475,8 +4459,7 @@ dependencies = [ [[package]] name = "parquet" version = "59.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5302d4da74d6596a1f11f9928767995b53bca657cbeea1e4e8c5074f8a1157dd" +source = "git+https://github.com/pydantic/arrow-rs.git?branch=claude%2Fpush-decoder-peek-59#fcd1bc55963cc1ac8fd70a76ca2fec88e7df4385" dependencies = [ "ahash", "arrow-array", diff --git a/Cargo.toml b/Cargo.toml index 87c23cc456651..f4f3cedd4f5a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -295,3 +295,26 @@ debug = false debug-assertions = false strip = "debuginfo" incremental = false + +# EXPERIMENT: pinned to the arrow-rs branch adding +# `parquet::arrow::push_decoder::plan_scan_ranges` (pydantic/arrow-rs, +# claude/push-decoder-peek-59 = apache/arrow-rs 59.1.0 + that one module). +# Every arrow crate must be patched together or two arrow versions collide. +[patch.crates-io] +arrow-arith = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-array = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-avro = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-buffer = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-cast = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-csv = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-data = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-flight = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-ipc = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-json = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-ord = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-row = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-schema = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-select = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow-string = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +arrow = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } +parquet = { git = "https://github.com/pydantic/arrow-rs.git", branch = "claude/push-decoder-peek-59" } diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 25b79a618830c..9ae77b35cc7cf 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -34,6 +34,7 @@ mod opener; mod page_filter; mod projection_read_plan; mod push_decoder; +pub use push_decoder::PEAK_STAGED_BYTES; mod reader; mod row_filter; mod row_group_filter; diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index af97a192fa7ce..a29ce6b010d6f 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -27,7 +27,8 @@ use crate::access_plan::PreparedAccessPlan; use crate::decoder_projection::DecoderProjection; use crate::page_filter::PagePruningAccessPlanFilter; use crate::push_decoder::{ - DecoderBuilderConfig, PushDecoderStreamState, RgPlanEntry, RowGroupPruner, + DecoderBuilderConfig, FetchPolicy, PushDecoderStreamState, ReaderSlot, RgPlanEntry, + RowGroupPruner, }; use crate::row_filter::RowFilterGenerator; use crate::row_group_filter::RowGroupAccessPlanFilter; @@ -550,10 +551,19 @@ impl ParquetOpenState { } ParquetOpenState::PruneWithStatistics(prepared) => { let prepared_row_groups = (*prepared).prune_row_groups()?; - if should_load_page_index( - prepared_row_groups.prepared.page_pruning_predicate.as_ref(), - &prepared_row_groups.row_groups, - ) { + // EXPERIMENT: the streaming policy plans at page granularity, + // which requires the offset index. Page-index loading is + // otherwise driven purely by whether *pruning* can use it, so + // without this the streaming path silently falls back to the + // push decoder on every scan that has no prunable predicate. + let streaming_needs_page_index = + matches!(FetchPolicy::from_env(), FetchPolicy::Streaming { .. }); + if streaming_needs_page_index + || should_load_page_index( + prepared_row_groups.prepared.page_pruning_predicate.as_ref(), + &prepared_row_groups.row_groups, + ) + { Ok(ParquetOpenState::LoadPageIndex( prepared_row_groups.load_page_index().boxed(), )) @@ -1386,6 +1396,79 @@ impl RowGroupsPrunedParquetOpen { prepared.virtual_state.as_deref(), )?; + // EXPERIMENT: streaming (batch-granular) scan path, selected via + // DF_FETCH_POLICY=streaming. A long-lived sync + // `ParquetRecordBatchReader` pulls from a shared buffer that the + // stream driver fills with exactly the page ranges each batch needs + // (plus bounded readahead). Falls back to the push-decoder path when + // row filters are active or the offset index is unavailable. + // + // Note this does not drive `ParquetPushDecoder`: only the *plan* — + // which pages this scan reads, in decode order — comes from arrow-rs + // (`plan_scan_ranges`, which merely lives in that crate's + // `push_decoder` module). The decoder itself cannot be used here + // because `NeedsData` resolves only at row-group granularity, so it + // cannot say what the next *batch* needs. Teaching it to would let + // this path drop `SharedBuffers` and the sync reader entirely and go + // back to being a pure scheduler; see apache/arrow-rs#10555. + if let FetchPolicy::Streaming { window } = FetchPolicy::from_env() { + let pushdown_active = + prepared.pushdown_filters && prepared.predicate.is_some(); + if !pushdown_active { + let streaming_access_plan = prepare_access_plan(access_plan.clone())?; + if let Some(streaming_plan) = crate::push_decoder::build_streaming_plan( + &file_metadata, + &streaming_access_plan.row_group_indexes, + decoder_projection.projection_mask(), + streaming_access_plan.row_selection.as_ref(), + ) { + let stream = crate::push_decoder::build_streaming_stream( + streaming_plan, + crate::push_decoder::StreamingScanConfig { + reader_metadata: reader_metadata.clone(), + row_group_indexes: streaming_access_plan.row_group_indexes, + row_selection: streaming_access_plan.row_selection, + decoder_projection, + batch_size: prepared.batch_size, + limit: prepared.limit, + reader: prepared.async_file_reader, + baseline_metrics: prepared.baseline_metrics, + window, + }, + )?; + let files_ranges_pruned_statistics = + prepared.file_metrics.files_ranges_pruned_statistics.clone(); + return match prepared.file_pruner { + Some(file_pruner) if file_pruner.is_watching() => { + Ok(EarlyStoppingStream::new( + stream, + file_pruner, + files_ranges_pruned_statistics, + ) + .boxed()) + } + _ => Ok(stream), + }; + } else { + // The streaming path needs page locations. Falling back + // silently makes "did it even run?" unanswerable without + // a benchmark round-trip, which is exactly how the + // no-page-index ClickBench files went unnoticed. + debug!( + "streaming fetch policy requested but {} has no offset index; \ + falling back to row-group-granular push decoding", + prepared.file_name + ); + } + } else { + debug!( + "streaming fetch policy requested but pushdown filters are \ + active for {}; falling back to row-group-granular push decoding", + prepared.file_name + ); + } + } + let (decoder, rg_plan) = { let pushdown_predicate = prepared .pushdown_filters @@ -1481,7 +1564,10 @@ impl RowGroupsPrunedParquetOpen { decoder: Some(decoder), active_reader: None, rg_plan, - reader: prepared.async_file_reader, + reader: ReaderSlot::Idle(prepared.async_file_reader), + fetch_policy: FetchPolicy::from_env(), + parquet_metadata: Arc::clone(reader_metadata.metadata()), + prefetched_row_groups: std::collections::HashSet::new(), decoder_projection, arrow_reader_metrics, predicate_cache_inner_records, diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 31bd365a4631d..36fc8eeff4af9 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -35,23 +35,31 @@ //! The opener constructs both halves and hands the state off to //! [`PushDecoderStreamState::into_stream`] for consumption. -use std::collections::VecDeque; +use std::collections::{HashSet, VecDeque}; +use std::ops::Range; use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::SchemaRef; +use bytes::Bytes; use futures::StreamExt; use futures::stream::BoxStream; use log::debug; use parquet::DecodeResult; use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::arrow_reader::RowSelection; use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; use parquet::arrow::arrow_reader::{ ArrowReaderMetadata, ParquetRecordBatchReader, RowSelectionPolicy, }; use parquet::arrow::async_reader::AsyncFileReader; -use parquet::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder}; +use parquet::arrow::push_decoder::{ + ParquetPushDecoder, ParquetPushDecoderBuilder, PlannedRange, plan_scan_ranges, +}; +use parquet::errors::ParquetError; use parquet::file::metadata::ParquetMetaData; +use parquet::file::reader::{ChunkReader, Length}; use datafusion_common::{DataFusionError, Result}; use datafusion_physical_expr::expressions::DynamicFilterTracking; @@ -111,6 +119,123 @@ pub(crate) struct RgPlanEntry { pub(crate) rg_index: usize, } +/// EXPERIMENT: peak bytes staged in the push decoder (buffered but not yet +/// handed to a reader), max over all streams since last reset. Benchmarks +/// reset and read this between runs. +pub static PEAK_STAGED_BYTES: std::sync::atomic::AtomicU64 = + std::sync::atomic::AtomicU64::new(0); + +fn note_staged_bytes(decoder: &ParquetPushDecoder) { + PEAK_STAGED_BYTES.fetch_max( + decoder.buffered_bytes(), + std::sync::atomic::Ordering::Relaxed, + ); +} + +/// EXPERIMENT: how the stream schedules I/O relative to decode. +/// +/// - `Off`: current main behavior — fetch exactly what the decoder asks +/// for, when it asks for it. I/O and decode strictly alternate. +/// - `Batched`: PR #23492 behavior — when the decoder asks for the current +/// row group's ranges, append the complete projected ranges of upcoming +/// row groups to the *same* blocking fetch, as long as +/// `buffered + staged <= budget`. Fewer round trips, but no overlap. +/// - `Pipelined`: when a row group's reader is handed over for decode, +/// spawn a *background* fetch for upcoming row groups' projected ranges +/// within the same byte budget. Decode of the current RG overlaps with +/// I/O for the next ones. +/// - `Streaming`: batch-granular readiness. One long-lived *sync* +/// [`ParquetRecordBatchReader`] pulls bytes through a shared in-memory +/// buffer; the stream driver computes, from the offset index, exactly +/// which page ranges the next batch needs, awaits their fetch, and keeps +/// a background readahead of up to `window` bytes in flight. Falls back +/// to `Off` when preconditions don't hold (no offset index, or row +/// filters are active). See [`StreamingScanState`]. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum FetchPolicy { + Off, + Batched { budget: u64 }, + Pipelined { budget: u64 }, + Streaming { window: u64 }, +} + +impl FetchPolicy { + pub(crate) fn from_env() -> Self { + let budget = std::env::var("DF_FETCH_BUDGET") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(20 * 1024 * 1024); + match std::env::var("DF_FETCH_POLICY").as_deref() { + Ok("batched") => FetchPolicy::Batched { budget }, + Ok("pipelined") => FetchPolicy::Pipelined { budget }, + Ok("streaming") => FetchPolicy::Streaming { window: budget }, + _ => FetchPolicy::Off, + } + } +} + +/// Payload returned by a background prefetch task: the lent reader, the +/// ranges it fetched, and the fetch result. +type PrefetchResult = ( + Box, + Vec>, + parquet::errors::Result>, +); + +/// The file reader is either available inline or lent out to a background +/// prefetch task (only under [`FetchPolicy::Pipelined`]). +pub(crate) enum ReaderSlot { + Idle(Box), + Busy(tokio::task::JoinHandle), + /// Transient state while ownership moves between the two above. + Empty, +} + +/// Compute the complete projected byte ranges of upcoming row groups that +/// fit in `budget` given `staged_bytes` already accounted for. Mirrors the +/// logic in PR #23492. `entries` must already exclude the row group +/// currently being fetched/decoded. +fn upcoming_row_group_ranges<'a>( + entries: impl Iterator, + projection: &ProjectionMask, + metadata: &ParquetMetaData, + prefetched_row_groups: &mut HashSet, + mut staged_bytes: u64, + budget: u64, +) -> Vec> { + let mut ranges = Vec::new(); + if staged_bytes >= budget { + return ranges; + } + for entry in entries { + if prefetched_row_groups.contains(&entry.rg_index) { + continue; + } + let row_group = metadata.row_group(entry.rg_index); + let row_group_ranges = row_group + .columns() + .iter() + .enumerate() + .filter(|(column_idx, _)| projection.leaf_included(*column_idx)) + .map(|(_, column)| { + let (start, len) = column.byte_range(); + start..start + len + }) + .collect::>(); + let row_group_bytes = row_group_ranges + .iter() + .map(|range| range.end - range.start) + .sum::(); + if staged_bytes.saturating_add(row_group_bytes) > budget { + break; + } + ranges.extend(row_group_ranges); + staged_bytes += row_group_bytes; + prefetched_row_groups.insert(entry.rg_index); + } + ranges +} + /// Runtime row-group pruner driven by a dynamic predicate (e.g. the /// threshold expression a `TopK` operator pushes down). /// @@ -243,7 +368,13 @@ pub(crate) struct PushDecoderStreamState { pub(crate) decoder: Option, pub(crate) active_reader: Option, pub(crate) rg_plan: VecDeque, - pub(crate) reader: Box, + pub(crate) reader: ReaderSlot, + /// EXPERIMENT: fetch scheduling policy (see [`FetchPolicy`]). + pub(crate) fetch_policy: FetchPolicy, + /// Parquet metadata used to compute projected ranges of upcoming RGs. + pub(crate) parquet_metadata: Arc, + /// Row groups whose projected ranges were already staged speculatively. + pub(crate) prefetched_row_groups: HashSet, /// Per-file projection: the mask installed on every decoder and the /// per-batch transform applied by [`Self::project_batch`]. pub(crate) decoder_projection: DecoderProjection, @@ -372,22 +503,77 @@ impl PushDecoderStreamState { // Step 3: drive the decoder. let decoder = self.decoder.as_mut().expect("decoder present"); match decoder.try_next_reader() { - Ok(DecodeResult::NeedsData(ranges)) => { - let data = self - .reader + Ok(DecodeResult::NeedsData(mut ranges)) => { + // EXPERIMENT: if a background prefetch is in flight, + // land it first — it may cover (part of) the request. + if matches!(self.reader, ReaderSlot::Busy(_)) { + let ReaderSlot::Busy(handle) = + std::mem::replace(&mut self.reader, ReaderSlot::Empty) + else { + unreachable!() + }; + let (reader, fetched_ranges, result) = match handle.await { + Ok(v) => v, + Err(e) => { + return Some(( + Err(DataFusionError::External(Box::new(e))), + self, + )); + } + }; + self.reader = ReaderSlot::Idle(reader); + match result { + Ok(data) => { + let decoder = + self.decoder.as_mut().expect("decoder present"); + if let Err(e) = decoder.push_ranges(fetched_ranges, data) + { + return Some((Err(DataFusionError::from(e)), self)); + } + note_staged_bytes(decoder); + } + Err(e) => { + return Some((Err(DataFusionError::from(e)), self)); + } + } + // Re-poll the decoder: the prefetched data may have + // satisfied this request entirely. + continue; + } + + // EXPERIMENT: batched policy (PR #23492) — extend the + // blocking fetch with upcoming row groups' ranges. + if let FetchPolicy::Batched { budget } = self.fetch_policy { + let buffered = self + .decoder + .as_ref() + .expect("decoder present") + .buffered_bytes(); + let required: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + ranges.extend(upcoming_row_group_ranges( + self.rg_plan.iter().skip(1), + self.decoder_projection.projection_mask(), + &self.parquet_metadata, + &mut self.prefetched_row_groups, + buffered.saturating_add(required), + budget, + )); + } + + let ReaderSlot::Idle(reader) = &mut self.reader else { + unreachable!("reader is idle here") + }; + let data = reader .get_byte_ranges(ranges.clone()) .await .map_err(DataFusionError::from); match data { Ok(data) => { - if let Err(e) = self - .decoder - .as_mut() - .expect("decoder present") - .push_ranges(ranges, data) - { + let decoder = self.decoder.as_mut().expect("decoder present"); + if let Err(e) = decoder.push_ranges(ranges, data) { return Some((Err(DataFusionError::from(e)), self)); } + note_staged_bytes(decoder); } Err(e) => return Some((Err(e), self)), } @@ -398,6 +584,74 @@ impl PushDecoderStreamState { // the decoder is about to read). self.rg_plan.pop_front(); self.active_reader = Some(reader); + + // EXPERIMENT: pipelined policy — while this RG decodes, + // fetch upcoming RGs' projected ranges in the background + // within the byte budget. + if let FetchPolicy::Pipelined { budget } = self.fetch_policy + && matches!(self.reader, ReaderSlot::Idle(_)) + { + let buffered = self + .decoder + .as_ref() + .expect("decoder present") + .buffered_bytes(); + // Hysteresis: without this, once the budget is + // mostly full of staged data every spawn only has + // ~one RG of headroom and the fetch degrades to + // one round trip per RG. Waiting until at least + // half the budget is free keeps gulps large. The + // final row groups are always eligible so the + // tail doesn't stall. + let headroom = budget.saturating_sub(buffered); + let remaining_bytes: u64 = self + .rg_plan + .iter() + .filter(|e| !self.prefetched_row_groups.contains(&e.rg_index)) + .map(|e| { + let rg = self.parquet_metadata.row_group(e.rg_index); + rg.columns() + .iter() + .enumerate() + .filter(|(i, _)| { + self.decoder_projection + .projection_mask() + .leaf_included(*i) + }) + .map(|(_, c)| c.byte_range().1) + .sum::() + }) + .sum(); + if headroom < (budget / 2).min(remaining_bytes) { + continue; + } + let prefetch_ranges = upcoming_row_group_ranges( + self.rg_plan.iter(), + self.decoder_projection.projection_mask(), + &self.parquet_metadata, + &mut self.prefetched_row_groups, + buffered, + budget, + ); + if !prefetch_ranges.is_empty() { + let ReaderSlot::Idle(mut reader) = + std::mem::replace(&mut self.reader, ReaderSlot::Empty) + else { + unreachable!() + }; + // POC: a raw tokio JoinHandle so the prefetch + // detaches (runs to completion) if the stream + // is dropped mid-fetch; a mergeable version + // would use `SpawnedTask` for cancel-safety. + #[expect(clippy::disallowed_methods)] + let handle = tokio::task::spawn(async move { + let result = + reader.get_byte_ranges(prefetch_ranges.clone()).await; + (reader, prefetch_ranges, result) + }); + self.reader = ReaderSlot::Busy(handle); + } + } } Ok(DecodeResult::Finished) => return None, Err(e) => { @@ -423,6 +677,489 @@ impl PushDecoderStreamState { } } +// =========================================================================== +// EXPERIMENT: `FetchPolicy::Streaming` — batch-granular readiness. +// +// One long-lived *sync* `ParquetRecordBatchReader` pulls bytes through +// [`SharedBuffers`] (an in-memory `ChunkReader`). The stream driver computes +// from the offset index exactly which page ranges the next batch needs, +// awaits their fetch (with up to `window` bytes of background readahead), +// then calls `next()` — which therefore never blocks on I/O. Dictionary +// pages are fetched and decoded once per row group (the reader persists), +// and page bytes are dropped as soon as the decode cursor passes them, so +// resident memory is bounded by the readahead window rather than row-group +// size. +// =========================================================================== + +/// In-memory byte store shared between the fetch side (inserts ranges as +/// they land) and the sync parquet reader (reads through `ChunkReader`). +/// Reads must be fully contained in a previously inserted range; the stream +/// driver guarantees this by construction, so a miss is a bug, not a wait. +#[derive(Clone)] +pub(crate) struct SharedBuffers { + inner: Arc>>, + file_len: u64, +} + +impl SharedBuffers { + fn new(file_len: u64) -> Self { + Self { + inner: Arc::new(std::sync::Mutex::new(Default::default())), + file_len, + } + } + + fn insert(&self, range: &Range, data: Bytes) { + self.inner.lock().unwrap().insert(range.start, data); + } + + fn remove(&self, start: u64) { + self.inner.lock().unwrap().remove(&start); + } + + /// Whether the page starting at `start` is already staged. + fn contains(&self, start: u64) -> bool { + self.inner.lock().unwrap().contains_key(&start) + } +} + +impl Length for SharedBuffers { + fn len(&self) -> u64 { + self.file_len + } +} + +pub(crate) struct SharedBuffersRead { + buffers: SharedBuffers, + pos: u64, +} + +impl std::io::Read for SharedBuffersRead { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let n = buf.len().min((self.buffers.file_len - self.pos) as usize); + if n == 0 { + return Ok(0); + } + let bytes = self + .buffers + .get_bytes(self.pos, n) + .map_err(std::io::Error::other)?; + buf[..n].copy_from_slice(&bytes); + self.pos += n as u64; + Ok(n) + } +} + +impl ChunkReader for SharedBuffers { + type T = SharedBuffersRead; + + fn get_read(&self, start: u64) -> parquet::errors::Result { + Ok(SharedBuffersRead { + buffers: self.clone(), + pos: start, + }) + } + + fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result { + let guard = self.inner.lock().unwrap(); + if let Some((&rstart, bytes)) = guard.range(..=start).next_back() { + let offset = start - rstart; + if offset as usize + length <= bytes.len() { + return Ok(bytes.slice(offset as usize..offset as usize + length)); + } + } + Err(ParquetError::General(format!( + "streaming scan buffer miss: {start}..{} not resident", + start + length as u64 + ))) + } +} + +/// A planned page plus the stream's eviction bookkeeping. +/// +/// The page itself (byte range + the span of selected rows it serves) comes +/// from arrow-rs's [`plan_scan_ranges`]; the only thing DataFusion adds is +/// whether the decode cursor has passed it and its bytes were released. +struct PlanPage { + planned: PlannedRange, + cleared: bool, +} + +impl PlanPage { + fn range(&self) -> &Range { + &self.planned.range + } +} + +/// Opaque prebuilt streaming fetch plan. +pub(crate) struct StreamingPlan { + pages: Vec, + total_selected: u64, + file_end: u64, +} + +/// Build the streaming fetch plan by asking arrow-rs which pages this scan +/// will read, in the order decoding needs them. +/// +/// Returns `None` when arrow-rs cannot plan at page granularity (no offset +/// index) — the caller falls back to the push-decoder path. Borrows only, so +/// callers can probe feasibility before committing resources. +pub(crate) fn build_streaming_plan( + metadata: &ParquetMetaData, + row_group_indexes: &[usize], + projection: &ProjectionMask, + selection: Option<&RowSelection>, +) -> Option { + let plan = plan_scan_ranges(metadata, row_group_indexes, projection, selection)?; + // `SharedBuffers` reports a file length to the sync reader; the end of the + // last projected column chunk is an upper bound on anything it will read. + let file_end = row_group_indexes + .iter() + .flat_map(|&rg| { + metadata.row_group(rg).columns().iter().map(|c| { + let (start, len) = c.byte_range(); + start + len + }) + }) + .max() + .unwrap_or(0); + Some(StreamingPlan { + total_selected: plan.total_selected_rows, + pages: plan + .ranges + .into_iter() + .map(|planned| PlanPage { + planned, + cleared: false, + }) + .collect(), + file_end, + }) +} + +pub(crate) struct StreamingScanConfig { + pub reader_metadata: ArrowReaderMetadata, + pub row_group_indexes: Vec, + pub row_selection: Option, + pub decoder_projection: DecoderProjection, + pub batch_size: usize, + pub limit: Option, + pub reader: Box, + pub baseline_metrics: BaselineMetrics, + pub window: u64, +} + +/// Build the streaming (batch-granular) scan stream from a prebuilt plan. +pub(crate) fn build_streaming_stream( + plan: StreamingPlan, + config: StreamingScanConfig, +) -> Result>> { + let StreamingScanConfig { + reader_metadata, + row_group_indexes, + row_selection, + decoder_projection, + batch_size, + limit, + reader, + baseline_metrics, + window, + } = config; + let StreamingPlan { + pages: plan, + total_selected, + file_end, + } = plan; + + let buffers = SharedBuffers::new(file_end); + let mut builder = ParquetRecordBatchReaderBuilder::new_with_metadata( + buffers.clone(), + reader_metadata, + ) + .with_projection(decoder_projection.projection_mask().clone()) + .with_batch_size(batch_size) + .with_row_groups(row_group_indexes); + if let Some(selection) = row_selection { + builder = builder.with_row_selection(selection); + } + if let Some(limit) = limit { + builder = builder.with_limit(limit); + } + let sync_reader = builder.build()?; + + let total_plan_bytes: u64 = plan.iter().map(|p| p.planned.len()).sum(); + let state = StreamingScanState { + plan, + total_plan_bytes, + total_selected, + batch_size: batch_size as u64, + window, + fetched_idx: 0, + inflight_start: 0, + clear_idx: 0, + resident_bytes: 0, + cursor: 0, + buffers, + slot: ReaderSlot::Idle(reader), + sync_reader, + decoder_projection, + baseline_metrics, + }; + Ok( + futures::stream::unfold(state, |state| async move { state.transition().await }) + .fuse() + .boxed(), + ) +} + +pub(crate) struct StreamingScanState { + plan: Vec, + /// Total bytes across all plan pages. Plans that fit the window are + /// fetched in one wave (see the inline-fetch site) — splitting a small + /// file's fetch into a required wave plus a readahead wave costs an + /// extra round trip per file, which dominates many-small-file workloads + /// (measured: TPC-DS under simulated latency). + total_plan_bytes: u64, + total_selected: u64, + batch_size: u64, + window: u64, + /// Plan pages `[0, fetched_idx)` have been requested (resident or in + /// the single in-flight background fetch). + fetched_idx: usize, + /// Start of the in-flight slice when the slot is `Busy`. + inflight_start: usize, + /// Scan start for dropping pages the cursor has passed. + clear_idx: usize, + resident_bytes: u64, + /// Selected rows emitted so far. + cursor: u64, + buffers: SharedBuffers, + slot: ReaderSlot, + sync_reader: ParquetRecordBatchReader, + decoder_projection: DecoderProjection, + baseline_metrics: BaselineMetrics, +} + +impl StreamingScanState { + /// First selected row not yet guaranteed decodable: pages whose + /// `sel_start` is below this must be resident before the next batch. + fn needed_end(&self) -> u64 { + (self.cursor + self.batch_size).min(self.total_selected) + } + + /// Whether any not-yet-landed plan page is required for the next batch. + fn required_pending(&self) -> bool { + let needed = self.needed_end(); + let first_unlanded = match self.slot { + ReaderSlot::Busy(_) => self.inflight_start, + _ => self.fetched_idx, + }; + self.plan + .get(first_unlanded) + .is_some_and(|p| p.planned.first_row < needed) + } + + /// Extent of the next fetch starting at `fetched_idx`. When + /// `required_only`, stop at the pages the next batch needs (keeps the + /// blocking inline fetch — and therefore time-to-first-batch — minimal); + /// otherwise extend with readahead while the window has room. + fn next_gulp_end(&self, required_only: bool) -> usize { + let needed = self.needed_end(); + let mut bytes = 0u64; + let mut end = self.fetched_idx; + while let Some(page) = self.plan.get(end) { + let len = page.planned.len(); + let required = page.planned.first_row < needed; + if !required + && (required_only || self.resident_bytes + bytes + len > self.window) + { + break; + } + bytes += len; + end += 1; + } + end + } + + /// Drop resident pages the decode cursor has fully passed. + fn clear_consumed(&mut self) { + let landed_end = match self.slot { + ReaderSlot::Busy(_) => self.inflight_start, + _ => self.fetched_idx, + }; + let mut idx = self.clear_idx; + while idx < landed_end { + let page = &mut self.plan[idx]; + if page.planned.first_row > self.cursor { + break; + } + if !page.cleared && page.planned.last_row <= self.cursor { + self.buffers.remove(page.range().start); + self.resident_bytes -= page.planned.len(); + page.cleared = true; + } + idx += 1; + } + while self + .plan + .get(self.clear_idx) + .is_some_and(|page| page.cleared) + { + self.clear_idx += 1; + } + } + + /// Byte ranges this wave still needs: the plan pages in + /// `[start_idx, end_idx)` not already staged. + /// + /// Deliberately no range merging. `ObjectStore::get_ranges` already + /// coalesces (1MB gap) for every store using the default implementation — + /// S3, GCS, Azure — while `LocalFileSystem` and friends override it and + /// coalesce not at all. A second pass here can therefore only raise the + /// effective threshold, never lower it, and it merges without knowing the + /// medium. Measurement agreed: a 4MB gap merged away 59 requests on + /// ClickBench but pulled 157MB of unprojected columns with them and ran + /// slower. The merge decision belongs to the layer that knows its own + /// round-trip cost. + fn wave_ranges(&self, start_idx: usize, end_idx: usize) -> Vec> { + let mut ranges: Vec> = self.plan[start_idx..end_idx] + .iter() + .filter(|p| !self.buffers.contains(p.range().start)) + .map(|p| p.range().clone()) + .collect(); + ranges.sort_by_key(|r| r.start); + ranges + } + + /// Stage a landed wave. Each fetched range is exactly one plan page. + fn install_wave(&mut self, ranges: &[Range], data: &[Bytes]) { + for (range, bytes) in ranges.iter().zip(data) { + self.resident_bytes += range.end - range.start; + self.buffers.insert(range, bytes.clone()); + } + PEAK_STAGED_BYTES + .fetch_max(self.resident_bytes, std::sync::atomic::Ordering::Relaxed); + } + + async fn transition(mut self) -> Option<(Result, Self)> { + loop { + // 1. Land the in-flight fetch when the next batch needs it (or + // when there is nothing left to decode without it). + if self.required_pending() { + match std::mem::replace(&mut self.slot, ReaderSlot::Empty) { + ReaderSlot::Busy(handle) => { + let (reader, ranges, result) = match handle.await { + Ok(v) => v, + Err(e) => { + return Some(( + Err(DataFusionError::External(Box::new(e))), + self, + )); + } + }; + self.slot = ReaderSlot::Idle(reader); + match result { + Ok(data) => { + self.install_wave(&ranges, &data); + } + Err(e) => { + return Some((Err(DataFusionError::from(e)), self)); + } + } + continue; + } + ReaderSlot::Idle(mut reader) => { + // When the whole plan fits the readahead window + // (small files), fetch it in a single wave — the + // extra round trip of a required-only wave would + // dominate. For plans larger than the window, fetch + // only what the next batch requires so the blocking + // wave — and therefore time-to-first-batch — stays + // small; readahead happens in the background + // (step 2). + // "Fill the cart": a blocking wave is a round trip + // we pay either way, so extend it with readahead up + // to the window — EXCEPT the file's very first wave + // of a larger-than-window plan, which stays + // required-only so time-to-first-batch tracks the + // first pages rather than the window. + let required_only = + self.fetched_idx == 0 && self.total_plan_bytes > self.window; + let end = self.next_gulp_end(required_only); + let ranges = self.wave_ranges(self.fetched_idx, end); + let result = reader.get_byte_ranges(ranges.clone()).await; + self.slot = ReaderSlot::Idle(reader); + match result { + Ok(data) => { + self.install_wave(&ranges, &data); + self.fetched_idx = end; + } + Err(e) => { + return Some((Err(DataFusionError::from(e)), self)); + } + } + continue; + } + ReaderSlot::Empty => unreachable!("slot never left empty"), + } + } + + // 2. Required data resident: start background readahead when the + // slot is idle and at least half the window is free (or the + // tail is all that remains). + if matches!(self.slot, ReaderSlot::Idle(_)) + && self.fetched_idx < self.plan.len() + { + let end = self.next_gulp_end(false); + let gulp_bytes: u64 = self.plan[self.fetched_idx..end] + .iter() + .map(|p| p.planned.len()) + .sum(); + let tail = end == self.plan.len(); + if end > self.fetched_idx && (gulp_bytes >= self.window / 2 || tail) { + let ranges = self.wave_ranges(self.fetched_idx, end); + let ReaderSlot::Idle(mut reader) = + std::mem::replace(&mut self.slot, ReaderSlot::Empty) + else { + unreachable!() + }; + self.inflight_start = self.fetched_idx; + self.fetched_idx = end; + // The repo's `SpawnedTask` aborts on drop; this POC + // documents detach-on-drop semantics instead. + #[expect(clippy::disallowed_methods)] + let handle = tokio::task::spawn(async move { + let result = reader.get_byte_ranges(ranges.clone()).await; + (reader, ranges, result) + }); + self.slot = ReaderSlot::Busy(handle); + } + } + + // 3. Decode one batch — never blocks: its pages are resident. + let timer = self.baseline_metrics.elapsed_compute().timer(); + let next = self.sync_reader.next(); + match next { + Some(Ok(batch)) => { + self.cursor += batch.num_rows() as u64; + let result = self.decoder_projection.map(&batch); + drop(timer); + self.clear_consumed(); + return Some((result, self)); + } + Some(Err(e)) => { + drop(timer); + return Some((Err(DataFusionError::from(e)), self)); + } + None => { + drop(timer); + return None; + } + } + } + } +} + #[cfg(test)] mod tests { use super::*;