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 4178002a feat(mosaic): read row-group prefetch depth and byte budget
from table options (#864)
4178002a is described below
commit 4178002ae77460a1f0518724e86c37eeb93ee6de
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Fri Sep 18 22:31:02 2026 +0800
feat(mosaic): read row-group prefetch depth and byte budget from table
options (#864)
---
.../datafusion/src/physical_plan/scan.rs | 10 +-
crates/paimon/src/arrow/format/mod.rs | 12 +-
crates/paimon/src/arrow/format/mosaic.rs | 238 +++++++++------------
crates/paimon/src/arrow/format/parquet.rs | 18 +-
crates/paimon/src/arrow/mod.rs | 4 +-
.../{parquet_read_budget.rs => read_budget.rs} | 205 +++++++++++++-----
crates/paimon/src/spec/core_options.rs | 88 +++++++-
crates/paimon/src/table/audit_log_table/read.rs | 2 +
crates/paimon/src/table/data_evolution_reader.rs | 24 ++-
crates/paimon/src/table/data_file_reader.rs | 19 +-
crates/paimon/src/table/format_read_builder.rs | 6 +-
crates/paimon/src/table/format_table_read.rs | 14 +-
crates/paimon/src/table/kv_file_reader.rs | 28 ++-
crates/paimon/src/table/read_builder.rs | 12 +-
crates/paimon/src/table/table_read.rs | 72 ++++++-
15 files changed, 496 insertions(+), 256 deletions(-)
diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs
b/crates/integrations/datafusion/src/physical_plan/scan.rs
index 432e3a2a..01a9909f 100644
--- a/crates/integrations/datafusion/src/physical_plan/scan.rs
+++ b/crates/integrations/datafusion/src/physical_plan/scan.rs
@@ -49,7 +49,7 @@ use datafusion::physical_plan::filter_pushdown::{
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning,
PlanProperties};
use futures::{FutureExt, StreamExt, TryStreamExt};
-use paimon::arrow::ParquetReadBudget;
+use paimon::arrow::ReadBudget;
use paimon::spec::{
CoreOptions, DataField, Datum, MergeEngine, Predicate, PredicateBuilder,
PredicateOperator,
};
@@ -804,7 +804,7 @@ pub struct PaimonTableScan {
/// Parquet row filter instead, avoiding duplicate decoder evaluation.
decoder_filters: Vec<Arc<dyn PhysicalExpr>>,
/// Query-wide budget shared by every DataFusion scan partition.
- parquet_read_budget: Arc<ParquetReadBudget>,
+ parquet_read_budget: Arc<ReadBudget>,
}
impl PaimonTableScan {
@@ -833,7 +833,7 @@ impl PaimonTableScan {
scan_trace,
pushed_variants,
case_sensitive,
- Arc::new(ParquetReadBudget::default()),
+ Arc::new(ReadBudget::default()),
)
}
@@ -852,7 +852,7 @@ impl PaimonTableScan {
) -> DFResult<Self> {
let options = table.schema().core_options();
let parquet_read_budget = Arc::new(
- ParquetReadBudget::new(
+ ReadBudget::new(
options
.parquet_row_group_parallelism()
.map_err(to_datafusion_error)?,
@@ -889,7 +889,7 @@ impl PaimonTableScan {
scan_trace: Option<ScanTrace>,
pushed_variants: Option<String>,
case_sensitive: bool,
- parquet_read_budget: Arc<ParquetReadBudget>,
+ parquet_read_budget: Arc<ReadBudget>,
) -> Self {
let plan_properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
diff --git a/crates/paimon/src/arrow/format/mod.rs
b/crates/paimon/src/arrow/format/mod.rs
index 17d21c9f..4fc9f93b 100644
--- a/crates/paimon/src/arrow/format/mod.rs
+++ b/crates/paimon/src/arrow/format/mod.rs
@@ -25,10 +25,11 @@ mod shredding;
#[cfg(feature = "vortex")]
mod vortex;
+pub(crate) use mosaic::MosaicPrefetchOptions;
#[cfg(test)]
pub(crate) use parquet::ParquetFormatWriter;
-use super::ParquetReadBudget;
+use super::ReadBudget;
use super::RowFilterFactory;
use crate::io::{FileIO, FileRead, OutputFile};
use crate::spec::stats::BinaryTableStats;
@@ -196,16 +197,19 @@ pub(crate) fn create_format_reader(
read_fields,
None,
blob::DEFAULT_BLOB_READ_PARALLELISM,
+ MosaicPrefetchOptions::default(),
)
}
-/// Create a format reader with a scan-shared Parquet resource budget.
+/// Create a format reader with a scan-shared Parquet resource budget and the
+/// per-file Mosaic prefetch settings.
pub(crate) fn create_format_reader_with_budget(
path: &str,
blob_as_descriptor: bool,
read_fields: &[DataField],
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
blob_parallelism: usize,
+ mosaic_prefetch: MosaicPrefetchOptions,
) -> crate::Result<Box<dyn FormatFileReader>> {
let lower = path.to_ascii_lowercase();
let reader: Box<dyn FormatFileReader> = if lower.ends_with(".parquet") {
@@ -227,7 +231,7 @@ pub(crate) fn create_format_reader_with_budget(
} else {
if lower.ends_with(".mosaic") {
return Ok(shredding::maybe_wrap_reader(
- Box::new(mosaic::MosaicFormatReader),
+
Box::new(mosaic::MosaicFormatReader::with_prefetch(mosaic_prefetch)),
read_fields,
));
}
diff --git a/crates/paimon/src/arrow/format/mosaic.rs
b/crates/paimon/src/arrow/format/mosaic.rs
index ae7c5011..fa0d8f93 100644
--- a/crates/paimon/src/arrow/format/mosaic.rs
+++ b/crates/paimon/src/arrow/format/mosaic.rs
@@ -20,8 +20,10 @@ use crate::arrow::build_target_arrow_schema;
use crate::arrow::filtering::{
predicates_may_match_with_schema, remap_predicates_to_file, StatsAccessor,
};
+use crate::arrow::read_budget::{ReadBudget, ReadPermit};
use crate::arrow::residual::{filter_record_batch_by_predicates,
widen_scan_fields};
use crate::io::FileRead;
+use crate::spec::MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION;
use crate::spec::{DataField, DataType as PaimonDataType, Datum, Predicate};
use crate::table::{ArrowRecordBatchStream, RowRange};
use crate::Error;
@@ -41,9 +43,37 @@ use std::collections::{HashMap, HashSet, VecDeque};
use std::io;
use std::ops::Range;
use std::panic::{catch_unwind, AssertUnwindSafe};
-use std::sync::{Arc, Condvar, Mutex, OnceLock};
+use std::sync::{Arc, Mutex, OnceLock};
+
+/// Row-group prefetch settings of one Mosaic file read, from the table options
+/// `mosaic.read.prefetch-row-groups` and `mosaic.read.prefetch-max-bytes`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct MosaicPrefetchOptions {
+ /// Row groups kept opening ahead of the consumer; `0` reads on demand.
+ pub row_groups: usize,
+ /// Estimated decoded bytes allowed in flight beyond the head row group.
+ pub max_bytes: usize,
+}
+
+impl Default for MosaicPrefetchOptions {
+ fn default() -> Self {
+ Self {
+ row_groups: DEFAULT_PREFETCH_ROW_GROUPS,
+ max_bytes: DEFAULT_PREFETCH_MAX_BYTES,
+ }
+ }
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct MosaicFormatReader {
+ prefetch: MosaicPrefetchOptions,
+}
-pub(crate) struct MosaicFormatReader;
+impl MosaicFormatReader {
+ pub(crate) fn with_prefetch(prefetch: MosaicPrefetchOptions) -> Self {
+ Self { prefetch }
+ }
+}
const DEFAULT_BATCH_SIZE: usize = 8192;
const DEFAULT_PREFETCH_ROW_GROUPS: usize = 8;
@@ -77,6 +107,7 @@ impl FormatFileReader for MosaicFormatReader {
file_fields: predicates.file_fields.clone(),
});
let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
+ let prefetch = self.prefetch;
let (mut batch_tx, mut batch_rx) = futures::channel::mpsc::channel(1);
let read_task = tokio::task::spawn_blocking(move || {
let result = read_mosaic_batches_blocking(
@@ -88,8 +119,8 @@ impl FormatFileReader for MosaicFormatReader {
batch_size,
row_selection,
handle,
- prefetch_row_groups: DEFAULT_PREFETCH_ROW_GROUPS,
- prefetch_max_bytes: DEFAULT_PREFETCH_MAX_BYTES,
+ prefetch_row_groups: prefetch.row_groups,
+ prefetch_max_bytes: prefetch.max_bytes,
},
|batch|
futures::executor::block_on(batch_tx.send(Ok(batch))).is_ok(),
);
@@ -135,112 +166,13 @@ struct PlannedRowGroup {
}
struct PendingRowGroup {
- permit: MosaicPrefetchPermit,
+ permit: ReadPermit,
task: MosaicTask<RecordBatch>,
}
struct BudgetedBatch {
batch: RecordBatch,
- permit: Option<Arc<MosaicPrefetchPermit>>,
-}
-
-#[derive(Default)]
-struct MosaicPrefetchBudgetState {
- bytes: usize,
- groups: usize,
-}
-
-struct MosaicPrefetchBudgetInner {
- max_bytes: usize,
- max_groups: usize,
- state: Mutex<MosaicPrefetchBudgetState>,
- released: Condvar,
-}
-
-#[derive(Clone)]
-struct MosaicPrefetchBudget {
- inner: Arc<MosaicPrefetchBudgetInner>,
-}
-
-impl MosaicPrefetchBudget {
- fn new(max_bytes: usize, max_groups: usize) -> Self {
- Self {
- inner: Arc::new(MosaicPrefetchBudgetInner {
- max_bytes,
- max_groups,
- state: Mutex::new(MosaicPrefetchBudgetState::default()),
- released: Condvar::new(),
- }),
- }
- }
-
- fn try_acquire(&self, estimated_bytes: usize) ->
Option<MosaicPrefetchPermit> {
- let mut state = self
- .inner
- .state
- .lock()
- .unwrap_or_else(|poisoned| poisoned.into_inner());
- if !self.can_acquire(&state, estimated_bytes) {
- return None;
- }
- Self::charge(&mut state, estimated_bytes);
- Some(MosaicPrefetchPermit {
- inner: Arc::clone(&self.inner),
- estimated_bytes,
- })
- }
-
- fn acquire(&self, estimated_bytes: usize) -> MosaicPrefetchPermit {
- let mut state = self
- .inner
- .state
- .lock()
- .unwrap_or_else(|poisoned| poisoned.into_inner());
- while !self.can_acquire(&state, estimated_bytes) {
- state = self
- .inner
- .released
- .wait(state)
- .unwrap_or_else(|poisoned| poisoned.into_inner());
- }
- Self::charge(&mut state, estimated_bytes);
- MosaicPrefetchPermit {
- inner: Arc::clone(&self.inner),
- estimated_bytes,
- }
- }
-
- fn can_acquire(&self, state: &MosaicPrefetchBudgetState, estimated_bytes:
usize) -> bool {
- if state.groups >= self.inner.max_groups {
- return false;
- }
- // An oversized head group may consume the whole budget so the reader
- // always makes progress, but it cannot overlap peers.
- state.groups == 0 || state.bytes.saturating_add(estimated_bytes) <=
self.inner.max_bytes
- }
-
- fn charge(state: &mut MosaicPrefetchBudgetState, estimated_bytes: usize) {
- state.groups += 1;
- state.bytes = state.bytes.saturating_add(estimated_bytes);
- }
-}
-
-struct MosaicPrefetchPermit {
- inner: Arc<MosaicPrefetchBudgetInner>,
- estimated_bytes: usize,
-}
-
-impl Drop for MosaicPrefetchPermit {
- fn drop(&mut self) {
- let mut state = self
- .inner
- .state
- .lock()
- .unwrap_or_else(|poisoned| poisoned.into_inner());
- state.groups -= 1;
- state.bytes = state.bytes.saturating_sub(self.estimated_bytes);
- self.inner.released.notify_all();
- }
+ permit: Option<Arc<ReadPermit>>,
}
type MosaicJob = Box<dyn FnOnce() + Send + 'static>;
@@ -476,12 +408,14 @@ fn read_mosaic_batches_blocking(
let projected_names = Arc::new(projected_names);
let executor = mosaic_executor()?;
let mut pending = VecDeque::<PendingRowGroup>::new();
- let prefetch_budget = MosaicPrefetchBudget::new(prefetch_max_bytes,
prefetch_row_groups);
+ let prefetch_budget =
+ ReadBudget::with_byte_granularity(prefetch_row_groups,
prefetch_max_bytes as u64, 1)?
+ .with_byte_option(MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION);
let mut outcome = Ok(());
'read: while !planned.is_empty() || !pending.is_empty() {
while let Some(next) = planned.front() {
- let Some(permit) =
prefetch_budget.try_acquire(next.estimated_bytes) else {
+ let Some(permit) =
prefetch_budget.try_acquire(next.estimated_bytes as u64) else {
break;
};
let plan = planned.pop_front().unwrap();
@@ -506,7 +440,13 @@ fn read_mosaic_batches_blocking(
};
// Earlier decoded batches can still own the budget downstream.
// With no task left to join, wait for their permits to be
released.
- let permit = prefetch_budget.acquire(next.estimated_bytes);
+ let permit = match
prefetch_budget.acquire_blocking(next.estimated_bytes as u64) {
+ Ok(permit) => permit,
+ Err(error) => {
+ outcome = Err(error);
+ break 'read;
+ }
+ };
let plan = planned.pop_front().unwrap();
let reader = Arc::clone(&mosaic_reader);
let names = Arc::clone(&projected_names);
@@ -600,7 +540,7 @@ fn apply_mosaic_residual(
fn send_mosaic_batch_chunks(
batch: RecordBatch,
batch_size: usize,
- permit: Option<Arc<MosaicPrefetchPermit>>,
+ permit: Option<Arc<ReadPermit>>,
send_batch: &mut impl FnMut(BudgetedBatch) -> bool,
) -> bool {
split_batch(batch, batch_size).into_iter().all(|batch| {
@@ -1308,7 +1248,7 @@ mod tests {
row_selection: Option<Vec<RowRange>>,
) -> crate::Result<Vec<RecordBatch>> {
let file_size = data.len() as u64;
- MosaicFormatReader
+ MosaicFormatReader::default()
.read_batch_stream(
Box::new(TestFileRead { data }),
file_size,
@@ -1330,7 +1270,7 @@ mod tests {
) -> crate::Result<Vec<Range<u64>>> {
let file_size = data.len() as u64;
let calls = Arc::new(Mutex::new(Vec::new()));
- let _: Vec<RecordBatch> = MosaicFormatReader
+ let _: Vec<RecordBatch> = MosaicFormatReader::default()
.read_batch_stream(
Box::new(TrackingFileRead {
data,
@@ -1356,7 +1296,7 @@ mod tests {
) -> crate::Result<Vec<Range<u64>>> {
let file_size = data.len() as u64;
let calls = Arc::new(Mutex::new(Vec::new()));
- let _: Vec<RecordBatch> = MosaicFormatReader
+ let _: Vec<RecordBatch> = MosaicFormatReader::default()
.read_batch_stream(
Box::new(TrackingFileRead {
data,
@@ -1525,6 +1465,55 @@ mod tests {
assert_eq!(ids.value(4), 5);
}
+ async fn read_via_format_reader_tracking(
+ data: Bytes,
+ fields: Vec<DataField>,
+ prefetch: MosaicPrefetchOptions,
+ ) -> (Vec<RecordBatch>, usize) {
+ let max_active = Arc::new(AtomicUsize::new(0));
+ let file_size = data.len() as u64;
+ let read = ConcurrentTrackingFileRead {
+ data,
+ active: Arc::new(AtomicUsize::new(0)),
+ max_active: Arc::clone(&max_active),
+ };
+ let stream = MosaicFormatReader::with_prefetch(prefetch)
+ .read_batch_stream(Box::new(read), file_size, &fields, None, None,
None)
+ .await
+ .unwrap();
+ let batches = stream.try_collect::<Vec<_>>().await.unwrap();
+ (batches, max_active.load(Ordering::SeqCst))
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_format_reader_prefetch_options_control_look_ahead() {
+ let fields = vec![data_fields()[0].clone()];
+
+ let (batches, max_active) = read_via_format_reader_tracking(
+ multi_row_group_mosaic(Vec::new()),
+ fields.clone(),
+ MosaicPrefetchOptions {
+ row_groups: 0,
+ max_bytes: usize::MAX,
+ },
+ )
+ .await;
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert_eq!(max_active, 1, "depth 0 must read row groups on demand");
+
+ let (batches, max_active) = read_via_format_reader_tracking(
+ multi_row_group_mosaic(Vec::new()),
+ fields,
+ MosaicPrefetchOptions::default(),
+ )
+ .await;
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert!(
+ max_active >= 2,
+ "default prefetch should overlap storage reads,
max_active={max_active}"
+ );
+ }
+
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_prefetch_overlaps_row_group_reads_and_preserves_order() {
let data = multi_row_group_mosaic(Vec::new());
@@ -1574,7 +1563,7 @@ mod tests {
}
#[test]
- fn
test_mosaic_prefetch_budget_is_cumulative_and_oversized_head_progresses() {
+ fn test_estimated_row_bytes_matches_java_prefetch_budget() {
assert_eq!(estimated_row_bytes(&[data_fields()[0].clone()]), 5);
assert_eq!(estimated_row_bytes(&[data_fields()[1].clone()]), 40);
assert_eq!(
@@ -1585,25 +1574,6 @@ mod tests {
)]),
128
);
-
- let budget = MosaicPrefetchBudget::new(20, 8);
- let first = budget.try_acquire(10).unwrap();
- let second = budget.try_acquire(10).unwrap();
- assert!(budget.try_acquire(10).is_none());
- drop(first);
- let third = budget.try_acquire(10).unwrap();
- drop((second, third));
-
- let oversized = budget.try_acquire(21).unwrap();
- assert!(budget.try_acquire(1).is_none());
- drop(oversized);
- assert!(budget.try_acquire(1).is_some());
-
- let group_limited = MosaicPrefetchBudget::new(usize::MAX, 2);
- let first = group_limited.try_acquire(1).unwrap();
- let second = group_limited.try_acquire(1).unwrap();
- assert!(group_limited.try_acquire(1).is_none());
- drop((first, second));
}
#[tokio::test]
@@ -1656,7 +1626,7 @@ mod tests {
let calls = Arc::new(Mutex::new(Vec::new()));
assert!(file_size > 64 * 1024);
- let batches = MosaicFormatReader
+ let batches = MosaicFormatReader::default()
.read_batch_stream(
Box::new(TrackingFileRead {
data,
@@ -1712,7 +1682,7 @@ mod tests {
.unwrap();
let batches = runtime.block_on(async move {
- MosaicFormatReader
+ MosaicFormatReader::default()
.read_batch_stream(
Box::new(RuntimeThreadFileRead {
data,
diff --git a/crates/paimon/src/arrow/format/parquet.rs
b/crates/paimon/src/arrow/format/parquet.rs
index 7b5ec592..758fb36e 100644
--- a/crates/paimon/src/arrow/format/parquet.rs
+++ b/crates/paimon/src/arrow/format/parquet.rs
@@ -18,10 +18,10 @@
use super::shredding::PhysicalFormatWriterFactory;
use super::{FilePredicates, FormatFileReader, FormatFileWriter,
FormatWriteResult};
use crate::arrow::filtering::{predicates_may_match_with_schema, StatsAccessor};
-use crate::arrow::parquet_read_budget::ParquetReadPermit;
+use crate::arrow::read_budget::ReadPermit;
use crate::arrow::shredding::map::MapShreddingReadPlan;
use crate::arrow::shredding::ShreddingReadPlan;
-use crate::arrow::{ParquetReadBudget, RowFilter, RowFilterContext};
+use crate::arrow::{ReadBudget, RowFilter, RowFilterContext};
use crate::io::{FileRead, OutputFile};
use crate::spec::stats::BinaryTableStats;
use crate::spec::{
@@ -56,11 +56,11 @@ use tokio::sync::mpsc;
#[derive(Default)]
pub(crate) struct ParquetFormatReader {
- read_budget: Option<Arc<ParquetReadBudget>>,
+ read_budget: Option<Arc<ReadBudget>>,
}
impl ParquetFormatReader {
- pub(crate) fn with_read_budget(read_budget: Arc<ParquetReadBudget>) ->
Self {
+ pub(crate) fn with_read_budget(read_budget: Arc<ReadBudget>) -> Self {
Self {
read_budget: Some(read_budget),
}
@@ -702,7 +702,7 @@ async fn read_row_group(
projection: ProjectionMask,
row_group_index: usize,
batch_size: Option<usize>,
- _permit: ParquetReadPermit,
+ _permit: ReadPermit,
sender: mpsc::Sender<ParquetRowGroupMessage>,
) {
let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata(
@@ -2300,7 +2300,7 @@ mod tests {
use crate::arrow::format::{
create_format_reader, create_format_writer, FormatFileReader,
FormatFileWriter,
};
- use crate::arrow::{build_target_arrow_schema, variant_arrow_type,
ParquetReadBudget};
+ use crate::arrow::{build_target_arrow_schema, variant_arrow_type,
ReadBudget};
use crate::io::FileIOBuilder;
use crate::spec::{
ArrayType, BigIntType, DataField, DataType, Datum, IntType, MapType,
PredicateBuilder,
@@ -2838,7 +2838,7 @@ mod tests {
DataType::Int(IntType::new()),
)];
let batches = ParquetFormatReader::with_read_budget(Arc::new(
- ParquetReadBudget::new(8, 256 * 1024 * 1024).unwrap(),
+ ReadBudget::new(8, 256 * 1024 * 1024).unwrap(),
))
.read_batch_stream(
Box::new(file_reader),
@@ -2910,7 +2910,7 @@ mod tests {
"id".to_string(),
DataType::Int(IntType::new()),
)];
- let budget = Arc::new(ParquetReadBudget::new(2, 256 * 1024 *
1024).unwrap());
+ let budget = Arc::new(ReadBudget::new(2, 256 * 1024 * 1024).unwrap());
let first = ParquetFormatReader::with_read_budget(Arc::clone(&budget))
.read_batch_stream(
Box::new(new_file_reader()),
@@ -2965,7 +2965,7 @@ mod tests {
#[tokio::test]
async fn test_parquet_diagnostics_include_reads_with_row_selection() {
let data = write_multi_row_group_parquet(32, 64,
EnabledStatistics::Chunk).await;
- let budget = Arc::new(ParquetReadBudget::new(8, 256 * 1024 *
1024).unwrap());
+ let budget = Arc::new(ReadBudget::new(8, 256 * 1024 * 1024).unwrap());
budget.enable_diagnostics();
let file_size = data.len() as u64;
let fields = vec![int_field("id")];
diff --git a/crates/paimon/src/arrow/mod.rs b/crates/paimon/src/arrow/mod.rs
index 1f017aee..4792cdfb 100644
--- a/crates/paimon/src/arrow/mod.rs
+++ b/crates/paimon/src/arrow/mod.rs
@@ -18,14 +18,14 @@
pub(crate) mod filtering;
pub(crate) mod format;
pub(crate) mod nested_evolution;
-mod parquet_read_budget;
pub(crate) mod partition;
+mod read_budget;
pub(crate) mod residual;
mod row_filter;
pub(crate) mod schema_evolution;
pub(crate) mod shredding;
-pub use parquet_read_budget::ParquetReadBudget;
+pub use read_budget::ReadBudget;
pub use row_filter::{RowFilter, RowFilterContext, RowFilterFactory};
use crate::spec::{
diff --git a/crates/paimon/src/arrow/parquet_read_budget.rs
b/crates/paimon/src/arrow/read_budget.rs
similarity index 58%
rename from crates/paimon/src/arrow/parquet_read_budget.rs
rename to crates/paimon/src/arrow/read_budget.rs
index bef9c6eb..334ac8e3 100644
--- a/crates/paimon/src/arrow/parquet_read_budget.rs
+++ b/crates/paimon/src/arrow/read_budget.rs
@@ -20,24 +20,31 @@ use std::sync::Arc;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use crate::spec::PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES_OPTION;
+
const BYTE_PERMIT_UNIT: u64 = 1024 * 1024;
+const DEFAULT_BYTE_OPTION: &str = "the row-group read budget";
const DEFAULT_PARALLELISM: usize = 8;
const DEFAULT_MAX_INFLIGHT_BYTES: u64 = 256 * 1024 * 1024;
-/// Shared resource budget for concurrent Parquet row-group reads.
+/// 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)]
-pub struct ParquetReadBudget {
+pub struct ReadBudget {
parallelism: usize,
row_groups: Arc<Semaphore>,
bytes: Arc<Semaphore>,
byte_permits: u32,
+ byte_permit_unit: u64,
max_inflight_bytes: u64,
+ byte_option: &'static str,
oversized_warning_logged: AtomicBool,
- diagnostics: Arc<ParquetReadDiagnostics>,
+ diagnostics: Arc<ReadBudgetDiagnostics>,
}
#[derive(Debug)]
-struct ParquetReadDiagnostics {
+struct ReadBudgetDiagnostics {
enabled: AtomicBool,
row_group_count: AtomicU64,
projected_bytes_min: AtomicU64,
@@ -47,7 +54,7 @@ struct ParquetReadDiagnostics {
peak_inflight: AtomicUsize,
}
-impl Default for ParquetReadDiagnostics {
+impl Default for ReadBudgetDiagnostics {
fn default() -> Self {
Self {
enabled: AtomicBool::new(false),
@@ -62,7 +69,7 @@ impl Default for ParquetReadDiagnostics {
}
#[derive(Debug, Default, PartialEq, Eq)]
-pub(crate) struct ParquetReadDiagnosticsSnapshot {
+pub(crate) struct ReadBudgetDiagnosticsSnapshot {
pub(crate) row_group_count: u64,
pub(crate) projected_bytes_min: u64,
pub(crate) projected_bytes_max: u64,
@@ -71,12 +78,28 @@ pub(crate) struct ParquetReadDiagnosticsSnapshot {
pub(crate) peak_inflight: usize,
}
-impl ParquetReadBudget {
+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.
pub fn new(parallelism: usize, max_inflight_bytes: u64) ->
crate::Result<Self> {
+ Ok(
+ Self::with_byte_granularity(parallelism, max_inflight_bytes,
BYTE_PERMIT_UNIT)?
+ .with_byte_option(PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES_OPTION),
+ )
+ }
+
+ /// 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.
+ pub(crate) fn with_byte_granularity(
+ parallelism: usize,
+ max_inflight_bytes: u64,
+ byte_permit_unit: u64,
+ ) -> crate::Result<Self> {
if parallelism == 0 || parallelism > Semaphore::MAX_PERMITS {
return Err(crate::Error::DataInvalid {
message: format!(
- "Parquet row-group parallelism must be between 1 and {},
got {parallelism}",
+ "Row-group read parallelism must be between 1 and {}, got
{parallelism}",
Semaphore::MAX_PERMITS
),
source: None,
@@ -84,13 +107,14 @@ impl ParquetReadBudget {
}
if max_inflight_bytes == 0 {
return Err(crate::Error::DataInvalid {
- message: "Parquet row-group max in-flight bytes must be
greater than 0".to_string(),
+ message: "Row-group read max in-flight bytes must be greater
than 0".to_string(),
source: None,
});
}
+ let byte_permit_unit = byte_permit_unit.max(1);
let max_byte_permits = Semaphore::MAX_PERMITS.min(u32::MAX as usize)
as u32;
let byte_permits = max_inflight_bytes
- .div_ceil(BYTE_PERMIT_UNIT)
+ .div_ceil(byte_permit_unit)
.min(u64::from(max_byte_permits)) as u32;
Ok(Self {
@@ -98,12 +122,20 @@ impl ParquetReadBudget {
row_groups: Arc::new(Semaphore::new(parallelism)),
bytes: Arc::new(Semaphore::new(byte_permits as usize)),
byte_permits,
+ byte_permit_unit,
max_inflight_bytes,
+ byte_option: DEFAULT_BYTE_OPTION,
oversized_warning_logged: AtomicBool::new(false),
- diagnostics: Arc::new(ParquetReadDiagnostics::default()),
+ diagnostics: Arc::new(ReadBudgetDiagnostics::default()),
})
}
+ /// Name the table option that set the byte budget, for the oversized
warning.
+ pub(crate) fn with_byte_option(mut self, byte_option: &'static str) ->
Self {
+ self.byte_option = byte_option;
+ self
+ }
+
pub fn parallelism(&self) -> usize {
self.parallelism
}
@@ -140,9 +172,9 @@ impl ParquetReadBudget {
);
}
- pub(crate) fn diagnostics(&self) -> ParquetReadDiagnosticsSnapshot {
+ pub(crate) fn diagnostics(&self) -> ReadBudgetDiagnosticsSnapshot {
let row_group_count =
self.diagnostics.row_group_count.load(Ordering::Relaxed);
- ParquetReadDiagnosticsSnapshot {
+ ReadBudgetDiagnosticsSnapshot {
row_group_count,
projected_bytes_min: if row_group_count == 0 {
0
@@ -159,39 +191,62 @@ impl ParquetReadBudget {
}
}
- pub(crate) async fn acquire(
- &self,
- projected_uncompressed_bytes: u64,
- ) -> crate::Result<ParquetReadPermit> {
+ /// Wait for a row-group slot and the estimated bytes. The Parquet reader
+ /// acquires from its async tasks.
+ 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(|error| crate::Error::UnexpectedError {
- message: "Parquet row-group read budget was
closed".to_string(),
- source: Some(Box::new(error)),
- })?;
- let requested = projected_uncompressed_bytes
- .max(1)
- .div_ceil(BYTE_PERMIT_UNIT)
- .min(u64::from(self.byte_permits)) as u32;
- if projected_uncompressed_bytes > self.max_inflight_bytes
+ .map_err(|_| Self::closed("row-group"))?;
+ let bytes = Arc::clone(&self.bytes)
+ .acquire_many_owned(self.byte_permits_for(estimated_bytes))
+ .await
+ .map_err(|_| Self::closed("byte"))?;
+ Ok(self.permit(row_group, bytes))
+ }
+
+ /// Take a row-group slot and the estimated bytes if both are free now. The
+ /// Mosaic reader refills its look-ahead without waiting.
+ 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)
+ .try_acquire_many_owned(self.byte_permits_for(estimated_bytes))
+ .ok()?;
+ Some(self.permit(row_group, bytes))
+ }
+
+ /// Block until a row-group slot and the estimated bytes are free. The
+ /// Mosaic reader decodes on its own threads, so it waits here.
+ pub(crate) fn acquire_blocking(&self, estimated_bytes: u64) ->
crate::Result<ReadPermit> {
+ futures::executor::block_on(self.acquire(estimated_bytes))
+ }
+
+ fn closed(resource: &str) -> crate::Error {
+ crate::Error::UnexpectedError {
+ message: format!("The {resource} read budget was closed"),
+ source: None,
+ }
+ }
+
+ fn byte_permits_for(&self, estimated_bytes: u64) -> u32 {
+ if estimated_bytes > self.max_inflight_bytes
&& !self.oversized_warning_logged.swap(true, Ordering::Relaxed)
{
log::warn!(
- "Parquet row group projected size
({projected_uncompressed_bytes} bytes) exceeds \
- read.parquet.row-group.max-inflight-bytes ({} bytes); it will
consume the entire \
- byte budget and may reduce row-group read parallelism;
increase the option if \
- memory allows",
+ "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; \
+ increase the option if memory allows",
+ self.byte_option,
self.max_inflight_bytes
);
}
- let bytes = Arc::clone(&self.bytes)
- .acquire_many_owned(requested)
- .await
- .map_err(|error| crate::Error::UnexpectedError {
- message: "Parquet byte read budget was closed".to_string(),
- source: Some(Box::new(error)),
- })?;
+ estimated_bytes
+ .max(1)
+ .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
@@ -203,15 +258,15 @@ impl ParquetReadBudget {
.fetch_max(current, Ordering::Relaxed);
Arc::clone(&self.diagnostics)
});
- Ok(ParquetReadPermit {
+ ReadPermit {
_row_group: row_group,
_bytes: bytes,
diagnostics,
- })
+ }
}
}
-impl Default for ParquetReadBudget {
+impl Default for ReadBudget {
fn default() -> Self {
Self::new(DEFAULT_PARALLELISM, DEFAULT_MAX_INFLIGHT_BYTES)
.expect("default Parquet read budget is valid")
@@ -219,13 +274,13 @@ impl Default for ParquetReadBudget {
}
#[derive(Debug)]
-pub(crate) struct ParquetReadPermit {
+pub(crate) struct ReadPermit {
_row_group: OwnedSemaphorePermit,
_bytes: OwnedSemaphorePermit,
- diagnostics: Option<Arc<ParquetReadDiagnostics>>,
+ diagnostics: Option<Arc<ReadBudgetDiagnostics>>,
}
-impl Drop for ParquetReadPermit {
+impl Drop for ReadPermit {
fn drop(&mut self) {
if let Some(diagnostics) = &self.diagnostics {
diagnostics.current_inflight.fetch_sub(1, Ordering::Relaxed);
@@ -240,7 +295,7 @@ mod tests {
#[tokio::test]
async fn shared_budget_blocks_until_permits_are_released() {
- let budget = Arc::new(ParquetReadBudget::new(2,
BYTE_PERMIT_UNIT).unwrap());
+ let budget = Arc::new(ReadBudget::new(2, BYTE_PERMIT_UNIT).unwrap());
let first = budget.acquire(2 * BYTE_PERMIT_UNIT).await.unwrap();
assert!(
@@ -259,7 +314,7 @@ mod tests {
#[tokio::test]
async fn diagnostics_aggregate_shared_row_group_reads() {
- let budget = Arc::new(ParquetReadBudget::new(2, 2 *
BYTE_PERMIT_UNIT).unwrap());
+ let budget = Arc::new(ReadBudget::new(2, 2 *
BYTE_PERMIT_UNIT).unwrap());
budget.enable_diagnostics();
budget.record_projected_row_groups(&[300, 100, 200]);
@@ -267,7 +322,7 @@ mod tests {
let second = budget.acquire(1).await.unwrap();
assert_eq!(
budget.diagnostics(),
- ParquetReadDiagnosticsSnapshot {
+ ReadBudgetDiagnosticsSnapshot {
row_group_count: 3,
projected_bytes_min: 100,
projected_bytes_max: 300,
@@ -285,18 +340,17 @@ mod tests {
#[test]
fn rejects_invalid_limits() {
- assert!(ParquetReadBudget::new(0, BYTE_PERMIT_UNIT).is_err());
- assert!(ParquetReadBudget::new(1, 0).is_err());
+ assert!(ReadBudget::new(0, BYTE_PERMIT_UNIT).is_err());
+ assert!(ReadBudget::new(1, 0).is_err());
assert!(
- ParquetReadBudget::new(Semaphore::MAX_PERMITS.saturating_add(1),
BYTE_PERMIT_UNIT)
- .is_err()
+ ReadBudget::new(Semaphore::MAX_PERMITS.saturating_add(1),
BYTE_PERMIT_UNIT).is_err()
);
}
#[tokio::test]
async fn oversized_row_group_consumes_the_budget() {
let max_inflight_bytes = 8 * BYTE_PERMIT_UNIT + 1;
- let budget = Arc::new(ParquetReadBudget::new(8,
max_inflight_bytes).unwrap());
+ let budget = Arc::new(ReadBudget::new(8, max_inflight_bytes).unwrap());
let first = budget.acquire(max_inflight_bytes).await.unwrap();
assert!(!budget.oversized_warning_logged.load(Ordering::Relaxed));
assert!(
@@ -314,7 +368,7 @@ mod tests {
#[tokio::test]
async fn small_row_groups_keep_exact_accounting() {
- let budget = Arc::new(ParquetReadBudget::new(4, 4 *
BYTE_PERMIT_UNIT).unwrap());
+ let budget = Arc::new(ReadBudget::new(4, 4 *
BYTE_PERMIT_UNIT).unwrap());
let mut permits = Vec::new();
for _ in 0..4 {
permits.push(budget.acquire(BYTE_PERMIT_UNIT).await.unwrap());
@@ -328,7 +382,7 @@ mod tests {
#[tokio::test]
async fn tiny_budget_still_admits_one_at_a_time() {
- let budget = Arc::new(ParquetReadBudget::new(8,
BYTE_PERMIT_UNIT).unwrap());
+ let budget = Arc::new(ReadBudget::new(8, BYTE_PERMIT_UNIT).unwrap());
let first = budget.acquire(100 * BYTE_PERMIT_UNIT).await.unwrap();
assert!(
tokio::time::timeout(Duration::from_millis(20), budget.acquire(1))
@@ -339,4 +393,49 @@ mod tests {
drop(first);
budget.acquire(1).await.unwrap();
}
+
+ #[test]
+ fn
try_acquire_is_cumulative_and_an_oversized_head_takes_the_whole_budget() {
+ // Byte-exact granularity, the way the Mosaic reader charges row
groups.
+ let budget = ReadBudget::with_byte_granularity(8, 20, 1).unwrap();
+ let first = budget.try_acquire(10).unwrap();
+ let second = budget.try_acquire(10).unwrap();
+ assert!(budget.try_acquire(10).is_none());
+ drop(first);
+ let third = budget.try_acquire(10).unwrap();
+ drop((second, third));
+
+ let oversized = budget.try_acquire(21).unwrap();
+ assert!(budget.try_acquire(1).is_none());
+ drop(oversized);
+ assert!(budget.try_acquire(1).is_some());
+
+ let slot_limited = ReadBudget::with_byte_granularity(2, u64::MAX,
1).unwrap();
+ let first = slot_limited.try_acquire(1).unwrap();
+ let second = slot_limited.try_acquire(1).unwrap();
+ assert!(slot_limited.try_acquire(1).is_none());
+ drop((first, second));
+ }
+
+ #[test]
+ fn acquire_blocking_waits_off_the_runtime_until_a_permit_is_released() {
+ let budget = Arc::new(ReadBudget::with_byte_granularity(1, 10,
1).unwrap());
+ let held = budget.try_acquire(10).unwrap();
+ let (tx, rx) = std::sync::mpsc::channel();
+ let waiter = std::thread::spawn({
+ let budget = Arc::clone(&budget);
+ move || {
+ let permit = budget.acquire_blocking(10);
+ tx.send(()).ok();
+ permit
+ }
+ });
+
+ assert!(
+ rx.recv_timeout(Duration::from_millis(20)).is_err(),
+ "the budget must hold the waiter until a permit is released"
+ );
+ drop(held);
+ waiter.join().unwrap().unwrap();
+ }
}
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index 51b51e55..a741a087 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -93,8 +93,10 @@ const MANIFEST_SORT_ENABLED_OPTION: &str =
"manifest-sort.enabled";
const WRITE_PARQUET_BUFFER_SIZE_OPTION: &str = "write.parquet-buffer-size";
const READ_BATCH_SIZE_OPTION: &str = "read.batch-size";
const PARQUET_ROW_GROUP_PARALLELISM_OPTION: &str =
"read.parquet.row-group.parallelism";
-const PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES_OPTION: &str =
+pub(crate) const PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES_OPTION: &str =
"read.parquet.row-group.max-inflight-bytes";
+const MOSAIC_READ_PREFETCH_ROW_GROUPS_OPTION: &str =
"mosaic.read.prefetch-row-groups";
+pub(crate) const MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION: &str =
"mosaic.read.prefetch-max-bytes";
pub(crate) const TABLE_READ_SEQUENCE_NUMBER_ENABLED_OPTION: &str =
"table-read.sequence-number.enabled";
pub(crate) const SEQUENCE_FIELD_OPTION: &str = "sequence.field";
@@ -142,6 +144,8 @@ const DEFAULT_WRITE_PARQUET_BUFFER_SIZE: i64 = 256 * 1024 *
1024;
const DEFAULT_READ_BATCH_SIZE: usize = 1024;
const DEFAULT_PARQUET_ROW_GROUP_PARALLELISM: usize = 8;
const DEFAULT_PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES: i64 = 256 * 1024 * 1024;
+const DEFAULT_MOSAIC_READ_PREFETCH_ROW_GROUPS: usize = 8;
+const DEFAULT_MOSAIC_READ_PREFETCH_MAX_BYTES: i64 = 64 * 1024 * 1024;
const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str =
"dynamic-bucket.target-row-num";
const DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM: i64 = 200_000;
const DEFAULT_GLOBAL_INDEX_ROW_COUNT_PER_SHARD: i64 = 100_000;
@@ -404,6 +408,41 @@ impl<'a> CoreOptions<'a> {
Ok(value)
}
+ /// Mosaic row groups kept opening ahead of the consumer per file; `0`
disables prefetch.
+ pub fn mosaic_read_prefetch_row_groups(&self) -> crate::Result<usize> {
+ let Some(raw) =
self.options.get(MOSAIC_READ_PREFETCH_ROW_GROUPS_OPTION) else {
+ return Ok(DEFAULT_MOSAIC_READ_PREFETCH_ROW_GROUPS);
+ };
+ raw.parse::<usize>().map_err(|error| crate::Error::DataInvalid {
+ message: format!(
+ "Option '{MOSAIC_READ_PREFETCH_ROW_GROUPS_OPTION}' must be a
non-negative integer, got: {raw}"
+ ),
+ source: Some(Box::new(error)),
+ })
+ }
+
+ /// Estimated decoded bytes a Mosaic file may hold in flight beyond the
head row group.
+ pub fn mosaic_read_prefetch_max_bytes(&self) -> crate::Result<u64> {
+ let value = match
self.options.get(MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION) {
+ Some(raw) => parse_memory_size(raw).ok_or_else(||
crate::Error::DataInvalid {
+ message: format!(
+ "Option '{MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION}' must be
a valid memory size, got: {raw}"
+ ),
+ source: None,
+ })?,
+ None => DEFAULT_MOSAIC_READ_PREFETCH_MAX_BYTES,
+ };
+ u64::try_from(value)
+ .ok()
+ .filter(|value| *value > 0)
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: format!(
+ "Option '{MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION}' must be
greater than 0, got: {value}"
+ ),
+ source: None,
+ })
+ }
+
/// Scan-wide projected uncompressed bytes for concurrent Parquet row
groups.
pub fn parquet_row_group_max_inflight_bytes(&self) -> crate::Result<u64> {
let value = match self
@@ -1655,6 +1694,53 @@ mod tests {
assert!(!CoreOptions::new(&options).file_index_read_enabled());
}
+ #[test]
+ fn test_mosaic_read_prefetch_options() {
+ let options = HashMap::new();
+ let core = CoreOptions::new(&options);
+ assert_eq!(core.mosaic_read_prefetch_row_groups().unwrap(), 8);
+ assert_eq!(
+ core.mosaic_read_prefetch_max_bytes().unwrap(),
+ 64 * 1024 * 1024
+ );
+
+ let options = HashMap::from([
+ (
+ MOSAIC_READ_PREFETCH_ROW_GROUPS_OPTION.to_string(),
+ "0".to_string(),
+ ),
+ (
+ MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION.to_string(),
+ "16 mb".to_string(),
+ ),
+ ]);
+ let core = CoreOptions::new(&options);
+ assert_eq!(core.mosaic_read_prefetch_row_groups().unwrap(), 0);
+ assert_eq!(
+ core.mosaic_read_prefetch_max_bytes().unwrap(),
+ 16 * 1024 * 1024
+ );
+
+ for value in ["-1", "invalid", "1.5"] {
+ let options = HashMap::from([(
+ MOSAIC_READ_PREFETCH_ROW_GROUPS_OPTION.to_string(),
+ value.to_string(),
+ )]);
+ assert!(CoreOptions::new(&options)
+ .mosaic_read_prefetch_row_groups()
+ .is_err());
+ }
+ for value in ["0", "-1", "invalid", "9223372036854775807 tb"] {
+ let options = HashMap::from([(
+ MOSAIC_READ_PREFETCH_MAX_BYTES_OPTION.to_string(),
+ value.to_string(),
+ )]);
+ assert!(CoreOptions::new(&options)
+ .mosaic_read_prefetch_max_bytes()
+ .is_err());
+ }
+ }
+
#[test]
fn test_parquet_row_group_read_budget_options() {
let options = HashMap::new();
diff --git a/crates/paimon/src/table/audit_log_table/read.rs
b/crates/paimon/src/table/audit_log_table/read.rs
index 625d3170..5b4aac43 100644
--- a/crates/paimon/src/table/audit_log_table/read.rs
+++ b/crates/paimon/src/table/audit_log_table/read.rs
@@ -28,6 +28,7 @@ use crate::spec::{
};
use crate::table::data_file_reader::DataFileReader;
use crate::table::kv_file_reader::{KeyValueFileReader, KeyValueReadConfig};
+use crate::table::table_read::configured_mosaic_prefetch;
use crate::table::ArrowRecordBatchStream;
use crate::DataSplit;
use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions, StringArray};
@@ -125,6 +126,7 @@ impl<'a> AuditLogRead<'a> {
merge_splits: false,
max_merge_input_streams: Some(MAX_MERGE_INPUT_STREAMS),
parquet_read_budget: Some(parquet_read_budget),
+ mosaic_prefetch:
configured_mosaic_prefetch(self.read.table)?,
},
)
.read_with_merge_function(
diff --git a/crates/paimon/src/table/data_evolution_reader.rs
b/crates/paimon/src/table/data_evolution_reader.rs
index c162c7e0..ce0ece4d 100644
--- a/crates/paimon/src/table/data_evolution_reader.rs
+++ b/crates/paimon/src/table/data_evolution_reader.rs
@@ -24,7 +24,8 @@ use super::data_file_reader::{
};
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
use crate::arrow::format::FilePredicates;
-use crate::arrow::{build_target_arrow_schema, ParquetReadBudget};
+use crate::arrow::format::MosaicPrefetchOptions;
+use crate::arrow::{build_target_arrow_schema, ReadBudget};
use crate::deletion_vector::{DeletionVector, DeletionVectorFactory};
use crate::io::FileIO;
use crate::spec::{
@@ -115,7 +116,8 @@ pub(crate) struct DataEvolutionReader {
blob_read_limiter: BlobReadLimiter,
blob_parallelism: usize,
batch_size: Option<usize>,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
+ mosaic_prefetch: MosaicPrefetchOptions,
read_timing: Option<Arc<DataFileReadTiming>>,
}
@@ -195,6 +197,7 @@ impl DataEvolutionReader {
blob_parallelism: DEFAULT_BLOB_READ_PARALLELISM,
batch_size: None,
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
read_timing: None,
})
}
@@ -213,12 +216,17 @@ impl DataEvolutionReader {
pub(crate) fn with_parquet_read_budget(
mut self,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
) -> Self {
self.parquet_read_budget = parquet_read_budget;
self
}
+ pub(crate) fn with_mosaic_prefetch(mut self, mosaic_prefetch:
MosaicPrefetchOptions) -> Self {
+ self.mosaic_prefetch = mosaic_prefetch;
+ self
+ }
+
pub(crate) fn with_read_timing(mut self, read_timing:
Option<Arc<DataFileReadTiming>>) -> Self {
self.read_timing = read_timing;
self
@@ -267,6 +275,7 @@ impl DataEvolutionReader {
.with_batch_size(self.batch_size)
.with_blob_parallelism(self.blob_parallelism)
.with_parquet_read_budget(self.parquet_read_budget.clone())
+ .with_mosaic_prefetch(self.mosaic_prefetch)
.with_read_timing(self.read_timing.clone());
for split in splits {
@@ -566,7 +575,8 @@ impl DataEvolutionReader {
None,
)?
.with_batch_size(self.batch_size)
- .with_parquet_read_budget(self.parquet_read_budget.clone());
+ .with_parquet_read_budget(self.parquet_read_budget.clone())
+ .with_mosaic_prefetch(self.mosaic_prefetch);
let mut stream = prescan.read(splits)?;
let mut view_structs = HashSet::new();
while let Some(batch) = stream.next().await {
@@ -631,6 +641,7 @@ impl DataEvolutionReader {
let blob_parallelism = self.blob_parallelism;
let batch_size = self.batch_size;
let parquet_read_budget = self.parquet_read_budget.clone();
+ let mosaic_prefetch = self.mosaic_prefetch;
let read_timing = self.read_timing.clone();
let anchor_deletion_vector = anchor_deletion_vector.clone();
// Batch size for column-merge output. Matches the default Parquet
reader batch size.
@@ -719,6 +730,7 @@ impl DataEvolutionReader {
blob_as_descriptor,
blob_parallelism,
source_parquet_read_budget.clone(),
+ mosaic_prefetch,
read_timing.clone(),
anchor_deletion_vector.as_ref(),
)
@@ -1258,7 +1270,8 @@ fn open_source_stream(
batch_size: Option<usize>,
blob_as_descriptor: bool,
blob_parallelism: usize,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
+ mosaic_prefetch: MosaicPrefetchOptions,
read_timing: Option<Arc<DataFileReadTiming>>,
anchor_deletion_vector: Option<&DeletionVectorContext>,
) -> crate::Result<ArrowRecordBatchStream> {
@@ -1328,6 +1341,7 @@ fn open_source_stream(
.with_blob_as_descriptor(blob_as_descriptor)
.with_blob_parallelism(blob_parallelism)
.with_parquet_read_budget(parquet_read_budget)
+ .with_mosaic_prefetch(mosaic_prefetch)
.with_read_timing(read_timing);
match source {
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index b24f03ad..7019e9bc 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -17,9 +17,9 @@
use crate::arrow::build_target_arrow_schema;
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
-use crate::arrow::format::create_format_reader_with_budget;
+use crate::arrow::format::{create_format_reader_with_budget,
MosaicPrefetchOptions};
use crate::arrow::schema_evolution::{create_index_mapping, NULL_FIELD_INDEX};
-use crate::arrow::ParquetReadBudget;
+use crate::arrow::ReadBudget;
use crate::deletion_vector::{DeletionVector, DeletionVectorFactory};
use crate::file_index::evaluator::evaluate_file_index;
use crate::file_index::file_index_result::FileIndexResult;
@@ -120,7 +120,8 @@ pub(crate) struct DataFileReader {
blob_as_descriptor: bool,
blob_parallelism: usize,
batch_size: Option<usize>,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
+ mosaic_prefetch: MosaicPrefetchOptions,
read_timing: Option<Arc<DataFileReadTiming>>,
}
@@ -146,6 +147,7 @@ impl DataFileReader {
blob_parallelism: DEFAULT_BLOB_READ_PARALLELISM,
batch_size: None,
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
read_timing: None,
}
}
@@ -173,12 +175,17 @@ impl DataFileReader {
pub(crate) fn with_parquet_read_budget(
mut self,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
) -> Self {
self.parquet_read_budget = parquet_read_budget;
self
}
+ pub(crate) fn with_mosaic_prefetch(mut self, mosaic_prefetch:
MosaicPrefetchOptions) -> Self {
+ self.mosaic_prefetch = mosaic_prefetch;
+ self
+ }
+
pub(crate) fn with_read_timing(mut self, read_timing:
Option<Arc<DataFileReadTiming>>) -> Self {
self.read_timing = read_timing;
self
@@ -460,6 +467,7 @@ impl DataFileReader {
let blob_parallelism = self.blob_parallelism;
let batch_size = self.batch_size;
let parquet_read_budget = self.parquet_read_budget.clone();
+ let mosaic_prefetch = self.mosaic_prefetch;
let read_timing = self.read_timing.clone();
let target_schema = build_target_arrow_schema(&read_type)?;
@@ -520,6 +528,7 @@ impl DataFileReader {
&format_read_fields,
parquet_read_budget,
blob_parallelism,
+ mosaic_prefetch,
)?;
let input_file = file_io.new_input(&path_to_read)?;
let open_start = read_timing.as_ref().map(|_| Instant::now());
@@ -729,6 +738,7 @@ impl DataFileReader {
let blob_as_descriptor = self.blob_as_descriptor;
let blob_parallelism = self.blob_parallelism;
let parquet_read_budget = self.parquet_read_budget.clone();
+ let mosaic_prefetch = self.mosaic_prefetch;
let target_schema = build_target_arrow_schema(&read_type)?;
let file_fields = data_fields.clone().unwrap_or_else(||
table_fields.clone());
@@ -794,6 +804,7 @@ impl DataFileReader {
&format_read_fields,
parquet_read_budget,
blob_parallelism,
+ mosaic_prefetch,
)?;
let input_file = file_io.new_input(&path_to_read)?;
let file_reader = input_file.reader().await?;
diff --git a/crates/paimon/src/table/format_read_builder.rs
b/crates/paimon/src/table/format_read_builder.rs
index f426e096..2e90822d 100644
--- a/crates/paimon/src/table/format_read_builder.rs
+++ b/crates/paimon/src/table/format_read_builder.rs
@@ -23,7 +23,7 @@ use super::read_builder::{resolve_projected_fields,
validate_projection_possible
use super::table_read::configured_parquet_read_budget;
use super::{Table, TableRead, TableScan};
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
-use crate::arrow::ParquetReadBudget;
+use crate::arrow::ReadBudget;
use crate::spec::{DataField, Predicate};
use crate::table::source::RowRange;
use crate::Result;
@@ -42,7 +42,7 @@ pub(crate) struct FormatReadBuilder<'a> {
limit: Option<usize>,
row_ranges: Option<Vec<RowRange>>,
case_sensitive: bool,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
blob_parallelism: usize,
}
@@ -113,7 +113,7 @@ impl<'a> FormatReadBuilder<'a> {
self
}
- pub(crate) fn with_parquet_read_budget(&mut self, budget:
Arc<ParquetReadBudget>) -> &mut Self {
+ pub(crate) fn with_parquet_read_budget(&mut self, budget: Arc<ReadBudget>)
-> &mut Self {
self.parquet_read_budget = Some(budget);
self
}
diff --git a/crates/paimon/src/table/format_table_read.rs
b/crates/paimon/src/table/format_table_read.rs
index c20c3839..d2281405 100644
--- a/crates/paimon/src/table/format_table_read.rs
+++ b/crates/paimon/src/table/format_table_read.rs
@@ -19,11 +19,11 @@
use super::data_file_reader::DataFileReader;
use super::read_builder::split_scan_predicates;
-use super::table_read::configured_parquet_read_budget;
+use super::table_read::{configured_mosaic_prefetch,
configured_parquet_read_budget};
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, ParquetReadBudget};
+use crate::arrow::{build_target_arrow_schema, ReadBudget};
use crate::spec::{DataField, Predicate};
use crate::{DataSplit, Error};
use arrow_array::{RecordBatch, RecordBatchOptions};
@@ -37,7 +37,7 @@ pub(crate) struct FormatTableRead<'a> {
read_type: Vec<DataField>,
data_predicates: Vec<Predicate>,
row_filter_factory: Option<Arc<dyn crate::arrow::RowFilterFactory>>,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
limit: Option<usize>,
blob_parallelism: usize,
}
@@ -85,7 +85,7 @@ impl<'a> FormatTableRead<'a> {
self
}
- pub(crate) fn with_parquet_read_budget(mut self, budget:
Arc<ParquetReadBudget>) -> Self {
+ pub(crate) fn with_parquet_read_budget(mut self, budget: Arc<ReadBudget>)
-> Self {
self.parquet_read_budget = Some(budget);
self
}
@@ -95,7 +95,7 @@ impl<'a> FormatTableRead<'a> {
self
}
- fn parquet_read_budget(&self) -> crate::Result<Arc<ParquetReadBudget>> {
+ 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),
@@ -127,6 +127,7 @@ impl<'a> FormatTableRead<'a> {
let schema_id = self.table.schema().id();
let mut remaining = self.limit;
let batch_size = Some(core_options.read_batch_size()?);
+ let mosaic_prefetch = configured_mosaic_prefetch(self.table)?;
let row_filter_factory = self.row_filter_factory.clone();
let parquet_read_budget = Some(self.parquet_read_budget()?);
let blob_parallelism = self.blob_parallelism;
@@ -147,7 +148,8 @@ impl<'a> FormatTableRead<'a> {
)
.with_batch_size(batch_size)
.with_blob_parallelism(blob_parallelism)
- .with_parquet_read_budget(parquet_read_budget.clone());
+ .with_parquet_read_budget(parquet_read_budget.clone())
+ .with_mosaic_prefetch(mosaic_prefetch);
if let Some(factory) = &row_filter_factory {
reader =
reader.with_row_filter_factory(Arc::clone(factory));
}
diff --git a/crates/paimon/src/table/kv_file_reader.rs
b/crates/paimon/src/table/kv_file_reader.rs
index 8bc665e1..0bcc97cd 100644
--- a/crates/paimon/src/table/kv_file_reader.rs
+++ b/crates/paimon/src/table/kv_file_reader.rs
@@ -30,7 +30,8 @@ use super::sort_merge::{
AggregateMergeFunction, DeduplicateMergeFunction, FirstRowMergeFunction,
MergeFunction,
PartialUpdateMergeFunction, SortMergeReaderBuilder,
};
-use crate::arrow::{build_target_arrow_schema, ParquetReadBudget};
+use crate::arrow::format::MosaicPrefetchOptions;
+use crate::arrow::{build_target_arrow_schema, ReadBudget};
use crate::deletion_vector::DeletionVectorFactory;
use crate::io::FileIO;
use crate::spec::{
@@ -81,7 +82,9 @@ pub(crate) struct KeyValueReadConfig {
/// This limits merge fan-in, not files: files within a run are opened
serially.
pub max_merge_input_streams: Option<usize>,
/// Scan-shared Parquet concurrency and projected-byte budget.
- pub parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ pub parquet_read_budget: Option<Arc<ReadBudget>>,
+ /// Per-file Mosaic row-group prefetch settings.
+ pub mosaic_prefetch: MosaicPrefetchOptions,
}
/// Keep only the conjuncts of `predicates` that reference primary-key columns,
@@ -588,7 +591,8 @@ impl KeyValueFileReader {
pushdown_predicates.clone(),
)
.with_batch_size(Some(config.read_batch_size))
-
.with_parquet_read_budget(group_parquet_read_budget.clone());
+
.with_parquet_read_budget(group_parquet_read_budget.clone())
+ .with_mosaic_prefetch(config.mosaic_prefetch);
let run_schema_manager = config.schema_manager.clone();
let run_file_io = file_io.clone();
let deletion_files_by_split =
deletion_files_by_split.clone();
@@ -1295,7 +1299,7 @@ mod tests {
file
}
- fn kv_reader_with_budget(table: &Table, budget: Arc<ParquetReadBudget>) ->
KeyValueFileReader {
+ fn kv_reader_with_budget(table: &Table, budget: Arc<ReadBudget>) ->
KeyValueFileReader {
let core_options = table.schema().core_options();
KeyValueFileReader::new(
table.file_io().clone(),
@@ -1314,6 +1318,7 @@ mod tests {
merge_splits: true,
max_merge_input_streams: None,
parquet_read_budget: Some(budget),
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
)
}
@@ -1433,6 +1438,7 @@ mod tests {
merge_splits: true,
max_merge_input_streams: Some(256),
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
);
@@ -1640,6 +1646,7 @@ mod tests {
merge_splits: false,
max_merge_input_streams: None,
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
)
.with_input_batch_sizes(input_batch_sizes.clone());
@@ -1710,7 +1717,8 @@ mod tests {
read_batch_size: core_options.read_batch_size().unwrap(),
merge_splits: false,
max_merge_input_streams: None,
- parquet_read_budget: Some(Arc::new(ParquetReadBudget::new(2,
256 << 20).unwrap())),
+ parquet_read_budget: Some(Arc::new(ReadBudget::new(2, 256 <<
20).unwrap())),
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
);
let batches = tokio::time::timeout(
@@ -1767,10 +1775,8 @@ mod tests {
assert_eq!(planned[0].len(), 1);
assert_eq!(planned[0][0].files.len(), 2);
- let reader = kv_reader_with_budget(
- &table,
- Arc::new(ParquetReadBudget::new(1, 256 << 20).unwrap()),
- );
+ let reader =
+ kv_reader_with_budget(&table, Arc::new(ReadBudget::new(1, 256 <<
20).unwrap()));
let batches = tokio::time::timeout(
std::time::Duration::from_secs(5),
reader
@@ -1826,7 +1832,7 @@ mod tests {
};
let first_split = split(vec![first_high, first_low]);
let second_split = split(vec![second_high, second_low]);
- let budget = Arc::new(ParquetReadBudget::new(1, 256 << 20).unwrap());
+ let budget = Arc::new(ReadBudget::new(1, 256 << 20).unwrap());
let first_reader = kv_reader_with_budget(&table, budget.clone());
let second_reader = kv_reader_with_budget(&table, budget);
@@ -1905,6 +1911,7 @@ mod tests {
merge_splits,
max_merge_input_streams: None,
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
)
.read(splits)
@@ -1972,6 +1979,7 @@ mod tests {
merge_splits: true,
max_merge_input_streams: Some(256),
parquet_read_budget: None,
+ mosaic_prefetch: MosaicPrefetchOptions::default(),
},
);
let batches = reader
diff --git a/crates/paimon/src/table/read_builder.rs
b/crates/paimon/src/table/read_builder.rs
index f5087883..89883efa 100644
--- a/crates/paimon/src/table/read_builder.rs
+++ b/crates/paimon/src/table/read_builder.rs
@@ -259,10 +259,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::ParquetReadBudget>,
- ) -> &mut Self {
+ pub fn with_parquet_read_budget(&mut self, budget:
Arc<crate::arrow::ReadBudget>) -> &mut Self {
match &mut self.0 {
ReadBuilderKind::Paimon(builder) => {
builder.with_parquet_read_budget(budget);
@@ -336,7 +333,7 @@ struct PaimonReadBuilder<'a> {
/// Kept apart so neither setter discards the other's constraint.
derived_row_ranges: Option<Vec<RowRange>>,
case_sensitive: bool,
- parquet_read_budget: Option<Arc<crate::arrow::ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<crate::arrow::ReadBudget>>,
blob_parallelism: usize,
}
@@ -485,10 +482,7 @@ impl<'a> PaimonReadBuilder<'a> {
self
}
- fn with_parquet_read_budget(
- &mut self,
- budget: Arc<crate::arrow::ParquetReadBudget>,
- ) -> &mut Self {
+ fn with_parquet_read_budget(&mut self, budget:
Arc<crate::arrow::ReadBudget>) -> &mut Self {
self.parquet_read_budget = Some(budget);
self
}
diff --git a/crates/paimon/src/table/table_read.rs
b/crates/paimon/src/table/table_read.rs
index 9d641841..47a0b0a7 100644
--- a/crates/paimon/src/table/table_read.rs
+++ b/crates/paimon/src/table/table_read.rs
@@ -24,7 +24,8 @@ use super::read_builder::split_scan_predicates;
use super::{ArrowRecordBatchStream, Table};
use crate::arrow::build_target_arrow_schema;
use crate::arrow::format::blob::DEFAULT_BLOB_READ_PARALLELISM;
-use crate::arrow::ParquetReadBudget;
+use crate::arrow::format::MosaicPrefetchOptions;
+use crate::arrow::ReadBudget;
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,
@@ -60,16 +61,22 @@ enum TableReadKind<'a> {
Format(FormatTableRead<'a>),
}
-pub(super) fn configured_parquet_read_budget(
- table: &Table,
-) -> crate::Result<Arc<ParquetReadBudget>> {
+pub(super) fn configured_parquet_read_budget(table: &Table) ->
crate::Result<Arc<ReadBudget>> {
let options = table.schema().core_options();
- Ok(Arc::new(ParquetReadBudget::new(
+ Ok(Arc::new(ReadBudget::new(
options.parquet_row_group_parallelism()?,
options.parquet_row_group_max_inflight_bytes()?,
)?))
}
+pub(crate) fn configured_mosaic_prefetch(table: &Table) ->
crate::Result<MosaicPrefetchOptions> {
+ let options = table.schema().core_options();
+ Ok(MosaicPrefetchOptions {
+ row_groups: options.mosaic_read_prefetch_row_groups()?,
+ max_bytes:
usize::try_from(options.mosaic_read_prefetch_max_bytes()?).unwrap_or(usize::MAX),
+ })
+}
+
impl<'a> TableRead<'a> {
/// Create a new TableRead with a specific read type (projected fields).
pub fn new(
@@ -170,7 +177,7 @@ impl<'a> TableRead<'a> {
/// Override the Parquet resource budget shared by this read.
#[doc(hidden)]
- pub fn with_parquet_read_budget(self, budget: Arc<ParquetReadBudget>) ->
Self {
+ pub fn with_parquet_read_budget(self, budget: Arc<ReadBudget>) -> Self {
match self.0 {
TableReadKind::Paimon(read) => {
Self(TableReadKind::Paimon(read.with_parquet_read_budget(budget)))
@@ -247,7 +254,7 @@ struct PaimonTableRead<'a> {
read_type: Vec<DataField>,
data_predicates: Vec<Predicate>,
row_filter_factory: Option<Arc<dyn crate::arrow::RowFilterFactory>>,
- parquet_read_budget: Option<Arc<ParquetReadBudget>>,
+ parquet_read_budget: Option<Arc<ReadBudget>>,
data_file_read_timing: Option<Arc<DataFileReadTiming>>,
blob_parallelism: usize,
}
@@ -305,7 +312,7 @@ impl<'a> PaimonTableRead<'a> {
self
}
- fn with_parquet_read_budget(mut self, budget: Arc<ParquetReadBudget>) ->
Self {
+ fn with_parquet_read_budget(mut self, budget: Arc<ReadBudget>) -> Self {
self.parquet_read_budget = Some(budget);
self
}
@@ -320,7 +327,7 @@ impl<'a> PaimonTableRead<'a> {
self
}
- fn parquet_read_budget(&self) -> crate::Result<Arc<ParquetReadBudget>> {
+ 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),
@@ -450,7 +457,8 @@ impl<'a> PaimonTableRead<'a> {
.with_file_index_read_enabled(core_options.file_index_read_enabled())
.with_batch_size(Some(core_options.read_batch_size()?))
.with_blob_parallelism(self.blob_parallelism)
- .with_parquet_read_budget(Some(self.parquet_read_budget()?));
+ .with_parquet_read_budget(Some(self.parquet_read_budget()?))
+ .with_mosaic_prefetch(configured_mosaic_prefetch(self.table)?);
let raw_stream = reader.read(&data_splits)?;
Ok(Box::pin(async_stream::try_stream! {
@@ -748,6 +756,7 @@ impl<'a> PaimonTableRead<'a> {
// a row-group permit across yielded batches can otherwise let
// the first side block the second side indefinitely.
parquet_read_budget: None,
+ mosaic_prefetch: configured_mosaic_prefetch(self.table)?,
},
);
reader.read(splits)
@@ -904,6 +913,7 @@ impl<'a> PaimonTableRead<'a> {
&& core_options.deletion_vectors_merge_on_read())
.then_some(MAX_MERGE_INPUT_STREAMS),
parquet_read_budget: Some(self.parquet_read_budget()?),
+ mosaic_prefetch: configured_mosaic_prefetch(self.table)?,
},
);
reader.read(splits)
@@ -931,6 +941,7 @@ impl<'a> PaimonTableRead<'a> {
.with_batch_size(Some(core_options.read_batch_size()?))
.with_blob_parallelism(self.blob_parallelism)
.with_parquet_read_budget(Some(self.parquet_read_budget()?))
+ .with_mosaic_prefetch(configured_mosaic_prefetch(self.table)?)
.with_read_timing(self.data_file_read_timing.clone());
reader.read(data_splits)
}
@@ -954,6 +965,7 @@ impl<'a> PaimonTableRead<'a> {
.with_batch_size(Some(core_options.read_batch_size()?))
.with_blob_parallelism(self.blob_parallelism)
.with_parquet_read_budget(Some(self.parquet_read_budget()?))
+ .with_mosaic_prefetch(configured_mosaic_prefetch(self.table)?)
.with_read_timing(self.data_file_read_timing.clone());
// The engine decoder filter is safe only on the plain append/raw path.
// This constructor is also used by raw-convertible primary-key splits,
@@ -1791,11 +1803,49 @@ mod tests {
));
let read = TableRead::new(&table, table.schema.fields().to_vec(),
Vec::new())
-
.with_parquet_read_budget(Arc::new(ParquetReadBudget::default()));
+ .with_parquet_read_budget(Arc::new(ReadBudget::default()));
assert!(read.to_arrow(&[]).is_ok());
}
}
+ fn table_with_mosaic_prefetch_options(row_groups: &str, max_bytes: &str)
-> Table {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("mosaic.read.prefetch-row-groups", row_groups)
+ .option("mosaic.read.prefetch-max-bytes", max_bytes);
+ Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "mosaic_prefetch_t"),
+ "memory:/mosaic_prefetch_t".to_string(),
+ TableSchema::new(0, &schema.build().unwrap()),
+ None,
+ )
+ }
+
+ #[test]
+ fn test_mosaic_prefetch_options_come_from_table_options() {
+ let table = table_with_mosaic_prefetch_options("0", "16 mb");
+ assert_eq!(
+ configured_mosaic_prefetch(&table).unwrap(),
+ MosaicPrefetchOptions {
+ row_groups: 0,
+ max_bytes: 16 * 1024 * 1024,
+ }
+ );
+ let read = TableRead::new(&table, table.schema.fields().to_vec(),
Vec::new());
+ assert!(read.to_arrow(&[]).is_ok());
+
+ for (row_groups, max_bytes) in [("-1", "64 mb"), ("8", "0")] {
+ let table = table_with_mosaic_prefetch_options(row_groups,
max_bytes);
+ let read = TableRead::new(&table, table.schema.fields().to_vec(),
Vec::new());
+ assert!(matches!(
+ read.to_arrow(&[]),
+ Err(crate::Error::DataInvalid { ref message, .. })
+ if message.contains("mosaic.read.prefetch")
+ ));
+ }
+ }
+
#[test]
fn test_direct_incremental_read_fails_closed_when_query_auth_enabled() {
let table = query_auth_table();
