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 2aeaf119 perf(mosaic): prefetch row groups with a Mosaic-owned budget
(#851)
2aeaf119 is described below
commit 2aeaf1193a0135633b630d13b05533493d2c7024
Author: jianguotian <[email protected]>
AuthorDate: Thu Sep 17 22:45:12 2026 +0800
perf(mosaic): prefetch row groups with a Mosaic-owned budget (#851)
---
crates/paimon/src/arrow/format/mosaic.rs | 630 +++++++++++++++++++++++++++++--
1 file changed, 599 insertions(+), 31 deletions(-)
diff --git a/crates/paimon/src/arrow/format/mosaic.rs
b/crates/paimon/src/arrow/format/mosaic.rs
index b10c756c..ae7c5011 100644
--- a/crates/paimon/src/arrow/format/mosaic.rs
+++ b/crates/paimon/src/arrow/format/mosaic.rs
@@ -31,19 +31,24 @@ use arrow_schema::{DataType as ArrowDataType, SchemaRef};
use async_stream::try_stream;
use async_trait::async_trait;
use bytes::Bytes;
+use crossbeam_channel::{Receiver, Sender};
use futures::{SinkExt, StreamExt};
use paimon_mosaic_core::reader::{InputFile, MosaicReader, ReaderAccess};
use paimon_mosaic_core::schema::MosaicSchema;
use paimon_mosaic_core::stats::ColumnStats;
use paimon_mosaic_core::values::Value as MosaicValue;
-use std::collections::{HashMap, HashSet};
+use std::collections::{HashMap, HashSet, VecDeque};
use std::io;
use std::ops::Range;
-use std::sync::Arc;
+use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::sync::{Arc, Condvar, Mutex, OnceLock};
pub(crate) struct MosaicFormatReader;
const DEFAULT_BATCH_SIZE: usize = 8192;
+const DEFAULT_PREFETCH_ROW_GROUPS: usize = 8;
+const DEFAULT_PREFETCH_MAX_BYTES: usize = 64 * 1024 * 1024;
+const MAX_EXECUTOR_WORKERS: usize = 32;
#[async_trait]
impl FormatFileReader for MosaicFormatReader {
@@ -72,7 +77,6 @@ impl FormatFileReader for MosaicFormatReader {
file_fields: predicates.file_fields.clone(),
});
let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
-
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(
@@ -84,6 +88,8 @@ impl FormatFileReader for MosaicFormatReader {
batch_size,
row_selection,
handle,
+ prefetch_row_groups: DEFAULT_PREFETCH_ROW_GROUPS,
+ prefetch_max_bytes: DEFAULT_PREFETCH_MAX_BYTES,
},
|batch|
futures::executor::block_on(batch_tx.send(Ok(batch))).is_ok(),
);
@@ -94,7 +100,11 @@ impl FormatFileReader for MosaicFormatReader {
Ok(try_stream! {
while let Some(batch) = batch_rx.next().await {
- yield batch?;
+ let BudgetedBatch { batch, permit } = batch?;
+ yield batch;
+ // Keep the Mosaic prefetch permit while the decoded batch is
+ // queued or yielded. Advancing or dropping the stream
releases it.
+ drop(permit);
}
read_task.await.map_err(|e| Error::DataInvalid {
message: format!("Mosaic read task failed: {e}"),
@@ -113,11 +123,227 @@ struct MosaicReadRequest {
batch_size: usize,
row_selection: Option<Vec<RowRange>>,
handle: tokio::runtime::Handle,
+ prefetch_row_groups: usize,
+ prefetch_max_bytes: usize,
+}
+
+struct PlannedRowGroup {
+ index: usize,
+ rows: usize,
+ selected_slices: Option<Vec<(usize, usize)>>,
+ estimated_bytes: usize,
+}
+
+struct PendingRowGroup {
+ permit: MosaicPrefetchPermit,
+ 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();
+ }
+}
+
+type MosaicJob = Box<dyn FnOnce() + Send + 'static>;
+
+struct MosaicExecutor {
+ sender: Sender<MosaicJob>,
+}
+
+struct MosaicTask<T> {
+ receiver: Receiver<crate::Result<T>>,
+}
+
+impl MosaicExecutor {
+ fn new() -> crate::Result<Self> {
+ let workers = std::thread::available_parallelism()
+ .map(|parallelism| parallelism.get().saturating_mul(2))
+ .unwrap_or(DEFAULT_PREFETCH_ROW_GROUPS)
+ .clamp(DEFAULT_PREFETCH_ROW_GROUPS, MAX_EXECUTOR_WORKERS);
+ let (sender, receiver): (Sender<MosaicJob>, Receiver<MosaicJob>) =
+ crossbeam_channel::bounded(workers.saturating_mul(2));
+ for worker in 0..workers {
+ let receiver = receiver.clone();
+ std::thread::Builder::new()
+ .name(format!("paimon-mosaic-{worker}"))
+ .spawn(move || {
+ while let Ok(job) = receiver.recv() {
+ job();
+ }
+ })
+ .map_err(|error| Error::UnexpectedError {
+ message: "failed to start Mosaic row-group
executor".to_string(),
+ source: Some(Box::new(error)),
+ })?;
+ }
+ Ok(Self { sender })
+ }
+
+ fn submit<T, F>(&self, task: F) -> crate::Result<MosaicTask<T>>
+ where
+ T: Send + 'static,
+ F: FnOnce() -> crate::Result<T> + Send + 'static,
+ {
+ let (sender, receiver) = crossbeam_channel::bounded(1);
+ self.sender
+ .send(Box::new(move || {
+ let result = match catch_unwind(AssertUnwindSafe(task)) {
+ Ok(result) => result,
+ Err(panic) => Err(Error::UnexpectedError {
+ message: format!(
+ "Mosaic row-group prefetch task panicked: {}",
+ panic_payload_message(panic)
+ ),
+ source: None,
+ }),
+ };
+ let _ = sender.send(result);
+ }))
+ .map_err(|_| Error::UnexpectedError {
+ message: "Mosaic row-group executor stopped before accepting a
task".to_string(),
+ source: None,
+ })?;
+ Ok(MosaicTask { receiver })
+ }
+}
+
+impl<T> MosaicTask<T> {
+ fn join(self) -> crate::Result<T> {
+ self.receiver
+ .recv()
+ .map_err(|error| Error::UnexpectedError {
+ message: "Mosaic row-group executor dropped a task
result".to_string(),
+ source: Some(Box::new(error)),
+ })?
+ }
+}
+
+fn get_or_try_init<'a, T, E>(
+ cell: &'a OnceLock<T>,
+ init_lock: &Mutex<()>,
+ initialize: impl FnOnce() -> Result<T, E>,
+) -> Result<&'a T, E> {
+ if let Some(value) = cell.get() {
+ return Ok(value);
+ }
+ let _guard = init_lock
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner());
+ if let Some(value) = cell.get() {
+ return Ok(value);
+ }
+ let value = initialize()?;
+ Ok(cell.get_or_init(|| value))
+}
+
+fn mosaic_executor() -> crate::Result<&'static MosaicExecutor> {
+ static EXECUTOR: OnceLock<MosaicExecutor> = OnceLock::new();
+ static INIT_LOCK: Mutex<()> = Mutex::new(());
+ get_or_try_init(&EXECUTOR, &INIT_LOCK, MosaicExecutor::new)
}
fn read_mosaic_batches_blocking(
request: MosaicReadRequest,
- mut send_batch: impl FnMut(RecordBatch) -> bool,
+ mut send_batch: impl FnMut(BudgetedBatch) -> bool,
) -> crate::Result<()> {
let MosaicReadRequest {
reader,
@@ -127,9 +353,13 @@ fn read_mosaic_batches_blocking(
batch_size,
row_selection,
handle,
+ prefetch_row_groups,
+ prefetch_max_bytes,
} = request;
- let mosaic_reader = MosaicReader::new(FileReadInputFile::new(reader,
handle), file_size)
- .map_err(mosaic_read_error)?;
+ let mosaic_reader = Arc::new(
+ MosaicReader::new(FileReadInputFile::new(reader, handle.clone()),
file_size)
+ .map_err(mosaic_read_error)?,
+ );
let file_column_names = mosaic_reader
.schema()
@@ -169,6 +399,7 @@ fn read_mosaic_batches_blocking(
build_file_column_indices(mosaic_reader.schema(),
&predicates.file_fields)
});
+ let mut planned = VecDeque::new();
let mut row_group_start = 0usize;
for row_group_index in 0..mosaic_reader.num_row_groups() {
let row_group_rows = mosaic_reader
@@ -210,36 +441,200 @@ fn read_mosaic_batches_blocking(
}
}
- let batch = if all_scan_columns_missing {
- let row_count = selected_slices
- .as_ref()
- .map_or(row_group_rows, |slices| selected_row_count(slices));
- empty_batch(read_schema.clone(), row_count)?
- } else {
- let names = projected_names
- .iter()
- .map(String::as_str)
- .collect::<Vec<_>>();
- let mut row_group_reader = mosaic_reader
- .row_group_reader_by_names(row_group_index, &names)
- .map_err(mosaic_read_error)?;
+ planned.push_back(PlannedRowGroup {
+ index: row_group_index,
+ rows: row_group_rows,
+ selected_slices,
+ estimated_bytes: row_group_rows
+ .saturating_mul(estimated_row_bytes(&existing_scan_fields)),
+ });
+ }
- let batch =
row_group_reader.read_columns().map_err(mosaic_read_error)?;
- take_row_slices(batch, selected_slices.as_deref(), &read_schema)?
+ if all_scan_columns_missing || prefetch_row_groups == 0 {
+ while let Some(plan) = planned.pop_front() {
+ let batch = if all_scan_columns_missing {
+ let row_count = plan
+ .selected_slices
+ .as_ref()
+ .map_or(plan.rows, |slices| selected_row_count(slices));
+ empty_batch(read_schema.clone(), row_count)?
+ } else {
+ read_planned_row_group(&mosaic_reader, &plan,
&projected_names, &read_schema)?
+ };
+ let batch = apply_mosaic_residual(batch, predicates.as_ref(),
&existing_scan_fields)?;
+ if !send_mosaic_batch_chunks(batch, batch_size, None, &mut
send_batch) {
+ return Ok(());
+ }
+ }
+ return Ok(());
+ }
+
+ if planned.is_empty() {
+ return Ok(());
+ }
+
+ 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 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 {
+ break;
+ };
+ let plan = planned.pop_front().unwrap();
+ let reader = Arc::clone(&mosaic_reader);
+ let names = Arc::clone(&projected_names);
+ let schema = Arc::clone(&read_schema);
+ let task = match executor
+ .submit(move || read_planned_row_group(&reader, &plan, &names,
&schema))
+ {
+ Ok(task) => task,
+ Err(error) => {
+ outcome = Err(error);
+ break 'read;
+ }
+ };
+ pending.push_back(PendingRowGroup { permit, task });
+ }
+
+ if pending.is_empty() {
+ let Some(next) = planned.front() else {
+ break;
+ };
+ // 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 plan = planned.pop_front().unwrap();
+ let reader = Arc::clone(&mosaic_reader);
+ let names = Arc::clone(&projected_names);
+ let schema = Arc::clone(&read_schema);
+ let task = match executor
+ .submit(move || read_planned_row_group(&reader, &plan, &names,
&schema))
+ {
+ Ok(task) => task,
+ Err(error) => {
+ outcome = Err(error);
+ break 'read;
+ }
+ };
+ pending.push_back(PendingRowGroup { permit, task });
+ }
+
+ let PendingRowGroup { permit, task } = pending.pop_front().unwrap();
+ let batch = match task.join() {
+ Ok(batch) => batch,
+ Err(error) => {
+ outcome = Err(error);
+ break;
+ }
};
- let batch = match predicates.as_ref() {
- Some(predicates) => {
- filter_record_batch_by_predicates(batch, predicates,
&existing_scan_fields)?
+ let batch = match apply_mosaic_residual(batch, predicates.as_ref(),
&existing_scan_fields) {
+ Ok(batch) => batch,
+ Err(error) => {
+ outcome = Err(error);
+ break;
}
- None => batch,
};
- for chunk in split_batch(batch, batch_size) {
- if !send_batch(chunk) {
- return Ok(());
+ if !send_mosaic_batch_chunks(batch, batch_size,
Some(Arc::new(permit)), &mut send_batch) {
+ planned.clear();
+ break;
+ }
+ }
+
+ // A detached row-group task would keep the reader and its async storage
+ // work alive after cancellation or an error. Join every scheduled task
+ // before returning; preserve the first user-visible failure.
+ while let Some(next) = pending.pop_front() {
+ if let Err(error) = next.task.join() {
+ if outcome.is_ok() {
+ outcome = Err(error);
}
}
}
- Ok(())
+ outcome
+}
+
+fn read_planned_row_group(
+ mosaic_reader: &MosaicReader<FileReadInputFile>,
+ plan: &PlannedRowGroup,
+ projected_names: &[String],
+ read_schema: &SchemaRef,
+) -> crate::Result<RecordBatch> {
+ let names = projected_names
+ .iter()
+ .map(String::as_str)
+ .collect::<Vec<_>>();
+ let mut row_group_reader = mosaic_reader
+ .row_group_reader_by_names(plan.index, &names)
+ .map_err(mosaic_read_error)?;
+ let batch = row_group_reader.read_columns().map_err(mosaic_read_error)?;
+ take_row_slices(batch, plan.selected_slices.as_deref(), read_schema)
+}
+
+fn panic_payload_message(panic: Box<dyn std::any::Any + Send>) -> String {
+ if let Some(message) = panic.downcast_ref::<&str>() {
+ (*message).to_string()
+ } else if let Some(message) = panic.downcast_ref::<String>() {
+ message.clone()
+ } else {
+ "unknown panic payload".to_string()
+ }
+}
+
+fn apply_mosaic_residual(
+ batch: RecordBatch,
+ predicates: Option<&FilePredicates>,
+ existing_scan_fields: &[DataField],
+) -> crate::Result<RecordBatch> {
+ match predicates {
+ Some(predicates) => {
+ filter_record_batch_by_predicates(batch, predicates,
existing_scan_fields)
+ }
+ None => Ok(batch),
+ }
+}
+
+fn send_mosaic_batch_chunks(
+ batch: RecordBatch,
+ batch_size: usize,
+ permit: Option<Arc<MosaicPrefetchPermit>>,
+ send_batch: &mut impl FnMut(BudgetedBatch) -> bool,
+) -> bool {
+ split_batch(batch, batch_size).into_iter().all(|batch| {
+ send_batch(BudgetedBatch {
+ batch,
+ permit: permit.clone(),
+ })
+ })
+}
+
+/// Rough decoded size of one projected row, matching Java Mosaic's prefetch
budget.
+fn estimated_row_bytes(fields: &[DataField]) -> usize {
+ fields
+ .iter()
+ .map(|field| match field.data_type() {
+ PaimonDataType::Boolean(_) | PaimonDataType::TinyInt(_) => 2,
+ PaimonDataType::SmallInt(_) => 3,
+ PaimonDataType::Int(_)
+ | PaimonDataType::Float(_)
+ | PaimonDataType::Date(_)
+ | PaimonDataType::Time(_) => 5,
+ PaimonDataType::BigInt(_)
+ | PaimonDataType::Double(_)
+ | PaimonDataType::Timestamp(_)
+ | PaimonDataType::LocalZonedTimestamp(_) => 9,
+ PaimonDataType::Decimal(_) => 17,
+ PaimonDataType::Char(_)
+ | PaimonDataType::VarChar(_)
+ | PaimonDataType::Binary(_)
+ | PaimonDataType::VarBinary(_) => 40,
+ _ => 128,
+ })
+ .sum::<usize>()
+ .max(1)
}
struct MosaicRowGroupStats<'a> {
@@ -637,7 +1032,37 @@ mod tests {
use paimon_mosaic_core::spec::COMPRESSION_NONE;
use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions};
use std::ops::Range;
- use std::sync::{Arc, Mutex};
+ use std::sync::atomic::{AtomicUsize, Ordering};
+ use std::sync::{Arc, Mutex, OnceLock};
+
+ #[test]
+ fn test_mosaic_executor_init_retries_after_transient_failure() {
+ let executor = OnceLock::new();
+ let init_lock = Mutex::new(());
+ let attempts = AtomicUsize::new(0);
+ let initialize = || {
+ if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
+ Err("transient thread creation failure")
+ } else {
+ Ok(42)
+ }
+ };
+
+ assert_eq!(
+ get_or_try_init(&executor, &init_lock, initialize),
+ Err("transient thread creation failure")
+ );
+ assert!(executor.get().is_none(), "a failed init must not be cached");
+ assert_eq!(
+ *get_or_try_init(&executor, &init_lock, initialize).unwrap(),
+ 42
+ );
+ assert_eq!(
+ *get_or_try_init(&executor, &init_lock, initialize).unwrap(),
+ 42
+ );
+ assert_eq!(attempts.load(Ordering::SeqCst), 2);
+ }
struct TestFileRead {
data: Bytes,
@@ -672,6 +1097,26 @@ mod tests {
runtime_thread: std::thread::ThreadId,
}
+ struct ConcurrentTrackingFileRead {
+ data: Bytes,
+ active: Arc<AtomicUsize>,
+ max_active: Arc<AtomicUsize>,
+ }
+
+ #[async_trait]
+ impl FileRead for ConcurrentTrackingFileRead {
+ async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
+ let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
+ self.max_active.fetch_max(active, Ordering::SeqCst);
+ tokio::time::sleep(std::time::Duration::from_millis(20)).await;
+ let start = usize::try_from(range.start).unwrap();
+ let end = usize::try_from(range.end).unwrap();
+ let result = self.data.slice(start..end);
+ self.active.fetch_sub(1, Ordering::SeqCst);
+ Ok(result)
+ }
+ }
+
#[async_trait]
impl FileRead for RuntimeThreadFileRead {
async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
@@ -967,6 +1412,48 @@ mod tests {
)
}
+ async fn read_with_prefetch_tracking(
+ data: Bytes,
+ fields: Vec<DataField>,
+ prefetch_row_groups: usize,
+ prefetch_max_bytes: usize,
+ ) -> (Vec<RecordBatch>, usize) {
+ let active = Arc::new(AtomicUsize::new(0));
+ let max_active = Arc::new(AtomicUsize::new(0));
+ let file_size = data.len() as u64;
+ let read = ConcurrentTrackingFileRead {
+ data,
+ active,
+ max_active: Arc::clone(&max_active),
+ };
+ let handle = tokio::runtime::Handle::current();
+ let batches = tokio::task::spawn_blocking(move || {
+ let mut batches = Vec::new();
+ read_mosaic_batches_blocking(
+ MosaicReadRequest {
+ reader: Box::new(read),
+ file_size,
+ read_fields: fields,
+ predicates: None,
+ batch_size: DEFAULT_BATCH_SIZE,
+ row_selection: None,
+ handle,
+ prefetch_row_groups,
+ prefetch_max_bytes,
+ },
+ |BudgetedBatch { batch, .. }| {
+ batches.push(batch);
+ true
+ },
+ )?;
+ Ok::<_, Error>(batches)
+ })
+ .await
+ .unwrap()
+ .unwrap();
+ (batches, max_active.load(Ordering::SeqCst))
+ }
+
fn timestamp_fields() -> Vec<DataField> {
vec![
DataField::new(
@@ -1038,6 +1525,87 @@ mod tests {
assert_eq!(ids.value(4), 5);
}
+ #[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());
+ let fields = vec![data_fields()[0].clone()];
+ let (batches, max_active) = read_with_prefetch_tracking(data, fields,
3, usize::MAX).await;
+
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert!(
+ max_active >= 2,
+ "row-group prefetch should overlap storage reads,
max_active={max_active}"
+ );
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_prefetch_byte_budget_limits_groups_ahead() {
+ let data = multi_row_group_mosaic(Vec::new());
+ let fields = vec![data_fields()[0].clone()];
+ let (batches, max_active) = read_with_prefetch_tracking(data, fields,
3, 1).await;
+
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert_eq!(
+ max_active, 1,
+ "a one-byte budget should admit only the mandatory head group"
+ );
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_disabled_prefetch_reads_row_groups_on_demand() {
+ let data = multi_row_group_mosaic(Vec::new());
+ let fields = vec![data_fields()[0].clone()];
+ let (batches, max_active) = read_with_prefetch_tracking(data, fields,
0, usize::MAX).await;
+
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert_eq!(max_active, 1);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_prefetch_byte_budget_is_cumulative() {
+ let data = multi_row_group_mosaic(Vec::new());
+ let fields = vec![data_fields()[0].clone()];
+ // Each two-row Int group is conservatively charged 10 bytes, so this
+ // admits two groups but not all three.
+ let (batches, max_active) = read_with_prefetch_tracking(data, fields,
3, 20).await;
+
+ assert_eq!(collect_i32_column(&batches, 0), vec![1, 2, 10, 11, 20,
21]);
+ assert_eq!(max_active, 2);
+ }
+
+ #[test]
+ fn
test_mosaic_prefetch_budget_is_cumulative_and_oversized_head_progresses() {
+ assert_eq!(estimated_row_bytes(&[data_fields()[0].clone()]), 5);
+ assert_eq!(estimated_row_bytes(&[data_fields()[1].clone()]), 40);
+ assert_eq!(
+ estimated_row_bytes(&[field(
+ 3,
+ "nested",
+ DataType::Array(ArrayType::new(DataType::Int(IntType::new())))
+ )]),
+ 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]
async fn test_row_id_predicate_is_rejected() {
let data = write_mosaic(&sample_batch());
