comphead commented on code in PR #6125:
URL: https://github.com/apache/datafusion-comet/pull/6125#discussion_r4084350178
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -217,6 +237,21 @@ impl ParquetFileReaderFactory for
EagerPageIndexReaderFactory {
metrics,
);
+ let reservation = self.memory_pool.as_ref().map(|pool| {
Review Comment:
**Registering a `MemoryConsumer` per reader shrinks the fair pool for the
whole task.**
`CometFairMemoryPool::try_grow` computes `limit = pool_size / state.num` and
compares it against the pool's *total* `used`, and `TaskSharedMemoryPool`
forwards `register`/`unregister`, so `num` counts every consumer in the task.
One registration per `create_reader` cuts every other consumer's ceiling by
`n/(n+1)` regardless of how many bytes the scan actually holds. In a q1-shaped
task (aggregate + sort + shuffle writer) that is roughly 25%.
DataFusion also calls `create_reader` a second time per file when
bloom-filter pruning runs, and keeps both readers alive across
`load_bloom_filters`
(`datafusion-datasource-parquet-55.1.0/src/opener/mod.rs:1234`), so `num`
churns and the decode estimate below is charged twice for that window. The
comment just below holds for `FileStream`, which does keep one reader per
partition, but not for that replacement reader.
This is accounting pressure rather than real pressure, and it is a plausible
contributor to the Spark disk-spill increase you measured for q9/q10.
Registering once on the factory, one consumer per scan operator, and handing
each reader a `new_empty()` handle keeps `num` stable and removes the churn.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -217,6 +237,21 @@ impl ParquetFileReaderFactory for
EagerPageIndexReaderFactory {
metrics,
);
+ let reservation = self.memory_pool.as_ref().map(|pool| {
+
MemoryConsumer::new(format!("ParquetScan[{partition_index}]")).register(pool)
+ });
+ // DataFusion opens one file per partition at a time and drops its
reader when the file
+ // is done, so holding the estimate on the reader charges it once per
running scan.
+ let decode_reservation = reservation.as_ref().map(|reservation| {
+ let decode = reservation.new_empty();
+ if decode.try_grow(self.decode_buffer_estimate).is_err() {
Review Comment:
**A refused reservation is not the only error `try_grow` can return here.**
`CometFairMemoryPool::try_grow` returns `resources_err!` only for its own
local-limit refusal and for a partial grant. Everything else comes out of
`self.acquire(additional)?`: a `SparkOutOfMemoryError` raised when another
consumer's spill throws, or the `InterruptedException` that
`ExecutionMemoryPool.acquireMemory` raises out of `lock.wait()` when the task
is killed. Those arrive as `CometError::JavaException` ->
`DataFusionError::External`, and `native/jni-bridge/src/errors.rs:229` keeps
them `External` deliberately so the original throwable survives back to the JNI
boundary.
`is_err()` erases that distinction, here and in `reserve_buffers`: the scan
keeps reading and the throwable is dropped. Matching on
`DataFusionError::ResourcesExhausted` for the keep-reading path and propagating
the rest would preserve it. Both call sites already return a `Result`.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -402,7 +524,8 @@ impl AsyncFileReader for EagerPageIndexReader {
/// retain delegation; native cloud wrappers observe coalescing. Requested
bytes update
/// `bytes_scanned` before I/O; successful
/// logical ranges update `data_bytes`, excluding gaps fetched by
cloud-store coalescing.
- /// The future borrows this reader and propagates store errors as external
Parquet errors.
+ /// With a memory pool, the returned buffers are charged to it until
dropped; see
+ /// [`ReservedBytes`]. The future borrows this reader and propagates store
errors as external Parquet errors.
Review Comment:
Nit: 103 chars against the crate's `max_width = 100`. rustfmt does not
rewrap comments, so it just reads inconsistently with the lines above it.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -240,6 +277,91 @@ struct EagerPageIndexReader {
metadata_cache: Arc<FileMetadataCache>,
metadata_size_hint: Option<usize>,
spark_variant_schema: bool,
+ /// Empty handle whose registration every fetched buffer's reservation
shares.
+ reservation: Option<MemoryReservation>,
+ /// Holds the decode-buffer estimate while this file is being read.
+ _decode_reservation: Option<MemoryReservation>,
+}
+
+/// Data page size limit that both parquet-mr and parquet-rs default to.
+const DEFAULT_PAGE_SIZE: usize = 1024 * 1024;
+
+/// Estimates what a scan holds while decoding beyond the compressed column
chunks, which
+/// [`ReservedBytes`] charges exactly. arrow-rs allocates the rest internally
with no hook to
+/// observe it: per projected leaf column, a decompressed data page bounded by
the writer's page
+/// size, plus the output batch being assembled. Dictionary pages are left
out: they are often
+/// absent or small, and counting them at the same bound doubled the estimate
against TPC-H
+/// measurements. Files written with larger pages or large dictionaries are
underestimated.
+pub(crate) fn decode_buffer_estimate(projected: &Schema, batch_size: usize) ->
usize {
Review Comment:
**The estimate never looks at the file it is opening.**
`leaves * 1 MiB + batch_size * row_width` is fixed at plan time and charged
in `create_reader` before any metadata is read, so it does not depend on the
file's size, row count, or row-group count. A 500-leaf projection charges 500
MiB per open reader whether the file holds 1,000 rows or 100M, and a scan over
many small files pays it on every open. Combined with the per-reader consumer
registration above, on a 2g pool one open reader can take most of the fair
limit. TPC-H's 7-to-16 column projections do not exercise this.
`partitioned_file.object_meta.size` is available in `create_reader`; scaling
by it, or deferring the grow until `get_metadata` has returned and sizing from
the row groups actually selected, would bound the worst case.
Two smaller things. This walks `projected.fields()` twice, once for
`row_width` and once for `columns`, where one `fold` does both. And
`decode_buffer_estimate` names this function, the factory field, and the
`with_memory_pool` parameter, which makes the doc link on `with_memory_pool`
ambiguous about which it points at. Worth a word in the doc comment that
`row_width` sums top-level fields while `columns` sums leaves, so a 10-leaf
struct contributes 32 bytes of width but 10 pages.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -240,6 +277,91 @@ struct EagerPageIndexReader {
metadata_cache: Arc<FileMetadataCache>,
metadata_size_hint: Option<usize>,
spark_variant_schema: bool,
+ /// Empty handle whose registration every fetched buffer's reservation
shares.
+ reservation: Option<MemoryReservation>,
Review Comment:
Two `Option<MemoryReservation>` that share a registration and a lifetime,
one of them permanently zero-sized. Growing the single handle by the estimate
in `create_reader` and keeping `new_empty()` per fetch collapses this to one
field and one closure, and the `_`-prefixed field goes away.
`reservation` also reads as "the scan's reservation" when what it actually
carries is the registration.
##########
docs/source/user-guide/latest/metrics.md:
##########
@@ -163,6 +163,7 @@ execution metrics. Counters accumulate per scan operator;
they do not instrument
| `scan_io_object_store_response_bytes_read` | Response bytes actually
consumed at that API, including bytes fetched between coalesced ranges. Not
HTTP wire bytes.
|
| `scan_io_metadata_cache_hits` | Successful, cache-eligible
metadata opens requiring no storage reads.
|
| `scan_io_metadata_cache_misses` | Successful, cache-eligible
metadata opens requiring storage reads. Failed opens and encrypted opens, which
bypass this shared cache, increment neither cache counter.
|
+| `scan_io_unreserved_bytes` | Bytes the scan held without a
memory reservation because the pool refused it: fetched data pages, plus the
per-reader decode-buffer estimate. The scan cannot spill, so it keeps reading
instead of failing. |
Review Comment:
Sharpening my earlier question: the number is cumulative, so it re-adds the
same working set on every refused fetch, and it also folds in the one-shot
decode estimate, which is not "bytes read" at all. Rendered as a size in the
SQL UI it looks like a live footprint, which is exactly what makes a large
value read as a leak.
Either split the two, or say in both this row and the `CometMetricNode`
label that it is cumulative refused bytes rather than a current figure.
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -824,7 +833,7 @@ mod tests {
.iter()
.filter(|metric| metric.value().name().starts_with("scan_io_"))
.collect::<Vec<_>>();
- assert_eq!(scan_io_metrics.len(), 9);
+ assert_eq!(scan_io_metrics.len(), 10);
Review Comment:
The Rust count goes 9 -> 10 but the Scala side did not follow.
`spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:2423-2432`
("Comet native metrics: scan") still enumerates the same nine names, so nothing
asserts the new metric actually reaches Spark's SQL metrics. Adding it to that
`Seq` is the only end-to-end check that the `CometMetricNode` entry is wired.
`spark/src/test/scala/org/apache/spark/sql/benchmark/CometReadBenchmark.scala:60-64`
carries the same list, plus a "zero or nine extra SQL accumulators" comment
that is now stale.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -240,6 +277,91 @@ struct EagerPageIndexReader {
metadata_cache: Arc<FileMetadataCache>,
metadata_size_hint: Option<usize>,
spark_variant_schema: bool,
+ /// Empty handle whose registration every fetched buffer's reservation
shares.
+ reservation: Option<MemoryReservation>,
+ /// Holds the decode-buffer estimate while this file is being read.
+ _decode_reservation: Option<MemoryReservation>,
+}
+
+/// Data page size limit that both parquet-mr and parquet-rs default to.
+const DEFAULT_PAGE_SIZE: usize = 1024 * 1024;
+
+/// Estimates what a scan holds while decoding beyond the compressed column
chunks, which
+/// [`ReservedBytes`] charges exactly. arrow-rs allocates the rest internally
with no hook to
+/// observe it: per projected leaf column, a decompressed data page bounded by
the writer's page
+/// size, plus the output batch being assembled. Dictionary pages are left
out: they are often
+/// absent or small, and counting them at the same bound doubled the estimate
against TPC-H
+/// measurements. Files written with larger pages or large dictionaries are
underestimated.
+pub(crate) fn decode_buffer_estimate(projected: &Schema, batch_size: usize) ->
usize {
+ fn leaves(data_type: &DataType) -> usize {
+ match data_type {
+ DataType::Struct(fields) => fields.iter().map(|f|
leaves(f.data_type())).sum(),
+ DataType::List(child)
+ | DataType::LargeList(child)
+ | DataType::FixedSizeList(child, _)
+ | DataType::ListView(child)
+ | DataType::LargeListView(child)
+ | DataType::Map(child, _) => leaves(child.data_type()),
+ _ => 1,
+ }
+ }
+ // Variable-width and nested values have no fixed width; assume a short
string.
+ let row_width: usize = projected
+ .fields()
+ .iter()
+ .map(|f| f.data_type().primitive_width().unwrap_or(32))
+ .sum();
+ let columns: usize = projected
+ .fields()
+ .iter()
+ .map(|f| leaves(f.data_type()))
+ .sum();
+ columns * DEFAULT_PAGE_SIZE + batch_size * row_width
+}
+
+/// A fetched data-page buffer that holds its share of the scan's memory
reservation.
+///
+/// The reader cannot see when the decoder is finished with the pages it
fetched: arrow-rs keeps
+/// the column chunks of the row group being decoded (and the pages sliced
from them) alive, and
+/// drops them once it moves on. Making the reservation the `Bytes` owner ties
the release to that
+/// last drop, so the pool sees the scan's input working set, roughly one row
group's projected
+/// column chunks per partition, without any change to the decoder.
+struct ReservedBytes {
+ bytes: Bytes,
+ _reservation: MemoryReservation,
+}
+
+impl AsRef<[u8]> for ReservedBytes {
+ fn as_ref(&self) -> &[u8] {
+ self.bytes.as_ref()
+ }
+}
+
+/// Charges `buffers` to a new reservation sharing `reservation`'s
registration, one pool call
+/// for the whole fetch. A scan cannot spill, so a refused grow does not fail
the read: the
+/// buffers are returned unaccounted and counted in `unreserved_bytes`. The
attempt still
+/// lets Spark ask the task's other consumers to spill before it refuses.
+fn reserve_buffers(
Review Comment:
Two things here.
**Coalesced cloud fetches are undercharged.** For
`ScanIoSource::ObjectStore`, `get_ranges` runs `coalesce_ranges(...,
OBJECT_STORE_COALESCE_DEFAULT)`, so each returned `Bytes` is a slice of one
merged allocation and any surviving slice pins the whole thing, gaps included.
Charging the sum of the logical slices leaves up to 1 MiB per gap resident and
unreserved, so "charged exactly" holds only for the delegating local path.
`scan_io_object_store_response_bytes_read` already measures the difference.
**One grow, N shrinks.** `split` per buffer gives every `Bytes` its own
reservation, so release is one `pool.shrink` -> one JNI `release_memory` per
column chunk, each taking the pool mutex. arrow-rs drops a row group's chunks
together anyway. A single reservation behind an `Arc`, shared by every
`ReservedBytes` of the fetch, gives 1 grow + 1 shrink and drops `split`
entirely.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -992,6 +1118,119 @@ mod tests {
);
}
+ async fn read_ranges_with_pool(pool: Arc<dyn MemoryPool>) -> (Vec<Bytes>,
usize) {
Review Comment:
This rebuilds the fixture `assert_range_read_contract` already has: same
`InMemory` store, same `b"0123456789"`, same factory -> `create_reader` ->
`get_byte_ranges` shape, with the ranges changed from `0..2, 4..6` to `0..2,
4..7` for no apparent reason.
`decode_estimate_is_held_while_the_reader_is_open` builds it a third time. One
helper parameterized by pool and decode estimate covers all three, and the
`assert_eq!(result, vec![...])` that both pool-based tests repeat belongs
inside it.
A `use datafusion::execution::memory_pool::GreedyMemoryPool;` in the test
module would also remove the two wrapped fully-qualified constructions below.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
