Bounded streaming: open streamed requests concurrently; scale windows by projected column count

- AsyncFileReader::get_bytes_stream now returns an Option of a 'static
  future so streamed requests can be opened concurrently with each other
  and with the materialized batch (previously opens were serialized on the
  reader borrow, paying one request latency per stream). None (the
  default) falls back to materialization.
- Decode windows now have a floor of min_window_bytes_per_column (default
  2MiB) per projected column so per-window fixed costs (array reader
  construction, boundary-page and dictionary-page re-decode) stay
  amortized; projections wide enough that one window covers the row group
  fall back to the non-windowed path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MGnT9oETcFMydc9cp95HLG

Author: adriangb

parquet/src/arrow/arrow_reader/mod.rs

@@ -185,6 +185,13 @@ pub struct BoundedStreamingOptions {
     ///
     /// [`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 {
@@ -193,6 +200,7 @@ impl Default for BoundedStreamingOptions {
             stream_threshold: 16 * 1024 * 1024,
             window_bytes: 8 * 1024 * 1024,
             coalesce_gap: 1024 * 1024,
+            min_window_bytes_per_column: 2 * 1024 * 1024,
         }
     }
 }

parquet/src/arrow/async_reader/adaptive.rs

@@ -315,7 +315,8 @@ where
     ///
     /// # Errors
     ///
-    /// Returns an error if called outside a tokio runtime context.
+    /// 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<Self> {
         let runtime = Handle::try_current().map_err(|e| {
             ParquetError::General(format!(
@@ -363,8 +364,9 @@ where
         .await
     }
 
-    /// `runtime` is used to spawn body-draining tasks
-    pub(crate) fn new(options: BoundedStreamingOptions, runtime: Handle, input: T) -> Self {
+    /// 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,
@@ -435,9 +437,11 @@ where
             return Ok(());
         }
 
-        // Can only issue requests while the input is idle. If a request is
-        // already in flight, plan again once it resolves (`poll_sources`
-        // returns progress when it does).
+        // 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(());
         }
@@ -519,6 +523,38 @@ where
             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));
         }
@@ -528,26 +564,14 @@ where
             }
         }
 
-        let runtime = self.runtime.clone();
+        if materialize.is_empty() {
+            // nothing to materialize: the input stays available
+            self.input = InputState::Idle(input);
+            return Ok(());
+        }
+
         let future = async move {
-            // Open the streamed requests first so their bodies start flowing
-            // while the materialized ranges download
-            for spec in specs {
-                let StreamSpec { span, tx } = spec;
-                match input.get_bytes_stream(span).await {
-                    Ok(body) => {
-                        runtime.spawn(pump_body(body, tx));
-                    }
-                    Err(e) => {
-                        let _ = tx.send(Err(e)).await;
-                    }
-                }
-            }
-            let data = if materialize.is_empty() {
-                vec![]
-            } else {
-                input.get_byte_ranges(materialize.clone()).await?
-            };
+            let data = input.get_byte_ranges(materialize.clone()).await?;
             Ok((input, materialize, data))
         }
         .boxed();

parquet/src/arrow/async_reader/mod.rs

@@ -52,7 +52,7 @@ mod metadata;
 pub use adaptive::AdaptiveFetcher;
 pub use metadata::*;
 
-use futures::stream::{BoxStream, StreamExt as _};
+use futures::stream::BoxStream;
 
 #[cfg(feature = "object_store")]
 mod store;
@@ -95,33 +95,28 @@ pub trait AsyncFileReader: Send {
     }
 
     /// Retrieve the bytes in `range` as a stream of chunks that can be
-    /// consumed incrementally, without buffering the whole range in memory.
+    /// 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. The stream is `'static` so the request can remain in flight
-    /// while the reader issues other requests.
+    /// 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 materializes the entire range via
-    /// [`Self::get_bytes`] and yields it as a single chunk, so implementations
-    /// that do not override this method behave exactly as before (correct, but
+    /// 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<u64>,
-    ) -> BoxFuture<'_, Result<BoxStream<'static, Result<Bytes>>>> {
-        let fut = self.get_bytes(range);
-        async move {
-            let bytes = fut.await?;
-            let stream: BoxStream<'static, Result<Bytes>> =
-                futures::stream::once(async move { Ok(bytes) }).boxed();
-            Ok(stream)
-        }
-        .boxed()
+    ) -> Option<BoxFuture<'static, Result<BoxStream<'static, Result<Bytes>>>>> {
+        let _ = range;
+        None
     }
 
     /// Return a future which results in the [`ParquetMetaData`] for this Parquet file.
@@ -159,7 +154,7 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
     fn get_bytes_stream(
         &mut self,
         range: Range<u64>,
-    ) -> BoxFuture<'_, Result<BoxStream<'static, Result<Bytes>>>> {
+    ) -> Option<BoxFuture<'static, Result<BoxStream<'static, Result<Bytes>>>>> {
         self.as_mut().get_bytes_stream(range)
     }
 
@@ -2193,15 +2188,15 @@ mod tests {
         fn get_bytes_stream(
             &mut self,
             range: Range<u64>,
-        ) -> BoxFuture<'_, Result<BoxStream<'static, Result<Bytes>>>> {
+        ) -> Option<BoxFuture<'static, Result<BoxStream<'static, Result<Bytes>>>>> {
             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<Result<Bytes>> = (0..data.len())
                 .step_by(chunk_size)
                 .map(|offset| Ok(data.slice(offset..(offset + chunk_size).min(data.len()))))
                 .collect();
-            futures::future::ready(Ok(futures::stream::iter(chunks).boxed())).boxed()
+            Some(futures::future::ready(Ok(futures::stream::iter(chunks).boxed())).boxed())
         }
 
         fn get_metadata<'a>(
@@ -2289,6 +2284,7 @@ mod tests {
             stream_threshold: 256 * 1024,
             window_bytes: 128 * 1024,
             coalesce_gap: 1024 * 1024,
+            min_window_bytes_per_column: 16 * 1024,
         }
     }
 
@@ -2403,6 +2399,7 @@ mod tests {
                 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(
@@ -2460,6 +2457,7 @@ mod tests {
                 stream_threshold: 64 * 1024,
                 window_bytes: 32 * 1024,
                 coalesce_gap: 1024 * 1024,
+                min_window_bytes_per_column: 8 * 1024,
             })
             .build()
             .unwrap();

parquet/src/arrow/async_reader/store.rs

@@ -199,23 +199,26 @@ impl AsyncFileReader for ParquetObjectReader {
     fn get_bytes_stream(
         &mut self,
         range: Range<u64>,
-    ) -> BoxFuture<'_, Result<futures::stream::BoxStream<'static, Result<Bytes>>>> {
+    ) -> Option<BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<Bytes>>>>>
+    {
         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()
         };
-        self.spawn(|store, path| {
+        Some(
             async move {
-                let result = store.get_opts(path, options).await?;
+                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()
-        })
+            .boxed(),
+        )
     }
 
     // This method doesn't directly call `self.spawn` because all of the IO that is done down the

parquet/src/arrow/push_decoder/reader_builder/mod.rs

@@ -746,8 +746,7 @@ impl RowGroupReaderBuilder {
                             &self.projection,
                             plan_builder.selection(),
                             self.row_selection_policy,
-                            streaming.stream_threshold,
-                            streaming.window_bytes,
+                            streaming,
                         ) {
                             return Ok(NextState::again(RowGroupDecoderState::DecodingWindows {
                                 windowed: Box::new(windowed),

parquet/src/arrow/push_decoder/reader_builder/windows.rs

@@ -133,9 +133,9 @@ pub(super) fn plan_windows(
     projection: &ProjectionMask,
     selection: Option<&RowSelection>,
     row_selection_policy: RowSelectionPolicy,
-    stream_threshold: u64,
-    window_bytes: u64,
+    options: &crate::arrow::arrow_reader::BoundedStreamingOptions,
 ) -> Option<WindowedRead> {
+    let stream_threshold = options.stream_threshold;
     if batch_size == 0 || row_count == 0 {
         return None;
     }
@@ -172,6 +172,14 @@ pub(super) fn plan_windows(
             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;
     }