From d38a4eef2563083e692ac60dbb5d413fc44f2b8d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:10:35 -0500 Subject: [PATCH 1/9] POC: pipelined (overlapped) row-group prefetch with a byte budget in the parquet push decoder Experimental alternative to #23492: instead of widening the blocking fetch, spawn a background fetch of upcoming row groups' projected ranges (same byte budget) while the current row group decodes, so I/O overlaps decode. Env-var switched (DF_FETCH_POLICY / DF_FETCH_BUDGET) for A/B benchmarking; not intended to merge in this form. Co-Authored-By: Claude Fable 5 --- datafusion/datasource-parquet/src/mod.rs | 1 + .../datasource-parquet/src/opener/mod.rs | 8 +- .../datasource-parquet/src/push_decoder.rs | 261 +++++++++++++++++- 3 files changed, 257 insertions(+), 13 deletions(-) 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..a81cbfa23cfb0 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; @@ -1481,7 +1482,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..82543c076e827 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -35,11 +35,13 @@ //! 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; @@ -111,6 +113,114 @@ 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. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum FetchPolicy { + Off, + Batched { budget: u64 }, + Pipelined { budget: 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 }, + _ => 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 +353,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 +488,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 +569,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) => { From 5ff7af62acc00fbfe33496798fd22245958be41b Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 08:08:47 -0500 Subject: [PATCH 2/9] POC: streaming (batch-granular) parquet scan via a persistent sync reader Adds DF_FETCH_POLICY=streaming: instead of buffering whole row groups, one long-lived sync ParquetRecordBatchReader pulls bytes through a shared in-memory ChunkReader. The stream driver computes from the offset index exactly which page ranges the next batch needs, awaits their fetch, and keeps up to DF_FETCH_BUDGET bytes of background readahead in flight; page bytes are dropped as soon as the decode cursor passes them. Because the reader persists for the whole file, dictionary pages are fetched and decoded exactly once, there is no per-row-group reader rebuild, and resident memory is bounded by the readahead window rather than row-group size. Supports projections, RowSelections (page skipping preserved via selected-row prefix sums), limits, and multi-RG files; falls back to the push-decoder path when row filters are active or the offset index is unavailable. Local simulated-latency results (300MB single-row-group file, 50ms latency, 100MB window, streaming scan): time-to-first-batch 112ms vs ~700ms for all row-group-granular policies, total 708ms vs ~845ms, peak staged 100MB vs 283MB. Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/opener/mod.rs | 48 ++ .../datasource-parquet/src/push_decoder.rs | 554 ++++++++++++++++++ 2 files changed, 602 insertions(+) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index a81cbfa23cfb0..935543b18db62 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1387,6 +1387,54 @@ impl RowGroupsPrunedParquetOpen { prepared.virtual_state.as_deref(), )?; + // EXPERIMENT: streaming (batch-granular) scan path, selected via + // DF_FETCH_POLICY=streaming. Bypasses the push decoder entirely: a + // long-lived sync reader 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. + 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), + }; + } + } + } + let (decoder, rg_plan) = { let pushdown_predicate = prepared .pushdown_filters diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 82543c076e827..3de52f5f97087 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -47,13 +47,17 @@ 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::errors::ParquetError; use parquet::file::metadata::ParquetMetaData; +use parquet::file::reader::{ChunkReader, Length}; use datafusion_common::{DataFusionError, Result}; use datafusion_physical_expr::expressions::DynamicFilterTracking; @@ -138,11 +142,19 @@ fn note_staged_bytes(decoder: &ParquetPushDecoder) { /// 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 { @@ -154,6 +166,7 @@ impl FetchPolicy { 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, } } @@ -662,6 +675,547 @@ 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); + } +} + +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 + ))) + } +} + +/// One fetchable unit of the streaming plan: a page (or a row group's +/// dictionary region), with its position in *selected-row* space so the +/// driver knows when it becomes needed and when it can be dropped. +struct PlanPage { + range: Range, + /// First selected row (in output order) this page contributes to. + sel_start: u64, + /// One past the last selected row this page contributes to. Dictionary + /// regions span their whole row group so they stay resident until the + /// row group is fully decoded. + sel_end: u64, + cleared: bool, +} + +/// Prefix-sum view over a `RowSelection`: how many rows are selected before +/// a given raw row index (raw = concatenated rows of the scanned row groups +/// in scan order). +struct SelectedPrefix { + /// (raw_start, selected_before, skip) per selector run. + runs: Vec<(u64, u64, bool)>, + total_raw: u64, + total_selected: u64, +} + +impl SelectedPrefix { + fn new(selection: Option<&RowSelection>, total_raw: u64) -> Self { + let Some(selection) = selection else { + return Self { + runs: vec![(0, 0, false)], + total_raw, + total_selected: total_raw, + }; + }; + let mut runs = Vec::new(); + let mut raw = 0u64; + let mut selected = 0u64; + for selector in selection.iter() { + runs.push((raw, selected, selector.skip)); + raw += selector.row_count as u64; + if !selector.skip { + selected += selector.row_count as u64; + } + } + // Rows past the end of the selection are not selected. + runs.push((raw, selected, true)); + Self { + runs, + total_raw, + total_selected: selected, + } + } + + fn selected_before(&self, raw: u64) -> u64 { + let raw = raw.min(self.total_raw); + let idx = self.runs.partition_point(|(start, _, _)| *start <= raw) - 1; + let (start, selected, skip) = self.runs[idx]; + if skip { + selected + } else { + selected + (raw - start) + } + } +} + +/// Opaque prebuilt streaming fetch plan (see [`build_streaming_plan`]). +pub(crate) struct StreamingPlan { + pages: Vec, + total_selected: u64, + file_end: u64, +} + +/// Build the streaming fetch plan: every projected page (plus per-RG +/// dictionary regions) in decode-need order. Returns `None` when the +/// offset index is unavailable — the caller falls back to the push-decoder +/// path. Borrows only, so callers can probe feasibility before committing +/// resources to the streaming path. +pub(crate) fn build_streaming_plan( + metadata: &ParquetMetaData, + row_group_indexes: &[usize], + projection: &ProjectionMask, + selection: Option<&RowSelection>, +) -> Option { + let offset_index = metadata.offset_index()?; + let total_raw: u64 = row_group_indexes + .iter() + .map(|&rg| metadata.row_group(rg).num_rows() as u64) + .sum(); + let prefix = SelectedPrefix::new(selection, total_raw); + + let mut plan: Vec = Vec::new(); + let mut file_end = 0u64; + let mut rg_raw_start = 0u64; + for &rg_idx in row_group_indexes { + let rg = metadata.row_group(rg_idx); + let rg_rows = rg.num_rows() as u64; + let rg_sel_start = prefix.selected_before(rg_raw_start); + let rg_sel_end = prefix.selected_before(rg_raw_start + rg_rows); + for (col_idx, column) in rg.columns().iter().enumerate() { + let (chunk_start, chunk_len) = column.byte_range(); + file_end = file_end.max(chunk_start + chunk_len); + if !projection.leaf_included(col_idx) { + continue; + } + let locations = offset_index + .get(rg_idx) + .and_then(|cols| cols.get(col_idx))? + .page_locations(); + if locations.is_empty() { + return None; + } + if rg_sel_start == rg_sel_end { + // No selected rows in this row group at all. + continue; + } + // Dictionary region: everything before the first data page. + let first_page = locations[0].offset as u64; + if first_page != chunk_start { + plan.push(PlanPage { + range: chunk_start..first_page, + sel_start: rg_sel_start, + sel_end: rg_sel_end, + cleared: false, + }); + } + for (i, loc) in locations.iter().enumerate() { + let raw_first = rg_raw_start + loc.first_row_index as u64; + let raw_end = locations + .get(i + 1) + .map(|next| rg_raw_start + next.first_row_index as u64) + .unwrap_or(rg_raw_start + rg_rows); + let sel_start = prefix.selected_before(raw_first); + let sel_end = prefix.selected_before(raw_end); + if sel_start == sel_end { + // Page contains no selected rows: never fetched (page + // skipping preserved). + continue; + } + let start = loc.offset as u64; + plan.push(PlanPage { + range: start..start + loc.compressed_page_size as u64, + sel_start, + sel_end, + cleared: false, + }); + } + } + rg_raw_start += rg_rows; + } + // Need order: by first selected row, dictionaries (wider spans) first + // among equals so they are resident before their data pages decode. + plan.sort_by_key(|p| (p.sel_start, std::cmp::Reverse(p.sel_end), p.range.start)); + Some(StreamingPlan { + pages: plan, + total_selected: prefix.total_selected, + 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 state = StreamingScanState { + plan, + 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_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.sel_start < 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.range.end - page.range.start; + let required = page.sel_start < 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.sel_start > self.cursor { + break; + } + if !page.cleared && page.sel_end <= self.cursor { + self.buffers.remove(page.range.start); + self.resident_bytes -= page.range.end - page.range.start; + page.cleared = true; + } + idx += 1; + } + while self + .plan + .get(self.clear_idx) + .is_some_and(|page| page.cleared) + { + self.clear_idx += 1; + } + } + + 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) => { + for (range, bytes) in ranges.iter().zip(data) { + self.buffers.insert(range, bytes); + self.resident_bytes += range.end - range.start; + } + PEAK_STAGED_BYTES.fetch_max( + self.resident_bytes, + std::sync::atomic::Ordering::Relaxed, + ); + } + Err(e) => { + return Some((Err(DataFusionError::from(e)), self)); + } + } + continue; + } + ReaderSlot::Idle(mut reader) => { + // Fetch only the pages the next batch requires — + // readahead happens in the background (step 2), so + // the blocking fetch stays small and TTFB low. + let end = self.next_gulp_end(true); + let ranges: Vec> = self.plan[self.fetched_idx..end] + .iter() + .map(|p| p.range.clone()) + .collect(); + let result = reader.get_byte_ranges(ranges.clone()).await; + self.slot = ReaderSlot::Idle(reader); + match result { + Ok(data) => { + for (range, bytes) in ranges.iter().zip(data) { + self.buffers.insert(range, bytes); + self.resident_bytes += range.end - range.start; + } + PEAK_STAGED_BYTES.fetch_max( + self.resident_bytes, + std::sync::atomic::Ordering::Relaxed, + ); + 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.range.end - p.range.start) + .sum(); + let tail = end == self.plan.len(); + if end > self.fetched_idx && (gulp_bytes >= self.window / 2 || tail) { + let ranges: Vec> = self.plan[self.fetched_idx..end] + .iter() + .map(|p| p.range.clone()) + .collect(); + 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::*; From bcbbda1ae52a5e77cfcd895e14ee9bc1c686c1bb Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 10:05:25 -0500 Subject: [PATCH 3/9] streaming: fetch small files in a single wave MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The required-only first fetch keeps TTFB low on large files, but for a file whose whole plan fits the readahead window it turns one fetch wave into two serial waves — an extra round trip per file, which dominates many-small-file workloads (TPC-DS under simulated latency regressed 1.3x). Fetch the entire plan in one wave when it fits the window; required-only kicks in only for plans larger than the window, preserving the giant-row-group TTFB and memory wins. Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/push_decoder.rs | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 3de52f5f97087..67877f08d89af 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -976,8 +976,10 @@ pub(crate) fn build_streaming_stream( } let sync_reader = builder.build()?; + let total_plan_bytes: u64 = plan.iter().map(|p| p.range.end - p.range.start).sum(); let state = StreamingScanState { plan, + total_plan_bytes, total_selected, batch_size: batch_size as u64, window, @@ -1001,6 +1003,12 @@ pub(crate) fn build_streaming_stream( 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, @@ -1125,10 +1133,16 @@ impl StreamingScanState { continue; } ReaderSlot::Idle(mut reader) => { - // Fetch only the pages the next batch requires — - // readahead happens in the background (step 2), so - // the blocking fetch stays small and TTFB low. - let end = self.next_gulp_end(true); + // 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). + let required_only = self.total_plan_bytes > self.window; + let end = self.next_gulp_end(required_only); let ranges: Vec> = self.plan[self.fetched_idx..end] .iter() .map(|p| p.range.clone()) From 11f6867b988f3277f41332b5a910b92a1d7193cf Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:10:21 -0500 Subject: [PATCH 4/9] streaming: fill blocking waves and coalesce scattered ranges Two cart-filling changes to close the gap with bulk row-group prefetch on scattered page plans: - A blocking wave is a round trip we pay either way, so extend it with readahead up to the window (except a large plan's first wave, which stays required-only to protect time-to-first-batch). - Merge fetch ranges whose file gap is at most DF_FETCH_COALESCE bytes (default 4MB): deliberate, bounded over-fetch of gap bytes that collapses the object-store GET count for page-precise plans. Gap bytes are dropped after slicing pages out of the fetched blobs. Local filtered-scan spot check (page-index selections, 50ms latency): 193 GETs -> 15, total 3817ms -> 1686ms, at 149MB -> 379MB fetched (the bytes-vs-round-trips dial; tune per store). Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/push_decoder.rs | 93 ++++++++++++++----- 1 file changed, 69 insertions(+), 24 deletions(-) diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 67877f08d89af..04c5945a490d4 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -987,6 +987,10 @@ pub(crate) fn build_streaming_stream( inflight_start: 0, clear_idx: 0, resident_bytes: 0, + coalesce_gap: std::env::var("DF_FETCH_COALESCE") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(4 * 1024 * 1024), cursor: 0, buffers, slot: ReaderSlot::Idle(reader), @@ -1020,6 +1024,9 @@ pub(crate) struct StreamingScanState { /// Scan start for dropping pages the cursor has passed. clear_idx: usize, resident_bytes: u64, + /// Merge fetch ranges whose gap is at most this many bytes (deliberate + /// over-fetch that collapses GET count for scattered page plans). + coalesce_gap: u64, /// Selected rows emitted so far. cursor: u64, buffers: SharedBuffers, @@ -1098,6 +1105,52 @@ impl StreamingScanState { } } + /// Fetch ranges for a wave of plan pages: merge page ranges whose file + /// gap is <= `coalesce_gap` into single requests ("buy the shelf + /// section") — deliberate over-fetch of small gaps that collapses the + /// object-store GET count for scattered page-precise plans. The gap + /// bytes are dropped after slicing (only page bytes are installed), so + /// window accounting stays page-based; the backing allocation lives + /// until its last page clears. + fn wave_ranges(&self, start_idx: usize, end_idx: usize) -> Vec> { + let mut sorted: Vec> = self.plan[start_idx..end_idx] + .iter() + .map(|p| p.range.clone()) + .collect(); + sorted.sort_by_key(|r| r.start); + let mut merged: Vec> = Vec::with_capacity(sorted.len()); + for r in sorted { + match merged.last_mut() { + Some(last) if r.start.saturating_sub(last.end) <= self.coalesce_gap => { + last.end = last.end.max(r.end); + } + _ => merged.push(r), + } + } + merged + } + + /// Install a fetched wave: slice each plan page's bytes out of the + /// merged fetch results and stage them in the shared buffers. + fn install_wave( + &mut self, + start_idx: usize, + end_idx: usize, + fetched: &[Range], + data: &[Bytes], + ) { + for page in &self.plan[start_idx..end_idx] { + let i = fetched.partition_point(|r| r.start <= page.range.start) - 1; + let offset = (page.range.start - fetched[i].start) as usize; + let len = (page.range.end - page.range.start) as usize; + self.buffers + .insert(&page.range, data[i].slice(offset..offset + len)); + self.resident_bytes += len as u64; + } + 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 @@ -1117,13 +1170,11 @@ impl StreamingScanState { self.slot = ReaderSlot::Idle(reader); match result { Ok(data) => { - for (range, bytes) in ranges.iter().zip(data) { - self.buffers.insert(range, bytes); - self.resident_bytes += range.end - range.start; - } - PEAK_STAGED_BYTES.fetch_max( - self.resident_bytes, - std::sync::atomic::Ordering::Relaxed, + self.install_wave( + self.inflight_start, + self.fetched_idx, + &ranges, + &data, ); } Err(e) => { @@ -1141,24 +1192,21 @@ impl StreamingScanState { // wave — and therefore time-to-first-batch — stays // small; readahead happens in the background // (step 2). - let required_only = self.total_plan_bytes > self.window; + // "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: Vec> = self.plan[self.fetched_idx..end] - .iter() - .map(|p| p.range.clone()) - .collect(); + 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) => { - for (range, bytes) in ranges.iter().zip(data) { - self.buffers.insert(range, bytes); - self.resident_bytes += range.end - range.start; - } - PEAK_STAGED_BYTES.fetch_max( - self.resident_bytes, - std::sync::atomic::Ordering::Relaxed, - ); + self.install_wave(self.fetched_idx, end, &ranges, &data); self.fetched_idx = end; } Err(e) => { @@ -1184,10 +1232,7 @@ impl StreamingScanState { .sum(); let tail = end == self.plan.len(); if end > self.fetched_idx && (gulp_bytes >= self.window / 2 || tail) { - let ranges: Vec> = self.plan[self.fetched_idx..end] - .iter() - .map(|p| p.range.clone()) - .collect(); + let ranges = self.wave_ranges(self.fetched_idx, end); let ReaderSlot::Idle(mut reader) = std::mem::replace(&mut self.slot, ReaderSlot::Empty) else { From 34f28eaf3b539645c3ef3cf3be6baf18f7d7523a Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 12:48:25 -0500 Subject: [PATCH 5/9] streaming: derive the fetch plan from arrow-rs instead of reimplementing it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming path previously reimplemented arrow-rs's page/dictionary/ selection walk to work out which bytes each batch needs, duplicating logic that `InMemoryRowGroup::fetch_ranges` already owns. That derivation belongs in the crate that must stay consistent with it, not here. Adds `parquet::arrow::push_decoder::plan_scan_ranges` upstream and uses it: it decomposes a scan into the pages it will read, in decode order, each tagged with the span of selected rows it serves. `build_streaming_plan` becomes a thin adapter, and DataFusion keeps only what is genuinely scheduling — readahead budget, wave sizing, range coalescing, eviction. Row tags rather than a fixed row or byte quantum keep the plan usable by callers budgeting in either unit; that matters for wide values, where a fixed row count implies an unbounded byte count. The remaining floor on resident bytes is one output batch's pages, which only a smaller `batch_size` (or byte-based batch sizing upstream) can lower. The API is unreleased, so this pins every arrow crate to the branch carrying it (pydantic/arrow-rs claude/push-decoder-peek-59 = apache 59.1.0 + that one module). That pin is why this POC cannot merge as-is; the intended path is upstreaming the arrow-rs side first. Verified against the same local benchmarks as the hand-rolled version: single 300MB row group 704ms total / 112ms TTFB / 99.9MB peak; filtered page-selection scan 15 GETs and 1689ms at 50ms latency; correctness checks (full scan, page-selection filter, LIMIT) identical. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 55 ++--- Cargo.toml | 23 ++ .../datasource-parquet/src/push_decoder.rs | 202 +++++------------- 3 files changed, 96 insertions(+), 184 deletions(-) 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/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 04c5945a490d4..1f5163917f4d3 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -54,7 +54,9 @@ 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}; @@ -768,160 +770,64 @@ impl ChunkReader for SharedBuffers { } } -/// One fetchable unit of the streaming plan: a page (or a row group's -/// dictionary region), with its position in *selected-row* space so the -/// driver knows when it becomes needed and when it can be dropped. +/// 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 { - range: Range, - /// First selected row (in output order) this page contributes to. - sel_start: u64, - /// One past the last selected row this page contributes to. Dictionary - /// regions span their whole row group so they stay resident until the - /// row group is fully decoded. - sel_end: u64, + planned: PlannedRange, cleared: bool, } -/// Prefix-sum view over a `RowSelection`: how many rows are selected before -/// a given raw row index (raw = concatenated rows of the scanned row groups -/// in scan order). -struct SelectedPrefix { - /// (raw_start, selected_before, skip) per selector run. - runs: Vec<(u64, u64, bool)>, - total_raw: u64, - total_selected: u64, -} - -impl SelectedPrefix { - fn new(selection: Option<&RowSelection>, total_raw: u64) -> Self { - let Some(selection) = selection else { - return Self { - runs: vec![(0, 0, false)], - total_raw, - total_selected: total_raw, - }; - }; - let mut runs = Vec::new(); - let mut raw = 0u64; - let mut selected = 0u64; - for selector in selection.iter() { - runs.push((raw, selected, selector.skip)); - raw += selector.row_count as u64; - if !selector.skip { - selected += selector.row_count as u64; - } - } - // Rows past the end of the selection are not selected. - runs.push((raw, selected, true)); - Self { - runs, - total_raw, - total_selected: selected, - } - } - - fn selected_before(&self, raw: u64) -> u64 { - let raw = raw.min(self.total_raw); - let idx = self.runs.partition_point(|(start, _, _)| *start <= raw) - 1; - let (start, selected, skip) = self.runs[idx]; - if skip { - selected - } else { - selected + (raw - start) - } +impl PlanPage { + fn range(&self) -> &Range { + &self.planned.range } } -/// Opaque prebuilt streaming fetch plan (see [`build_streaming_plan`]). +/// Opaque prebuilt streaming fetch plan. pub(crate) struct StreamingPlan { pages: Vec, total_selected: u64, file_end: u64, } -/// Build the streaming fetch plan: every projected page (plus per-RG -/// dictionary regions) in decode-need order. Returns `None` when the -/// offset index is unavailable — the caller falls back to the push-decoder -/// path. Borrows only, so callers can probe feasibility before committing -/// resources to the streaming path. +/// 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 offset_index = metadata.offset_index()?; - let total_raw: u64 = row_group_indexes + 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() - .map(|&rg| metadata.row_group(rg).num_rows() as u64) - .sum(); - let prefix = SelectedPrefix::new(selection, total_raw); - - let mut plan: Vec = Vec::new(); - let mut file_end = 0u64; - let mut rg_raw_start = 0u64; - for &rg_idx in row_group_indexes { - let rg = metadata.row_group(rg_idx); - let rg_rows = rg.num_rows() as u64; - let rg_sel_start = prefix.selected_before(rg_raw_start); - let rg_sel_end = prefix.selected_before(rg_raw_start + rg_rows); - for (col_idx, column) in rg.columns().iter().enumerate() { - let (chunk_start, chunk_len) = column.byte_range(); - file_end = file_end.max(chunk_start + chunk_len); - if !projection.leaf_included(col_idx) { - continue; - } - let locations = offset_index - .get(rg_idx) - .and_then(|cols| cols.get(col_idx))? - .page_locations(); - if locations.is_empty() { - return None; - } - if rg_sel_start == rg_sel_end { - // No selected rows in this row group at all. - continue; - } - // Dictionary region: everything before the first data page. - let first_page = locations[0].offset as u64; - if first_page != chunk_start { - plan.push(PlanPage { - range: chunk_start..first_page, - sel_start: rg_sel_start, - sel_end: rg_sel_end, - cleared: false, - }); - } - for (i, loc) in locations.iter().enumerate() { - let raw_first = rg_raw_start + loc.first_row_index as u64; - let raw_end = locations - .get(i + 1) - .map(|next| rg_raw_start + next.first_row_index as u64) - .unwrap_or(rg_raw_start + rg_rows); - let sel_start = prefix.selected_before(raw_first); - let sel_end = prefix.selected_before(raw_end); - if sel_start == sel_end { - // Page contains no selected rows: never fetched (page - // skipping preserved). - continue; - } - let start = loc.offset as u64; - plan.push(PlanPage { - range: start..start + loc.compressed_page_size as u64, - sel_start, - sel_end, - cleared: false, - }); - } - } - rg_raw_start += rg_rows; - } - // Need order: by first selected row, dictionaries (wider spans) first - // among equals so they are resident before their data pages decode. - plan.sort_by_key(|p| (p.sel_start, std::cmp::Reverse(p.sel_end), p.range.start)); + .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 { - pages: plan, - total_selected: prefix.total_selected, + total_selected: plan.total_selected_rows, + pages: plan + .ranges + .into_iter() + .map(|planned| PlanPage { + planned, + cleared: false, + }) + .collect(), file_end, }) } @@ -976,7 +882,7 @@ pub(crate) fn build_streaming_stream( } let sync_reader = builder.build()?; - let total_plan_bytes: u64 = plan.iter().map(|p| p.range.end - p.range.start).sum(); + let total_plan_bytes: u64 = plan.iter().map(|p| p.planned.len()).sum(); let state = StreamingScanState { plan, total_plan_bytes, @@ -1052,7 +958,7 @@ impl StreamingScanState { }; self.plan .get(first_unlanded) - .is_some_and(|p| p.sel_start < needed) + .is_some_and(|p| p.planned.first_row < needed) } /// Extent of the next fetch starting at `fetched_idx`. When @@ -1064,8 +970,8 @@ impl StreamingScanState { let mut bytes = 0u64; let mut end = self.fetched_idx; while let Some(page) = self.plan.get(end) { - let len = page.range.end - page.range.start; - let required = page.sel_start < needed; + let len = page.planned.len(); + let required = page.planned.first_row < needed; if !required && (required_only || self.resident_bytes + bytes + len > self.window) { @@ -1086,12 +992,12 @@ impl StreamingScanState { let mut idx = self.clear_idx; while idx < landed_end { let page = &mut self.plan[idx]; - if page.sel_start > self.cursor { + if page.planned.first_row > self.cursor { break; } - if !page.cleared && page.sel_end <= self.cursor { - self.buffers.remove(page.range.start); - self.resident_bytes -= page.range.end - page.range.start; + 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; @@ -1115,7 +1021,7 @@ impl StreamingScanState { fn wave_ranges(&self, start_idx: usize, end_idx: usize) -> Vec> { let mut sorted: Vec> = self.plan[start_idx..end_idx] .iter() - .map(|p| p.range.clone()) + .map(|p| p.range().clone()) .collect(); sorted.sort_by_key(|r| r.start); let mut merged: Vec> = Vec::with_capacity(sorted.len()); @@ -1140,11 +1046,11 @@ impl StreamingScanState { data: &[Bytes], ) { for page in &self.plan[start_idx..end_idx] { - let i = fetched.partition_point(|r| r.start <= page.range.start) - 1; - let offset = (page.range.start - fetched[i].start) as usize; - let len = (page.range.end - page.range.start) as usize; + let i = fetched.partition_point(|r| r.start <= page.range().start) - 1; + let offset = (page.range().start - fetched[i].start) as usize; + let len = (page.planned.len()) as usize; self.buffers - .insert(&page.range, data[i].slice(offset..offset + len)); + .insert(page.range(), data[i].slice(offset..offset + len)); self.resident_bytes += len as u64; } PEAK_STAGED_BYTES @@ -1228,7 +1134,7 @@ impl StreamingScanState { let end = self.next_gulp_end(false); let gulp_bytes: u64 = self.plan[self.fetched_idx..end] .iter() - .map(|p| p.range.end - p.range.start) + .map(|p| p.planned.len()) .sum(); let tail = end == self.plan.len(); if end > self.fetched_idx && (gulp_bytes >= self.window / 2 || tail) { From 075e816e8d8c666dda45b76a4c1df36b87cd8ed5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:09:46 -0500 Subject: [PATCH 6/9] streaming: load the page index when the streaming policy is active MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Page-index loading is driven entirely by whether *pruning* can use it (should_load_page_index requires a page-pruning predicate that can still eliminate row groups). The streaming path plans at page granularity and returns None without an offset index, so on any scan with no prunable predicate it was silently falling back to the push decoder. That is the likely explanation for the benchmark bot reporting "no change" on every clickbench and tpcds query while tpch (whose queries carry selective range predicates) saw 1.69x and 6.0x: the streaming path was probably never executing on those suites. Force the load when the policy is on, which is also just correct — the offset index is required infrastructure for this path, not an optional pruning aid. This costs an extra metadata read per file where pruning would not have needed one, so many-small-file workloads may pay for it; the point of the change is to make the comparison meaningful either way. Co-Authored-By: Claude Fable 5 --- datafusion/datasource-parquet/src/opener/mod.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 935543b18db62..922b3091d7b1b 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -551,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(), )) From ee1db7ba7cc44b0e45f48dba4300cada2adfd782 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:20:26 -0500 Subject: [PATCH 7/9] streaming: log fallbacks, install covered pages, lower the coalesce default Three fixes from investigating why clickbench showed no change: - Log when the streaming policy is requested but cannot run (no offset index, or pushdown filters active). The fallback was silent, which is why it took three benchmark rounds to notice that the published ClickBench files carry no page index at all (0 of 105 columns) and the path was never executing. - Install every planned page a fetched blob covers, not just the current wave's. Coalescing over-fetches the gaps between pages and those gaps can contain pages planned for a later wave. Measured worth only ~7MB of 157MB on clickbench, so this was not the source of the over-fetch I suspected, but re-fetching bytes already in hand is still wrong. - Lower the default coalesce gap from 4MB to 1MB, matching object_store's OBJECT_STORE_COALESCE_DEFAULT. 4MB was a guess and measurement contradicts it: on the page-index ClickBench copy it merged away 59 requests but pulled 157MB of unprojected columns with them and ran slower (q22 204ms/236MB vs 170ms/212MB at 1MB). With several partitions fetching concurrently a saved round trip is worth far less than latency x bandwidth suggests. At 1MB, bytes fetched return to parity with the unscheduled path (1655MB vs 1650MB) and the speedup holds at 1.53x. Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/opener/mod.rs | 16 ++++++++ .../datasource-parquet/src/push_decoder.rs | 40 +++++++++++++++++-- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 922b3091d7b1b..e1fa7caff2dd8 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1440,7 +1440,23 @@ impl RowGroupsPrunedParquetOpen { } _ => 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 + ); } } diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 1f5163917f4d3..006b51c3cec14 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -716,6 +716,11 @@ impl SharedBuffers { 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 { @@ -893,10 +898,17 @@ pub(crate) fn build_streaming_stream( inflight_start: 0, clear_idx: 0, resident_bytes: 0, + // 1MB matches object_store's own OBJECT_STORE_COALESCE_DEFAULT, and + // measurement agrees: on ClickBench (page-index copy, 50ms latency, + // 8 partitions) a 4MB gap merged away 59 requests but pulled 157MB of + // unprojected columns with them, and ran *slower* — with several + // partitions fetching concurrently, a saved round trip is worth much + // less than the naive latency x bandwidth break-even suggests. At 1MB + // bytes fetched return to parity with the unscheduled path. coalesce_gap: std::env::var("DF_FETCH_COALESCE") .ok() .and_then(|v| v.parse::().ok()) - .unwrap_or(4 * 1024 * 1024), + .unwrap_or(1024 * 1024), cursor: 0, buffers, slot: ReaderSlot::Idle(reader), @@ -1019,8 +1031,10 @@ impl StreamingScanState { /// window accounting stays page-based; the backing allocation lives /// until its last page clears. fn wave_ranges(&self, start_idx: usize, end_idx: usize) -> Vec> { + // Skip pages an earlier wave's gap fill already staged. let mut sorted: Vec> = self.plan[start_idx..end_idx] .iter() + .filter(|p| !self.buffers.contains(p.range().start)) .map(|p| p.range().clone()) .collect(); sorted.sort_by_key(|r| r.start); @@ -1045,10 +1059,28 @@ impl StreamingScanState { fetched: &[Range], data: &[Bytes], ) { - for page in &self.plan[start_idx..end_idx] { - let i = fetched.partition_point(|r| r.start <= page.range().start) - 1; + // Install every planned page the fetched blobs cover — not just this + // wave's. Coalescing deliberately over-fetches the gaps between + // pages, and those gaps routinely contain pages planned for a *later* + // wave (a different column, or later rows). Slicing them out now is + // free; discarding them means paying to fetch the same bytes twice, + // which measured as ~10% of all bytes read on ClickBench. + let _ = end_idx; + for idx in start_idx..self.plan.len() { + let page = &self.plan[idx]; + if page.cleared || self.buffers.contains(page.range().start) { + continue; + } + let i = fetched.partition_point(|r| r.start <= page.range().start); + if i == 0 || page.range().end > fetched[i - 1].end { + // Not covered by this fetch. Pages are ordered by decode + // need rather than file offset, so this says nothing about + // whether later pages are covered — keep scanning. + continue; + } + let i = i - 1; let offset = (page.range().start - fetched[i].start) as usize; - let len = (page.planned.len()) as usize; + let len = page.planned.len() as usize; self.buffers .insert(page.range(), data[i].slice(offset..offset + len)); self.resident_bytes += len as u64; From 7f3e44840d3cc56a9d4158e86280cc29d28add93 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:29:08 -0500 Subject: [PATCH 8/9] streaming: drop scheduler-side range coalescing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ObjectStore::get_ranges` already coalesces ranges within 1MB for every store using the default implementation (S3, GCS, Azure), while LocalFileSystem and the other overriding stores coalesce not at all. A second pass in the scheduler can therefore only raise the effective threshold, never lower it, and it makes the merge decision without knowing the medium — so DF_FETCH_COALESCE was both redundant on real object stores and unable to express "don't merge" there. It also never earned its place. Removing it entirely is marginally faster than any setting of it on the page-index ClickBench copy (50ms latency, 8 partitions, medians of 2): 1.54x vs 1.53x at 1MB and 1.51x at 4MB, with bytes fetched back near parity with the unscheduled path (1680MB vs 1650MB, against 1807MB at 4MB). Per query: 35 faster, 0 slower, 1 neutral. The scan's own wins do not depend on it: the tpch 1.69x and tpch10 6.4x results were measured at bcbbda1, before this code existed. What remains is the actual idea — batch-granular readiness, so decode starts on the first pages instead of waiting for a whole row group. If a scheduler-side merge is wanted later, it should be driven by characteristics the store reports rather than a constant, so that the two layers stop duplicating a decision neither can make alone. Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/push_decoder.rs | 97 +++++-------------- 1 file changed, 22 insertions(+), 75 deletions(-) diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 006b51c3cec14..36fc8eeff4af9 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -898,17 +898,6 @@ pub(crate) fn build_streaming_stream( inflight_start: 0, clear_idx: 0, resident_bytes: 0, - // 1MB matches object_store's own OBJECT_STORE_COALESCE_DEFAULT, and - // measurement agrees: on ClickBench (page-index copy, 50ms latency, - // 8 partitions) a 4MB gap merged away 59 requests but pulled 157MB of - // unprojected columns with them, and ran *slower* — with several - // partitions fetching concurrently, a saved round trip is worth much - // less than the naive latency x bandwidth break-even suggests. At 1MB - // bytes fetched return to parity with the unscheduled path. - coalesce_gap: std::env::var("DF_FETCH_COALESCE") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(1024 * 1024), cursor: 0, buffers, slot: ReaderSlot::Idle(reader), @@ -942,9 +931,6 @@ pub(crate) struct StreamingScanState { /// Scan start for dropping pages the cursor has passed. clear_idx: usize, resident_bytes: u64, - /// Merge fetch ranges whose gap is at most this many bytes (deliberate - /// over-fetch that collapses GET count for scattered page plans). - coalesce_gap: u64, /// Selected rows emitted so far. cursor: u64, buffers: SharedBuffers, @@ -1023,67 +1009,33 @@ impl StreamingScanState { } } - /// Fetch ranges for a wave of plan pages: merge page ranges whose file - /// gap is <= `coalesce_gap` into single requests ("buy the shelf - /// section") — deliberate over-fetch of small gaps that collapses the - /// object-store GET count for scattered page-precise plans. The gap - /// bytes are dropped after slicing (only page bytes are installed), so - /// window accounting stays page-based; the backing allocation lives - /// until its last page clears. + /// 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> { - // Skip pages an earlier wave's gap fill already staged. - let mut sorted: Vec> = self.plan[start_idx..end_idx] + 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(); - sorted.sort_by_key(|r| r.start); - let mut merged: Vec> = Vec::with_capacity(sorted.len()); - for r in sorted { - match merged.last_mut() { - Some(last) if r.start.saturating_sub(last.end) <= self.coalesce_gap => { - last.end = last.end.max(r.end); - } - _ => merged.push(r), - } - } - merged + ranges.sort_by_key(|r| r.start); + ranges } - /// Install a fetched wave: slice each plan page's bytes out of the - /// merged fetch results and stage them in the shared buffers. - fn install_wave( - &mut self, - start_idx: usize, - end_idx: usize, - fetched: &[Range], - data: &[Bytes], - ) { - // Install every planned page the fetched blobs cover — not just this - // wave's. Coalescing deliberately over-fetches the gaps between - // pages, and those gaps routinely contain pages planned for a *later* - // wave (a different column, or later rows). Slicing them out now is - // free; discarding them means paying to fetch the same bytes twice, - // which measured as ~10% of all bytes read on ClickBench. - let _ = end_idx; - for idx in start_idx..self.plan.len() { - let page = &self.plan[idx]; - if page.cleared || self.buffers.contains(page.range().start) { - continue; - } - let i = fetched.partition_point(|r| r.start <= page.range().start); - if i == 0 || page.range().end > fetched[i - 1].end { - // Not covered by this fetch. Pages are ordered by decode - // need rather than file offset, so this says nothing about - // whether later pages are covered — keep scanning. - continue; - } - let i = i - 1; - let offset = (page.range().start - fetched[i].start) as usize; - let len = page.planned.len() as usize; - self.buffers - .insert(page.range(), data[i].slice(offset..offset + len)); - self.resident_bytes += len as u64; + /// 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); @@ -1108,12 +1060,7 @@ impl StreamingScanState { self.slot = ReaderSlot::Idle(reader); match result { Ok(data) => { - self.install_wave( - self.inflight_start, - self.fetched_idx, - &ranges, - &data, - ); + self.install_wave(&ranges, &data); } Err(e) => { return Some((Err(DataFusionError::from(e)), self)); @@ -1144,7 +1091,7 @@ impl StreamingScanState { self.slot = ReaderSlot::Idle(reader); match result { Ok(data) => { - self.install_wave(self.fetched_idx, end, &ranges, &data); + self.install_wave(&ranges, &data); self.fetched_idx = end; } Err(e) => { From 2400652c60f9027ef4d914fb36f8812fe0bedd1b Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:55:39 -0500 Subject: [PATCH 9/9] docs: clarify that the streaming path does not drive the push decoder The comment said the streaming path "bypasses the push decoder entirely", which is true of the decoder but reads as contradictory now that the path imports `plan_scan_ranges` from arrow-rs's `push_decoder` module. Only the plan comes from there; the decoder is not constructed, because `NeedsData` resolves at row-group granularity and so cannot express what the next batch needs. Co-Authored-By: Claude Fable 5 --- .../datasource-parquet/src/opener/mod.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index e1fa7caff2dd8..a29ce6b010d6f 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1397,11 +1397,20 @@ impl RowGroupsPrunedParquetOpen { )?; // EXPERIMENT: streaming (batch-granular) scan path, selected via - // DF_FETCH_POLICY=streaming. Bypasses the push decoder entirely: a - // long-lived sync reader 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. + // 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();