viirya commented on code in PR #6125:
URL: https://github.com/apache/datafusion-comet/pull/6125#discussion_r4085630425
##########
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.
Review Comment:
This holds only in a narrow case. `NativeMemoryConsumer.spill()` always
returns 0, so native consumers never give anything back. `CometFairMemoryPool`
refuses on its local `pool_size / num` check without calling Spark at all. So
only JVM consumers in the same task can be asked to spill.
In that case there is also a cost worth writing down. Say
`acquireExecutionMemory` spills a JVM consumer, such as the JVM columnar
shuffle's `CometShuffleMemoryAllocator`, and still can only grant part of the
request. `CometUnifiedMemoryPool::try_grow` releases the whole grant and
returns an error, and the scan reads unaccounted anyway. The spill it caused
bought nothing. Could the comment say exactly when a refusal triggers a spill?
Could the PR description mention this trade-off, since it currently presents
the spill as a benefit?
##########
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(
+ reservation: &MemoryReservation,
+ buffers: Vec<Bytes>,
+ scan_io_metrics: &ScanIoMetrics,
+) -> Vec<Bytes> {
+ let total = buffers.iter().map(Bytes::len).sum();
+ let fetch = reservation.new_empty();
+ if fetch.try_grow(total).is_err() {
Review Comment:
Under `greedy_unified`, every refused grow here goes to
`CometTaskMemoryManager.acquireMemory`. Whenever `acquired < size`, that method
logs a `warn` and calls `internal.showMemoryUsage()`, which prints every
consumer in the task. A spillable operator only hits that a bounded number of
times before it spills. This path doesn't stop. The next row group asks again,
and so does every file open in `create_reader`. On a long scan under pressure,
that means a full memory dump per row group per task.
Could the reader stop asking once it has been refused, for example for the
rest of the file? Or could it back off in some other way? `fair_unified` mostly
avoids this because its local-limit refusal never reaches Spark, but the greedy
pool does not.
##########
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) {
+ let store = Arc::new(InMemory::new());
+ let location = Path::from("reserved.parquet");
+ store
+ .put(&location, Bytes::from_static(b"0123456789").into())
+ .await
+ .unwrap();
+ let runtime =
datafusion::execution::runtime_env::RuntimeEnv::default();
+ let metrics = ExecutionPlanMetricsSet::new();
+ let factory = EagerPageIndexReaderFactory::new(
+ store,
+ runtime.cache_manager.get_file_metadata_cache(),
+ ScanIoSource::Local,
+ &metrics,
+ )
+ .with_memory_pool(pool, 0);
+ let mut reader = factory
+ .create_reader(
+ 0,
+ PartitionedFile::new(location.to_string(), 10),
+ None,
+ &metrics,
+ )
+ .unwrap();
+ let result = reader.get_byte_ranges(vec![0..2, 4..7]).await.unwrap();
+ drop(reader);
+ let unreserved = metrics
+ .clone_inner()
+ .sum_by_name("scan_io_unreserved_bytes")
+ .unwrap()
+ .as_usize();
+ (result, unreserved)
+ }
+
+ #[tokio::test]
+ async fn fetched_data_pages_hold_a_reservation_until_dropped() {
+ let pool: Arc<dyn MemoryPool> = Arc::new(
+ datafusion::execution::memory_pool::GreedyMemoryPool::new(1024),
+ );
+ let (mut result, unreserved) =
read_ranges_with_pool(Arc::clone(&pool)).await;
+ assert_eq!(
+ result,
+ vec![Bytes::from_static(b"01"), Bytes::from_static(b"456")]
+ );
+ assert_eq!(unreserved, 0);
+ // The reservation outlives the reader and follows each buffer,
including slices of it.
+ assert_eq!(pool.reserved(), 5);
+ let slice = result[1].slice(1..2);
+ result.remove(1);
+ assert_eq!(pool.reserved(), 5);
+ drop(slice);
+ assert_eq!(pool.reserved(), 2);
+ drop(result);
+ assert_eq!(pool.reserved(), 0);
+ }
+
+ #[test]
+ fn decode_estimate_counts_leaf_columns_and_one_batch() {
+ use arrow::datatypes::{Field, Fields};
+ let schema = Schema::new(vec![
+ Field::new("a", DataType::Int64, true),
+ Field::new("s", DataType::Utf8, true),
+ Field::new(
+ "st",
+ DataType::Struct(Fields::from(vec![
+ Field::new("x", DataType::Int32, true),
+ Field::new("y", DataType::Int32, true),
+ ])),
+ true,
+ ),
+ ]);
+ // Four leaves; row width 8 + 32 + 32 (struct has no primitive width).
+ assert_eq!(
+ decode_buffer_estimate(&schema, 100),
+ 4 * DEFAULT_PAGE_SIZE + 100 * 72
+ );
+ }
+
+ #[tokio::test]
+ async fn decode_estimate_is_held_while_the_reader_is_open() {
+ let pool: Arc<dyn MemoryPool> = Arc::new(
+ datafusion::execution::memory_pool::GreedyMemoryPool::new(1024),
+ );
+ let runtime =
datafusion::execution::runtime_env::RuntimeEnv::default();
+ let metrics = ExecutionPlanMetricsSet::new();
+ let factory = EagerPageIndexReaderFactory::new(
+ Arc::new(InMemory::new()),
+ runtime.cache_manager.get_file_metadata_cache(),
+ ScanIoSource::Local,
+ &metrics,
+ )
+ .with_memory_pool(Arc::clone(&pool), 100);
+ let reader = factory
+ .create_reader(0, PartitionedFile::new("f.parquet", 10), None,
&metrics)
+ .unwrap();
+ assert_eq!(pool.reserved(), 100);
+ drop(reader);
+ assert_eq!(pool.reserved(), 0);
+ }
+
+ #[tokio::test]
+ async fn refused_reservation_returns_data_unaccounted() {
Review Comment:
Once the refusal path tells `ResourcesExhausted` apart from other errors
(per @comphead's comment), could you add a test that uses a small `MemoryPool`
whose `try_grow` returns `DataFusionError::External`? It would check that
`get_byte_ranges` and `create_reader` propagate the error instead of counting
it as unreserved. All the tests here use `GreedyMemoryPool`, which only ever
returns `ResourcesExhausted`, so nothing covers the difference between refused
and failed today.
##########
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 {
Review Comment:
Adding one thing to @comphead's note about view arrays aliasing these
buffers. If that invariant ever breaks, the problem is bigger than a delayed
release. The final `shrink` would then run from the FFI release callback when
the JVM closes the batch. Both `CometFairMemoryPool::shrink` and
`CometUnifiedMemoryPool::shrink` `panic!` if the JNI release fails, and that
panic would happen inside a release callback. That seems worth spelling out in
the comment that records the invariant, so whoever turns on `Utf8View` later
knows what depends on it.
##########
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:
Following up on @comphead's question about blocking in `create_reader`. I
think the concern is wider than this call. When a plan has no JVM sources,
`jni_api.rs:1058` spawns it onto the shared tokio runtime. That means the grow
in `reserve_buffers` runs on a runtime worker as well.
`ExecutionMemoryPool.acquireMemory` parks in `lock.wait()` while a task
holds less than its minimum share (`curMem + toGrant < minMemoryPerTask`). That
is most likely at the very start of a task, and that is exactly when this scan
now makes its first request. Other operators already grow from spawned streams,
but they do it later in the task. This change makes every scan-rooted task do
it first thing.
There is a related effect on @comphead's `InterruptedException` point. On a
worker thread, the task-kill interrupt is never delivered, so a killed task's
scan can stay parked until some other task releases memory. Is that acceptable
here? If not, should the scan avoid a blocking acquire for its first grows?
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -194,7 +196,14 @@ pub(crate) fn init_datasource_exec(
scan_io_source,
parquet_source.metrics(),
)
- .with_spark_variant_schema(projects_variant),
+ .with_spark_variant_schema(projects_variant)
+ .with_memory_pool(
Review Comment:
This turns on reservation accounting for every native Parquet scan and there
is no way to turn it off. The benchmarks are encouraging. But your 2g run shows
q9's Spark disk spill going from 8.0 to 15.0 GiB and q10's from 5.6 to 6.9 GiB.
The decode estimate is also still an open question in @comphead's review.
Would you consider adding a `CometConf` flag, following
`config_conventions.md`, that skips `with_memory_pool`? That would let users
who regress disable it while the estimate is refined. A default of `true` seems
fine to me.
--
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]
