This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 86f0cae8 feat: integrate reader memory reservations with DataFusion
(#914)
86f0cae8 is described below
commit 86f0cae85e0ce2d88def5f8401821b0338c40d9f
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Sep 23 17:59:03 2026 +0800
feat: integrate reader memory reservations with DataFusion (#914)
---
.github/workflows/ci.yml | 8 +
crates/integrations/datafusion/README.md | 14 +
crates/integrations/datafusion/src/error.rs | 7 +-
crates/integrations/datafusion/src/lib.rs | 1 +
crates/integrations/datafusion/src/memory.rs | 57 ++
.../datafusion/src/physical_plan/audit_log.rs | 4 +-
.../datafusion/src/physical_plan/scan.rs | 159 +++++-
crates/integrations/datafusion/src/sql_context.rs | 6 +-
crates/paimon/src/arrow/format/parquet.rs | 582 ++++++++++++++++++---
crates/paimon/src/arrow/read_budget.rs | 199 +++++--
crates/paimon/src/error.rs | 3 +
crates/paimon/src/lib.rs | 1 +
crates/paimon/src/resource/memory.rs | 150 ++++++
crates/paimon/src/resource/mod.rs | 112 ++++
crates/paimon/src/resource/tests.rs | 163 ++++++
crates/paimon/src/table/data_evolution_reader.rs | 85 ++-
crates/paimon/src/table/format_table_read.rs | 19 +-
crates/paimon/src/table/kv_file_reader.rs | 103 ++--
crates/paimon/src/table/read_builder.rs | 58 +-
crates/paimon/src/table/table_read.rs | 38 +-
crates/paimon/tests/reader_resources_test.rs | 332 ++++++++++++
21 files changed, 1922 insertions(+), 179 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 1b84a707..fb37e7f0 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -185,6 +185,14 @@ jobs:
steps:
- uses: actions/checkout@v7
+ - name: Limit Windows test build memory
+ if: runner.os == 'Windows'
+ shell: pwsh
+ run: |
+ # Full-feature test binaries can exhaust the runner's paging file.
+ "CARGO_BUILD_JOBS=2" >> $env:GITHUB_ENV
+ "CARGO_PROFILE_TEST_DEBUG=line-tables-only" >> $env:GITHUB_ENV
+
- name: Test
run: cargo test --locked -p paimon --all-targets --features
fulltext,vortex
env:
diff --git a/crates/integrations/datafusion/README.md
b/crates/integrations/datafusion/README.md
index 94515721..b9583726 100644
--- a/crates/integrations/datafusion/README.md
+++ b/crates/integrations/datafusion/README.md
@@ -24,6 +24,20 @@
This crate contains the integration of [Apache
DataFusion](https://datafusion.apache.org/) and [Apache
Paimon](https://paimon.apache.org/).
+## Reader memory
+
+Paimon Parquet scans use the memory pool from the executing DataFusion
`TaskContext`.
+Configure it through a DataFusion `RuntimeEnv`, including with
+`SQLContext::builder().with_runtime_env(runtime_env)`. Scan partitions share
that
+pool with downstream operators and reserve projected row-group working
estimates
+before data I/O. A scan consumer cannot spill; exhaustion returns
+`DataFusionError::ResourcesExhausted` after speculative prefetch has been
deferred.
+
+Downstream consumers account for output batches they retain. The reader leaves
+Arrow buffers unchanged and releases its working reservation when its decoder
is
+dropped, including cancellation. Metadata, merge state, transient batches and
+allocation overhead are not fully accounted, so the pool is not an RSS limit.
+
## REST Catalog views and SQL functions
`SQLContext` can read, execute, create, and drop persistent views and can
create SQL scalar functions in a Paimon REST Catalog:
diff --git a/crates/integrations/datafusion/src/error.rs
b/crates/integrations/datafusion/src/error.rs
index 92b27281..bb876f8c 100644
--- a/crates/integrations/datafusion/src/error.rs
+++ b/crates/integrations/datafusion/src/error.rs
@@ -19,5 +19,10 @@ use datafusion::common::error::GenericError;
/// Converts a Paimon error into a DataFusion error.
pub fn to_datafusion_error(error: paimon::Error) ->
datafusion::error::DataFusionError {
- datafusion::error::DataFusionError::External(GenericError::from(error))
+ match error {
+ paimon::Error::ResourceExhausted { message } => {
+ datafusion::error::DataFusionError::ResourcesExhausted(message)
+ }
+ other =>
datafusion::error::DataFusionError::External(GenericError::from(other)),
+ }
}
diff --git a/crates/integrations/datafusion/src/lib.rs
b/crates/integrations/datafusion/src/lib.rs
index 5d581137..7043303e 100644
--- a/crates/integrations/datafusion/src/lib.rs
+++ b/crates/integrations/datafusion/src/lib.rs
@@ -51,6 +51,7 @@ mod format_table_truncate;
mod full_text_search;
mod hybrid_search;
mod lateral_vector_search;
+mod memory;
mod merge_into;
mod partition_count_pushdown;
mod physical_plan;
diff --git a/crates/integrations/datafusion/src/memory.rs
b/crates/integrations/datafusion/src/memory.rs
new file mode 100644
index 00000000..7419b218
--- /dev/null
+++ b/crates/integrations/datafusion/src/memory.rs
@@ -0,0 +1,57 @@
+// 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.
+
+//! Connect reader working-memory reservations to the executing DataFusion
task.
+
+use std::sync::Arc;
+
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use paimon::resource::{MemoryPool, ResourceContext};
+
+#[derive(Debug)]
+struct DataFusionMemoryPool {
+ reservation: MemoryReservation,
+}
+
+impl MemoryPool for DataFusionMemoryPool {
+ fn try_reserve(&self, bytes: usize) -> paimon::Result<()> {
+ self.reservation
+ .try_grow(bytes)
+ .map_err(|error| paimon::Error::ResourceExhausted {
+ message: error.to_string(),
+ })
+ }
+
+ fn release(&self, bytes: usize) {
+ self.reservation.shrink(bytes);
+ }
+}
+
+pub(crate) fn reader_resources(
+ context: &TaskContext,
+ partition: usize,
+) -> paimon::Result<ResourceContext> {
+ // Register per execution, not on the reusable plan. Reader working state
+ // cannot spill; downstream operators account for batches they retain.
+ let reservation =
MemoryConsumer::new(format!("PaimonTableScan[{partition}]"))
+ .with_can_spill(false)
+ .register(context.memory_pool());
+ ResourceContext::builder()
+ .memory_pool(Arc::new(DataFusionMemoryPool { reservation }))
+ .build()
+}
diff --git a/crates/integrations/datafusion/src/physical_plan/audit_log.rs
b/crates/integrations/datafusion/src/physical_plan/audit_log.rs
index 5fdfeccb..072779b5 100644
--- a/crates/integrations/datafusion/src/physical_plan/audit_log.rs
+++ b/crates/integrations/datafusion/src/physical_plan/audit_log.rs
@@ -62,9 +62,9 @@ impl ExecutionPlan for PaimonAuditLogScan {
fn execute(
&self,
partition: usize,
- _context: Arc<TaskContext>,
+ context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
- self.inner.execute_with(partition, |read, splits| {
+ self.inner.execute_with(partition, context, |read, splits| {
AuditLogRead::new(read)?.to_arrow(splits)
})
}
diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs
b/crates/integrations/datafusion/src/physical_plan/scan.rs
index f75f0f69..531c4fb2 100644
--- a/crates/integrations/datafusion/src/physical_plan/scan.rs
+++ b/crates/integrations/datafusion/src/physical_plan/scan.rs
@@ -1000,6 +1000,7 @@ impl PaimonTableScan {
pub(crate) fn execute_with(
&self,
partition: usize,
+ context: Arc<TaskContext>,
read_splits: impl FnOnce(TableRead<'_>, &[DataSplit]) ->
paimon::Result<ArrowRecordBatchStream>
+ Send
+ 'static,
@@ -1021,7 +1022,10 @@ impl PaimonTableScan {
let parquet_read_budget = Arc::clone(&self.parquet_read_budget);
let fut = async move {
+ let resources = crate::memory::reader_resources(&context,
partition)
+ .map_err(to_datafusion_error)?;
let mut read_builder = table.new_read_builder();
+ read_builder.with_resources(resources);
let runtime_filter_plan = partition_runtime_decoder_filters(
&decoder_filters,
table.schema().fields(),
@@ -1172,9 +1176,9 @@ impl ExecutionPlan for PaimonTableScan {
fn execute(
&self,
partition: usize,
- _context: Arc<TaskContext>,
+ context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
- self.execute_with(partition, |read, splits| read.to_arrow(splits))
+ self.execute_with(partition, context, |read, splits|
read.to_arrow(splits))
}
fn partition_statistics(&self, partition: Option<usize>) ->
DFResult<Arc<Statistics>> {
@@ -2308,6 +2312,157 @@ mod tests {
assert!(result.updated_node.is_none());
}
+ fn memory_scan_fixture() -> (tempfile::TempDir, PaimonTableScan, usize) {
+ use parquet::file::reader::{FileReader, SerializedFileReader};
+ let dir = tempdir().unwrap();
+ let bucket = dir.path().join("bucket-0");
+ fs::create_dir_all(&bucket).unwrap();
+ let file = bucket.join("data.parquet");
+ write_int_parquet_file(&file, vec![("id", (0..8).collect())], None);
+ let metadata =
SerializedFileReader::new(fs::File::open(&file).unwrap()).unwrap();
+ let estimate = metadata
+ .metadata()
+ .row_group(0)
+ .column(0)
+ .uncompressed_size() as usize;
+ let schema = PaimonSchema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("read.batch-size", "2")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ FileIOBuilder::new("file").build().unwrap(),
+ Identifier::new("default", "memory_test"),
+ local_file_path(dir.path()),
+ TableSchema::new(0, &schema),
+ None,
+ );
+ let split = paimon::DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(local_file_path(&bucket))
+ .with_total_buckets(1)
+ .with_data_files(vec![test_data_file(
+ "data.parquet",
+ 8,
+ fs::metadata(file).unwrap().len() as i64,
+ )])
+ .build()
+ .unwrap();
+ let scan = PaimonTableScan::new(
+ test_schema(),
+ table,
+ test_read_type(),
+ None,
+ vec![Arc::from(vec![split.clone()]), Arc::from(vec![split])],
+ None,
+ false,
+ None,
+ None,
+ true,
+ );
+ (dir, scan, estimate)
+ }
+
+ fn memory_session(
+ pool: Arc<dyn datafusion::execution::memory_pool::MemoryPool>,
+ ) -> SessionContext {
+ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+ SessionContext::new_with_config_rt(
+ datafusion::prelude::SessionConfig::new(),
+ RuntimeEnvBuilder::new()
+ .with_memory_pool(pool)
+ .build_arc()
+ .unwrap(),
+ )
+ }
+
+ #[tokio::test]
+ async fn scan_memory_honors_each_execution_pool_and_audit_reads() {
+ use datafusion::error::DataFusionError;
+ use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool};
+ for audit in [false, true] {
+ let (_dir, scan, estimate) = memory_scan_fixture();
+ let plan: Arc<dyn ExecutionPlan> = if audit {
+
Arc::new(super::super::audit_log::PaimonAuditLogScan::new(scan))
+ } else {
+ Arc::new(scan)
+ };
+ let rejected = Arc::new(GreedyMemoryPool::new(0));
+ let ctx = memory_session(rejected.clone());
+ let mut stream = plan.execute(0, ctx.task_ctx()).unwrap();
+ assert!(matches!(
+ stream.try_next().await,
+ Err(DataFusionError::ResourcesExhausted(_))
+ ));
+ assert!(stream.try_next().await.unwrap().is_none());
+ assert_eq!(rejected.reserved(), 0);
+
+ // Reusing a plan must bind to the current execution, never an
earlier pool.
+ let admitted = Arc::new(GreedyMemoryPool::new(estimate));
+ let ctx = memory_session(admitted.clone());
+ let output = plan
+ .execute(0, ctx.task_ctx())
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(collect_ids(&output), (0..8).collect::<Vec<_>>());
+ assert_eq!(
+ admitted.reserved(),
+ 0,
+ "returned batches belong to their consumer"
+ );
+ assert_eq!(rejected.reserved(), 0);
+ }
+ }
+
+ #[tokio::test]
+ async fn
scan_memory_shares_capacity_with_partitions_and_downstream_consumers() {
+ use datafusion::error::DataFusionError;
+ use datafusion::execution::memory_pool::{FairSpillPool,
MemoryConsumer, MemoryPool};
+ let (_dir, scan, estimate) = memory_scan_fixture();
+ let pool: Arc<dyn MemoryPool> = Arc::new(FairSpillPool::new(2 *
estimate));
+ let ctx = memory_session(pool.clone());
+ // A spillable peer must not reduce the reader's share: its decoder
+ // cannot spill and must be registered as a non-spillable consumer.
+ let _spillable_peer = MemoryConsumer::new("spillable peer")
+ .with_can_spill(true)
+ .register(&pool);
+ let downstream = MemoryConsumer::new("downstream
state").register(&pool);
+ downstream.try_grow(estimate).unwrap();
+ let mut first = scan.execute(0, ctx.task_ctx()).unwrap();
+ let retained = first.try_next().await.unwrap().unwrap();
+ assert_eq!(pool.reserved(), 2 * estimate);
+ let mut second = scan.execute(1, ctx.task_ctx()).unwrap();
+ assert!(matches!(
+ second.try_next().await,
+ Err(DataFusionError::ResourcesExhausted(_))
+ ));
+ assert!(second.try_next().await.unwrap().is_none());
+ assert_eq!(
+ pool.reserved(),
+ 2 * estimate,
+ "failed admission must roll back"
+ );
+
+ // Cancelling the decoder releases its estimate even if output
survives.
+ drop(first);
+ assert_eq!(pool.reserved(), estimate);
+ assert_eq!(collect_ids(&[retained]), vec![0, 1]);
+ drop(downstream);
+ assert_eq!(pool.reserved(), 0);
+ let output = scan
+ .execute(1, ctx.task_ctx())
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(collect_ids(&output), (0..8).collect::<Vec<_>>());
+ assert_eq!(pool.reserved(), 0);
+ }
+
#[tokio::test]
async fn test_scan_applies_retained_runtime_filter_by_default() {
let tempdir = tempdir().unwrap();
diff --git a/crates/integrations/datafusion/src/sql_context.rs
b/crates/integrations/datafusion/src/sql_context.rs
index 34c33035..d13bb93b 100644
--- a/crates/integrations/datafusion/src/sql_context.rs
+++ b/crates/integrations/datafusion/src/sql_context.rs
@@ -109,7 +109,11 @@ pub struct SQLContext {
///
/// The builder preserves Paimon's session configuration while allowing callers
/// to customize DataFusion runtime resources such as memory pools, temporary
-/// directories, and object store registries.
+/// directories, and object store registries. Parquet scans reserve projected
+/// row-group working estimates from the execution's memory pool. These readers
+/// cannot spill; supported downstream DataFusion operators may spill their
state.
+/// Output batches are accounted by consumers that retain them, not by the
reader.
+/// The pool does not cover every allocation or bound process RSS.
///
/// # Example
///
diff --git a/crates/paimon/src/arrow/format/parquet.rs
b/crates/paimon/src/arrow/format/parquet.rs
index c1c8be80..22d02ce8 100644
--- a/crates/paimon/src/arrow/format/parquet.rs
+++ b/crates/paimon/src/arrow/format/parquet.rs
@@ -54,8 +54,8 @@ use parquet::file::statistics::Statistics as
ParquetStatistics;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::ops::Range;
-use std::sync::Arc;
-use tokio::sync::mpsc;
+use std::sync::{Arc, Mutex};
+use tokio::sync::{mpsc, oneshot};
pub(crate) struct ParquetFormatReader {
read_budget: Option<Arc<ReadBudget>>,
@@ -723,6 +723,21 @@ impl FormatFileReader for ParquetFormatReader {
}
}
+ let mut memory_mask = mask.clone();
+ for predicate in &decoder_predicates {
+ memory_mask.union(predicate.projection());
+ }
+ // Resource-aware sequential reads build one row-group stream at a
time,
+ // so admission precedes its data I/O. Preserve stateful decoder
filters
+ // across those streams instead of recreating the engine's factory.
+ let shared_filters = if self.read_budget.as_ref().is_some_and(|b|
b.has_resources()) {
+ std::mem::take(&mut decoder_predicates)
+ .into_iter()
+ .map(SharedParquetPredicate::new)
+ .collect::<Vec<_>>()
+ } else {
+ Vec::new()
+ };
if !decoder_predicates.is_empty() {
batch_stream_builder =
batch_stream_builder.with_row_filter(ParquetRowFilter::new(decoder_predicates));
@@ -766,7 +781,8 @@ impl FormatFileReader for ParquetFormatReader {
//
// Row-group receivers are consumed in order and buffer one batch each,
// preserving positional `_ROW_ID`, sort order, and batch backpressure.
- // Reads with predicates retain the original single-stream path.
+ // Reads with predicates stay sequential; resource-aware reads reserve
+ // each row group before opening its decoder.
let read_budget = self
.read_budget
.as_ref()
@@ -781,7 +797,9 @@ impl FormatFileReader for ParquetFormatReader {
let selected_row_groups = self
.read_budget
.as_ref()
- .filter(|budget| row_group_parallelism > 1 ||
budget.diagnostics_enabled())
+ .filter(|budget| {
+ row_group_parallelism > 1 || budget.diagnostics_enabled() ||
budget.has_resources()
+ })
.map(|budget| {
let mut row_group_selection = combined_selection;
let selected_row_groups = batch_stream_builder
@@ -799,7 +817,12 @@ impl FormatFileReader for ParquetFormatReader {
{
return None;
}
- let projected_bytes =
projected_row_group_bytes(row_group, &mask);
+ let projection = if budget.has_resources() {
+ &memory_mask
+ } else {
+ &mask
+ };
+ let projected_bytes =
projected_row_group_bytes(row_group, projection);
Some((row_group_index, selection, projected_bytes))
})
.collect::<Vec<_>>();
@@ -825,17 +848,38 @@ impl FormatFileReader for ParquetFormatReader {
let Ok(slot) = row_group_tx.reserve().await else {
return;
};
- let permit = match tokio::select! {
+ let (batch_tx, batch_rx) = mpsc::channel(1);
+ let (demand_tx, demand_rx) = oneshot::channel();
+ slot.send((batch_rx, demand_tx));
+ let admitted = tokio::select! {
_ = row_group_tx.closed() => return,
permit = read_budget.acquire(projected_bytes) =>
permit,
- } {
+ };
+ let admitted = match admitted {
+ Err(Error::ResourceExhausted { .. }) => {
+ // Speculative prefetch must not fail a read merely
+ // because an earlier group is still decoding.
Retry
+ // once this group is actually requested, after the
+ // consumer has drained its predecessors. Never
wait
+ // for the caller to drop retained output buffers.
+ tokio::select! {
+ _ = row_group_tx.closed() => return,
+ _ = demand_rx => {},
+ }
+ tokio::select! {
+ _ = row_group_tx.closed() => return,
+ permit = read_budget.acquire(projected_bytes)
=> permit,
+ }
+ }
+ other => other,
+ };
+ let permit = match admitted {
Ok(permit) => permit,
Err(error) => {
- slot.send(Err(error));
+ let _ =
batch_tx.send(ParquetRowGroupMessage::Error(error)).await;
return;
}
};
- let (batch_tx, batch_rx) = mpsc::channel(1);
let row_group_reader = Arc::clone(&shared_reader);
let row_group_metadata = reader_metadata.clone();
let row_group_mask = mask.clone();
@@ -850,17 +894,17 @@ impl FormatFileReader for ParquetFormatReader {
permit,
batch_tx,
));
- slot.send(Ok(batch_rx));
}
});
let stream = async_stream::try_stream! {
for _ in 0..row_group_count {
- let mut batches = row_group_rx.recv().await.ok_or_else(|| {
+ let (mut batches, demand) =
row_group_rx.recv().await.ok_or_else(|| {
Error::UnexpectedError {
message: "Parquet row-group coordinator stopped
early".to_string(),
source: None,
}
- })??;
+ })?;
+ let _ = demand.send(());
let mut completed = false;
while let Some(message) = batches.recv().await {
match message {
@@ -888,6 +932,46 @@ impl FormatFileReader for ParquetFormatReader {
return Ok(stream.boxed());
}
+ if let Some(budget) = self.read_budget.as_ref().filter(|b|
b.has_resources()) {
+ let budget = Arc::clone(budget);
+ let selected = selected_row_groups.expect("resource-aware reads
need a selection plan");
+ let metadata = ArrowReaderMetadata::try_new(
+ batch_stream_builder.metadata().clone(),
+ ArrowReaderOptions::new(),
+ )?;
+ let residual = (!all_enforced).then(|| FilePredicates {
+ predicates: preds.to_vec(),
+ row_filter_factory: None,
+ file_fields: file_fields.to_vec(),
+ });
+ return Ok(async_stream::try_stream! {
+ for (index, selection, estimated_bytes) in selected {
+ // Foreground and merge reads only try memory admission.
They
+ // never hold a shared scheduling slot while awaiting
another
+ // merge input, nor wait for output buffers owned
downstream.
+ let _permit = budget.reserve_memory(estimated_bytes)?;
+ let mut stream = build_row_group_stream(
+ Arc::clone(&shared_reader), file_size,
metadata.clone(), mask.clone(),
+ index, batch_size, selection, shared_filters.clone(),
+ )?;
+ while let Some(batch) = stream.next().await {
+ let batch = batch.map_err(Error::from)?;
+ let batch = match &map_read_plan {
+ Some(plan) => plan.assemble_batch(&batch)?,
+ None => batch,
+ };
+ let batch = match &residual {
+ Some(predicates) =>
crate::arrow::residual::filter_record_batch_by_predicates(
+ batch, predicates, &scan_fields,
+ )?,
+ None => batch,
+ };
+ yield batch;
+ }
+ }
+ }.boxed());
+ }
+
let batch_stream = batch_stream_builder.build()?;
if all_enforced {
@@ -977,58 +1061,139 @@ async fn read_row_group(
row_group_index: usize,
batch_size: Option<usize>,
selection: Option<RowSelection>,
- _permit: ReadPermit,
+ permit: ReadPermit,
sender: mpsc::Sender<ParquetRowGroupMessage>,
) {
+ let stream = match build_row_group_stream(
+ reader,
+ file_size,
+ reader_metadata,
+ projection,
+ row_group_index,
+ batch_size,
+ selection,
+ Vec::new(),
+ ) {
+ Ok(stream) => stream,
+ Err(error) => {
+ drop(permit);
+ let _ = sender.send(ParquetRowGroupMessage::Error(error)).await;
+ return;
+ }
+ };
+ forward_row_group_batches(stream, sender, Some(permit)).await;
+}
+
+#[derive(Clone)]
+struct SharedParquetPredicate {
+ projection: ProjectionMask,
+ inner: Arc<Mutex<Box<dyn ArrowPredicate>>>,
+}
+
+impl SharedParquetPredicate {
+ fn new(predicate: Box<dyn ArrowPredicate>) -> Self {
+ Self {
+ projection: predicate.projection().clone(),
+ inner: Arc::new(Mutex::new(predicate)),
+ }
+ }
+}
+
+impl ArrowPredicate for SharedParquetPredicate {
+ fn projection(&self) -> &ProjectionMask {
+ &self.projection
+ }
+
+ fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray,
arrow_schema::ArrowError> {
+ self.inner
+ .lock()
+ .map_err(|_| {
+ arrow_schema::ArrowError::ComputeError(
+ "Parquet row predicate mutex was poisoned".into(),
+ )
+ })?
+ .evaluate(batch)
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn build_row_group_stream(
+ reader: Arc<dyn FileRead>,
+ file_size: u64,
+ metadata: ArrowReaderMetadata,
+ projection: ProjectionMask,
+ index: usize,
+ batch_size: Option<usize>,
+ selection: Option<RowSelection>,
+ predicates: Vec<SharedParquetPredicate>,
+) ->
crate::Result<parquet::arrow::async_reader::ParquetRecordBatchStream<ArrowFileReader>>
{
let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata(
ArrowFileReader::new(file_size, reader),
- reader_metadata,
+ metadata,
)
.with_projection(projection)
- .with_row_groups(vec![row_group_index]);
+ .with_row_groups(vec![index]);
if let Some(selection) = selection {
builder = builder.with_row_selection(selection);
}
if let Some(size) = batch_size {
builder = builder.with_batch_size(size);
}
- let mut stream = match builder.build() {
- Ok(stream) => stream,
- Err(error) => {
- let _ = sender
- .send(ParquetRowGroupMessage::Error(error.into()))
- .await;
- return;
- }
- };
-
- forward_row_group_batches(&mut stream, sender).await;
+ if !predicates.is_empty() {
+ builder = builder.with_row_filter(ParquetRowFilter::new(
+ predicates
+ .into_iter()
+ .map(|p| Box::new(p) as Box<dyn ArrowPredicate>)
+ .collect(),
+ ));
+ }
+ builder.build().map_err(Error::from)
}
-async fn forward_row_group_batches<S, E>(
- mut stream: S,
+fn forward_row_group_batches<S, E>(
+ stream: S,
sender: mpsc::Sender<ParquetRowGroupMessage>,
-) where
+ permit: Option<ReadPermit>,
+) -> impl std::future::Future<Output = ()>
+where
S: futures::Stream<Item = std::result::Result<RecordBatch, E>> + Unpin,
E: Into<Error>,
{
- loop {
- let Ok(slot) = sender.reserve().await else {
- return;
- };
- let next = tokio::select! {
- _ = sender.closed() => return,
- next = stream.next() => next,
- };
- match next {
- Some(Ok(batch)) => slot.send(ParquetRowGroupMessage::Batch(batch)),
- Some(Err(error)) => {
- slot.send(ParquetRowGroupMessage::Error(error.into()));
- return;
- }
- None => {
- slot.send(ParquetRowGroupMessage::Done);
+ // Struct fields drop in declaration order, even if the future is cancelled
+ // before its first poll. Keep the working reservation until the decoder
drops.
+ struct ReservedDecoder<S> {
+ stream: S,
+ _permit: Option<ReadPermit>,
+ }
+ let mut decoder = ReservedDecoder {
+ stream,
+ _permit: permit,
+ };
+ async move {
+ loop {
+ let Ok(slot) = sender.reserve().await else {
return;
+ };
+ let next = tokio::select! {
+ _ = sender.closed() => return,
+ next = decoder.stream.next() => next,
+ };
+ match next {
+ Some(Ok(batch)) => {
+ slot.send(ParquetRowGroupMessage::Batch(batch));
+ }
+ Some(Err(error)) => {
+ drop(decoder);
+ slot.send(ParquetRowGroupMessage::Error(error.into()));
+ return;
+ }
+ None => {
+ // Release decoder buffers and their working estimate
before
+ // the next demanded group retries.
+ drop(decoder);
+ slot.send(ParquetRowGroupMessage::Done);
+ return;
+ }
}
}
}
@@ -3401,7 +3566,7 @@ mod tests {
tracked_polls.fetch_add(1, AtomicOrdering::SeqCst);
});
let (tx, mut rx) = mpsc::channel(1);
- let task = tokio::spawn(forward_row_group_batches(stream, tx));
+ let task = tokio::spawn(forward_row_group_batches(stream, tx, None));
tokio::time::sleep(Duration::from_millis(20)).await;
assert_eq!(
@@ -3420,6 +3585,62 @@ mod tests {
task.await.unwrap();
}
+ #[tokio::test]
+ async fn cancelled_row_group_drops_decoder_before_returning_memory() {
+ use crate::resource::ResourceContext;
+
+ struct Decoder {
+ resources: ResourceContext,
+ charge_at_drop: Arc<AtomicUsize>,
+ }
+ impl futures::Stream for Decoder {
+ type Item = Result<RecordBatch, Error>;
+
+ fn poll_next(
+ self: std::pin::Pin<&mut Self>,
+ _: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Option<Self::Item>> {
+ std::task::Poll::Pending
+ }
+ }
+ impl Drop for Decoder {
+ fn drop(&mut self) {
+ self.charge_at_drop.store(
+ self.resources.metrics().reserved_memory_bytes,
+ AtomicOrdering::SeqCst,
+ );
+ }
+ }
+
+ for mode in 0..3 {
+ let resources =
ResourceContext::builder().memory_limit(8).build().unwrap();
+ let budget =
ReadBudget::default().with_resources(resources.clone());
+ let permit = budget.reserve_memory(8).unwrap();
+ let charge_at_drop = Arc::new(AtomicUsize::new(usize::MAX));
+ let decoder = Decoder {
+ resources: resources.clone(),
+ charge_at_drop: charge_at_drop.clone(),
+ };
+ let (tx, rx) = mpsc::channel(1);
+ let mut forwarding = Box::pin(forward_row_group_batches(decoder,
tx, Some(permit)));
+ if mode > 0 {
+ assert!(futures::poll!(&mut forwarding).is_pending());
+ }
+ if mode == 2 {
+ drop(rx);
+ forwarding.as_mut().await;
+ }
+ // Cover an unpolled future, cancellation during I/O, and a closed
receiver.
+ drop(forwarding);
+ assert_eq!(
+ charge_at_drop.load(AtomicOrdering::SeqCst),
+ 8,
+ "mode {mode}"
+ );
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
+
#[tokio::test]
async fn test_row_group_batch_forwarding_stops_during_pending_io() {
let (polled_tx, polled_rx) = tokio::sync::oneshot::channel();
@@ -3431,7 +3652,7 @@ mod tests {
std::task::Poll::Pending::<Option<Result<RecordBatch, Error>>>
});
let (tx, rx) = mpsc::channel(1);
- let task = tokio::spawn(forward_row_group_batches(stream, tx));
+ let task = tokio::spawn(forward_row_group_batches(stream, tx, None));
polled_rx.await.unwrap();
drop(rx);
@@ -4333,6 +4554,228 @@ mod tests {
}
}
+ #[tokio::test]
+ async fn resource_admission_precedes_data_reads_on_all_parquet_paths() {
+ use crate::resource::ResourceContext;
+ for (group_rows, parallelism, filtered) in [
+ (16, 4, false),
+ (64, 4, false),
+ (16, 1, false),
+ (16, 4, true),
+ ] {
+ let data = Bytes::from(
+ write_multi_row_group_parquet(group_rows, 64,
EnabledStatistics::None, true).await,
+ );
+ let tracker = TrackingFileRead::new(data.clone());
+ let resources =
ResourceContext::builder().memory_limit(0).build().unwrap();
+ let budget = Arc::new(
+ ReadBudget::new(parallelism, 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone()),
+ );
+ let fields = vec![int_field("id"), int_field("value")];
+ let predicates = filtered.then(|| FilePredicates {
+ predicates: vec![PredicateBuilder::new(&fields)
+ .greater_than("value", Datum::Int(10))
+ .unwrap()],
+ row_filter_factory: None,
+ file_fields: fields.clone(),
+ });
+ let mut stream = ParquetFormatReader::with_read_budget(budget)
+ .read_batch_stream(
+ Box::new(tracker.clone()),
+ data.len() as u64,
+ &fields[..1],
+ predicates.as_ref(),
+ Some(16),
+ None,
+ )
+ .await
+ .unwrap();
+ tracker.reset();
+ assert!(matches!(
+ stream.next().await.unwrap(),
+ Err(Error::ResourceExhausted { .. })
+ ));
+ assert!(stream.next().await.is_none());
+ assert_eq!(
+ tracker.bytes_read(),
+ 0,
+ "rejected work must not start data I/O"
+ );
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
+
+ #[tokio::test]
+ async fn
shared_memory_pressure_reduces_prefetch_without_failing_the_read() {
+ use crate::resource::ResourceContext;
+ let data = Bytes::from(
+ write_multi_row_group_parquet(64, 256, EnabledStatistics::Chunk,
false).await,
+ );
+ let metadata = load_metadata_with_page_index(&data, true);
+ let mask =
super::ProjectionMask::roots(metadata.file_metadata().schema_descr(), [0]);
+ let limit = metadata
+ .row_groups()
+ .iter()
+ .map(|rg| super::projected_row_group_bytes(rg, &mask))
+ .max()
+ .unwrap() as usize;
+ // Only one row group fits; the consumer owns its output batches.
+ let resources = ResourceContext::builder()
+ .memory_limit(limit)
+ .build()
+ .unwrap();
+ let budget = Arc::new(
+ ReadBudget::new(8, 256 * 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone()),
+ );
+ budget.enable_diagnostics();
+ let mut stream = ParquetFormatReader::with_read_budget(budget.clone())
+ .read_batch_stream(
+ Box::new(TrackingFileRead::new(data.clone())),
+ data.len() as u64,
+ &[int_field("id")],
+ None,
+ Some(16),
+ None,
+ )
+ .await
+ .unwrap();
+ let mut ids = Vec::new();
+ while let Some(batch) = tokio::time::timeout(Duration::from_secs(2),
stream.try_next())
+ .await
+ .unwrap()
+ .unwrap()
+ {
+ ids.extend_from_slice(
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values(),
+ );
+ }
+ assert_eq!(ids, (0..256).collect::<Vec<_>>());
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ assert!(resources.metrics().peak_reserved_memory_bytes <= limit);
+ assert_eq!(budget.diagnostics().peak_inflight, 1);
+ }
+
+ #[tokio::test]
+ async fn
cancelling_shared_budget_prefetch_releases_memory_with_escaped_output() {
+ use crate::resource::ResourceContext;
+ let data = Bytes::from(
+ write_multi_row_group_parquet(64, 256, EnabledStatistics::Chunk,
false).await,
+ );
+ let resources = ResourceContext::builder()
+ .memory_limit(1024 * 1024)
+ .build()
+ .unwrap();
+ let budget = Arc::new(
+ ReadBudget::new(4, 256 * 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone()),
+ );
+ budget.enable_diagnostics();
+ let mut stream = ParquetFormatReader::with_read_budget(budget.clone())
+ .read_batch_stream(
+ Box::new(TrackingFileRead::new(data.clone())),
+ data.len() as u64,
+ &[int_field("id")],
+ None,
+ Some(16),
+ None,
+ )
+ .await
+ .unwrap();
+ let batch = stream.try_next().await.unwrap().unwrap();
+ let escaped = batch.column(0).slice(0, 1);
+ drop(batch);
+ drop(stream);
+ tokio::time::timeout(Duration::from_secs(2), async {
+ while budget.diagnostics().current_inflight != 0 {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .unwrap();
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ assert_eq!(
+ escaped
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 0
+ );
+ drop(escaped);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+
+ #[tokio::test]
+ async fn
filtered_resource_reads_reserve_predicate_columns_and_preserve_row_selection() {
+ use crate::resource::ResourceContext;
+ let data =
+ Bytes::from(write_multi_row_group_parquet(16, 64,
EnabledStatistics::None, true).await);
+ let metadata = load_metadata_with_page_index(&data, false);
+ let mask =
super::ProjectionMask::roots(metadata.file_metadata().schema_descr(), [0]);
+ let id_only =
super::projected_row_group_bytes(&metadata.row_groups()[0], &mask) as usize;
+ let fields = vec![int_field("id"), int_field("value")];
+ let predicates = FilePredicates {
+ predicates: vec![PredicateBuilder::new(&fields)
+ .greater_or_equal("value", Datum::Int(250))
+ .unwrap()],
+ row_filter_factory: None,
+ file_fields: fields.clone(),
+ };
+ for limit in [id_only, 1024 * 1024] {
+ let resources = ResourceContext::builder()
+ .memory_limit(limit)
+ .build()
+ .unwrap();
+ let budget = Arc::new(
+ ReadBudget::new(4, 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone()),
+ );
+ let mut stream = ParquetFormatReader::with_read_budget(budget)
+ .read_batch_stream(
+ Box::new(TrackingFileRead::new(data.clone())),
+ data.len() as u64,
+ &fields[..1],
+ Some(&predicates),
+ Some(8),
+ Some(vec![RowRange::new(10, 54)]),
+ )
+ .await
+ .unwrap();
+ if limit == id_only {
+ assert!(matches!(
+ stream.next().await.unwrap(),
+ Err(Error::ResourceExhausted { .. })
+ ));
+ assert!(stream.next().await.is_none());
+ } else {
+ let mut ids = Vec::new();
+ while let Some(batch) = stream.try_next().await.unwrap() {
+ ids.extend_from_slice(
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values(),
+ );
+ }
+ assert_eq!(ids, (25..=54).collect::<Vec<_>>());
+ }
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
+
#[tokio::test]
async fn test_sparse_read_buffer_owners_and_cancellation() {
use crate::io::FileRead;
@@ -5013,6 +5456,7 @@ mod tests {
.unwrap();
let props = parquet::file::properties::WriterProperties::builder()
.set_compression(parquet::basic::Compression::UNCOMPRESSED)
+ .set_max_row_group_row_count(Some(16))
.build();
let mut bytes = Vec::new();
{
@@ -5045,22 +5489,38 @@ mod tests {
file_fields: fields.clone(),
};
- ParquetFormatReader::default()
- .read_batch_stream(
- Box::new(input.reader().await.unwrap()),
- file_size,
- &fields[..1],
- Some(&predicates),
- None,
- None,
- )
- .await
- .unwrap()
- .try_collect::<Vec<_>>()
- .await
+ let resources = crate::resource::ResourceContext::builder()
+ .memory_limit(1024 * 1024)
+ .build()
.unwrap();
+ for reader in [
+ ParquetFormatReader::default(),
+ ParquetFormatReader::with_read_budget(Arc::new(
+ ReadBudget::default().with_resources(resources.clone()),
+ )),
+ ] {
+ order.lock().unwrap().clear();
+ reader
+ .read_batch_stream(
+ Box::new(input.reader().await.unwrap()),
+ file_size,
+ &fields[..1],
+ Some(&predicates),
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
- assert_eq!(*order.lock().unwrap(), vec!["small", "large"]);
+ assert_eq!(
+ *order.lock().unwrap(),
+ vec!["small", "large", "small", "large"]
+ );
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
}
/// Read `[name]` from the `(id, name, age)` parquet file under
`predicates`
diff --git a/crates/paimon/src/arrow/read_budget.rs
b/crates/paimon/src/arrow/read_budget.rs
index 334ac8e3..fc6531bd 100644
--- a/crates/paimon/src/arrow/read_budget.rs
+++ b/crates/paimon/src/arrow/read_budget.rs
@@ -20,6 +20,8 @@ use std::sync::Arc;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use crate::resource::{MemoryReservation, ResourceContext};
+
use crate::spec::PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES_OPTION;
const BYTE_PERMIT_UNIT: u64 = 1024 * 1024;
@@ -27,19 +29,23 @@ const DEFAULT_BYTE_OPTION: &str = "the row-group read
budget";
const DEFAULT_PARALLELISM: usize = 8;
const DEFAULT_MAX_INFLIGHT_BYTES: u64 = 256 * 1024 * 1024;
-/// Resource budget for concurrent row-group reads, bounding both the number of
-/// row groups in flight and their estimated bytes. Parquet holds one per scan;
-/// Mosaic builds one per open file from its own prefetch options.
-#[derive(Debug)]
+/// Scheduling window for concurrent row-group reads, with memory reservations
+/// backed by a [`ResourceContext`]. Slot and prefetch-byte limits control how
+/// much work may start; the context admits the full estimate without clamping.
+/// Without an explicit context, memory reservations have no additional limit
+/// and the existing oversized-row-group scheduling behavior is preserved.
+#[derive(Debug, Clone)]
pub struct ReadBudget {
parallelism: usize,
row_groups: Arc<Semaphore>,
- bytes: Arc<Semaphore>,
+ // A look-ahead window, independent of shared memory accounting.
+ prefetch: Arc<Semaphore>,
+ resources: Option<ResourceContext>,
byte_permits: u32,
byte_permit_unit: u64,
max_inflight_bytes: u64,
byte_option: &'static str,
- oversized_warning_logged: AtomicBool,
+ oversized_warning_logged: Arc<AtomicBool>,
diagnostics: Arc<ReadBudgetDiagnostics>,
}
@@ -79,8 +85,8 @@ pub(crate) struct ReadBudgetDiagnosticsSnapshot {
}
impl ReadBudget {
- /// Scan-wide budget with MiB-granular byte permits, as the Parquet reader
- /// takes it; its oversized-row-group warning names the Parquet option.
+ /// Scan-wide scheduling window with MiB-granular prefetch permits.
+ /// Its oversized-row-group warning names the Parquet option.
pub fn new(parallelism: usize, max_inflight_bytes: u64) ->
crate::Result<Self> {
Ok(
Self::with_byte_granularity(parallelism, max_inflight_bytes,
BYTE_PERMIT_UNIT)?
@@ -88,9 +94,9 @@ impl ReadBudget {
)
}
- /// Budget whose byte permits count `byte_permit_unit` bytes each. Mosaic
- /// charges the exact estimated bytes of a row group, so it passes `1`; at
- /// most `u32::MAX` permits are tracked, which caps that unit at 4 GiB.
+ /// Scheduling window whose permits count `byte_permit_unit` bytes each.
+ /// Mosaic passes `1`; at most `u32::MAX` permits are tracked, which caps
+ /// that window at 4 GiB. Memory reservations always use the full estimate.
pub(crate) fn with_byte_granularity(
parallelism: usize,
max_inflight_bytes: u64,
@@ -120,12 +126,13 @@ impl ReadBudget {
Ok(Self {
parallelism,
row_groups: Arc::new(Semaphore::new(parallelism)),
- bytes: Arc::new(Semaphore::new(byte_permits as usize)),
+ prefetch: Arc::new(Semaphore::new(byte_permits as usize)),
+ resources: None,
byte_permits,
byte_permit_unit,
max_inflight_bytes,
byte_option: DEFAULT_BYTE_OPTION,
- oversized_warning_logged: AtomicBool::new(false),
+ oversized_warning_logged: Arc::new(AtomicBool::new(false)),
diagnostics: Arc::new(ReadBudgetDiagnostics::default()),
})
}
@@ -191,28 +198,88 @@ impl ReadBudget {
}
}
- /// Wait for a row-group slot and the estimated bytes. The Parquet reader
- /// acquires from its async tasks.
+ /// Bind a shared memory budget while preserving this scheduling window.
+ /// Working estimates draw from this context. An
+ /// oversized row group must still fit the context's memory limit.
+ pub fn with_resources(&self, resources: ResourceContext) -> Self {
+ Self {
+ resources: Some(resources),
+ ..self.clone()
+ }
+ }
+
+ pub(crate) fn has_resources(&self) -> bool {
+ self.resources.is_some()
+ }
+
+ /// Merge inputs need to advance in lockstep. Keep memory admission, while
+ /// disabling background row-group prefetch that could hold slots they
need.
+ pub(crate) fn without_prefetch(&self) -> Self {
+ Self {
+ parallelism: 1,
+ ..self.clone()
+ }
+ }
+
+ /// Wait for scheduling capacity, then try memory admission. Memory
rejection
+ /// is immediate: waiting for a caller's retained buffers could deadlock.
pub(crate) async fn acquire(&self, estimated_bytes: u64) ->
crate::Result<ReadPermit> {
let row_group = Arc::clone(&self.row_groups)
.acquire_owned()
.await
.map_err(|_| Self::closed("row-group"))?;
- let bytes = Arc::clone(&self.bytes)
+ let prefetch = Arc::clone(&self.prefetch)
.acquire_many_owned(self.byte_permits_for(estimated_bytes))
.await
- .map_err(|_| Self::closed("byte"))?;
- Ok(self.permit(row_group, bytes))
+ .map_err(|_| Self::closed("prefetch"))?;
+ let mut permit = self.reserve_memory(estimated_bytes)?;
+ permit.prefetch = Some((row_group, prefetch));
+ Ok(permit)
}
- /// Take a row-group slot and the estimated bytes if both are free now. The
- /// Mosaic reader refills its look-ahead without waiting.
+ /// Take scheduling capacity and memory if both are available now.
pub(crate) fn try_acquire(&self, estimated_bytes: u64) ->
Option<ReadPermit> {
let row_group = Arc::clone(&self.row_groups).try_acquire_owned().ok()?;
- let bytes = Arc::clone(&self.bytes)
+ let prefetch = Arc::clone(&self.prefetch)
.try_acquire_many_owned(self.byte_permits_for(estimated_bytes))
.ok()?;
- Some(self.permit(row_group, bytes))
+ let mut permit = self.reserve_memory(estimated_bytes).ok()?;
+ permit.prefetch = Some((row_group, prefetch));
+ Some(permit)
+ }
+
+ /// Foreground reads, including merge inputs, must not wait while another
+ /// input holds a scheduling slot. They still reserve from the shared pool.
+ pub(crate) fn reserve_memory(&self, estimated_bytes: u64) ->
crate::Result<ReadPermit> {
+ let bytes =
+ usize::try_from(estimated_bytes).map_err(|_|
crate::Error::ResourceExhausted {
+ message: format!("Row-group estimate {estimated_bytes} exceeds
addressable memory"),
+ })?;
+ let memory = self
+ .resources
+ .as_ref()
+ .map(|resources| {
+ let mut memory = resources.reservation();
+ memory.try_grow(bytes)?;
+ Ok::<_, crate::Error>(memory)
+ })
+ .transpose()?;
+ let diagnostics = self.diagnostics_enabled().then(|| {
+ let current = self
+ .diagnostics
+ .current_inflight
+ .fetch_add(1, Ordering::Relaxed)
+ + 1;
+ self.diagnostics
+ .peak_inflight
+ .fetch_max(current, Ordering::Relaxed);
+ Arc::clone(&self.diagnostics)
+ });
+ Ok(ReadPermit {
+ _memory: memory,
+ prefetch: None,
+ diagnostics,
+ })
}
/// Block until a row-group slot and the estimated bytes are free. The
@@ -234,7 +301,7 @@ impl ReadBudget {
{
log::warn!(
"A row group's estimated size ({estimated_bytes} bytes)
exceeds {} ({} bytes); it \
- will consume the entire byte budget and may reduce row-group
read parallelism; \
+ will occupy the entire prefetch window and may reduce
row-group read parallelism; \
increase the option if memory allows",
self.byte_option,
self.max_inflight_bytes
@@ -245,25 +312,6 @@ impl ReadBudget {
.div_ceil(self.byte_permit_unit)
.min(u64::from(self.byte_permits)) as u32
}
-
- fn permit(&self, row_group: OwnedSemaphorePermit, bytes:
OwnedSemaphorePermit) -> ReadPermit {
- let diagnostics = self.diagnostics_enabled().then(|| {
- let current = self
- .diagnostics
- .current_inflight
- .fetch_add(1, Ordering::Relaxed)
- + 1;
- self.diagnostics
- .peak_inflight
- .fetch_max(current, Ordering::Relaxed);
- Arc::clone(&self.diagnostics)
- });
- ReadPermit {
- _row_group: row_group,
- _bytes: bytes,
- diagnostics,
- }
- }
}
impl Default for ReadBudget {
@@ -275,8 +323,9 @@ impl Default for ReadBudget {
#[derive(Debug)]
pub(crate) struct ReadPermit {
- _row_group: OwnedSemaphorePermit,
- _bytes: OwnedSemaphorePermit,
+ // Return memory before waking scheduling waiters.
+ _memory: Option<MemoryReservation>,
+ prefetch: Option<(OwnedSemaphorePermit, OwnedSemaphorePermit)>,
diagnostics: Option<Arc<ReadBudgetDiagnostics>>,
}
@@ -438,4 +487,66 @@ mod tests {
drop(held);
waiter.join().unwrap().unwrap();
}
+
+ #[tokio::test]
+ async fn row_groups_and_other_consumers_share_one_memory_limit() {
+ let resources =
ResourceContext::builder().memory_limit(32).build().unwrap();
+ let first = ReadBudget::new(2, 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone());
+ let second = ReadBudget::new(2, 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone());
+ let read = first.acquire(20).await.unwrap();
+ let mut consumer = resources.reservation();
+ consumer.try_grow(8).unwrap();
+ assert!(matches!(
+ second.acquire(5).await,
+ Err(crate::Error::ResourceExhausted { .. })
+ ));
+ assert_eq!(resources.metrics().reserved_memory_bytes, 28);
+ // Rejected memory admission must release both scheduling resources.
+ let other = second.acquire(4).await.unwrap();
+ assert_eq!(resources.metrics().reserved_memory_bytes, 32);
+ drop((read, other, consumer));
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ assert_eq!(resources.metrics().peak_reserved_memory_bytes, 32);
+ }
+
+ #[tokio::test]
+ async fn oversized_prefetch_admission_cannot_bypass_the_memory_limit() {
+ let resources =
ResourceContext::builder().memory_limit(3).build().unwrap();
+ let budget = ReadBudget::with_byte_granularity(2, 1, 1)
+ .unwrap()
+ .with_resources(resources.clone());
+ assert!(matches!(
+ budget.acquire(4).await,
+ Err(crate::Error::ResourceExhausted { .. })
+ ));
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ let permit = budget.acquire(3).await.unwrap();
+ // The prefetch window is clamped to one; the memory charge is all
three.
+ assert_eq!(resources.metrics().reserved_memory_bytes, 3);
+ drop(permit);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+
+ #[tokio::test]
+ async fn cancelling_a_prefetch_waiter_does_not_hold_slots_or_memory() {
+ let resources =
ResourceContext::builder().memory_limit(32).build().unwrap();
+ let budget = ReadBudget::new(2, 1024 * 1024)
+ .unwrap()
+ .with_resources(resources.clone());
+ let first = budget.acquire(16).await.unwrap();
+ let mut waiting = Box::pin(budget.acquire(16));
+ assert!(futures::poll!(&mut waiting).is_pending());
+ assert_eq!(resources.metrics().reserved_memory_bytes, 16);
+ drop(waiting);
+ drop(first);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ let first = budget.acquire(32).await.unwrap();
+ drop(first);
+ assert_eq!(budget.row_groups.available_permits(), 2);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
}
diff --git a/crates/paimon/src/error.rs b/crates/paimon/src/error.rs
index 18696cdd..44aa4a74 100644
--- a/crates/paimon/src/error.rs
+++ b/crates/paimon/src/error.rs
@@ -23,6 +23,9 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
/// Error type for paimon.
#[derive(Debug, Snafu)]
pub enum Error {
+ /// A configured resource budget could not admit the requested reservation.
+ #[snafu(display("Paimon resource exhausted: {}", message))]
+ ResourceExhausted { message: String },
#[snafu(whatever, display("Paimon data invalid for {}: {:?}", message,
source))]
DataInvalid {
message: String,
diff --git a/crates/paimon/src/lib.rs b/crates/paimon/src/lib.rs
index 9e6b1340..12c67eb5 100644
--- a/crates/paimon/src/lib.rs
+++ b/crates/paimon/src/lib.rs
@@ -38,6 +38,7 @@ pub mod full_text;
pub mod io;
pub mod lumina;
mod predicate_stats;
+pub mod resource;
pub mod spec;
pub mod table;
pub mod variant;
diff --git a/crates/paimon/src/resource/memory.rs
b/crates/paimon/src/resource/memory.rs
new file mode 100644
index 00000000..4be2feea
--- /dev/null
+++ b/crates/paimon/src/resource/memory.rs
@@ -0,0 +1,150 @@
+// 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.
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+use crate::{Error, Result};
+
+/// Optional adapter to an embedding engine's memory budget.
+///
+/// Calls may run concurrently on any runtime thread. Failed reservations must
+/// leave the pool unchanged. Methods must not panic; `release` is infallible
and
+/// may run during unwinding. The pool must not wait for, or call back into, a
+/// reader to free memory. Reclamation belongs to the consumer, outside the
pool.
+pub trait MemoryPool: std::fmt::Debug + Send + Sync + 'static {
+ fn try_reserve(&self, bytes: usize) -> Result<()>;
+ fn release(&self, bytes: usize);
+}
+
+/// Exact reservation counters, not measurements of the process's allocations.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct ResourceMetrics {
+ pub reserved_memory_bytes: usize,
+ pub peak_reserved_memory_bytes: usize,
+}
+
+#[derive(Debug)]
+pub(super) struct MemoryAccount {
+ limit: usize,
+ external: Option<Arc<dyn MemoryPool>>,
+ // Includes requests awaiting approval from the external pool. This makes
+ // concurrent admission obey the local cap even before external approval.
+ admitted: AtomicUsize,
+ reserved: AtomicUsize,
+ peak: AtomicUsize,
+}
+
+impl MemoryAccount {
+ pub(super) fn new(limit: Option<usize>, external: Option<Arc<dyn
MemoryPool>>) -> Self {
+ Self {
+ limit: limit.unwrap_or(usize::MAX),
+ external,
+ admitted: AtomicUsize::new(0),
+ reserved: AtomicUsize::new(0),
+ peak: AtomicUsize::new(0),
+ }
+ }
+
+ fn try_reserve(&self, bytes: usize) -> Result<()> {
+ if bytes == 0 {
+ return Ok(());
+ }
+ self.admitted
+ .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
+ used.checked_add(bytes).filter(|&next| next <= self.limit)
+ })
+ .map_err(|used| Error::ResourceExhausted {
+ message: format!(
+ "Cannot reserve {bytes} bytes: {used} bytes already
reserved or pending, limit {}",
+ self.limit
+ ),
+ })?;
+ if let Some(pool) = &self.external {
+ if let Err(error) = pool.try_reserve(bytes) {
+ self.admitted.fetch_sub(bytes, Ordering::Relaxed);
+ return Err(error);
+ }
+ }
+ let used = self.reserved.fetch_add(bytes, Ordering::Relaxed) + bytes;
+ self.peak.fetch_max(used, Ordering::Relaxed);
+ Ok(())
+ }
+
+ fn release(&self, bytes: usize) {
+ if bytes == 0 {
+ return;
+ }
+ self.reserved.fetch_sub(bytes, Ordering::Relaxed);
+ if let Some(pool) = &self.external {
+ pool.release(bytes);
+ }
+ self.admitted.fetch_sub(bytes, Ordering::Relaxed);
+ }
+
+ pub(super) fn metrics(&self) -> ResourceMetrics {
+ ResourceMetrics {
+ reserved_memory_bytes: self.reserved.load(Ordering::Relaxed),
+ peak_reserved_memory_bytes: self.peak.load(Ordering::Relaxed),
+ }
+ }
+}
+
+/// Owned reservation. Moving it transfers accounting; dropping it releases it.
+///
+/// This type deliberately does not implement `Clone`. Shared allocations
should
+/// retain one reservation in their shared owner rather than charge every
alias.
+#[derive(Debug)]
+pub struct MemoryReservation {
+ account: Arc<MemoryAccount>,
+ size: usize,
+}
+
+impl MemoryReservation {
+ pub(super) fn new(account: Arc<MemoryAccount>) -> Self {
+ Self { account, size: 0 }
+ }
+
+ pub fn size(&self) -> usize {
+ self.size
+ }
+
+ /// Grow atomically with respect to admission. Failure preserves this
guard.
+ pub fn try_grow(&mut self, bytes: usize) -> Result<()> {
+ self.account.try_reserve(bytes)?;
+ // Account-wide checked admission also guarantees this sum fits.
+ self.size += bytes;
+ Ok(())
+ }
+
+ /// Reserve additional bytes, or release surplus bytes, to reach `size`.
+ pub fn try_resize(&mut self, size: usize) -> Result<()> {
+ if size > self.size {
+ self.try_grow(size - self.size)?;
+ } else {
+ self.account.release(self.size - size);
+ self.size = size;
+ }
+ Ok(())
+ }
+}
+
+impl Drop for MemoryReservation {
+ fn drop(&mut self) {
+ self.account.release(self.size);
+ }
+}
diff --git a/crates/paimon/src/resource/mod.rs
b/crates/paimon/src/resource/mod.rs
new file mode 100644
index 00000000..24e99f18
--- /dev/null
+++ b/crates/paimon/src/resource/mod.rs
@@ -0,0 +1,112 @@
+// 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.
+
+//! Shared memory reservations for readers and embedding engines.
+//!
+//! Parquet readers reserve each selected row group's projected uncompressed
+//! column size before data I/O, including columns needed by decoder
predicates.
+//! The decoder owns that reservation and releases it when dropped. Concurrency
+//! slots and the prefetch window remain separate scheduling controls.
+//!
+//! Consumers are responsible for reserving memory they retain. Output batches
+//! are ordinary Arrow batches: a downstream consumer that holds them must
reserve
+//! its own memory. The reader does not wrap their buffers or charge their
aliases.
+//!
+//! Reservations are accounting, not an allocator or an RSS limit. Row-group
+//! charges are estimates; metadata, merge state, transient batches and
allocation
+//! overhead are not fully covered, and actual memory can exceed the estimates.
+
+mod memory;
+pub use memory::{MemoryPool, MemoryReservation, ResourceMetrics};
+
+use std::sync::Arc;
+
+use self::memory::MemoryAccount;
+use crate::Result;
+
+/// Memory budget shared by consumers in one logical operation.
+///
+/// Clones share admission and metrics. Each consumer owns its reservations and
+/// releases them when its working state is dropped. Retained output batches
are
+/// the caller's responsibility, including when they outlive the reader.
+///
+/// ```
+/// use paimon::resource::ResourceContext;
+///
+/// let resources = ResourceContext::builder()
+/// .memory_limit(256 * 1024 * 1024)
+/// .build()?;
+/// // Pass resources.clone() to ReadBuilder::with_resources.
+/// // Other consumers reserve from the same budget for their own retained
state.
+/// let mut reservation = resources.reservation();
+/// reservation.try_grow(1024)?;
+/// assert_eq!(resources.metrics().reserved_memory_bytes, 1024);
+/// drop(reservation);
+/// assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+/// # Ok::<(), paimon::Error>(())
+/// ```
+#[derive(Clone, Debug)]
+pub struct ResourceContext {
+ memory: Arc<MemoryAccount>,
+}
+
+impl ResourceContext {
+ pub fn builder() -> ResourceContextBuilder {
+ ResourceContextBuilder::default()
+ }
+
+ /// Create an initially empty reservation owned by a consumer.
+ pub fn reservation(&self) -> MemoryReservation {
+ MemoryReservation::new(Arc::clone(&self.memory))
+ }
+
+ /// Read accounting counters. The two counters are sampled independently.
+ pub fn metrics(&self) -> ResourceMetrics {
+ self.memory.metrics()
+ }
+}
+
+/// Configure a local limit and an optional embedding-engine pool.
+#[derive(Default, Debug)]
+pub struct ResourceContextBuilder {
+ memory_limit: Option<usize>,
+ memory_pool: Option<Arc<dyn MemoryPool>>,
+}
+
+impl ResourceContextBuilder {
+ /// Limit outstanding reservations. Zero allows only zero-byte
reservations.
+ /// Omitting this setting leaves the context without a local limit.
+ pub fn memory_limit(mut self, bytes: usize) -> Self {
+ self.memory_limit = Some(bytes);
+ self
+ }
+
+ /// Also reserve from an external pool. Both limits must permit each
request.
+ pub fn memory_pool(mut self, pool: Arc<dyn MemoryPool>) -> Self {
+ self.memory_pool = Some(pool);
+ self
+ }
+
+ pub fn build(self) -> Result<ResourceContext> {
+ Ok(ResourceContext {
+ memory: Arc::new(MemoryAccount::new(self.memory_limit,
self.memory_pool)),
+ })
+ }
+}
+
+#[cfg(test)]
+mod tests;
diff --git a/crates/paimon/src/resource/tests.rs
b/crates/paimon/src/resource/tests.rs
new file mode 100644
index 00000000..7db2d122
--- /dev/null
+++ b/crates/paimon/src/resource/tests.rs
@@ -0,0 +1,163 @@
+// 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.
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Barrier};
+
+use super::{MemoryPool, ResourceContext};
+use crate::{Error, Result};
+
+fn used(context: &ResourceContext) -> usize {
+ context.metrics().reserved_memory_bytes
+}
+
+#[test]
+fn reservations_share_the_limit_and_failed_growth_is_atomic() {
+ let resources =
ResourceContext::builder().memory_limit(8).build().unwrap();
+ let mut first = resources.reservation();
+ let mut second = resources.clone().reservation();
+ first.try_grow(8).unwrap();
+ assert!(matches!(
+ second.try_grow(1),
+ Err(Error::ResourceExhausted { .. })
+ ));
+ assert_eq!(second.size(), 0);
+ first.try_resize(3).unwrap();
+ second.try_grow(5).unwrap();
+ assert!(first.try_resize(9).is_err());
+ assert_eq!(first.size(), 3);
+ assert_eq!(used(&resources), 8);
+ drop(first);
+ assert_eq!(used(&resources), 5);
+ drop(second);
+ assert_eq!(used(&resources), 0);
+ assert_eq!(resources.metrics().peak_reserved_memory_bytes, 8);
+}
+
+#[test]
+fn zero_limit_and_address_space_overflow_are_rejected() {
+ let zero = ResourceContext::builder().memory_limit(0).build().unwrap();
+ let mut reservation = zero.reservation();
+ reservation.try_grow(0).unwrap();
+ assert!(reservation.try_grow(1).is_err());
+ assert_eq!(used(&zero), 0);
+
+ let unlimited = ResourceContext::builder().build().unwrap();
+ let mut reservation = unlimited.reservation();
+ // Reservations are accounting; this does not allocate memory.
+ reservation.try_grow(usize::MAX).unwrap();
+ assert!(reservation.try_grow(1).is_err());
+ assert!(unlimited.reservation().try_grow(1).is_err());
+ assert_eq!(used(&unlimited), usize::MAX);
+ drop(reservation);
+ assert_eq!(used(&unlimited), 0);
+}
+
+#[derive(Debug)]
+struct EnginePool {
+ limit: usize,
+ reserved: AtomicUsize,
+ released: AtomicUsize,
+}
+
+impl MemoryPool for EnginePool {
+ fn try_reserve(&self, bytes: usize) -> Result<()> {
+ self.reserved
+ .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
+ used.checked_add(bytes).filter(|next| *next <= self.limit)
+ })
+ .map_err(|_| Error::ResourceExhausted {
+ message: "engine budget".to_string(),
+ })?;
+ Ok(())
+ }
+
+ fn release(&self, bytes: usize) {
+ assert!(self.reserved.fetch_sub(bytes, Ordering::Relaxed) >= bytes);
+ self.released.fetch_add(bytes, Ordering::Relaxed);
+ }
+}
+
+#[test]
+fn external_pool_failure_rolls_back_local_admission() {
+ let pool = Arc::new(EnginePool {
+ limit: 12,
+ reserved: AtomicUsize::new(0),
+ released: AtomicUsize::new(0),
+ });
+ let left = ResourceContext::builder()
+ .memory_limit(10)
+ .memory_pool(pool.clone())
+ .build()
+ .unwrap();
+ let right = ResourceContext::builder()
+ .memory_limit(10)
+ .memory_pool(pool.clone())
+ .build()
+ .unwrap();
+ let mut first = left.reservation();
+ let mut second = right.reservation();
+ first.try_grow(8).unwrap();
+ let error = second.try_grow(5).unwrap_err();
+ assert!(matches!(error, Error::ResourceExhausted { message } if message ==
"engine budget"));
+ assert_eq!(used(&right), 0);
+ assert_eq!(right.metrics().peak_reserved_memory_bytes, 0);
+ second.try_grow(4).unwrap();
+ assert!(first.try_grow(3).is_err());
+ drop(first);
+ // This reaches the local limit, proving the rejected request was rolled
back.
+ second.try_grow(6).unwrap();
+ assert_eq!(used(&right), 10);
+ assert_eq!(pool.reserved.load(Ordering::Relaxed), 10);
+ drop(second);
+ assert_eq!(pool.reserved.load(Ordering::Relaxed), 0);
+ assert_eq!(pool.released.load(Ordering::Relaxed), 18);
+}
+
+#[test]
+fn concurrent_reservations_cannot_overcommit() {
+ const THREADS: usize = 16;
+ let resources =
ResourceContext::builder().memory_limit(32).build().unwrap();
+ let start = Arc::new(Barrier::new(THREADS + 1));
+ let ready = Arc::new(Barrier::new(THREADS + 1));
+ let release = Arc::new(Barrier::new(THREADS + 1));
+ let threads: Vec<_> = (0..THREADS)
+ .map(|_| {
+ let resources = resources.clone();
+ let (start, ready, release) = (start.clone(), ready.clone(),
release.clone());
+ std::thread::spawn(move || {
+ let mut reservation = resources.reservation();
+ start.wait();
+ let admitted = reservation.try_grow(8).is_ok();
+ ready.wait();
+ release.wait();
+ admitted
+ })
+ })
+ .collect();
+ start.wait();
+ ready.wait();
+ assert_eq!(used(&resources), 32);
+ release.wait();
+ let admitted = threads
+ .into_iter()
+ .map(|thread| usize::from(thread.join().unwrap()))
+ .sum::<usize>();
+ assert_eq!(admitted, 4);
+ assert_eq!(used(&resources), 0);
+ assert_eq!(resources.metrics().peak_reserved_memory_bytes, 32);
+}
diff --git a/crates/paimon/src/table/data_evolution_reader.rs
b/crates/paimon/src/table/data_evolution_reader.rs
index fb775ae0..d9e4ad6f 100644
--- a/crates/paimon/src/table/data_evolution_reader.rs
+++ b/crates/paimon/src/table/data_evolution_reader.rs
@@ -991,7 +991,9 @@ impl DataEvolutionReader {
let source_parquet_read_budget = if active_source_indices.len() ==
1 {
parquet_read_budget.clone()
} else {
- None
+ parquet_read_budget.as_ref()
+ .filter(|budget| budget.has_resources())
+ .map(|budget| Arc::new(budget.without_prefetch()))
};
let mut source_streams: Vec<Option<ArrowRecordBatchStream>> =
source_plan
.sources
@@ -7605,6 +7607,87 @@ mod tests {
}
}
+ #[tokio::test]
+ async fn test_evolution_multi_source_merge_preserves_memory_admission() {
+ let tempdir = tempdir().unwrap();
+ let table_path = local_file_path(tempdir.path());
+ let bucket_dir = tempdir.path().join("bucket-0");
+ fs::create_dir_all(&bucket_dir).unwrap();
+
+ let id_path = bucket_dir.join("id.parquet");
+ write_int_parquet_file(&id_path, vec![("id", vec![1, 2, 3, 4])],
Some(2));
+ let value_path = bucket_dir.join("value.parquet");
+ write_int_parquet_file(&value_path, vec![("value", vec![10, 20, 30,
40])], Some(2));
+
+ let table = two_col_evolution_table(table_path);
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(local_file_path(&bucket_dir))
+ .with_total_buckets(1)
+ .with_data_files(vec![
+ data_file_meta_with_path(
+ "id.parquet",
+ 0,
+ 4,
+ 1,
+ id_path.metadata().unwrap().len() as i64,
+ Some(vec!["id"]),
+ ),
+ data_file_meta_with_path(
+ "value.parquet",
+ 0,
+ 4,
+ 2,
+ value_path.metadata().unwrap().len() as i64,
+ Some(vec!["value"]),
+ ),
+ ])
+ .build()
+ .unwrap();
+
+ for limit in [None, Some(0), Some(1024 * 1024)] {
+ let resources = limit.map(|limit| {
+ crate::resource::ResourceContext::builder()
+ .memory_limit(limit)
+ .build()
+ .unwrap()
+ });
+ let mut builder = table.new_read_builder();
+ builder.with_parquet_read_budget(Arc::new(ReadBudget::new(2,
1).unwrap()));
+ if let Some(resources) = &resources {
+ builder.with_resources(resources.clone());
+ }
+ let result = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ builder
+ .new_read()
+ .unwrap()
+ .to_arrow(std::slice::from_ref(&split))
+ .unwrap()
+ .try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("column merges must not wait on shared prefetch permits");
+
+ if limit == Some(0) {
+ assert!(matches!(result, Err(Error::ResourceExhausted { ..
})));
+ } else {
+ let batches = result.unwrap();
+ assert_eq!(collect_int_values(&batches, "id"), vec![1, 2, 3,
4]);
+ assert_eq!(collect_int_values(&batches, "value"), vec![10, 20,
30, 40]);
+ if let Some(resources) = &resources {
+ let peak = resources.metrics().peak_reserved_memory_bytes;
+ assert!(peak > 0 && peak <= limit.unwrap());
+ }
+ }
+ if let Some(resources) = &resources {
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
+ }
+
#[tokio::test]
async fn
test_evolution_input_decode_honors_read_batch_size_on_all_file_paths() {
let tempdir = tempdir().unwrap();
diff --git a/crates/paimon/src/table/format_table_read.rs
b/crates/paimon/src/table/format_table_read.rs
index e419e78e..125a2a07 100644
--- a/crates/paimon/src/table/format_table_read.rs
+++ b/crates/paimon/src/table/format_table_read.rs
@@ -24,6 +24,7 @@ use super::{ArrowRecordBatchStream, Table};
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
use crate::arrow::partition::partition_array;
use crate::arrow::{build_target_arrow_schema, ReadBudget};
+use crate::resource::ResourceContext;
use crate::spec::{DataField, Predicate};
use crate::{DataSplit, Error};
use arrow_array::{RecordBatch, RecordBatchOptions};
@@ -38,6 +39,7 @@ pub(crate) struct FormatTableRead<'a> {
data_predicates: Vec<Predicate>,
row_filter_factory: Option<Arc<dyn crate::arrow::RowFilterFactory>>,
parquet_read_budget: Option<Arc<ReadBudget>>,
+ resources: Option<ResourceContext>,
limit: Option<usize>,
blob_parallelism: usize,
}
@@ -55,6 +57,7 @@ impl<'a> FormatTableRead<'a> {
data_predicates,
row_filter_factory: None,
parquet_read_budget: None,
+ resources: None,
limit,
blob_parallelism: DEFAULT_BLOB_READ_PARALLELISM,
}
@@ -95,11 +98,19 @@ impl<'a> FormatTableRead<'a> {
self
}
+ pub(crate) fn with_resources(&mut self, resources: ResourceContext) {
+ self.resources = Some(resources);
+ }
+
fn parquet_read_budget(&self) -> crate::Result<Arc<ReadBudget>> {
- match &self.parquet_read_budget {
- Some(budget) => Ok(Arc::clone(budget)),
- None => configured_parquet_read_budget(self.table),
- }
+ let budget = match &self.parquet_read_budget {
+ Some(budget) => Arc::clone(budget),
+ None => configured_parquet_read_budget(self.table)?,
+ };
+ Ok(match &self.resources {
+ Some(resources) =>
Arc::new(budget.with_resources(resources.clone())),
+ None => budget,
+ })
}
pub(crate) fn to_arrow(
diff --git a/crates/paimon/src/table/kv_file_reader.rs
b/crates/paimon/src/table/kv_file_reader.rs
index 36ca7ed8..174b6542 100644
--- a/crates/paimon/src/table/kv_file_reader.rs
+++ b/crates/paimon/src/table/kv_file_reader.rs
@@ -616,7 +616,9 @@ impl KeyValueFileReader {
let group_parquet_read_budget = if input_stream_count == 1
{
config.parquet_read_budget.clone()
} else {
- None
+ config.parquet_read_budget.as_ref()
+ .filter(|budget| budget.has_resources())
+ .map(|budget| Arc::new(budget.without_prefetch()))
};
let mut file_streams: Vec<ArrowRecordBatchStream> =
Vec::new();
@@ -2035,43 +2037,70 @@ mod tests {
let planned = plan_merge_groups(std::slice::from_ref(&split),
Some(&comparator), false);
assert_eq!(planned.len(), 1);
assert_eq!(planned[0].len(), 2);
- let core_options = table.schema().core_options();
- let reader = KeyValueFileReader::new(
- table.file_io().clone(),
- KeyValueReadConfig {
- table_name: table.identifier().full_name(),
- table_options: table.schema().options().clone(),
- schema_manager: table.schema_manager().clone(),
- table_schema_id: table.schema().id(),
- table_fields: table.schema().fields().to_vec(),
- read_type: table.schema().fields().to_vec(),
- predicates: Vec::new(),
- primary_keys: table.schema().trimmed_primary_keys(),
- table_primary_keys: table.schema().primary_keys().to_vec(),
- merge_engine: core_options.merge_engine().unwrap(),
- sequence_fields: Vec::new(),
- read_batch_size: core_options.read_batch_size().unwrap(),
- merge_splits: false,
- max_merge_input_streams: None,
- parquet_read_budget: Some(Arc::new(ReadBudget::new(2, 256 <<
20).unwrap())),
- mosaic_prefetch: MosaicPrefetchOptions::default(),
- },
- );
- let batches = tokio::time::timeout(
- std::time::Duration::from_secs(5),
- reader
- .read(std::slice::from_ref(split.as_ref()))
- .unwrap()
- .try_collect::<Vec<_>>(),
- )
- .await
- .expect("multiple sorted-run inputs must not deadlock on shared
Parquet permits")
- .unwrap();
+ for limit in [None, Some(0), Some(1024 * 1024)] {
+ let core_options = table.schema().core_options();
+ let budget = ReadBudget::new(2, 256 << 20).unwrap();
+ let resources = limit.map(|limit| {
+ crate::resource::ResourceContext::builder()
+ .memory_limit(limit)
+ .build()
+ .unwrap()
+ });
+ let budget = match &resources {
+ Some(resources) => budget.with_resources(resources.clone()),
+ None => budget,
+ };
+ let reader = KeyValueFileReader::new(
+ table.file_io().clone(),
+ KeyValueReadConfig {
+ table_name: table.identifier().full_name(),
+ table_options: table.schema().options().clone(),
+ schema_manager: table.schema_manager().clone(),
+ table_schema_id: table.schema().id(),
+ table_fields: table.schema().fields().to_vec(),
+ read_type: table.schema().fields().to_vec(),
+ predicates: Vec::new(),
+ primary_keys: table.schema().trimmed_primary_keys(),
+ table_primary_keys: table.schema().primary_keys().to_vec(),
+ merge_engine: core_options.merge_engine().unwrap(),
+ sequence_fields: Vec::new(),
+ read_batch_size: core_options.read_batch_size().unwrap(),
+ merge_splits: false,
+ max_merge_input_streams: None,
+ parquet_read_budget: Some(Arc::new(budget)),
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
+ },
+ );
+ let result = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ reader
+ .read(std::slice::from_ref(split.as_ref()))
+ .unwrap()
+ .try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("multiple sorted-run inputs must not deadlock on shared
Parquet permits");
+ if limit == Some(0) {
+ assert!(matches!(
+ result,
+ Err(crate::Error::ResourceExhausted { .. })
+ ));
+ } else {
+ let batches = result.unwrap();
- assert_eq!(
- batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
- 128
- );
+ assert_eq!(
+ batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
+ 128
+ );
+ drop(batches);
+ if let Some(resources) = &resources {
+ assert!(resources.metrics().peak_reserved_memory_bytes >
0);
+ }
+ }
+ if let Some(resources) = resources {
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
}
#[tokio::test]
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index 063b898c..3561616a 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -27,6 +27,7 @@ use super::partition_filter::PartitionFilter;
use super::table_read::{configured_parquet_read_budget, TableRead};
use super::{Table, TableScan};
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
+use crate::resource::ResourceContext;
use crate::spec::{CoreOptions, DataField, Predicate};
use crate::table::source::RowRange;
use crate::{Error, Result};
@@ -112,7 +113,10 @@ fn normalize_filter(table: &Table, filter: Predicate) ->
NormalizedFilter {
/// Rust keeps a names-based projection API for ergonomics, while aligning the
/// resulting read semantics with Java Paimon's order-preserving projection.
#[derive(Debug, Clone)]
-pub struct ReadBuilder<'a>(ReadBuilderKind<'a>);
+pub struct ReadBuilder<'a> {
+ kind: ReadBuilderKind<'a>,
+ resources: Option<ResourceContext>,
+}
#[derive(Debug, Clone)]
enum ReadBuilderKind<'a> {
@@ -122,13 +126,27 @@ enum ReadBuilderKind<'a> {
impl<'a> ReadBuilder<'a> {
pub(crate) fn new(table: &'a Table) -> Self {
- if table.is_format_table() {
- Self(ReadBuilderKind::Format(FormatReadBuilder::new(table)))
+ let kind = if table.is_format_table() {
+ ReadBuilderKind::Format(FormatReadBuilder::new(table))
} else {
- Self(ReadBuilderKind::Paimon(PaimonReadBuilder::new(table)))
+ ReadBuilderKind::Paimon(PaimonReadBuilder::new(table))
+ };
+ Self {
+ kind,
+ resources: None,
}
}
+ /// Share Parquet working-memory reservations across reads.
+ ///
+ /// Row-group data reads draw from this context; callers account for
retained outputs.
+ /// See [`ResourceContext`] for the estimate-based accounting contract.
+ /// When admission fails, the stream returns a resource-exhausted error
and ends.
+ pub fn with_resources(&mut self, resources: ResourceContext) -> &mut Self {
+ self.resources = Some(resources);
+ self
+ }
+
/// Set column projection by name. Output order follows the
caller-specified order.
/// An empty list is a valid zero-column projection.
///
@@ -139,7 +157,7 @@ impl<'a> ReadBuilder<'a> {
/// case-insensitively, or a case-fold duplicate/ambiguity — surface from
/// [`new_read`](Self::new_read) using the effective case sensitivity.
pub fn with_projection(&mut self, columns: &[&str]) -> Result<&mut Self> {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_projection(columns)?;
}
@@ -165,7 +183,7 @@ impl<'a> ReadBuilder<'a> {
/// retroactively change a predicate already passed to
/// [`with_filter`](Self::with_filter).
pub fn with_case_sensitive(&mut self, case_sensitive: bool) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_case_sensitive(case_sensitive);
}
@@ -179,7 +197,7 @@ impl<'a> ReadBuilder<'a> {
/// Set the full read type, including nested field pruning or
connector-defined
/// logical read types such as Variant extractions.
pub fn with_read_type(&mut self, read_type: Vec<DataField>) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_read_type(read_type);
}
@@ -192,7 +210,7 @@ impl<'a> ReadBuilder<'a> {
/// Set a filter predicate for scan planning and conservative read pruning.
pub fn with_filter(&mut self, filter: Predicate) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_filter(filter);
}
@@ -205,7 +223,7 @@ impl<'a> ReadBuilder<'a> {
/// Whether a translated predicate is exact at the table-provider boundary.
pub fn is_exact_filter_pushdown(&self, filter: &Predicate) -> bool {
- match &self.0 {
+ match &self.kind {
ReadBuilderKind::Paimon(builder) =>
builder.is_exact_filter_pushdown(filter),
ReadBuilderKind::Format(builder) =>
builder.is_exact_filter_pushdown(filter),
}
@@ -214,7 +232,7 @@ impl<'a> ReadBuilder<'a> {
/// Set Data Evolution row ID ranges `[from, to]` (inclusive).
/// An empty vector selects no rows. Format tables are not supported.
pub fn with_row_ranges(&mut self, ranges: Vec<RowRange>) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_row_ranges(ranges);
}
@@ -228,7 +246,7 @@ impl<'a> ReadBuilder<'a> {
/// Push a row-limit hint down to scan planning. Data-evolution reads also
/// enforce this limit before resolving BLOB payloads.
pub fn with_limit(&mut self, limit: usize) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_limit(limit);
}
@@ -247,7 +265,7 @@ impl<'a> ReadBuilder<'a> {
source: None,
});
}
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_blob_parallelism(blob_parallelism);
}
@@ -261,7 +279,7 @@ impl<'a> ReadBuilder<'a> {
/// Inject a Parquet budget shared with sibling scan partitions.
#[doc(hidden)]
pub fn with_parquet_read_budget(&mut self, budget:
Arc<crate::arrow::ReadBudget>) -> &mut Self {
- match &mut self.0 {
+ match &mut self.kind {
ReadBuilderKind::Paimon(builder) => {
builder.with_parquet_read_budget(budget);
}
@@ -274,7 +292,7 @@ impl<'a> ReadBuilder<'a> {
/// Create a table scan. Call [TableScan::plan] to get splits.
pub fn new_scan(&self) -> TableScan<'a> {
- match &self.0 {
+ match &self.kind {
ReadBuilderKind::Paimon(builder) => builder.new_scan(),
ReadBuilderKind::Format(builder) => builder.new_scan(),
}
@@ -293,7 +311,7 @@ impl<'a> ReadBuilder<'a> {
start_exclusive: i64,
end_inclusive: i64,
) -> IncrementalScan<'a> {
- match &self.0 {
+ match &self.kind {
ReadBuilderKind::Paimon(builder) => IncrementalScan::new(
builder.table,
builder.new_scan(),
@@ -310,10 +328,14 @@ impl<'a> ReadBuilder<'a> {
/// Create a table read for consuming splits (e.g. from a scan plan).
pub fn new_read(&self) -> Result<TableRead<'a>> {
- match &self.0 {
+ let read = match &self.kind {
ReadBuilderKind::Paimon(builder) => builder.new_read(),
ReadBuilderKind::Format(builder) => builder.new_read(),
- }
+ }?;
+ Ok(match &self.resources {
+ Some(resources) => read.with_resources(resources.clone()),
+ None => read,
+ })
}
}
@@ -761,7 +783,7 @@ mod tests {
use test_utils::{local_file_path, test_data_file, write_int_parquet_file};
fn paimon_builder<'a, 'b>(builder: &'b ReadBuilder<'a>) -> &'b
PaimonReadBuilder<'a> {
- match &builder.0 {
+ match &builder.kind {
ReadBuilderKind::Paimon(inner) => inner,
ReadBuilderKind::Format(_) => panic!("expected Paimon read
builder"),
}
diff --git a/crates/paimon/src/table/table_read.rs
b/crates/paimon/src/table/table_read.rs
index 70576d30..3313df0e 100644
--- a/crates/paimon/src/table/table_read.rs
+++ b/crates/paimon/src/table/table_read.rs
@@ -26,6 +26,7 @@ use crate::arrow::build_target_arrow_schema;
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
use crate::arrow::format::MosaicPrefetchOptions;
use crate::arrow::ReadBudget;
+use crate::resource::ResourceContext;
use crate::spec::{
BigIntType, CoreOptions, DataField, DataType, MergeEngine, Predicate,
TinyIntType,
ROW_KIND_FIELD_ID, ROW_KIND_FIELD_NAME, SEQUENCE_NUMBER_FIELD_ID,
SEQUENCE_NUMBER_FIELD_NAME,
@@ -133,6 +134,18 @@ impl<'a> TableRead<'a> {
}
}
+ /// Share Parquet working-memory reservations with other consumers.
+ ///
+ /// Callers account for output batches they retain. Budget exhaustion is a
+ /// terminal stream error. See [`ResourceContext`] for the accounting
scope.
+ pub fn with_resources(mut self, resources: ResourceContext) -> Self {
+ match &mut self.0 {
+ TableReadKind::Paimon(read) => read.resources = Some(resources),
+ TableReadKind::Format(read) => read.with_resources(resources),
+ }
+ self
+ }
+
/// Set a filter predicate.
pub fn with_filter(self, filter: Predicate) -> Self {
match self.0 {
@@ -287,6 +300,7 @@ struct PaimonTableRead<'a> {
data_predicates: Vec<Predicate>,
row_filter_factory: Option<Arc<dyn crate::arrow::RowFilterFactory>>,
parquet_read_budget: Option<Arc<ReadBudget>>,
+ resources: Option<ResourceContext>,
data_file_read_timing: Option<Arc<DataFileReadTiming>>,
blob_parallelism: usize,
limit: Option<usize>,
@@ -305,6 +319,7 @@ impl<'a> PaimonTableRead<'a> {
data_predicates,
row_filter_factory: None,
parquet_read_budget: None,
+ resources: None,
data_file_read_timing: None,
blob_parallelism: DEFAULT_BLOB_READ_PARALLELISM,
limit: None,
@@ -362,10 +377,14 @@ impl<'a> PaimonTableRead<'a> {
}
fn parquet_read_budget(&self) -> crate::Result<Arc<ReadBudget>> {
- match &self.parquet_read_budget {
- Some(budget) => Ok(Arc::clone(budget)),
- None => configured_parquet_read_budget(self.table),
- }
+ let budget = match &self.parquet_read_budget {
+ Some(budget) => Arc::clone(budget),
+ None => configured_parquet_read_budget(self.table)?,
+ };
+ Ok(match &self.resources {
+ Some(resources) =>
Arc::new(budget.with_resources(resources.clone())),
+ None => budget,
+ })
}
/// Returns an [`ArrowRecordBatchStream`] for an incremental scan plan.
@@ -808,6 +827,10 @@ impl<'a> PaimonTableRead<'a> {
});
}
}
+ let budget = self.parquet_read_budget()?;
+ let parquet_read_budget = budget
+ .has_resources()
+ .then(|| Arc::new(budget.without_prefetch()));
let reader = KeyValueFileReader::new(
self.table.file_io.clone(),
KeyValueReadConfig {
@@ -829,10 +852,9 @@ impl<'a> PaimonTableRead<'a> {
read_batch_size: core_options.read_batch_size()?,
merge_splits: true,
max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS),
- // Diff primes the before and after streams in sequence.
Keeping
- // a row-group permit across yielded batches can otherwise let
- // the first side block the second side indefinitely.
- parquet_read_budget: None,
+ // Diff advances before/after in lockstep. Disable prefetch so
+ // neither side waits on shared slots, but keep memory
admission.
+ parquet_read_budget,
mosaic_prefetch: configured_mosaic_prefetch(self.table)?,
},
);
diff --git a/crates/paimon/tests/reader_resources_test.rs
b/crates/paimon/tests/reader_resources_test.rs
new file mode 100644
index 00000000..db4a7f88
--- /dev/null
+++ b/crates/paimon/tests/reader_resources_test.rs
@@ -0,0 +1,332 @@
+// 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.
+
+mod common;
+
+use std::sync::Arc;
+
+use arrow_array::Int32Array;
+use common::incremental_helpers::{
+ make_batch, memory_table, persist_table_schema, setup_dirs, write_batch,
+};
+use futures::{StreamExt, TryStreamExt};
+use paimon::resource::ResourceContext;
+use paimon::spec::{DataType, Datum, IntType, PredicateBuilder, Schema,
TableSchema};
+use paimon::table::{ArrowRecordBatchStream, AuditLogRead, IncrementalScanMode,
Table};
+use paimon::Error;
+
+async fn parquet_table(primary_key: bool) -> Table {
+ let mut schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("file.format", "parquet");
+ if primary_key {
+ schema = schema.primary_key(["id"]).option("bucket", "1");
+ }
+ let path = "memory:/reader_resources";
+ let (io, table) = memory_table(path, TableSchema::new(0,
&schema.build().unwrap()));
+ setup_dirs(&io, path).await;
+ persist_table_schema(&io, path, table.schema()).await;
+ write_batch(&table, &make_batch(vec![1, 2, 3], vec![10, 20, 30])).await;
+ table
+}
+
+#[tokio::test]
+async fn parquet_reader_releases_working_memory_while_output_survives() {
+ for primary_key in [false, true] {
+ let table = parquet_table(primary_key).await;
+ let resources = ResourceContext::builder()
+ .memory_limit(1024 * 1024)
+ .build()
+ .unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_resources(resources.clone());
+ builder.with_projection(&["id"]).unwrap();
+ let plan = builder.new_scan().plan().await.unwrap();
+ let predicate = PredicateBuilder::new(table.schema().fields())
+ .greater_or_equal("id", Datum::Int(2))
+ .unwrap();
+ // Builder cloning and subsequent consuming setters must preserve the
context.
+ let read = builder
+ .clone()
+ .new_read()
+ .unwrap()
+ .with_filter(predicate)
+ .with_blob_parallelism(2)
+ .unwrap()
+ .with_parquet_read_budget(Arc::new(
+ paimon::arrow::ReadBudget::new(2, 1024 * 1024).unwrap(),
+ ));
+ let mut stream = read.to_arrow(plan.splits()).unwrap();
+ let batch = stream.next().await.unwrap().unwrap();
+ let ids = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ids.values().as_ref(), &[2, 3]);
+ let bytes = resources.metrics().reserved_memory_bytes;
+ assert!(bytes > 0);
+ let escaped = batch.column(0).slice(1, 1);
+ drop(batch);
+ drop(stream);
+ drop(read);
+ drop(builder);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ assert_eq!(
+ escaped
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 3
+ );
+ drop(escaped);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+}
+
+async fn assert_exhausted(mut stream: ArrowRecordBatchStream, resources:
&ResourceContext) {
+ assert!(matches!(
+ stream.next().await.unwrap(),
+ Err(Error::ResourceExhausted { .. })
+ ));
+ assert!(
+ stream.next().await.is_none(),
+ "exhaustion must end the stream"
+ );
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+}
+
+#[tokio::test]
+async fn every_read_output_mode_honors_the_budget() {
+ let table = parquet_table(false).await;
+ let resources =
ResourceContext::builder().memory_limit(0).build().unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_resources(resources.clone());
+ let plan = builder.new_scan().plan().await.unwrap();
+ let incremental = builder
+ .new_incremental_scan(IncrementalScanMode::Delta, 0, 1)
+ .plan()
+ .await
+ .unwrap();
+ let read = builder.new_read().unwrap();
+ assert_exhausted(read.to_arrow(plan.splits()).unwrap(), &resources).await;
+ assert_exhausted(
+ read.to_arrow_with_row_kind(plan.splits()).unwrap(),
+ &resources,
+ )
+ .await;
+ assert_exhausted(read.to_incremental_arrow(&incremental).unwrap(),
&resources).await;
+ assert_exhausted(read.to_audit_log_arrow(&incremental).unwrap(),
&resources).await;
+ assert_exhausted(
+ AuditLogRead::new(read)
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap(),
+ &resources,
+ )
+ .await;
+ // No buffers are needed for a zero-column projection, even under a zero
budget.
+ builder.with_projection(&[]).unwrap();
+ let batches: Vec<_> = builder
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(
+ batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+ 3
+ );
+ assert!(batches.iter().all(|batch| batch.num_columns() == 0));
+}
+
+#[tokio::test]
+async fn readers_and_downstream_consumers_share_explicit_reservations() {
+ let table = parquet_table(false).await;
+ let plan = table.new_read_builder().new_scan().plan().await.unwrap();
+ let resources = ResourceContext::builder()
+ .memory_limit(1024 * 1024)
+ .build()
+ .unwrap();
+ let read = table
+ .new_read_builder()
+ .new_read()
+ .unwrap()
+ .with_resources(resources.clone());
+ let mut stream = read.to_arrow(plan.splits()).unwrap();
+ let retained = stream.next().await.unwrap().unwrap();
+ drop(stream);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ // The downstream consumer explicitly accounts for batches it retains.
+ let retained_bytes = retained.get_array_memory_size();
+ let mut downstream = resources.reservation();
+ downstream.try_grow(retained_bytes).unwrap();
+ // Occupy the rest so another reader cannot start its row-group data read.
+ let mut other_consumer = resources.reservation();
+ other_consumer
+ .try_grow(1024 * 1024 - retained_bytes)
+ .unwrap();
+ let mut second = read.clone().to_arrow(plan.splits()).unwrap();
+ assert!(matches!(
+ second.next().await.unwrap(),
+ Err(Error::ResourceExhausted { .. })
+ ));
+ assert!(second.next().await.is_none());
+ drop(second);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 1024 * 1024);
+ drop(other_consumer);
+ drop(retained);
+ drop(downstream);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ // A fresh read can use capacity released by the failed read's siblings.
+ let batches: Vec<_> = read
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(
+ batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
+ 3
+ );
+ drop(batches);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+}
+
+#[tokio::test]
+async fn format_table_reader_uses_the_same_resource_context() {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .option("type", "format-table")
+ .option("file.format", "parquet")
+ .build()
+ .unwrap();
+ let path = "memory:/format_reader_resources";
+ let (io, table) = memory_table(path, TableSchema::new(0, &schema));
+ io.mkdirs(path).await.unwrap();
+ let input = make_batch(vec![1, 2], vec![10, 20]);
+ let mut bytes = Vec::new();
+ let mut writer =
+ parquet::arrow::ArrowWriter::try_new(&mut bytes, input.schema(),
None).unwrap();
+ writer.write(&input).unwrap();
+ writer.close().unwrap();
+ io.new_output(&format!("{path}/data.parquet"))
+ .unwrap()
+ .write(bytes.into())
+ .await
+ .unwrap();
+ let resources =
ResourceContext::builder().memory_limit(0).build().unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_resources(resources.clone());
+ let plan = builder.new_scan().plan().await.unwrap();
+ let read = builder
+ .clone()
+ .new_read()
+ .unwrap()
+ .with_blob_parallelism(1)
+ .unwrap();
+ assert_exhausted(read.to_arrow(plan.splits()).unwrap(), &resources).await;
+ assert_exhausted(
+ read.to_arrow_with_row_kind(plan.splits()).unwrap(),
+ &resources,
+ )
+ .await;
+
+ let resources = ResourceContext::builder()
+ .memory_limit(1024)
+ .build()
+ .unwrap();
+ let output: Vec<_> = read
+ .with_resources(resources.clone())
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(output.len(), 1);
+ assert_eq!(output[0].column(0).as_ref(), input.column(0).as_ref());
+ assert!(resources.metrics().peak_reserved_memory_bytes > 0);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ drop(output);
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+}
+
+#[tokio::test]
+async fn primary_key_diff_and_audit_diff_share_working_memory() {
+ let table = parquet_table(true).await;
+ write_batch(&table, &make_batch(vec![2], vec![200])).await;
+ for limit in [0, 1024 * 1024] {
+ let resources = ResourceContext::builder()
+ .memory_limit(limit)
+ .build()
+ .unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_resources(resources.clone());
+ let plan = builder
+ .new_incremental_scan(IncrementalScanMode::Diff, 1, 2)
+ .plan()
+ .await
+ .unwrap();
+ let read = builder.new_read().unwrap();
+ for audit in [false, true] {
+ let stream = if audit {
+ read.to_audit_log_arrow(&plan).unwrap()
+ } else {
+ read.to_incremental_arrow(&plan).unwrap()
+ };
+ let result = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ stream.try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("before/after streams must not deadlock on shared
capacity");
+ if limit == 0 {
+ assert!(matches!(result, Err(Error::ResourceExhausted { ..
})));
+ } else {
+ let batches = result.unwrap();
+ let mut values = Vec::new();
+ for batch in &batches {
+ let ids = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert!(ids.values().iter().all(|id| *id == 2));
+ values.extend_from_slice(
+ batch
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values(),
+ );
+ }
+ values.sort_unstable();
+ assert_eq!(values, if audit { vec![20, 200] } else { vec![200]
});
+ assert!(resources.metrics().peak_reserved_memory_bytes > 0);
+ }
+ assert_eq!(resources.metrics().reserved_memory_bytes, 0);
+ }
+ }
+}
