diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 621c1c812afd..23f8a510ed3f 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -145,6 +145,64 @@ pub struct ArrowReaderBuilder { pub(crate) metrics: ArrowReaderMetrics, pub(crate) max_predicate_cache_size: usize, + + pub(crate) bounded_streaming: Option, +} + +/// Options controlling adaptive bounded streaming of large column chunks. +/// +/// When enabled (see [`ArrowReaderBuilder::with_bounded_streaming`]), the +/// async reader decodes row groups whose projected column chunks exceed +/// [`Self::stream_threshold`] in bounded *windows* instead of materializing +/// every projected byte before decoding: +/// +/// * Column chunks whose requested bytes are below `stream_threshold` keep the +/// existing behavior: their ranges are coalesced and fully materialized with +/// the minimum number of requests. +/// * Larger chunks are fetched with a single ranged request each whose body is +/// consumed incrementally ([`AsyncFileReader::get_bytes_stream`]), feeding +/// the decoder through a buffer bounded by roughly `window_bytes` (plus one +/// decode window). Decode of earlier windows overlaps transfer of later +/// ones. +/// +/// Requires the Parquet offset index (page locations); row groups without one, +/// or with active row filters, fall back to the materialized path. +/// +/// [`AsyncFileReader::get_bytes_stream`]: https://docs.rs/parquet/latest/parquet/arrow/async_reader/trait.AsyncFileReader.html +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundedStreamingOptions { + /// Minimum number of requested bytes in a single column chunk for it to be + /// streamed rather than materialized. Default: 16 MiB. + pub stream_threshold: u64, + /// Target compressed bytes per decode window, and the approximate + /// per-stream fetch-ahead bound. Default: 8 MiB. + pub window_bytes: u64, + /// Ranges closer together than this are assumed to be coalesced into a + /// single request by the underlying reader. Must match the coalescing of + /// the [`AsyncFileReader::get_byte_ranges`] implementation for the + /// request-count guarantees to be exact. Default: 1 MiB (the + /// `object_store` default). + /// + /// [`AsyncFileReader::get_byte_ranges`]: https://docs.rs/parquet/latest/parquet/arrow/async_reader/trait.AsyncFileReader.html + pub coalesce_gap: u64, + /// A decode window must hold at least this many bytes per projected + /// column, so that per-window fixed costs (rebuilding one array reader + /// per column, re-decompressing one boundary page per column) stay + /// amortized. Wide projections therefore get proportionally larger + /// windows; projections so wide that one window would cover the whole + /// row group fall back to the non-windowed path. Default: 2 MiB. + pub min_window_bytes_per_column: u64, +} + +impl Default for BoundedStreamingOptions { + fn default() -> Self { + Self { + stream_threshold: 16 * 1024 * 1024, + window_bytes: 8 * 1024 * 1024, + coalesce_gap: 1024 * 1024, + min_window_bytes_per_column: 2 * 1024 * 1024, + } + } } impl Debug for ArrowReaderBuilder { @@ -184,6 +242,7 @@ impl ArrowReaderBuilder { offset: None, metrics: ArrowReaderMetrics::Disabled, max_predicate_cache_size: 100 * 1024 * 1024, // 100MB default cache size + bounded_streaming: None, } } @@ -440,6 +499,20 @@ impl ArrowReaderBuilder { ..self } } + + /// Enable adaptive bounded streaming of large column chunks for the async + /// decoder. See [`BoundedStreamingOptions`] for details. + /// + /// Defaults to disabled. Only used by [`ParquetRecordBatchStream`] and the + /// push decoder; the sync reader ignores this option. + /// + /// [`ParquetRecordBatchStream`]: https://docs.rs/parquet/latest/parquet/arrow/async_reader/struct.ParquetRecordBatchStream.html + pub fn with_bounded_streaming(self, options: BoundedStreamingOptions) -> Self { + Self { + bounded_streaming: Some(options), + ..self + } + } } /// Options that control how [`ParquetMetaData`] is read when constructing @@ -1198,6 +1271,8 @@ impl ParquetRecordBatchReaderBuilder { metrics, // Not used for the sync reader, see https://github.com/apache/arrow-rs/issues/8000 max_predicate_cache_size: _, + // Only used by the async / push decoders + bounded_streaming: _, } = self; // Try to avoid allocate large buffer diff --git a/parquet/src/arrow/async_reader/adaptive.rs b/parquet/src/arrow/async_reader/adaptive.rs new file mode 100644 index 000000000000..a8e79173af91 --- /dev/null +++ b/parquet/src/arrow/async_reader/adaptive.rs @@ -0,0 +1,690 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adaptive fetching for [`ParquetRecordBatchStream`]: small ranges are +//! coalesced and materialized exactly as before, large column chunks are +//! fetched with a single ranged request each whose body feeds the decoder +//! incrementally through a bounded buffer. +//! +//! See [`BoundedStreamingOptions`] for the user-facing description. +//! +//! [`ParquetRecordBatchStream`]: super::ParquetRecordBatchStream + +use crate::arrow::arrow_reader::BoundedStreamingOptions; +use crate::arrow::async_reader::AsyncFileReader; +use crate::arrow::push_decoder::UpcomingFetchPlan; +use crate::errors::{ParquetError, Result}; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::stream::{BoxStream, StreamExt}; +use std::collections::{HashSet, VecDeque}; +use std::ops::Range; +use std::task::{Context, Poll}; +use tokio::runtime::Handle; +use tokio::sync::mpsc; + +/// Target size of a single message from a body-draining task to the fetcher. +/// Small network chunks are aggregated up to roughly this size so the +/// channel's message-count capacity corresponds to a byte bound. +const TARGET_MESSAGE_BYTES: usize = 256 * 1024; + +/// Returns a sorted list of ranges that cover `ranges`, merging ranges whose +/// gap is at most `gap`. +/// +/// This replicates `object_store::coalesce_ranges` (`merge_ranges`) so the +/// fetch plan matches what a `get_byte_ranges` implementation backed by +/// `ObjectStore::get_ranges` would issue. +fn merge_ranges(ranges: &[Range], gap: u64) -> Vec> { + if ranges.is_empty() { + return vec![]; + } + + let mut ranges = ranges.to_vec(); + ranges.sort_unstable_by_key(|range| range.start); + + let mut ret = Vec::with_capacity(ranges.len()); + let mut start_idx = 0; + let mut end_idx = 1; + + while start_idx != ranges.len() { + let mut range_end = ranges[start_idx].end; + + while end_idx != ranges.len() + && ranges[end_idx] + .start + .checked_sub(range_end) + .map(|delta| delta <= gap) + .unwrap_or(true) + { + range_end = range_end.max(ranges[end_idx].end); + end_idx += 1; + } + + ret.push(ranges[start_idx].start..range_end); + + start_idx = end_idx; + end_idx += 1; + } + + ret +} + +fn key(range: &Range) -> (u64, u64) { + (range.start, range.end) +} + +/// Drains `body` into `tx`, aggregating small chunks up to +/// [`TARGET_MESSAGE_BYTES`]. Exits when the body ends, errors, or the +/// receiver is dropped (checked even while awaiting the body, so dropping the +/// fetcher promptly drops the body and releases its connection). The bounded +/// channel provides backpressure: when the fetcher stops consuming, this task +/// parks on `send` and, transitively, the server stalls via flow control. +async fn pump_body(mut body: BoxStream<'static, Result>, tx: mpsc::Sender>) { + let mut pending: Vec = Vec::new(); + let mut pending_bytes = 0usize; + loop { + let item = tokio::select! { + biased; + _ = tx.closed() => return, + item = body.next() => item, + }; + match item { + Some(Ok(bytes)) => { + pending_bytes += bytes.len(); + pending.push(bytes); + if pending_bytes >= TARGET_MESSAGE_BYTES { + let message = concat_bytes(std::mem::take(&mut pending), pending_bytes); + pending_bytes = 0; + if tx.send(Ok(message)).await.is_err() { + return; + } + } + } + Some(Err(e)) => { + let _ = tx.send(Err(e)).await; + return; + } + None => break, + } + } + if !pending.is_empty() { + let message = concat_bytes(pending, pending_bytes); + let _ = tx.send(Ok(message)).await; + } +} + +fn concat_bytes(mut parts: Vec, total: usize) -> Bytes { + if parts.len() == 1 { + return parts.pop().expect("checked len"); + } + let mut out = Vec::with_capacity(total); + for part in parts { + out.extend_from_slice(&part); + } + out.into() +} + +/// A single in-flight streamed request +#[derive(Debug)] +struct ActiveStream { + /// Requested ranges this stream will satisfy, ascending, disjoint. + /// Members are removed once delivered to the decoder. + members: VecDeque>, + /// Members ending at or below this offset (small ranges preceding the + /// streamed chunk) are delivered as soon as their bytes arrive; members + /// above it are delivered only once the decoder asks for them. + eager_end: u64, + /// Absolute offset of the next byte the body will produce + watermark: u64, + /// Received, not-yet-delivered bytes: (absolute start offset, data) + segments: VecDeque<(u64, Bytes)>, + /// Dropping this receiver makes the pump task exit promptly (it selects + /// on the channel being closed), dropping the body mid-transfer + rx: mpsc::Receiver>, + /// The body has ended + finished: bool, + /// The body errored; surfaced when this stream's data is next needed + error: Option, +} + +impl ActiveStream { + /// True if an undelivered, currently-deliverable member still needs bytes + /// beyond the watermark + fn needs_progress(&self, wanted: &HashSet<(u64, u64)>) -> bool { + !self.finished + && self.error.is_none() + && self.members.iter().any(|m| { + m.end > self.watermark && (m.end <= self.eager_end || wanted.contains(&key(m))) + }) + } + + /// True if the decoder is currently waiting on data from this stream + fn is_needed_now(&self, wanted: &HashSet<(u64, u64)>) -> bool { + self.members.iter().any(|m| wanted.contains(&key(m))) + } + + /// Drop buffered bytes below `pos` (gap bytes between requested ranges) + fn drop_below(&mut self, pos: u64) { + while let Some((start, bytes)) = self.segments.front_mut() { + let end = *start + bytes.len() as u64; + if end <= pos { + self.segments.pop_front(); + } else if *start < pos { + let cut = (pos - *start) as usize; + *bytes = bytes.slice(cut..); + *start = pos; + break; + } else { + break; + } + } + } + + /// Assemble the bytes for `member`, which must be fully covered by the + /// buffered segments + fn extract(&mut self, member: &Range) -> Result { + self.drop_below(member.start); + let len = (member.end - member.start) as usize; + + // single segment covering the whole member: zero copy + if let Some((start, bytes)) = self.segments.front() { + if *start <= member.start && *start + bytes.len() as u64 >= member.end { + let offset = (member.start - *start) as usize; + let out = bytes.slice(offset..offset + len); + self.drop_below(member.end); + return Ok(out); + } + } + + let mut out = Vec::with_capacity(len); + for (start, bytes) in self.segments.iter() { + if *start >= member.end { + break; + } + let seg_end = *start + bytes.len() as u64; + let from = member.start.max(*start); + let to = member.end.min(seg_end); + if from < to { + out.extend_from_slice(&bytes[(from - *start) as usize..(to - *start) as usize]); + } + } + if out.len() != len { + return Err(ParquetError::General(format!( + "streamed request ended before covering range {}..{} ({} of {} bytes available)", + member.start, + member.end, + out.len(), + len + ))); + } + self.drop_below(member.end); + Ok(out.into()) + } +} + +/// Specification of a streamed request to open +struct StreamSpec { + span: Range, + tx: mpsc::Sender>, +} + +/// Resolves to (input, materialized ranges, their bytes) +type IssueFuture = BoxFuture<'static, Result<(T, Vec>, Vec)>>; + +/// Ownership of the underlying [`AsyncFileReader`]. It is moved into a future +/// while requests are being issued (`get_bytes_stream` / `get_byte_ranges` +/// borrow the reader), and returned when the future resolves. +enum InputState { + Idle(T), + Busy(IssueFuture), +} + +impl std::fmt::Debug for InputState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Idle(_) => write!(f, "InputState::Idle"), + Self::Busy(_) => write!(f, "InputState::Busy"), + } + } +} + +/// Adaptive fetcher for the bounded streaming mode: fetches small coalesced +/// ranges exactly like [`AsyncFileReader::get_byte_ranges`], and consumes +/// large spans incrementally via [`AsyncFileReader::get_bytes_stream`] with +/// bounded buffering. +/// +/// Used internally by [`ParquetRecordBatchStream`] when +/// [`BoundedStreamingOptions`] are set. External drivers of +/// [`ParquetPushDecoder`] (systems that implement their own fetch loops) can +/// use it directly via [`Self::try_new`] and [`Self::fetch_more`]: +/// +/// ```ignore +/// let mut fetcher = AdaptiveFetcher::try_new(options, reader)?; +/// loop { +/// match decoder.try_decode()? { +/// DecodeResult::NeedsData(ranges) => { +/// let plan = decoder.upcoming_fetch_plan(); +/// let (ranges, data) = fetcher.fetch_more(ranges, plan).await?; +/// decoder.push_ranges(ranges, data)?; +/// } +/// DecodeResult::Data(batch) => { /* ... */ } +/// DecodeResult::Finished => break, +/// } +/// } +/// ``` +/// +/// [`ParquetRecordBatchStream`]: super::ParquetRecordBatchStream +/// [`ParquetPushDecoder`]: crate::arrow::push_decoder::ParquetPushDecoder +#[derive(Debug)] +pub struct AdaptiveFetcher { + options: BoundedStreamingOptions, + runtime: Handle, + input: InputState, + streams: Vec, + /// Ranges assigned to an in-flight fetch (stream member or materialize + /// batch); removed on delivery + planned: HashSet<(u64, u64)>, + /// Ranges the decoder needs to make progress now + wanted: HashSet<(u64, u64)>, + /// Assembled data not yet pushed to the decoder + ready_ranges: Vec>, + ready_data: Vec, +} + +impl AdaptiveFetcher +where + T: AsyncFileReader + Send + 'static, +{ + /// Create a new fetcher over `input`, spawning body-draining tasks on the + /// current tokio runtime. + /// + /// # Errors + /// + /// Returns an error if called outside a tokio runtime context. To supply + /// an explicit runtime, use [`Self::new`]. + pub fn try_new(options: BoundedStreamingOptions, input: T) -> Result { + let runtime = Handle::try_current().map_err(|e| { + ParquetError::General(format!( + "AdaptiveFetcher requires a tokio runtime to drive streamed request bodies: {e}" + )) + })?; + Ok(Self::new(options, runtime, input)) + } + + /// Fetch (at least) the data the decoder needs to make progress now. + /// + /// `needs` are the ranges from [`DecodeResult::NeedsData`] and `plan` the + /// corresponding [`ParquetPushDecoder::upcoming_fetch_plan`]. Returns + /// (ranges, data) pairs ready to be passed to + /// [`ParquetPushDecoder::push_ranges`]; call again with the decoder's next + /// request if it still needs more. + /// + /// [`DecodeResult::NeedsData`]: crate::DecodeResult::NeedsData + /// [`ParquetPushDecoder`]: crate::arrow::push_decoder::ParquetPushDecoder + /// [`ParquetPushDecoder::push_ranges`]: crate::arrow::push_decoder::ParquetPushDecoder::push_ranges + /// [`ParquetPushDecoder::upcoming_fetch_plan`]: crate::arrow::push_decoder::ParquetPushDecoder::upcoming_fetch_plan + pub async fn fetch_more( + &mut self, + needs: Vec>, + plan: UpcomingFetchPlan, + ) -> Result<(Vec>, Vec)> { + self.want(&needs); + std::future::poll_fn(|cx| { + loop { + if let Err(e) = self.plan_and_issue(&plan) { + return Poll::Ready(Err(e)); + } + let progressed = match self.poll_sources(cx) { + Ok(progressed) => progressed, + Err(e) => return Poll::Ready(Err(e)), + }; + if self.has_ready() { + return Poll::Ready(Ok(self.take_ready())); + } + if !progressed { + return Poll::Pending; + } + } + }) + .await + } + + /// Create a new fetcher over `input`. `runtime` is used to spawn + /// body-draining tasks. + pub fn new(options: BoundedStreamingOptions, runtime: Handle, input: T) -> Self { + Self { + options, + runtime, + input: InputState::Idle(input), + streams: Vec::new(), + planned: HashSet::new(), + wanted: HashSet::new(), + ready_ranges: Vec::new(), + ready_data: Vec::new(), + } + } + + pub(crate) fn has_ready(&self) -> bool { + !self.ready_ranges.is_empty() + } + + /// Take the assembled (range, bytes) pairs, to push into the decoder + pub(crate) fn take_ready(&mut self) -> (Vec>, Vec) { + ( + std::mem::take(&mut self.ready_ranges), + std::mem::take(&mut self.ready_data), + ) + } + + /// Record the ranges the decoder needs to make progress now + pub(crate) fn want(&mut self, ranges: &[Range]) { + for range in ranges { + self.wanted.insert(key(range)); + } + } + + /// Plan and issue requests for any not-yet-planned upcoming ranges. + /// + /// The plan first reproduces the coalesced request set that + /// `get_byte_ranges` would issue (see [`merge_ranges`]). Coalesced groups + /// that contain no streamable span are materialized through a single + /// `get_byte_ranges` call, exactly as the non-adaptive path. Groups + /// containing streamable spans (advertised by the decoder via + /// [`UpcomingFetchPlan::streamable`]) are split into one streamed request + /// per streamable span; each stream also carries the small ranges + /// preceding its span, and any ranges after the last span go to the + /// materialize batch. + pub(crate) fn plan_and_issue(&mut self, plan: &UpcomingFetchPlan) -> Result<()> { + // Drop streams whose remaining members are no longer needed (e.g. the + // decoder hit its row limit and skipped the rest of a row group). + // Aborts their body tasks. + let upcoming: HashSet<(u64, u64)> = plan.ranges.iter().map(key).collect(); + for stream in &mut self.streams { + let wanted = &self.wanted; + let planned = &mut self.planned; + stream.members.retain(|m| { + let keep = upcoming.contains(&key(m)) || wanted.contains(&key(m)); + if !keep { + planned.remove(&key(m)); + } + keep + }); + } + self.streams.retain(|s| !s.members.is_empty()); + + let unplanned: Vec> = plan + .ranges + .iter() + .filter(|r| !self.planned.contains(&key(r))) + .cloned() + .collect(); + if unplanned.is_empty() { + return Ok(()); + } + + // Streamed requests can be opened at any time (they own their own + // connections), but the materialize batch borrows the input, so it + // can only be issued while the input is idle. If a materialize + // request is already in flight, plan again once it resolves + // (`poll_sources` returns progress when it does). + if matches!(self.input, InputState::Busy(_)) { + return Ok(()); + } + + let mut materialize: Vec> = Vec::new(); + let mut specs: Vec = Vec::new(); + + let capacity = usize::try_from(self.options.window_bytes) + .unwrap_or(usize::MAX) + .div_ceil(TARGET_MESSAGE_BYTES) + .clamp(2, 64); + + for group in merge_ranges(&unplanned, self.options.coalesce_gap) { + // members of this coalesced group, ascending + let mut members: VecDeque> = unplanned + .iter() + .filter(|r| r.start >= group.start && r.end <= group.end) + .cloned() + .collect(); + members.make_contiguous().sort_unstable_by_key(|r| r.start); + + // streamable spans intersecting this group, large enough to be + // worth a dedicated request + let mut cuts: Vec> = plan + .streamable + .iter() + .filter_map(|s| { + let start = s.start.max(group.start); + let end = s.end.min(group.end); + (end > start && end - start >= self.options.stream_threshold) + .then_some(start..end) + }) + .collect(); + cuts.sort_unstable_by_key(|r| r.start); + + if cuts.is_empty() { + materialize.extend(members); + continue; + } + + for cut in cuts { + let mut stream_members = VecDeque::new(); + while let Some(front) = members.front() { + if front.start < cut.end { + stream_members.push_back(members.pop_front().expect("front exists")); + } else { + break; + } + } + if stream_members.is_empty() { + continue; + } + let span = stream_members.front().expect("non-empty").start + ..stream_members.back().expect("non-empty").end; + let (tx, rx) = mpsc::channel(capacity); + let watermark = span.start; + specs.push(StreamSpec { + span: span.clone(), + tx, + }); + self.streams.push(ActiveStream { + members: stream_members, + eager_end: cut.start, + watermark, + segments: VecDeque::new(), + rx, + finished: false, + error: None, + }); + } + // anything after the last streamable span + materialize.extend(members); + } + + let InputState::Idle(mut input) = std::mem::replace( + &mut self.input, + InputState::Busy(futures::future::pending().boxed()), + ) else { + unreachable!("checked idle above"); + }; + + // Open all streamed requests concurrently: each open future is + // 'static (owns its connection), so it is spawned immediately and + // runs in parallel with the other opens and with the materialize + // batch below. Readers that do not support streamed requests + // (`get_bytes_stream` returns None) fall back to materialization. + for spec in specs { + let StreamSpec { span, tx } = spec; + match input.get_bytes_stream(span.clone()) { + Some(open) => { + self.runtime.spawn(async move { + match open.await { + Ok(body) => pump_body(body, tx).await, + Err(e) => { + let _ = tx.send(Err(e)).await; + } + } + }); + } + None => { + // Unsupported: drop the just-registered stream and + // materialize its members instead + let idx = self + .streams + .iter() + .position(|s| s.members.front().map(|m| m.start) == Some(span.start)) + .expect("stream registered above"); + let stream = self.streams.swap_remove(idx); + materialize.extend(stream.members); + } + } + } + + for range in &materialize { + self.planned.insert(key(range)); + } + for stream in &self.streams { + for member in &stream.members { + self.planned.insert(key(member)); + } + } + + if materialize.is_empty() { + // nothing to materialize: the input stays available + self.input = InputState::Idle(input); + return Ok(()); + } + + let future = async move { + let data = input.get_byte_ranges(materialize.clone()).await?; + Ok((input, materialize, data)) + } + .boxed(); + self.input = InputState::Busy(future); + Ok(()) + } + + /// Poll all in-flight requests, assembling any data the decoder is + /// waiting for. Returns `true` if any progress was made; `false` means + /// everything is pending (wakers registered). + pub(crate) fn poll_sources(&mut self, cx: &mut Context<'_>) -> Result { + let mut progressed = false; + + // The materialize batch (and stream opening) + if let InputState::Busy(future) = &mut self.input { + match future.poll_unpin(cx) { + Poll::Ready(Ok((input, ranges, data))) => { + for range in &ranges { + self.planned.remove(&key(range)); + self.wanted.remove(&key(range)); + } + self.ready_ranges.extend(ranges); + self.ready_data.extend(data); + self.input = InputState::Idle(input); + progressed = true; + } + Poll::Ready(Err(e)) => return Err(e), + Poll::Pending => {} + } + } + + // Streamed bodies + let mut failed: Option = None; + for stream in &mut self.streams { + // Surface a stream failure once the decoder actually needs this + // stream's data + if stream.error.is_some() && stream.is_needed_now(&self.wanted) { + failed = Some(stream.error.take().expect("checked above")); + break; + } + + while stream.needs_progress(&self.wanted) { + match stream.rx.poll_recv(cx) { + Poll::Ready(Some(Ok(bytes))) => { + let start = stream.watermark; + stream.watermark += bytes.len() as u64; + stream.segments.push_back((start, bytes)); + progressed = true; + } + Poll::Ready(Some(Err(e))) => { + stream.error = Some(e); + break; + } + Poll::Ready(None) => { + stream.finished = true; + break; + } + Poll::Pending => break, + } + } + + // Deliver members in offset order as they become deliverable + while let Some(member) = stream.members.front().cloned() { + let deliverable = member.end <= stream.watermark + && (member.end <= stream.eager_end || self.wanted.contains(&key(&member))); + if !deliverable { + break; + } + let bytes = stream.extract(&member)?; + stream.members.pop_front(); + self.planned.remove(&key(&member)); + self.wanted.remove(&key(&member)); + self.ready_ranges.push(member); + self.ready_data.push(bytes); + progressed = true; + } + + // A finished body that cannot cover a still-needed member is an + // error surfaced on demand + if stream.finished + && stream.error.is_none() + && stream + .members + .front() + .is_some_and(|m| m.end > stream.watermark) + { + stream.error = Some(ParquetError::EOF( + "streamed request body ended before all requested bytes arrived".to_string(), + )); + } + } + if let Some(e) = failed { + return Err(e); + } + self.streams + .retain(|s| !s.members.is_empty() || s.error.is_some()); + + Ok(progressed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_merge_ranges_matches_object_store_semantics() { + // gap of 1: 0..10 and 11..20 merge, 22..30 does not + let merged = merge_ranges(&[22..30, 0..10, 11..20], 1); + assert_eq!(merged, vec![0..20, 22..30]); + // overlapping ranges merge + let merged = merge_ranges(&[0..10, 5..15], 0); + assert_eq!(merged, vec![0..15]); + assert_eq!(merge_ranges(&[], 10), Vec::>::new()); + } +} diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index 5a0083b7164d..522a08d15b18 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -47,9 +47,13 @@ use crate::bloom_filter::{ use crate::errors::{ParquetError, Result}; use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +mod adaptive; mod metadata; +pub use adaptive::AdaptiveFetcher; pub use metadata::*; +use futures::stream::BoxStream; + #[cfg(feature = "object_store")] mod store; @@ -90,6 +94,31 @@ pub trait AsyncFileReader: Send { .boxed() } + /// Retrieve the bytes in `range` as a stream of chunks that can be + /// consumed incrementally, without buffering the whole range in memory, + /// or `None` if this reader does not support streamed requests. + /// + /// Used by the bounded streaming mode (see + /// [`BoundedStreamingOptions`](crate::arrow::arrow_reader::BoundedStreamingOptions)) + /// to consume large column chunks incrementally: the returned future + /// issues the request, and the resulting stream yields the body as it + /// arrives. Both the future and the stream are `'static` (owning their + /// connection) so several streamed requests can be opened concurrently + /// and remain in flight while the reader issues other requests. + /// + /// The default implementation returns `None`, which makes callers fall + /// back to [`Self::get_byte_ranges`] (existing implementors keep working, + /// without the memory benefit). Implementations backed by a network store + /// should return the response body stream directly; see the + /// `ParquetObjectReader` implementation. + fn get_bytes_stream( + &mut self, + range: Range, + ) -> Option>>>> { + let _ = range; + None + } + /// Return a future which results in the [`ParquetMetaData`] for this Parquet file. /// /// This is an asynchronous operation as it may involve reading the file @@ -122,6 +151,13 @@ impl AsyncFileReader for Box { self.as_mut().get_byte_ranges(ranges) } + fn get_bytes_stream( + &mut self, + range: Range, + ) -> Option>>>> { + self.as_mut().get_bytes_stream(range) + } + fn get_metadata<'a>( &'a mut self, options: Option<&'a ArrowReaderOptions>, @@ -589,6 +625,7 @@ impl ParquetRecordBatchStreamBuilder { offset, metrics, max_predicate_cache_size, + bounded_streaming, } = self; // Ensure schema of ParquetRecordBatchStream respects projection, and does @@ -614,15 +651,27 @@ impl ParquetRecordBatchStreamBuilder { offset, metrics, max_predicate_cache_size, + bounded_streaming: bounded_streaming.clone(), } .build()?; - let request_state = RequestState::None { input: input.0 }; + // Adaptive streamed fetching requires a tokio runtime to drive request + // bodies in the background; without one, fall back to the materialized + // path (identical behavior to the option being off) + let (request_state, adaptive) = + match bounded_streaming.zip(tokio::runtime::Handle::try_current().ok()) { + Some((options, runtime)) => ( + RequestState::Done, + Some(AdaptiveFetcher::new(options, runtime, input.0)), + ), + None => (RequestState::None { input: input.0 }, None), + }; Ok(ParquetRecordBatchStream { schema: projected_schema, decoder, request_state, + adaptive, }) } } @@ -709,10 +758,14 @@ impl std::fmt::Debug for RequestState { pub struct ParquetRecordBatchStream { /// Output schema of the stream schema: SchemaRef, - /// Input and Outstanding IO request, if any + /// Input and Outstanding IO request, if any (owns the input unless + /// adaptive streaming is active) request_state: RequestState, /// Decoding state machine (no IO) decoder: ParquetPushDecoder, + /// When bounded streaming is enabled, owns the input and all in-flight + /// requests instead of `request_state` + adaptive: Option>, } impl std::fmt::Debug for ParquetRecordBatchStream { @@ -750,7 +803,38 @@ where /// - `Ok(None)` if the stream has ended. /// - `Err(error)` if the stream has errored. All subsequent calls will return `Ok(None)`. /// - `Ok(Some(reader))` which holds all the data for the row group. + /// + /// Note: when bounded streaming is enabled (see + /// [`BoundedStreamingOptions`](crate::arrow::arrow_reader::BoundedStreamingOptions)), + /// row groups with large column chunks are decoded in bounded windows and + /// this method returns one reader per *window* rather than per row group. + /// Each returned reader still owns all the data it needs. pub async fn next_row_group(&mut self) -> Result> { + if self.adaptive.is_some() { + let result = std::future::poll_fn(|cx| { + let Self { + decoder, adaptive, .. + } = self; + let fetcher = adaptive.as_mut().expect("checked above"); + match Self::poll_adaptive(decoder, fetcher, cx, |d| d.try_next_reader()) { + Ok(poll) => poll.map(Ok), + Err(e) => Poll::Ready(Err(e)), + } + }) + .await; + return match result { + Ok(DecodeResult::Data(reader)) => Ok(Some(reader)), + Ok(DecodeResult::Finished) => Ok(None), + Ok(DecodeResult::NeedsData(_)) => Err(ParquetError::General( + "Internal Error: adaptive poll returned NeedsData".to_string(), + )), + Err(e) => { + // drop in-flight requests and stop + self.adaptive = None; + Err(e) + } + }; + } loop { // Take ownership of request state to process, leaving self in a // valid state @@ -800,6 +884,8 @@ where } Err(e) => { self.request_state = RequestState::Done; + // drop any in-flight adaptive requests + self.adaptive = None; Poll::Ready(Some(Err(e))) } } @@ -810,11 +896,65 @@ impl ParquetRecordBatchStream where T: AsyncFileReader + Unpin + Send + 'static, { + /// Drive the decoder with the adaptive fetcher until it produces a result + /// or everything is pending. + /// + /// `decode` is the decoder entry point to drive ([`ParquetPushDecoder::try_decode`] + /// or [`ParquetPushDecoder::try_next_reader`]). Never returns + /// `Poll::Ready(DecodeResult::NeedsData)`. + fn poll_adaptive( + decoder: &mut ParquetPushDecoder, + fetcher: &mut AdaptiveFetcher, + cx: &mut Context<'_>, + decode: impl Fn(&mut ParquetPushDecoder) -> Result>, + ) -> Result>> { + loop { + if fetcher.has_ready() { + let (ranges, data) = fetcher.take_ready(); + decoder.push_ranges(ranges, data)?; + } + match decode(decoder)? { + DecodeResult::NeedsData(now) => { + fetcher.want(&now); + loop { + // (Re)plan: a no-op when everything needed is already + // in flight, otherwise issues the new requests (which + // may have to wait for the input to become idle) + let plan = decoder.upcoming_fetch_plan(); + fetcher.plan_and_issue(&plan)?; + let progressed = fetcher.poll_sources(cx)?; + if fetcher.has_ready() { + break; // outer loop pushes the data and decodes + } + if !progressed { + return Ok(Poll::Pending); + } + } + } + other => return Ok(Poll::Ready(other)), + } + } + } + /// Inner state machine /// /// Note this is separate from poll_next so we can use ? operator to check for errors /// as it returns `Result>>` fn poll_next_inner(&mut self, cx: &mut Context<'_>) -> Result>> { + if self.adaptive.is_some() { + let Self { + decoder, adaptive, .. + } = self; + let fetcher = adaptive.as_mut().expect("checked above"); + return match Self::poll_adaptive(decoder, fetcher, cx, |d| d.try_decode())? { + Poll::Ready(DecodeResult::Data(batch)) => Ok(Poll::Ready(Some(batch))), + Poll::Ready(DecodeResult::Finished) => Ok(Poll::Ready(None)), + Poll::Ready(DecodeResult::NeedsData(_)) => Err(ParquetError::General( + "Internal Error: adaptive poll returned NeedsData".to_string(), + )), + Poll::Pending => Ok(Poll::Pending), + }; + } loop { let request_state = std::mem::replace(&mut self.request_state, RequestState::Done); match request_state { @@ -1996,4 +2136,338 @@ mod tests { Ok(()) } + + /// [`AsyncFileReader`] over in-memory data with a chunked + /// [`AsyncFileReader::get_bytes_stream`], recording all requests + #[derive(Clone)] + struct StreamingTestReader { + data: Bytes, + /// chunk size for streamed responses + chunk_size: usize, + /// each get_byte_ranges call's ranges + range_requests: Arc>>>>, + /// each get_bytes_stream span + stream_requests: Arc>>>, + } + + impl StreamingTestReader { + fn new(data: Bytes) -> Self { + Self { + data, + chunk_size: 1024, + range_requests: Default::default(), + stream_requests: Default::default(), + } + } + } + + impl AsyncFileReader for StreamingTestReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.range_requests + .lock() + .unwrap() + .push(vec![range.clone()]); + futures::future::ready(Ok(self + .data + .slice(range.start as usize..range.end as usize))) + .boxed() + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, Result>> { + self.range_requests.lock().unwrap().push(ranges.clone()); + let result = ranges + .iter() + .map(|r| self.data.slice(r.start as usize..r.end as usize)) + .collect(); + futures::future::ready(Ok(result)).boxed() + } + + fn get_bytes_stream( + &mut self, + range: Range, + ) -> Option>>>> { + self.stream_requests.lock().unwrap().push(range.clone()); + let data = self.data.slice(range.start as usize..range.end as usize); + let chunk_size = self.chunk_size; + let chunks: Vec> = (0..data.len()) + .step_by(chunk_size) + .map(|offset| Ok(data.slice(offset..(offset + chunk_size).min(data.len())))) + .collect(); + Some(futures::future::ready(Ok(futures::stream::iter(chunks).boxed())).boxed()) + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let mut metadata_reader = ParquetMetaDataReader::new(); + if let Some(opts) = options { + metadata_reader = metadata_reader + .with_column_index_policy(opts.column_index_policy()) + .with_offset_index_policy(opts.offset_index_policy()); + } + let metadata = Arc::new(metadata_reader.parse_and_finish(&self.data).unwrap()); + futures::future::ready(Ok(metadata)).boxed() + } + } + + /// A file with skewed column sizes: tiny int, small strings, large binary + /// values; several row groups, small pages + fn skewed_test_file() -> Bytes { + let mut rand = rng(); + let mut buf = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(500)) + .set_write_batch_size(50) + .set_data_page_size_limit(16 * 1024) + .build(); + + let ids = Int32Array::from_iter(0..1500); + let smalls = StringArray::from_iter_values((0..1500).map(|i| format!("string-value-{i}"))); + let bigs = arrow_array::BinaryArray::from_iter_values((0..1500).map(|i| { + let len = 2048 + (i % 7) * 331; + (0..len) + .map(|j| ((i * 31 + j * 7 + rand.random_range(0..13)) % 251) as u8) + .collect::>() + })); + let batch = RecordBatch::try_from_iter([ + ("i", Arc::new(ids) as ArrayRef), + ("s", Arc::new(smalls) as ArrayRef), + ("b", Arc::new(bigs) as ArrayRef), + ]) + .unwrap(); + + let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + buf.into() + } + + async fn read_all_with_options( + data: Bytes, + streaming: Option, + projection: Option>, + batch_size: usize, + selection: Option, + limit_offset: Option<(usize, usize)>, + ) -> (Vec, StreamingTestReader) { + let reader = StreamingTestReader::new(data); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let mut builder = + ParquetRecordBatchStreamBuilder::new_with_options(reader.clone(), options) + .await + .unwrap() + .with_batch_size(batch_size); + if let Some(streaming) = streaming { + builder = builder.with_bounded_streaming(streaming); + } + if let Some(projection) = projection { + let mask = ProjectionMask::leaves(builder.parquet_schema(), projection); + builder = builder.with_projection(mask); + } + if let Some(selection) = selection { + builder = builder.with_row_selection(selection); + } + if let Some((offset, limit)) = limit_offset { + builder = builder.with_offset(offset).with_limit(limit); + } + let stream = builder.build().unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + (batches, reader) + } + + fn test_streaming_options() -> crate::arrow::arrow_reader::BoundedStreamingOptions { + crate::arrow::arrow_reader::BoundedStreamingOptions { + stream_threshold: 256 * 1024, + window_bytes: 128 * 1024, + coalesce_gap: 1024 * 1024, + min_window_bytes_per_column: 16 * 1024, + } + } + + #[tokio::test] + async fn test_bounded_streaming_byte_identical_full_scan() { + let data = skewed_test_file(); + + let (baseline, baseline_reader) = + read_all_with_options(data.clone(), None, None, 100, None, None).await; + let (streamed, streamed_reader) = read_all_with_options( + data.clone(), + Some(test_streaming_options()), + None, + 100, + None, + None, + ) + .await; + + // exact same batches, including batch boundaries + assert_eq!(baseline, streamed); + + // the large binary chunks were streamed: one streamed request per row + // group (columns are adjacent so each row group is a single coalesced + // group ending in the large chunk) + let streams = streamed_reader.stream_requests.lock().unwrap(); + assert_eq!(streams.len(), 3, "streams: {streams:?}"); + + // request count: baseline issues one call per row group; adaptive + // replaces each with a single streamed request and no materialized + // requests + let baseline_calls = baseline_reader.range_requests.lock().unwrap().len(); + let streamed_calls = streamed_reader.range_requests.lock().unwrap().len(); + let streamed_materialized: usize = streamed_reader + .range_requests + .lock() + .unwrap() + .iter() + .map(|c| c.len()) + .sum(); + assert_eq!(baseline_calls, 3); + assert_eq!(streamed_materialized, 0, "everything rides the streams"); + assert_eq!(streams.len() + streamed_calls, 3); + } + + #[tokio::test] + async fn test_bounded_streaming_mixed_projection_with_gap() { + let data = skewed_test_file(); + // project only the tiny and huge columns; use a small coalesce gap so + // the unprojected string column splits them into separate groups + let options = crate::arrow::arrow_reader::BoundedStreamingOptions { + coalesce_gap: 1024, + ..test_streaming_options() + }; + + let (baseline, _) = + read_all_with_options(data.clone(), None, Some(vec![0, 2]), 333, None, None).await; + let (streamed, reader) = read_all_with_options( + data.clone(), + Some(options), + Some(vec![0, 2]), + 333, + None, + None, + ) + .await; + assert_eq!(baseline, streamed); + + // the tiny int column materializes, the big binary column streams + let streams = reader.stream_requests.lock().unwrap(); + assert_eq!(streams.len(), 3); + let materialized: Vec>> = reader.range_requests.lock().unwrap().clone(); + assert_eq!(materialized.len(), 3, "one materialize call per row group"); + // no materialized range overlaps a streamed span + for call in &materialized { + for range in call { + for span in streams.iter() { + assert!( + range.end <= span.start || range.start >= span.end, + "materialized {range:?} overlaps streamed {span:?}" + ); + } + } + } + } + + #[tokio::test] + async fn test_bounded_streaming_with_row_selection_fuzz() { + let data = skewed_test_file(); + let mut rand = rng(); + + for _ in 0..10 { + let mut selectors = vec![]; + let mut total_rows = 0; + let mut skip = false; + while total_rows < 1500 { + let row_count: usize = rand.random_range(1..320).min(1500 - total_rows); + selectors.push(RowSelector { row_count, skip }); + total_rows += row_count; + skip = !skip; + } + let selection = RowSelection::from(selectors); + let batch_size = rand.random_range(16..600); + let projection: Vec = match rand.random_range(0..3) { + 0 => vec![0, 1, 2], + 1 => vec![0, 2], + _ => vec![2], + }; + // small enough that streaming engages for at least full-ish + // selections + let options = crate::arrow::arrow_reader::BoundedStreamingOptions { + stream_threshold: 64 * 1024, + window_bytes: 32 * 1024, + coalesce_gap: 4 * 1024, + min_window_bytes_per_column: 8 * 1024, + }; + + let (baseline, _) = read_all_with_options( + data.clone(), + None, + Some(projection.clone()), + batch_size, + Some(selection.clone()), + None, + ) + .await; + let (streamed, _) = read_all_with_options( + data.clone(), + Some(options), + Some(projection), + batch_size, + Some(selection), + None, + ) + .await; + assert_eq!(baseline, streamed); + } + } + + #[tokio::test] + async fn test_bounded_streaming_offset_limit() { + let data = skewed_test_file(); + for (offset, limit) in [(0, 100), (700, 600), (499, 2), (1400, 1000)] { + let (baseline, _) = + read_all_with_options(data.clone(), None, None, 128, None, Some((offset, limit))) + .await; + let (streamed, _) = read_all_with_options( + data.clone(), + Some(test_streaming_options()), + None, + 128, + None, + Some((offset, limit)), + ) + .await; + assert_eq!(baseline, streamed, "offset={offset} limit={limit}"); + } + } + + #[tokio::test] + async fn test_bounded_streaming_cancel_mid_stream() { + let data = skewed_test_file(); + let reader = StreamingTestReader::new(data); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let mut stream = ParquetRecordBatchStreamBuilder::new_with_options(reader.clone(), options) + .await + .unwrap() + .with_batch_size(100) + .with_bounded_streaming(crate::arrow::arrow_reader::BoundedStreamingOptions { + stream_threshold: 64 * 1024, + window_bytes: 32 * 1024, + coalesce_gap: 1024 * 1024, + min_window_bytes_per_column: 8 * 1024, + }) + .build() + .unwrap(); + + // read one batch then drop mid row group + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 100); + assert!(!reader.stream_requests.lock().unwrap().is_empty()); + drop(stream); + // give the pump tasks a chance to observe the closed channels + tokio::task::yield_now().await; + } } diff --git a/parquet/src/arrow/async_reader/store.rs b/parquet/src/arrow/async_reader/store.rs index d47ca744d8f6..bd550ddd895a 100644 --- a/parquet/src/arrow/async_reader/store.rs +++ b/parquet/src/arrow/async_reader/store.rs @@ -196,6 +196,31 @@ impl AsyncFileReader for ParquetObjectReader { self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed()) } + fn get_bytes_stream( + &mut self, + range: Range, + ) -> Option>>>> + { + use futures::{StreamExt as _, TryStreamExt}; + let store = Arc::clone(&self.store); + let path = self.path.clone(); + let options = GetOptions { + range: Some(GetRange::Bounded(range)), + ..Default::default() + }; + Some( + async move { + let result = store.get_opts(&path, options).await?; + // Note: for file-backed stores this streams in small (8KiB) + // chunks; wrap the store in `object_store::chunked::ChunkedStore` + // for larger chunks if that matters + let stream = result.into_stream().map_err(ParquetError::from).boxed(); + Ok::<_, ParquetError>(stream) + } + .boxed(), + ) + } + // This method doesn't directly call `self.spawn` because all of the IO that is done down the // line due to this method call is done through `self.get_bytes` and/or `self.get_byte_ranges`. // When `self` is passed into `ParquetMetaDataReader::load_and_finish`, it treats it as diff --git a/parquet/src/arrow/in_memory_row_group.rs b/parquet/src/arrow/in_memory_row_group.rs index 6c5f013159d5..92cd2d3ef80c 100644 --- a/parquet/src/arrow/in_memory_row_group.rs +++ b/parquet/src/arrow/in_memory_row_group.rs @@ -18,14 +18,102 @@ use crate::arrow::ProjectionMask; use crate::arrow::array_reader::RowGroups; use crate::arrow::arrow_reader::RowSelection; -use crate::column::page::{PageIterator, PageReader}; +use crate::column::page::{Page, PageIterator, PageMetadata, PageReader}; use crate::errors::ParquetError; +use crate::errors::Result; use crate::file::metadata::{ParquetMetaData, RowGroupMetaData}; use crate::file::page_index::offset_index::OffsetIndexMetaData; use crate::file::reader::{ChunkReader, Length, SerializedPageReader}; use bytes::{Buf, Bytes}; +use std::collections::HashMap; use std::ops::Range; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; + +/// Decompressed dictionary pages shared between the decode windows of one row +/// group (see `BoundedStreamingOptions`), keyed by column index. +/// +/// Each window builds fresh column readers, whose page readers would otherwise +/// re-read and re-decompress the chunk's dictionary page for every window. +/// The first window to read a column's dictionary page stores the decompressed +/// [`Page`] here; later windows skip the page in the underlying reader (free +/// when page locations are available) and reuse the cached one. Dropped with +/// the row group's windowed-read state, so at most one decompressed dictionary +/// page per projected column is retained at a time. +#[derive(Debug, Default)] +pub(crate) struct DictionaryPageCache { + pages: Mutex>, +} + +impl DictionaryPageCache { + fn get(&self, column: usize) -> Option { + self.pages.lock().unwrap().get(&column).cloned() + } + + fn put(&self, column: usize, page: Page) { + self.pages.lock().unwrap().insert(column, page); + } +} + +/// A [`PageReader`] adapter that reuses a cached decompressed dictionary page +/// (see [`DictionaryPageCache`]) instead of re-reading it from the chunk data. +struct CachedDictionaryPageReader { + inner: R, + cache: Arc, + column: usize, + /// Whether the (potential) dictionary page at the start of the chunk has + /// been handled + dictionary_handled: bool, +} + +impl Iterator for CachedDictionaryPageReader { + type Item = Result; + + fn next(&mut self) -> Option { + self.get_next_page().transpose() + } +} + +impl PageReader for CachedDictionaryPageReader { + fn get_next_page(&mut self) -> Result> { + if !self.dictionary_handled { + self.dictionary_handled = true; + let next_is_dictionary = matches!( + self.inner.peek_next_page()?, + Some(PageMetadata { is_dict: true, .. }) + ); + if next_is_dictionary { + if let Some(page) = self.cache.get(self.column) { + // skip the underlying dictionary page without reading it + // (with page locations this only drops its metadata) + self.inner.skip_next_page()?; + return Ok(Some(page)); + } + let page = self.inner.get_next_page()?; + if let Some(page) = &page { + self.cache.put(self.column, page.clone()); + } + return Ok(page); + } + } + self.inner.get_next_page() + } + + fn peek_next_page(&mut self) -> Result> { + self.inner.peek_next_page() + } + + fn skip_next_page(&mut self) -> Result<()> { + if !self.dictionary_handled { + // a skipped dictionary page needs no caching + self.dictionary_handled = true; + } + self.inner.skip_next_page() + } + + fn at_record_boundary(&mut self) -> Result { + self.inner.at_record_boundary() + } +} /// An in-memory collection of column chunks #[derive(Debug)] @@ -36,6 +124,9 @@ pub(crate) struct InMemoryRowGroup<'a> { pub(crate) row_count: usize, pub(crate) row_group_idx: usize, pub(crate) metadata: &'a ParquetMetaData, + /// If set, dictionary pages are cached and shared across the decode + /// windows of this row group + pub(crate) dictionary_page_cache: Option>, } /// What ranges to fetch for the columns in this row group @@ -218,7 +309,15 @@ impl RowGroups for InMemoryRowGroup<'_> { column_chunk_metadata, )?; - let page_reader: Box = Box::new(page_reader); + let page_reader: Box = match &self.dictionary_page_cache { + Some(cache) => Box::new(CachedDictionaryPageReader { + inner: page_reader, + cache: Arc::clone(cache), + column: i, + dictionary_handled: false, + }), + None => Box::new(page_reader), + }; Ok(Box::new(ColumnChunkIterator { reader: Some(Ok(page_reader)), diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs index ba075b4c08b6..1253a14ee47a 100644 --- a/parquet/src/arrow/push_decoder/mod.rs +++ b/parquet/src/arrow/push_decoder/mod.rs @@ -266,6 +266,7 @@ impl ParquetPushDecoderBuilder { metrics, row_selection_policy, max_predicate_cache_size, + bounded_streaming, } = self; // If no row groups were specified, read all of them @@ -288,6 +289,7 @@ impl ParquetPushDecoderBuilder { max_predicate_cache_size, buffers, row_selection_policy, + bounded_streaming, ); // Initialize the decoder with the configured options @@ -331,6 +333,7 @@ fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderB max_predicate_cache_size, metrics, row_selection_policy, + bounded_streaming, buffers, } = reader_builder; @@ -352,6 +355,7 @@ fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderB offset, metrics, max_predicate_cache_size, + bounded_streaming, } // Carry the decoder's already-fetched bytes across the rebuild so the new // decoder does not re-request them. @@ -509,6 +513,12 @@ impl ParquetPushDecoder { self.state.buffered_bytes() } + /// Returns everything the active row group will still need, plus which + /// byte spans will be consumed incrementally. See [`UpcomingFetchPlan`]. + pub fn upcoming_fetch_plan(&self) -> UpcomingFetchPlan { + self.state.upcoming_fetch_plan() + } + /// Clear any staged byte ranges currently buffered for future decode work. /// /// This clears byte ranges still owned by the decoder's internal @@ -617,6 +627,30 @@ impl ParquetPushDecoder { } } +/// The full remaining data need of a [`ParquetPushDecoder`], plus which byte +/// spans are suitable for incremental (streamed) consumption. +/// +/// Returned by [`ParquetPushDecoder::upcoming_fetch_plan`]. Unlike +/// [`DecodeResult::NeedsData`], which reports only what the decoder needs to +/// make progress *now*, `ranges` reports everything the active row group will +/// still request, so callers can plan coalesced fetches up front. +/// +/// `streamable` lists the full requested byte spans of column chunks that the +/// decoder will consume incrementally (in windows) when bounded streaming is +/// enabled (see +/// [`BoundedStreamingOptions`](crate::arrow::arrow_reader::BoundedStreamingOptions)). +/// Ranges inside a streamable span will be requested gradually, in increasing +/// offset order, so a caller may satisfy them from a single incrementally +/// consumed ranged request. When empty, all ranges will be needed at once to +/// make progress. +#[derive(Debug, Clone, Default)] +pub struct UpcomingFetchPlan { + /// All byte ranges the active row group still needs + pub ranges: Vec>, + /// Requested spans that will be consumed incrementally, in offset order + pub streamable: Vec>, +} + /// Internal state machine for the [`ParquetPushDecoder`] #[derive(Debug)] enum ParquetDecoderState { @@ -824,6 +858,20 @@ impl ParquetDecoderState { } } + /// See [`ParquetPushDecoder::upcoming_fetch_plan`] + fn upcoming_fetch_plan(&self) -> UpcomingFetchPlan { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.upcoming_fetch_plan(), + ParquetDecoderState::DecodingRowGroup { + remaining_row_groups, + .. + } => remaining_row_groups.upcoming_fetch_plan(), + ParquetDecoderState::Finished => UpcomingFetchPlan::default(), + } + } + /// Clear any staged ranges currently buffered in the decoder. fn clear_all_ranges(&mut self) { match self { diff --git a/parquet/src/arrow/push_decoder/reader_builder/data.rs b/parquet/src/arrow/push_decoder/reader_builder/data.rs index 6fbc2090b06e..bf0e00c922db 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/data.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/data.rs @@ -54,6 +54,11 @@ impl DataRequest { .collect() } + /// All ranges this request will read + pub fn ranges(&self) -> &[Range] { + &self.ranges + } + /// Returns the chunks from the buffers that satisfy this request fn get_chunks(&self, buffers: &PushBuffers) -> Result, ParquetError> { self.ranges @@ -83,6 +88,29 @@ impl DataRequest { parquet_metadata: &'a ParquetMetaData, projection: &ProjectionMask, buffers: &mut PushBuffers, + ) -> Result, ParquetError> { + self.try_into_in_memory_row_group_retaining( + row_group_idx, + row_count, + parquet_metadata, + projection, + buffers, + &|_| false, + ) + } + + /// Like [`Self::try_into_in_memory_row_group`], but ranges for which + /// `retain` returns true are kept in `buffers` instead of being cleared + /// (used by windowed decoding, where dictionary pages and window-boundary + /// pages are shared between successive windows). + pub fn try_into_in_memory_row_group_retaining<'a>( + self, + row_group_idx: usize, + row_count: usize, + parquet_metadata: &'a ParquetMetaData, + projection: &ProjectionMask, + buffers: &mut PushBuffers, + retain: &dyn Fn(&Range) -> bool, ) -> Result, ParquetError> { let chunks = self.get_chunks(buffers)?; @@ -101,12 +129,14 @@ impl DataRequest { offset_index: get_offset_index(parquet_metadata, row_group_idx), row_group_idx, metadata: parquet_metadata, + dictionary_page_cache: None, }; in_memory_row_group.fill_column_chunks(projection, page_start_offsets, chunks); - // Clear the ranges that were explicitly requested - buffers.clear_ranges(&ranges); + // Clear the ranges that were explicitly requested, except retained ones + let ranges_to_clear: Vec> = ranges.into_iter().filter(|r| !retain(r)).collect(); + buffers.clear_ranges(&ranges_to_clear); Ok(in_memory_row_group) } @@ -205,6 +235,7 @@ impl<'a> DataRequestBuilder<'a> { offset_index: get_offset_index(parquet_metadata, row_group_idx), row_group_idx, metadata: parquet_metadata, + dictionary_page_cache: None, }; let FetchRanges { diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index b773f95fd55e..a902b80a041f 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -17,18 +17,21 @@ mod data; mod filter; +mod windows; use crate::arrow::ProjectionMask; use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions, RowGroupCache}; use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; use crate::arrow::arrow_reader::selection::{LoadedRowRanges, RowSelectionStrategy}; use crate::arrow::arrow_reader::{ - ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, RowFilter, RowSelection, - RowSelectionPolicy, + BoundedStreamingOptions, ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, + RowFilter, RowSelection, RowSelectionPolicy, }; use crate::arrow::in_memory_row_group::ColumnChunkData; +use crate::arrow::push_decoder::UpcomingFetchPlan; use crate::arrow::push_decoder::reader_builder::data::DataRequestBuilder; use crate::arrow::push_decoder::reader_builder::filter::CacheInfo; +use crate::arrow::push_decoder::reader_builder::windows::{WindowPlan, WindowedRead}; use crate::arrow::schema::ParquetField; use crate::errors::ParquetError; use crate::file::metadata::ParquetMetaData; @@ -38,6 +41,7 @@ use bytes::Bytes; use data::DataRequest; use filter::AdvanceResult; use filter::FilterInfo; +use std::collections::HashSet; use std::ops::Range; use std::sync::{Arc, RwLock}; @@ -84,6 +88,15 @@ enum RowGroupDecoderState { /// Any cached filter results cache_info: Option, }, + /// Decoding the row group in bounded windows (see + /// [`BoundedStreamingOptions`]): one reader is produced per window, and + /// only the front window's data needs to be resident to make progress. + DecodingWindows { + windowed: Box, + /// Budget remaining after this row group (already applied to the + /// windows' selections) + remaining_budget: RowBudget, + }, /// Finished (or not yet started) reading this group Finished, } @@ -263,6 +276,9 @@ pub(crate) struct RowGroupReaderBuilder { /// Strategy for materialising row selections row_selection_policy: RowSelectionPolicy, + /// If set, decode row groups with large column chunks in bounded windows + bounded_streaming: Option, + /// Current state of the decoder. /// /// It is taken when processing, and must be put back before returning @@ -287,6 +303,7 @@ pub(crate) struct RowGroupReaderBuilderParts { pub max_predicate_cache_size: usize, pub metrics: ArrowReaderMetrics, pub row_selection_policy: RowSelectionPolicy, + pub bounded_streaming: Option, /// Bytes already pushed into the decoder, carried across a rebuild so they /// are not re-requested. pub buffers: PushBuffers, @@ -305,6 +322,7 @@ impl RowGroupReaderBuilder { max_predicate_cache_size: usize, buffers: PushBuffers, row_selection_policy: RowSelectionPolicy, + bounded_streaming: Option, ) -> Self { Self { batch_size, @@ -315,6 +333,7 @@ impl RowGroupReaderBuilder { metrics, max_predicate_cache_size, row_selection_policy, + bounded_streaming, state: Some(RowGroupDecoderState::Finished), buffers, } @@ -335,6 +354,7 @@ impl RowGroupReaderBuilder { max_predicate_cache_size, metrics, row_selection_policy, + bounded_streaming, state: _, buffers, } = self; @@ -346,6 +366,7 @@ impl RowGroupReaderBuilder { max_predicate_cache_size, metrics, row_selection_policy, + bounded_streaming, buffers, } } @@ -710,6 +731,31 @@ impl RowGroupReaderBuilder { )); } + // If bounded streaming is enabled and this row group's + // projection requests any chunk large enough to stream, decode + // it in bounded windows. Not applied when filters ran for this + // row group (`cache_info` / previously read `column_chunks`): + // the predicate cache assumes whole-row-group decode. + if let Some(streaming) = self.bounded_streaming.as_ref() { + if cache_info.is_none() && column_chunks.is_none() { + if let Some(windowed) = windows::plan_windows( + row_group_idx, + row_count, + self.batch_size, + &self.metadata, + &self.projection, + plan_builder.selection(), + self.row_selection_policy, + streaming, + ) { + return Ok(NextState::again(RowGroupDecoderState::DecodingWindows { + windowed: Box::new(windowed), + remaining_budget, + })); + } + } + } + let data_request = DataRequestBuilder::new( row_group_idx, row_count, @@ -805,6 +851,79 @@ impl RowGroupReaderBuilder { }, ) } + RowGroupDecoderState::DecodingWindows { + mut windowed, + remaining_budget, + } => { + let Some(front) = windowed.windows.front() else { + return Ok(NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Finished { remaining_budget }, + )); + }; + + // Only the *front* window's data must be resident to make + // progress; this is the decoder's backpressure signal. The + // full remaining need is available via + // [`RowGroupReaderBuilder::upcoming_fetch_plan`]. + let missing = front.data_request.needed_ranges(&self.buffers); + if !missing.is_empty() { + return Ok(NextState::result( + RowGroupDecoderState::DecodingWindows { + windowed, + remaining_budget, + }, + RowGroupBuildResult::NeedsData(missing), + )); + } + + let WindowPlan { + data_request, + mut plan_builder, + } = windowed.windows.pop_front().expect("checked above"); + + let ranges: Vec> = data_request.ranges().to_vec(); + // Ranges shared with later windows (dictionary pages and + // window-boundary pages) stay buffered until their last use + let retained = windowed.finish_front_window(&ranges); + + let mut row_group = data_request.try_into_in_memory_row_group_retaining( + windowed.row_group_idx, + windowed.row_count, + &self.metadata, + &self.projection, + &mut self.buffers, + &|range| retained.contains_key(&(range.start, range.end)), + )?; + // dictionary pages are decompressed once and shared across + // this row group's windows + row_group.dictionary_page_cache = Some(Arc::clone(&windowed.dictionary_page_cache)); + + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + &self.projection, + self.row_group_offset_index(windowed.row_group_idx), + windowed.row_count, + ); + let plan = plan_builder.build(); + + let array_reader = ArrayReaderBuilder::new(&row_group, &self.metrics) + .with_batch_size(self.batch_size) + .with_parquet_metadata(&self.metadata) + .build_array_reader(self.fields.as_deref(), &self.projection)?; + + let batch_reader = ParquetRecordBatchReader::new(array_reader, plan); + NextState::result( + RowGroupDecoderState::DecodingWindows { + windowed, + remaining_budget, + }, + RowGroupBuildResult::Data { + batch_reader, + remaining_budget, + }, + ) + } RowGroupDecoderState::Finished => { return Err(ParquetError::General(String::from( "Internal Error: try_build called without an active row group", @@ -814,6 +933,40 @@ impl RowGroupReaderBuilder { Ok(result) } + /// Returns everything the active row group will still need, plus which + /// byte spans are suitable for incremental (streamed) consumption. + /// + /// Unlike the `NeedsData` result, which only reports what is needed to + /// make progress *now*, this reports the full remaining need so callers + /// can plan coalesced fetches up front. + pub(crate) fn upcoming_fetch_plan(&self) -> UpcomingFetchPlan { + match self.state.as_ref() { + Some(RowGroupDecoderState::DecodingWindows { windowed, .. }) => { + let mut seen = HashSet::new(); + let mut ranges = Vec::new(); + for window in &windowed.windows { + for range in window.data_request.needed_ranges(&self.buffers) { + if seen.insert((range.start, range.end)) { + ranges.push(range); + } + } + } + UpcomingFetchPlan { + ranges, + streamable: windowed.streamable.clone(), + } + } + Some(RowGroupDecoderState::WaitingOnData { data_request, .. }) + | Some(RowGroupDecoderState::WaitingOnFilterData { data_request, .. }) => { + UpcomingFetchPlan { + ranges: data_request.needed_ranges(&self.buffers), + streamable: Vec::new(), + } + } + _ => UpcomingFetchPlan::default(), + } + } + /// Which columns should be cached? /// /// Returns the columns that are used by the filters *and* then used in the diff --git a/parquet/src/arrow/push_decoder/reader_builder/windows.rs b/parquet/src/arrow/push_decoder/reader_builder/windows.rs new file mode 100644 index 000000000000..bff3656bd742 --- /dev/null +++ b/parquet/src/arrow/push_decoder/reader_builder/windows.rs @@ -0,0 +1,352 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Window planning for bounded streaming decode. +//! +//! A *window* is a contiguous row span of a row group whose boundaries fall on +//! selected-row multiples of `batch_size`. Each window is decoded with the +//! existing sparse-selection machinery (its own [`DataRequest`] and +//! [`ReadPlanBuilder`]), so the sequence of emitted batches is identical to +//! decoding the whole row group at once, while only one window's bytes (plus +//! pinned dictionary/boundary pages) need to be resident at a time. + +use crate::arrow::ProjectionMask; +use crate::arrow::arrow_reader::{ReadPlanBuilder, RowSelection, RowSelectionPolicy, RowSelector}; +use crate::arrow::in_memory_row_group::DictionaryPageCache; +use crate::arrow::push_decoder::reader_builder::data::{DataRequest, DataRequestBuilder}; +use crate::file::metadata::ParquetMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use std::collections::{HashMap, VecDeque}; +use std::ops::Range; +use std::sync::Arc; + +/// One decode window: the data it needs and the plan to decode it. +#[derive(Debug)] +pub(super) struct WindowPlan { + pub data_request: DataRequest, + pub plan_builder: ReadPlanBuilder, +} + +/// State for decoding a row group in bounded windows. +#[derive(Debug)] +pub(super) struct WindowedRead { + pub row_group_idx: usize, + pub row_count: usize, + /// Remaining windows, front is next to decode + pub windows: VecDeque, + /// How many remaining windows still need each range. Ranges shared by + /// multiple windows (dictionary pages, window-boundary pages) are kept in + /// the push buffers until their last use. + range_uses: HashMap<(u64, u64), usize>, + /// Full requested byte spans of the column chunks large enough to be + /// streamed. Advertised to the fetch layer via + /// `ParquetPushDecoder::upcoming_fetch_plan` so it can decide which + /// coalesced requests to consume incrementally. + pub streamable: Vec>, + /// Decompressed dictionary pages shared across this row group's windows, + /// so each is decompressed once rather than once per window + pub dictionary_page_cache: Arc, +} + +impl WindowedRead { + /// Called after the front window's reader has been built: decrement the + /// use count of each of its ranges. Returns a set of ranges that are + /// still needed by later windows (and so must be retained in the push + /// buffers). + pub(super) fn finish_front_window(&mut self, ranges: &[Range]) -> HashMap<(u64, u64), ()> { + let mut retained = HashMap::new(); + for range in ranges { + let key = (range.start, range.end); + if let Some(uses) = self.range_uses.get_mut(&key) { + *uses -= 1; + if *uses == 0 { + self.range_uses.remove(&key); + } else { + retained.insert(key, ()); + } + } + } + retained + } +} + +/// Per-column byte estimator over the offset index. +struct ColumnPages { + /// First row index of each page, ascending, starting at 0 + first_rows: Vec, + /// Prefix sums of compressed page sizes: `prefix[i]` = bytes of pages `0..i` + prefix: Vec, +} + +impl ColumnPages { + fn new(offset_index: &OffsetIndexMetaData) -> Self { + let locations = offset_index.page_locations(); + let mut first_rows = Vec::with_capacity(locations.len()); + let mut prefix = Vec::with_capacity(locations.len() + 1); + prefix.push(0); + let mut total = 0u64; + for location in locations { + first_rows.push(location.first_row_index as u64); + total += location.compressed_page_size as u64; + prefix.push(total); + } + Self { first_rows, prefix } + } + + /// Compressed bytes of the pages overlapping rows `lo..hi`. + /// + /// This intentionally ignores the row selection within the span (an + /// overestimate), which can only make windows smaller than the target. + fn bytes_for_span(&self, lo: u64, hi: u64) -> u64 { + if self.first_rows.is_empty() || lo >= hi { + return 0; + } + // last page whose first row is <= lo + let idx_lo = self.first_rows.partition_point(|&fr| fr <= lo) - 1; + // last page whose first row is < hi + let idx_hi = self.first_rows.partition_point(|&fr| fr < hi) - 1; + self.prefix[idx_hi + 1] - self.prefix[idx_lo] + } +} + +/// Attempt to plan a windowed read of a row group. +/// +/// `selection` is the final (budget-applied) selection for the row group, or +/// `None` if all rows are selected. Returns `None` when windowing does not +/// apply: no offset index, no projected chunk requests at least +/// `stream_threshold` bytes, or fewer than two windows would result. +#[expect(clippy::too_many_arguments)] +pub(super) fn plan_windows( + row_group_idx: usize, + row_count: usize, + batch_size: usize, + parquet_metadata: &ParquetMetaData, + projection: &ProjectionMask, + selection: Option<&RowSelection>, + row_selection_policy: RowSelectionPolicy, + options: &crate::arrow::arrow_reader::BoundedStreamingOptions, +) -> Option { + let stream_threshold = options.stream_threshold; + if batch_size == 0 || row_count == 0 { + return None; + } + let offset_index = parquet_metadata + .offset_index() + .filter(|index| !index.is_empty()) + .map(|index| index[row_group_idx].as_slice())?; + + let row_group_meta = parquet_metadata.row_group(row_group_idx); + let num_columns = row_group_meta.columns().len(); + if offset_index.len() != num_columns { + return None; + } + + // Determine which projected chunks are large enough to stream, based on + // the bytes the current selection will actually request from them. + let mut streamable: Vec> = Vec::new(); + let mut total_requested: u64 = 0; + for (idx, column_offset_index) in offset_index.iter().enumerate() { + if !projection.leaf_included(idx) { + continue; + } + let (start, len) = row_group_meta.column(idx).byte_range(); + let requested: u64 = match selection { + Some(selection) => selection + .scan_ranges(column_offset_index.page_locations()) + .iter() + .map(|r| r.end - r.start) + .sum(), + None => len, + }; + total_requested += requested; + if requested >= stream_threshold { + streamable.push(start..start + len); + } + } + // See BoundedStreamingOptions::min_window_bytes_per_column: windows must + // amortize per-window, per-column fixed costs + let projected_columns = (0..num_columns) + .filter(|idx| projection.leaf_included(*idx)) + .count() as u64; + let window_bytes = options + .window_bytes + .max(projected_columns * options.min_window_bytes_per_column); + if streamable.is_empty() || total_requested <= window_bytes { + return None; + } + + // Per-column byte estimators for the projected columns + let estimators: Vec = (0..num_columns) + .filter(|idx| projection.leaf_included(*idx)) + .map(|idx| ColumnPages::new(&offset_index[idx])) + .collect(); + + // Row-space positions after each full batch of selected rows + let batch_ends = batch_end_positions(selection, row_count, batch_size); + if batch_ends.len() < 2 { + return None; + } + + // Greedily grow windows batch-by-batch until the estimated bytes exceed + // the target + let mut window_spans: Vec> = Vec::new(); + let mut window_start = 0u64; + let mut prev_end: Option = None; + for &batch_end in &batch_ends { + let estimate: u64 = estimators + .iter() + .map(|e| e.bytes_for_span(window_start, batch_end)) + .sum(); + if estimate > window_bytes { + if let Some(prev) = prev_end { + window_spans.push(window_start..prev); + window_start = prev; + } + } + prev_end = Some(batch_end); + } + // final window runs to the end of the row group + window_spans.push(window_start..row_count as u64); + + if window_spans.len() < 2 { + return None; + } + + // Build the per-window selections, data requests and read plans + let mut working = selection + .cloned() + .unwrap_or_else(|| RowSelection::from(vec![RowSelector::select(row_count)])); + let mut windows = VecDeque::with_capacity(window_spans.len()); + let mut range_uses: HashMap<(u64, u64), usize> = HashMap::new(); + for span in &window_spans { + let front = working.split_off((span.end - span.start) as usize); + let mut selectors = Vec::new(); + if span.start > 0 { + selectors.push(RowSelector::skip(span.start as usize)); + } + selectors.extend(front.iter().filter(|s| s.row_count > 0).cloned()); + if (span.end as usize) < row_count { + selectors.push(RowSelector::skip(row_count - span.end as usize)); + } + let window_selection = RowSelection::from(selectors); + + let data_request = DataRequestBuilder::new( + row_group_idx, + row_count, + batch_size, + parquet_metadata, + projection, + ) + .with_selection(Some(&window_selection)) + .build(); + + for range in data_request.ranges() { + *range_uses.entry((range.start, range.end)).or_insert(0) += 1; + } + + let plan_builder = ReadPlanBuilder::new(batch_size) + .with_selection(Some(window_selection)) + .with_row_selection_policy(row_selection_policy); + + windows.push_back(WindowPlan { + data_request, + plan_builder, + }); + } + + Some(WindowedRead { + row_group_idx, + row_count, + windows, + range_uses, + streamable, + dictionary_page_cache: Arc::new(DictionaryPageCache::default()), + }) +} + +/// Returns the row-space end position of each `batch_size`-selected-rows +/// batch, including a final (possibly partial) batch ending after the last +/// selected row. +fn batch_end_positions( + selection: Option<&RowSelection>, + row_count: usize, + batch_size: usize, +) -> Vec { + let mut ends = Vec::new(); + let mut row_pos = 0usize; + let mut selected_in_batch = 0usize; + let mut any_selected_since_last = false; + + let all_rows; + let selectors: Box> = match selection { + Some(selection) => Box::new(selection.iter().cloned()), + None => { + all_rows = [RowSelector::select(row_count)]; + Box::new(all_rows.iter().cloned()) + } + }; + + for selector in selectors { + if selector.skip { + row_pos += selector.row_count; + continue; + } + let mut remaining = selector.row_count; + while remaining > 0 { + let take = remaining.min(batch_size - selected_in_batch); + selected_in_batch += take; + row_pos += take; + remaining -= take; + any_selected_since_last = true; + if selected_in_batch == batch_size { + ends.push(row_pos as u64); + selected_in_batch = 0; + any_selected_since_last = false; + } + } + } + if any_selected_since_last { + ends.push(row_pos as u64); + } + ends +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_batch_end_positions_no_selection() { + let ends = batch_end_positions(None, 10, 4); + assert_eq!(ends, vec![4, 8, 10]); + } + + #[test] + fn test_batch_end_positions_with_selection() { + // skip 2, select 3, skip 1, select 4 => selected rows at 2,3,4, 6,7,8,9 + let selection = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(1), + RowSelector::select(4), + ]); + let ends = batch_end_positions(Some(&selection), 10, 2); + // batches of 2 selected rows end after rows 3 (2,3), 6 (4,6), 8 (7,8), + // and the final partial batch after row 9 + assert_eq!(ends, vec![4, 7, 9, 10]); + } +} diff --git a/parquet/src/arrow/push_decoder/remaining.rs b/parquet/src/arrow/push_decoder/remaining.rs index d2658189b7c1..d6d9c1f1142a 100644 --- a/parquet/src/arrow/push_decoder/remaining.rs +++ b/parquet/src/arrow/push_decoder/remaining.rs @@ -299,6 +299,13 @@ impl RemainingRowGroups { self.row_group_reader_builder.buffered_bytes() } + /// See [`ParquetPushDecoder::upcoming_fetch_plan`] + /// + /// [`ParquetPushDecoder::upcoming_fetch_plan`]: crate::arrow::push_decoder::ParquetPushDecoder::upcoming_fetch_plan + pub fn upcoming_fetch_plan(&self) -> crate::arrow::push_decoder::UpcomingFetchPlan { + self.row_group_reader_builder.upcoming_fetch_plan() + } + /// Clear any staged ranges currently buffered for future decode work pub fn clear_all_ranges(&mut self) { self.row_group_reader_builder.clear_all_ranges();