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 05fc0c67 [core] Add bounded parallel table read execution (#1183)
05fc0c67 is described below
commit 05fc0c6781a1bc6bd3aaa77dfacf8ac8e1070047
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Oct 20 17:18:10 2026 +0800
[core] Add bounded parallel table read execution (#1183)
---
bindings/python/README.md | 11 +
.../python/python/pypaimon_rust/datafusion.pyi | 17 +-
bindings/python/src/read.rs | 125 ++++-
bindings/python/tests/test_read.py | 101 ++++
crates/paimon/src/table/audit_log_table/read.rs | 2 +-
crates/paimon/src/table/mod.rs | 1 +
crates/paimon/src/table/read_execution.rs | 619 +++++++++++++++++++++
crates/paimon/src/table/table_read.rs | 328 +++++++++--
crates/paimon/tests/parallel_table_read_test.rs | 589 ++++++++++++++++++++
crates/paimon/tests/rest_catalog_test.rs | 103 ++++
10 files changed, 1824 insertions(+), 72 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index adebdc1e..4e5ae99b 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -132,6 +132,17 @@ scan = read_builder.new_scan()
plan = scan.plan()
batches = read_builder.new_read().read(plan.splits())
+# Core schedules a bounded set of split readers and owns their global limit.
+# Set ordered=True to retain the supplied split order. The default emits
+# whichever split group is ready first. Ordered parallel reads disable
background
+# Parquet row-group prefetch; split workers still read concurrently.
+reader = read_builder.new_read().read_arrow(plan.splits(), parallelism=4,
ordered=False)
+try:
+ for batch in reader:
+ print(batch)
+finally:
+ reader.close()
+
print(f"\nRead: {batches[0].num_rows} rows")
print(batches[0])
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index badf0a44..c5b66362 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -209,10 +209,21 @@ class RecordBatchReader:
...
class TableRead:
- def read_arrow(self, splits: Sequence[Split]) -> RecordBatchReader:
- """Lazily read splits as PyArrow RecordBatches."""
+ def read_arrow(self, splits: Sequence[Split], *, parallelism: int = 1,
+ ordered: bool = False) -> RecordBatchReader:
+ """Lazily read splits as PyArrow RecordBatches.
+
+ parallelism must be positive and bounds concurrent split groups.
+ ordered=True emits groups in the supplied split order; the default
+ emits ready groups first. The row limit applies across the whole read.
+ Ordered parallel reads disable background Parquet row-group prefetch.
+ close() cancels all pending groups.
+ """
+ ...
+ def read(self, splits: Sequence[Split], *, parallelism: int = 1,
+ ordered: bool = False) -> List[pyarrow.RecordBatch]:
+ """Collect the same stream returned by read_arrow()."""
...
- def read(self, splits: Sequence[Split]) -> List[pyarrow.RecordBatch]: ...
class ReadBuilder:
def with_read_type(self, read_type_json: str) -> "ReadBuilder":
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index dbe21716..ed7759f3 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -904,6 +904,8 @@ impl PyTableRead {
&self,
py: Python<'_>,
splits: &Bound<'_, PyAny>,
+ parallelism: usize,
+ ordered: bool,
) -> PyResult<ArrowRecordBatchStream> {
let splits = extract_splits(splits)?;
py.detach(|| {
@@ -926,7 +928,12 @@ impl PyTableRead {
// Validate config (e.g. projection) before the empty-splits fast
// path so an invalid projection fails consistently regardless of
// how many splits are passed.
- let read = builder.new_read().map_err(to_py_err)?;
+ let read = builder
+ .new_read()
+ .map_err(to_py_err)?
+ .with_parallelism(parallelism)
+ .map_err(to_py_err)?
+ .with_ordered(ordered);
let splits = runtime()
.block_on(read.authorize_splits(&splits))
.map_err(to_py_err)?;
@@ -941,20 +948,35 @@ impl PyTableRead {
#[pymethods]
impl PyTableRead {
- /// Lazily read the given splits as an iterator of PyArrow RecordBatches.
+ /// Lazily read the given splits as PyArrow RecordBatches.
+ /// Core owns split scheduling, the global row limit and cancellation.
+ /// ordered=True retains the supplied split order; the default emits ready
groups first.
+ #[pyo3(signature = (splits, *, parallelism = 1, ordered = false))]
fn read_arrow(
&self,
py: Python<'_>,
splits: &Bound<'_, PyAny>,
+ parallelism: usize,
+ ordered: bool,
) -> PyResult<PyRecordBatchReader> {
- Ok(PyRecordBatchReader::from_stream(
- self.read_stream(py, splits)?,
- ))
+ Ok(PyRecordBatchReader::from_stream(self.read_stream(
+ py,
+ splits,
+ parallelism,
+ ordered,
+ )?))
}
/// Read the given splits into a list of PyArrow RecordBatches.
- fn read(&self, py: Python<'_>, splits: &Bound<'_, PyAny>) ->
PyResult<Vec<Py<PyAny>>> {
- let stream = self.read_stream(py, splits)?;
+ #[pyo3(signature = (splits, *, parallelism = 1, ordered = false))]
+ fn read(
+ &self,
+ py: Python<'_>,
+ splits: &Bound<'_, PyAny>,
+ parallelism: usize,
+ ordered: bool,
+ ) -> PyResult<Vec<Py<PyAny>>> {
+ let stream = self.read_stream(py, splits, parallelism, ordered)?;
let rt = runtime();
let batches = py.detach(|| {
rt.block_on(stream.try_collect::<Vec<_>>())
@@ -1181,8 +1203,97 @@ mod tests {
use std::task::Poll;
use std::time::Duration;
+ use paimon::catalog::Identifier;
+ use paimon::io::FileIOBuilder;
+ use paimon::spec::{IntType, Schema, TableSchema};
+ use pyo3::types::PyList;
+
use super::*;
+ #[test]
+ fn table_read_execution_keywords_validate_before_empty_splits() {
+ Python::attach(|py| {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "read_execution_keywords"),
+ "memory:/read-execution-keywords".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let read = Py::new(py,
PyReadBuilder::new(table).new_read()).unwrap();
+ for method in ["read", "read_arrow"] {
+ for ordered in [false, true] {
+ let kwargs = PyDict::new(py);
+ kwargs.set_item("parallelism", 3).unwrap();
+ kwargs.set_item("ordered", ordered).unwrap();
+ let result = read
+ .bind(py)
+ .call_method(method, (PyList::empty(py),),
Some(&kwargs))
+ .unwrap();
+ if method == "read" {
+ assert_eq!(result.len().unwrap(), 0);
+ } else {
+ let reader: PyRef<PyRecordBatchReader> =
result.extract().unwrap();
+ assert!(reader.next_record_batch().unwrap().is_none());
+ reader.close_reader();
+ reader.close_reader();
+ }
+
+ kwargs.set_item("parallelism", 0).unwrap();
+ let error = read
+ .bind(py)
+ .call_method(method, (PyList::empty(py),),
Some(&kwargs))
+ .unwrap_err();
+ assert!(error.is_instance_of::<PyValueError>(py));
+ assert!(error.to_string().contains("parallelism"));
+ }
+
+ let signature = py
+ .import("inspect")
+ .unwrap()
+ .getattr("signature")
+ .unwrap()
+ .call1((read.bind(py).getattr(method).unwrap(),))
+ .unwrap();
+ let ordered = signature
+ .getattr("parameters")
+ .unwrap()
+ .get_item("ordered")
+ .unwrap();
+ assert!(!ordered
+ .getattr("default")
+ .unwrap()
+ .extract::<bool>()
+ .unwrap());
+ assert_eq!(
+ ordered
+ .getattr("kind")
+ .unwrap()
+ .getattr("name")
+ .unwrap()
+ .extract::<String>()
+ .unwrap(),
+ "KEYWORD_ONLY"
+ );
+
+ // Keep the existing no-keyword entry point and reject
positional
+ // execution options, which would be ambiguous in future APIs.
+ read.bind(py)
+ .call_method1(method, (PyList::empty(py),))
+ .unwrap();
+ let error = read
+ .bind(py)
+ .call_method1(method, (PyList::empty(py), 3))
+ .unwrap_err();
+ assert!(error.is_instance_of::<PyTypeError>(py));
+ }
+ });
+ }
+
#[test]
fn record_batch_reader_releases_stream_on_completion_or_error() {
for fail in [false, true] {
diff --git a/bindings/python/tests/test_read.py
b/bindings/python/tests/test_read.py
index 4e7d199d..4b56326e 100644
--- a/bindings/python/tests/test_read.py
+++ b/bindings/python/tests/test_read.py
@@ -1202,6 +1202,79 @@ def
test_read_arrow_close_is_idempotent_and_stops_iteration():
reader.close()
[email protected]("method", ["read", "read_arrow"])
[email protected]("ordered", [False, True])
+def test_parallel_read_preserves_rows_and_requested_split_order(method,
ordered):
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_partitioned_table(warehouse)
+ builder = table.new_read_builder().with_projection(["id", "dt"])
+ splits = list(reversed(builder.new_scan().plan().splits()))
+ assert len(splits) >= 3
+ read = builder.new_read()
+ expected = [row for split in splits for batch in read.read([split])
+ for row in batch.to_pylist()]
+ # Generators follow the same iterable contract as the original reader.
+ batches = getattr(read, method)((split for split in splits),
parallelism=3, ordered=ordered)
+ actual = [row for batch in batches for row in batch.to_pylist()]
+ assert sorted(actual, key=lambda row: row["id"]) == sorted(expected,
key=lambda row: row["id"])
+ if ordered:
+ assert actual == expected
+
+
[email protected]("method", ["read", "read_arrow"])
[email protected]("ordered", [False, True])
[email protected]("limit", [0, 1, 3])
+def test_parallel_read_shares_filter_and_limit_across_splits(method, ordered,
limit):
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_partitioned_table(warehouse)
+ splits = table.new_read_builder().new_scan().plan().splits()
+ builder = (table.new_read_builder().with_projection(["id"])
+ .with_filter({"method": "greaterThan", "field": "id",
"literals": [1]}))
+ read = builder.new_read()
+ expected = [row for split in splits for batch in read.read([split])
+ for row in batch.to_pylist()]
+ limited = builder.with_limit(limit).new_read()
+ batches = getattr(limited, method)(splits, parallelism=3,
ordered=ordered)
+ actual = [row for batch in batches for row in batch.to_pylist()]
+ assert len(actual) == min(limit, len(expected))
+ assert all(row in expected for row in actual)
+ assert len({row["id"] for row in actual}) == len(actual)
+ if ordered:
+ assert actual == expected[:limit]
+
+
[email protected]("method", ["read", "read_arrow"])
+def test_parallel_read_keeps_zero_column_row_counts(method):
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_partitioned_table(warehouse)
+ splits = table.new_read_builder().new_scan().plan().splits()
+ read =
table.new_read_builder().with_projection([]).with_limit(3).new_read()
+ batches = list(getattr(read, method)(splits, parallelism=3,
ordered=True))
+ assert sum(batch.num_rows for batch in batches) == 3
+ assert all(batch.num_columns == 0 for batch in batches)
+
+
[email protected]("method", ["read", "read_arrow"])
+def test_parallel_read_rejects_zero_parallelism_before_empty_splits(method):
+ with tempfile.TemporaryDirectory() as warehouse:
+ read = _make_table_with_data(warehouse).new_read_builder().new_read()
+ with pytest.raises(ValueError, match="parallelism"):
+ getattr(read, method)([], parallelism=0)
+
+
+def test_parallel_read_arrow_close_after_first_batch_stops_iteration():
+ with tempfile.TemporaryDirectory() as warehouse:
+ table = _make_partitioned_table(warehouse)
+ builder = table.new_read_builder()
+ splits = builder.new_scan().plan().splits()
+ reader = builder.new_read().read_arrow(splits, parallelism=3)
+ assert next(reader).num_rows > 0
+ reader.close()
+ assert reader.read_next_batch() is None
+ assert list(reader) == []
+ reader.close()
+
+
def test_read_empty_splits():
with tempfile.TemporaryDirectory() as warehouse:
table = _make_table_with_data(warehouse)
@@ -1626,6 +1699,34 @@ def
test_incremental_changelog_and_auto_plan_physical_changelog_files():
assert empty.splits() == []
[email protected]("method", ["read", "read_arrow"])
[email protected]("ordered", [False, True])
+def test_parallel_read_keeps_streaming_row_kinds(method, ordered):
+ with tempfile.TemporaryDirectory() as warehouse:
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": warehouse})
+ ctx.sql("CREATE SCHEMA paimon.parallel_changes")
+ ctx.sql("""CREATE TABLE paimon.parallel_changes.t (
+ dt STRING, id INT, value STRING, PRIMARY KEY (dt, id)
+ ) PARTITIONED BY (dt) WITH ('bucket' = '1', 'changelog-producer' =
'input')""")
+ ctx.sql("INSERT INTO paimon.parallel_changes.t VALUES ('p1', 1, 'a'),
('p2', 2, 'b')")
+ ctx.sql("INSERT INTO paimon.parallel_changes.t VALUES ('p3', 3, 'c')")
+ table = PaimonCatalog({"warehouse":
warehouse}).get_table("parallel_changes.t")
+ builder =
table.new_read_builder().with_projection(["id"]).with_include_row_kind(True)
+ splits = list(reversed(builder.new_incremental_scan(0, 2,
"changelog").plan().splits()))
+ assert len(splits) == 3
+ assert all(split.is_streaming() for split in splits)
+ read = builder.new_read()
+ expected = [row for split in splits for batch in read.read([split])
+ for row in batch.to_pylist()]
+ batches = getattr(read, method)(splits, parallelism=3, ordered=ordered)
+ actual = [row for batch in batches for row in batch.to_pylist()]
+ assert sorted(actual, key=lambda row: row["id"]) == [
+ {"rowkind": "+I", "id": 1}, {"rowkind": "+I", "id": 2},
{"rowkind": "+I", "id": 3}]
+ if ordered:
+ assert actual == expected
+
+
@pytest.mark.parametrize("mode", ["changelog", "auto"])
def test_batch_incremental_changelog_packing_sharding_and_limits(mode):
with tempfile.TemporaryDirectory() as warehouse:
diff --git a/crates/paimon/src/table/audit_log_table/read.rs
b/crates/paimon/src/table/audit_log_table/read.rs
index c5e55fc3..4719b39b 100644
--- a/crates/paimon/src/table/audit_log_table/read.rs
+++ b/crates/paimon/src/table/audit_log_table/read.rs
@@ -46,7 +46,7 @@ impl<'a> AuditLogRead<'a> {
// Query-auth is decided per split in `to_arrow`; only the type is
known here.
crate::spec::CoreOptions::new(read.table().schema().options())
.ensure_type_paimon_served(&read.table().identifier().full_name())?;
- match read.0 {
+ match read.kind {
TableReadKind::Paimon(read) => Ok(Self { read }),
TableReadKind::Format(_) => Err(crate::Error::Unsupported {
message: "Format tables do not support audit log batch
read".to_string(),
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index e811ca07..a8b35eb1 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -139,6 +139,7 @@ pub use system_table::SystemRead;
mod key_normalization;
mod read_builder;
+mod read_execution;
mod read_limit;
pub mod referenced_files;
pub(crate) mod rest_env;
diff --git a/crates/paimon/src/table/read_execution.rs
b/crates/paimon/src/table/read_execution.rs
new file mode 100644
index 00000000..b7a4ba3b
--- /dev/null
+++ b/crates/paimon/src/table/read_execution.rs
@@ -0,0 +1,619 @@
+// 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.
+
+//! Bounded split execution shared by table readers and language bindings.
+
+use super::{read_limit::take_limited_batch, ArrowRecordBatchStream};
+use crate::DataSplit;
+use arrow_array::RecordBatch;
+use futures::future::BoxFuture;
+use futures::{FutureExt, StreamExt};
+use std::collections::BTreeMap;
+use std::sync::Arc;
+use tokio::task::JoinSet;
+
+pub(super) type SplitReader =
+ Arc<dyn Fn(&[DataSplit], Option<usize>) ->
crate::Result<ArrowRecordBatchStream> + Send + Sync>;
+
+type NextBatch = (
+ usize,
+ Option<crate::Result<RecordBatch>>,
+ ArrowRecordBatchStream,
+);
+
+/// A contiguous minimax partition balances file bytes without changing split
order.
+fn split_groups(splits: Vec<DataSplit>, parallelism: usize) ->
Vec<Vec<DataSplit>> {
+ let weights = splits
+ .iter()
+ .map(|split| {
+ split
+ .data_files()
+ .iter()
+ .fold(0_u64, |sum, file| {
+
sum.saturating_add(u64::try_from(file.file_size).unwrap_or(0))
+ })
+ .max(1)
+ })
+ .collect::<Vec<_>>();
+ let ranges = group_ranges(&weights, parallelism);
+ let mut splits = splits.into_iter();
+ ranges
+ .into_iter()
+ .map(|(start, end)| splits.by_ref().take(end - start).collect())
+ .collect()
+}
+
+fn group_ranges(weights: &[u64], parallelism: usize) -> Vec<(usize, usize)> {
+ let workers = parallelism.min(weights.len());
+ if workers == 0 {
+ return Vec::new();
+ }
+ let mut low = *weights.iter().max().unwrap();
+ let mut high = weights.iter().copied().fold(0_u64, u64::saturating_add);
+ while low < high {
+ let capacity = low + (high - low) / 2;
+ let mut needed = 1;
+ let mut current = 0_u64;
+ for &weight in weights {
+ if current.saturating_add(weight) > capacity {
+ needed += 1;
+ current = weight;
+ } else {
+ current = current.saturating_add(weight);
+ }
+ }
+ if needed <= workers {
+ high = capacity;
+ } else {
+ low = capacity + 1;
+ }
+ }
+ let mut ranges = Vec::with_capacity(workers);
+ let mut start = 0;
+ for worker in 0..workers {
+ let left = workers - worker;
+ if left == 1 {
+ ranges.push((start, weights.len()));
+ break;
+ }
+ let mut end = start;
+ let mut weight = 0_u64;
+ while end < weights.len() - (left - 1) {
+ let next = weight.saturating_add(weights[end]);
+ if end > start && next > low {
+ break;
+ }
+ weight = next;
+ end += 1;
+ }
+ ranges.push((start, end));
+ start = end;
+ }
+ ranges
+}
+
+pub(super) fn execute(
+ splits: Vec<DataSplit>,
+ parallelism: usize,
+ ordered: bool,
+ limit: Option<usize>,
+ read: SplitReader,
+) -> ArrowRecordBatchStream {
+ let streams = split_groups(splits, parallelism)
+ .into_iter()
+ .map(|group| {
+ let read = Arc::clone(&read);
+ Box::pin(async_stream::try_stream! {
+ let mut remaining = limit;
+ for split in group {
+ if remaining == Some(0) { break; }
+ // Open later splits only on demand, with the remaining
quota.
+ let mut stream = read(std::slice::from_ref(&split),
remaining)?;
+ while remaining != Some(0) {
+ let Some(batch) = stream.next().await else { break; };
+ let batch = batch?;
+ if batch.num_rows() == 0 { continue; }
+ yield take_limited_batch(batch, &mut remaining);
+ }
+ }
+ }) as ArrowRecordBatchStream
+ })
+ .collect();
+ merge_streams(streams, ordered, limit)
+}
+
+fn next_batch(index: usize, mut stream: ArrowRecordBatchStream) ->
BoxFuture<'static, NextBatch> {
+ async move {
+ let batch = stream.next().await;
+ (index, batch, stream)
+ }
+ .boxed()
+}
+
+/// Every group has one pending read or one ready batch. Multiple groups run
+/// on owned runtime tasks so synchronous decoder work can use separate CPUs.
+/// Dropping the output aborts every task through JoinSet's drop
implementation.
+fn merge_streams(
+ mut streams: Vec<ArrowRecordBatchStream>,
+ ordered: bool,
+ limit: Option<usize>,
+) -> ArrowRecordBatchStream {
+ if limit == Some(0) {
+ return Box::pin(futures::stream::empty());
+ }
+ // Ordinary one-worker reads retain the original executor context and avoid
+ // requiring a Tokio runtime just to advance a stream.
+ if streams.len() <= 1 {
+ return limit_single_stream(
+ streams
+ .pop()
+ .unwrap_or_else(|| Box::pin(futures::stream::empty())),
+ limit,
+ );
+ }
+ Box::pin(async_stream::try_stream! {
+ let handle = match tokio::runtime::Handle::try_current() {
+ Ok(handle) => handle,
+ Err(error) => {
+ drop(streams);
+ Err(crate::Error::UnexpectedError {
+ message: "Parallel table reads require a Tokio
runtime".into(),
+ source: Some(Box::new(error)),
+ })?;
+ return;
+ }
+ };
+ let count = streams.len();
+ let mut pending = JoinSet::new();
+ for (index, stream) in streams.into_iter().enumerate() {
+ pending.spawn_on(next_batch(index, stream), &handle);
+ }
+ let mut ready = BTreeMap::new();
+ let mut current = 0;
+ let mut remaining = limit;
+ while remaining != Some(0) && current < count {
+ let next = if ordered {
+ loop {
+ if let Some(next) = ready.remove(¤t) { break
Ok(Some(next)); }
+ let next = match pending.join_next().await {
+ Some(Ok(next)) => next,
+ Some(Err(error)) => break Err(worker_error(error)),
+ None => break Ok(None),
+ };
+ // Surface a failed worker even while an earlier group is
+ // blocked; cleanup happens before exposing the error.
+ if next.0 == current ||
next.1.as_ref().is_some_and(Result::is_err) {
+ break Ok(Some(next));
+ }
+ ready.insert(next.0, next);
+ }
+ } else {
+ pending.join_next().await.transpose().map_err(worker_error)
+ };
+ let (index, batch, stream) = match next {
+ Ok(Some(next)) => next,
+ Ok(None) => break,
+ Err(error) => {
+ drop(ready);
+ pending.shutdown().await;
+ Err(error)?;
+ break;
+ }
+ };
+ match batch {
+ Some(Err(error)) => {
+ // async_stream's `?` yields before returning. Release
every
+ // worker first, even if the caller retains the errored
stream.
+ drop(stream);
+ drop(ready);
+ pending.shutdown().await;
+ Err(error)?;
+ break;
+ }
+ Some(Ok(batch)) => {
+ if batch.num_rows() > 0 {
+ let batch = take_limited_batch(batch, &mut remaining);
+ if remaining == Some(0) {
+ // Reaching the quota cancels unused work before
the
+ // final batch is handed to the caller.
+ drop(stream);
+ drop(ready);
+ pending.shutdown().await;
+ yield batch;
+ break;
+ }
+ yield batch;
+ }
+ // Schedule another decode only after the caller resumes.
+ pending.spawn_on(next_batch(index, stream), &handle);
+ }
+ None if ordered => current += 1,
+ None => {},
+ }
+ }
+ })
+}
+
+fn worker_error(error: tokio::task::JoinError) -> crate::Error {
+ crate::Error::UnexpectedError {
+ message: format!("Table read worker failed: {error}"),
+ source: Some(Box::new(error)),
+ }
+}
+
+fn limit_single_stream(
+ mut stream: ArrowRecordBatchStream,
+ limit: Option<usize>,
+) -> ArrowRecordBatchStream {
+ Box::pin(async_stream::try_stream! {
+ let mut remaining = limit;
+ while remaining != Some(0) {
+ let Some(batch) = stream.next().await else { break; };
+ match batch {
+ Err(error) => {
+ drop(stream);
+ Err(error)?;
+ break;
+ }
+ Ok(batch) => {
+ if batch.num_rows() == 0 { continue; }
+ let batch = take_limited_batch(batch, &mut remaining);
+ if remaining == Some(0) {
+ drop(stream);
+ yield batch;
+ break;
+ }
+ yield batch;
+ }
+ }
+ }
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::Int32Array;
+ use futures::channel::oneshot;
+ use futures::TryStreamExt;
+ use std::sync::atomic::{AtomicUsize, Ordering};
+ use std::time::Duration;
+
+ async fn wait_for_count(count: &AtomicUsize, expected: usize) {
+ tokio::time::timeout(Duration::from_secs(5), async {
+ while count.load(Ordering::SeqCst) != expected {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("worker did not reach its expected lifecycle state");
+ }
+
+ fn batch(value: i32) -> RecordBatch {
+ RecordBatch::try_from_iter([(
+ "id",
+ Arc::new(Int32Array::from(vec![value])) as arrow_array::ArrayRef,
+ )])
+ .unwrap()
+ }
+
+ fn value(batch: &RecordBatch) -> i32 {
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0)
+ }
+
+ struct DropGuard(Arc<AtomicUsize>);
+ impl Drop for DropGuard {
+ fn drop(&mut self) {
+ self.0.fetch_add(1, Ordering::SeqCst);
+ }
+ }
+
+ fn controlled_stream(
+ values: Vec<i32>,
+ first: Option<oneshot::Receiver<()>>,
+ reads: Arc<AtomicUsize>,
+ drops: Arc<AtomicUsize>,
+ ) -> ArrowRecordBatchStream {
+ let guard = DropGuard(drops);
+ Box::pin(async_stream::try_stream! {
+ let _guard = guard;
+ if let Some(first) = first { let _ = first.await; }
+ for item in values {
+ reads.fetch_add(1, Ordering::SeqCst);
+ yield batch(item);
+ }
+ })
+ }
+
+ #[test]
+ fn contiguous_groups_balance_skewed_files_without_reordering() {
+ assert_eq!(
+ group_ranges(&[100, 1, 1, 1, 1, 1], 3),
+ vec![(0, 1), (1, 5), (5, 6)]
+ );
+ assert_eq!(group_ranges(&[1, 1, 1, 1], 2), vec![(0, 2), (2, 4)]);
+ assert_eq!(group_ranges(&[1, 1], 8), vec![(0, 1), (1, 2)]);
+ assert!(group_ranges(&[], 2).is_empty());
+ }
+
+ #[tokio::test]
+ async fn ordered_execution_prefetches_one_batch_per_group() {
+ let (release, wait) = oneshot::channel();
+ let reads = Arc::new(AtomicUsize::new(0));
+ let later_reads = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let mut output = merge_streams(
+ vec![
+ controlled_stream(vec![1, 2], Some(wait), reads.clone(),
drops.clone()),
+ controlled_stream(vec![3, 4], None, later_reads.clone(),
drops.clone()),
+ ],
+ true,
+ None,
+ );
+ assert!(output.next().now_or_never().is_none());
+ wait_for_count(&later_reads, 1).await;
+ assert_eq!(later_reads.load(Ordering::SeqCst), 1);
+ release.send(()).unwrap();
+ assert_eq!(value(&output.next().await.unwrap().unwrap()), 1);
+ assert_eq!(later_reads.load(Ordering::SeqCst), 1);
+ let rest = output
+ .map_ok(|batch| value(&batch))
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(rest, vec![2, 3, 4]);
+ assert_eq!(reads.load(Ordering::SeqCst), 2);
+ assert_eq!(drops.load(Ordering::SeqCst), 2);
+ }
+
+ #[tokio::test]
+ async fn
unordered_execution_yields_ready_groups_and_close_cancels_pending_reads() {
+ let (_release, wait) = oneshot::channel();
+ let reads = Arc::new(AtomicUsize::new(0));
+ let later_reads = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let mut output = merge_streams(
+ vec![
+ controlled_stream(vec![1], Some(wait), reads.clone(),
drops.clone()),
+ controlled_stream(vec![2, 3], None, later_reads.clone(),
drops.clone()),
+ ],
+ false,
+ None,
+ );
+ assert_eq!(value(&output.next().await.unwrap().unwrap()), 2);
+ assert_eq!(reads.load(Ordering::SeqCst), 0);
+ assert_eq!(later_reads.load(Ordering::SeqCst), 1);
+ drop(output);
+ wait_for_count(&drops, 2).await;
+ assert_eq!(later_reads.load(Ordering::SeqCst), 1);
+ }
+
+ #[tokio::test]
+ async fn worker_errors_cancel_other_readers_even_in_ordered_mode() {
+ let (_release, wait) = oneshot::channel();
+ let reads = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let ready_reads = Arc::new(AtomicUsize::new(0));
+ let ready_before_failure = ready_reads.clone();
+ let failing = Box::pin(futures::stream::once(async move {
+ wait_for_count(&ready_before_failure, 1).await;
+ Err(crate::Error::DataInvalid {
+ message: "worker failed".into(),
+ source: None,
+ })
+ })) as ArrowRecordBatchStream;
+ let mut output = merge_streams(
+ vec![
+ controlled_stream(vec![1], Some(wait), reads, drops.clone()),
+ controlled_stream(vec![2, 3], None, ready_reads.clone(),
drops.clone()),
+ failing,
+ ],
+ true,
+ None,
+ );
+ assert!(output
+ .next()
+ .await
+ .unwrap()
+ .unwrap_err()
+ .to_string()
+ .contains("worker failed"));
+ assert_eq!(ready_reads.load(Ordering::SeqCst), 1);
+ assert_eq!(drops.load(Ordering::SeqCst), 2);
+ // Keeping the failed stream alive must not retain another worker.
+ assert!(output.next().await.is_none());
+ }
+
+ #[tokio::test]
+ async fn global_limit_stops_before_requesting_another_batch() {
+ let reads = (0..3)
+ .map(|_| Arc::new(AtomicUsize::new(0)))
+ .collect::<Vec<_>>();
+ let drops = Arc::new(AtomicUsize::new(0));
+ let streams = reads
+ .iter()
+ .enumerate()
+ .map(|(i, reads)| {
+ controlled_stream(
+ vec![i as i32, i as i32 + 10],
+ None,
+ reads.clone(),
+ drops.clone(),
+ )
+ })
+ .collect();
+ let mut output = merge_streams(streams, false, Some(1));
+ let last = output.next().await.unwrap().unwrap();
+ assert_eq!(last.num_rows(), 1);
+ let started = reads
+ .iter()
+ .map(|count| count.load(Ordering::SeqCst))
+ .collect::<Vec<_>>();
+ assert!(started.iter().all(|count| *count <= 1));
+ assert!((1..=3).contains(&started.iter().sum::<usize>()));
+ assert_eq!(drops.load(Ordering::SeqCst), 3);
+ // The caller may keep the final batch and stream without polling
again.
+ assert!(output.next().await.is_none());
+ assert_eq!(
+ reads
+ .iter()
+ .map(|count| count.load(Ordering::SeqCst))
+ .collect::<Vec<_>>(),
+ started
+ );
+ }
+
+ #[test]
+ fn one_worker_and_empty_reads_do_not_require_a_tokio_runtime() {
+ let reads = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ futures::executor::block_on(async {
+ let mut output = merge_streams(
+ vec![controlled_stream(
+ vec![1, 2],
+ None,
+ reads.clone(),
+ drops.clone(),
+ )],
+ false,
+ Some(1),
+ );
+ assert_eq!(value(&output.next().await.unwrap().unwrap()), 1);
+ assert_eq!(reads.load(Ordering::SeqCst), 1);
+ assert_eq!(drops.load(Ordering::SeqCst), 1);
+ assert!(output.next().await.is_none());
+ assert!(merge_streams(Vec::new(), true,
None).next().await.is_none());
+ let guard = DropGuard(drops.clone());
+ let failing = Box::pin(async_stream::stream! {
+ let _guard = guard;
+ yield Err(crate::Error::DataInvalid {
+ message: "serial worker failed".into(), source: None,
+ });
+ std::future::pending::<()>().await;
+ }) as ArrowRecordBatchStream;
+ let mut output = merge_streams(vec![failing], false, None);
+ assert!(output.next().await.unwrap().is_err());
+ assert_eq!(drops.load(Ordering::SeqCst), 2);
+ assert!(output.next().await.is_none());
+ let zero_drops = Arc::new(AtomicUsize::new(0));
+ let streams = (0..2)
+ .map(|i| controlled_stream(vec![i], None, reads.clone(),
zero_drops.clone()))
+ .collect();
+ let mut output = merge_streams(streams, false, Some(0));
+ assert_eq!(zero_drops.load(Ordering::SeqCst), 2);
+ assert!(output.next().await.is_none());
+ let parallel_drops = Arc::new(AtomicUsize::new(0));
+ let streams = (0..2)
+ .map(|i| controlled_stream(vec![i], None, reads.clone(),
parallel_drops.clone()))
+ .collect();
+ let mut output = merge_streams(streams, false, None);
+ let error = output.next().await.unwrap().unwrap_err();
+ assert!(matches!(&error, crate::Error::UnexpectedError { .. }));
+ assert!(error.to_string().contains("Tokio runtime"));
+ assert_eq!(parallel_drops.load(Ordering::SeqCst), 2);
+ assert!(output.next().await.is_none());
+ });
+ }
+
+ #[tokio::test]
+ async fn
a_panicking_worker_cancels_other_workers_before_returning_the_error() {
+ let (_release, wait) = oneshot::channel();
+ let drops = Arc::new(AtomicUsize::new(0));
+ let panicking = Box::pin(futures::stream::once(async {
+ panic!("decoder panic");
+ #[allow(unreachable_code)]
+ Ok(batch(2))
+ })) as ArrowRecordBatchStream;
+ let mut output = merge_streams(
+ vec![
+ controlled_stream(
+ vec![1],
+ Some(wait),
+ Arc::new(AtomicUsize::new(0)),
+ drops.clone(),
+ ),
+ panicking,
+ ],
+ true,
+ None,
+ );
+ let error = output.next().await.unwrap().unwrap_err();
+ assert!(matches!(error, crate::Error::UnexpectedError { .. }));
+ assert_eq!(drops.load(Ordering::SeqCst), 1);
+ assert!(output.next().await.is_none());
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn
independent_workers_poll_synchronous_decoders_on_separate_threads() {
+ #[derive(Default)]
+ struct State {
+ entered: usize,
+ active: usize,
+ peak: usize,
+ threads: std::collections::HashSet<std::thread::ThreadId>,
+ }
+ for ordered in [false, true] {
+ let shared = Arc::new((
+ std::sync::Mutex::new(State::default()),
+ std::sync::Condvar::new(),
+ ));
+ let streams = (0..2)
+ .map(|value| {
+ let shared = shared.clone();
+ Box::pin(futures::stream::once(async move {
+ // This synchronous portion represents a decoder's CPU
work.
+ // The bounded rendezvous lets the former inline
executor
+ // finish rather than hanging forever: it records
peak=1.
+ let (lock, changed) = &*shared;
+ let mut state = lock.lock().unwrap();
+ state.entered += 1;
+ state.active += 1;
+ state.peak = state.peak.max(state.active);
+ state.threads.insert(std::thread::current().id());
+ changed.notify_all();
+ let deadline = std::time::Instant::now() +
Duration::from_secs(2);
+ while state.entered < 2 {
+ let remaining =
+
deadline.saturating_duration_since(std::time::Instant::now());
+ let (next, timeout) = changed.wait_timeout(state,
remaining).unwrap();
+ state = next;
+ if timeout.timed_out() {
+ break;
+ }
+ }
+ state.active -= 1;
+ drop(state);
+ Ok(batch(value))
+ })) as ArrowRecordBatchStream
+ })
+ .collect();
+ let batches = merge_streams(streams, ordered, None)
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(batches.len(), 2);
+ let state = shared.0.lock().unwrap();
+ assert_eq!(state.peak, 2, "synchronous decoder work was
serialized");
+ assert_eq!(state.threads.len(), 2);
+ }
+ }
+}
diff --git a/crates/paimon/src/table/table_read.rs
b/crates/paimon/src/table/table_read.rs
index 39efafd6..47d38e5e 100644
--- a/crates/paimon/src/table/table_read.rs
+++ b/crates/paimon/src/table/table_read.rs
@@ -50,12 +50,17 @@ mod audit;
pub use audit::AuditLogRead;
const MAX_MERGE_INPUT_STREAMS: usize = 256;
+const MAX_TOTAL_BLOB_READS: usize = 64;
/// Table read: reads data from splits (e.g. produced by [TableScan::plan]).
///
/// Reference:
[pypaimon.read.table_read.TableRead](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/read/table_read.py)
#[derive(Debug, Clone)]
-pub struct TableRead<'a>(TableReadKind<'a>);
+pub struct TableRead<'a> {
+ kind: TableReadKind<'a>,
+ parallelism: usize,
+ ordered: bool,
+}
#[derive(Debug, Clone)]
enum TableReadKind<'a> {
@@ -84,7 +89,7 @@ impl<'a> TableRead<'a> {
/// original rows. This does not authorize a read: every split must still
/// carry a whole-table unrestricted grant issued to this exact handle.
pub(super) fn for_update(mut self) -> Self {
- if let TableReadKind::Paimon(read) = &mut self.0 {
+ if let TableReadKind::Paimon(read) = &mut self.kind {
read.update_read = true;
}
self
@@ -99,11 +104,15 @@ impl<'a> TableRead<'a> {
if table.is_format_table() {
Self::new_format(table, read_type, data_predicates, None)
} else {
- Self(TableReadKind::Paimon(PaimonTableRead::new(
- table,
- read_type,
- data_predicates,
- )))
+ Self {
+ kind: TableReadKind::Paimon(PaimonTableRead::new(
+ table,
+ read_type,
+ data_predicates,
+ )),
+ parallelism: 1,
+ ordered: false,
+ }
}
}
@@ -113,17 +122,21 @@ impl<'a> TableRead<'a> {
data_predicates: Vec<Predicate>,
limit: Option<usize>,
) -> Self {
- Self(TableReadKind::Format(FormatTableRead::new(
- table,
- read_type,
- data_predicates,
- limit,
- )))
+ Self {
+ kind: TableReadKind::Format(FormatTableRead::new(
+ table,
+ read_type,
+ data_predicates,
+ limit,
+ )),
+ parallelism: 1,
+ ordered: false,
+ }
}
/// Schema (fields) that this read will produce.
pub fn read_type(&self) -> &[DataField] {
- match &self.0 {
+ match &self.kind {
TableReadKind::Paimon(read) => read.read_type(),
TableReadKind::Format(read) => read.read_type(),
}
@@ -131,7 +144,7 @@ impl<'a> TableRead<'a> {
/// Data predicates for read-side pruning.
pub fn data_predicates(&self) -> &[Predicate] {
- match &self.0 {
+ match &self.kind {
TableReadKind::Paimon(read) => read.data_predicates(),
TableReadKind::Format(read) => read.data_predicates(),
}
@@ -139,7 +152,7 @@ impl<'a> TableRead<'a> {
/// Table for this read.
pub fn table(&self) -> &Table {
- match &self.0 {
+ match &self.kind {
TableReadKind::Paimon(read) => read.table(),
TableReadKind::Format(read) => read.table(),
}
@@ -150,31 +163,68 @@ impl<'a> TableRead<'a> {
/// 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 {
+ match &mut self.kind {
TableReadKind::Paimon(read) => read.resources = Some(resources),
TableReadKind::Format(read) => read.with_resources(resources),
}
self
}
+ /// Read independent splits concurrently, with at most one ready batch per
worker.
+ ///
+ /// Applies to [`Self::to_arrow`] and [`Self::to_arrow_with_row_kind`] on
+ /// Paimon tables. Incremental diff and audit readers keep their own
execution.
+ /// Multiple groups run on Tokio workers; poll parallel reads in a Tokio
runtime.
+ /// The default is one worker. All workers share the same Parquet resource
+ /// budget. Reads that resolve BLOB values before a finite limit remain
+ /// sequential so unused payloads are never fetched speculatively.
+ pub fn with_parallelism(mut self, parallelism: usize) ->
crate::Result<Self> {
+ if parallelism == 0 {
+ return Err(crate::Error::DataInvalid {
+ message: "Read parallelism must be greater than zero".into(),
+ source: None,
+ });
+ }
+ self.parallelism = parallelism;
+ Ok(self)
+ }
+
+ /// Control whether output follows the supplied split order.
+ ///
+ /// When true, batches follow the supplied split order while reads remain
+ /// concurrent. When false (the default), ready split groups can emit
first.
+ /// Rows within a split keep that split reader's existing order in both
modes.
+ /// Ordered parallel reads disable background Parquet row-group prefetch so
+ /// a later group cannot retain scheduling permits needed by an earlier
one.
+ /// Foreground reads still share memory admission from
[`Self::with_resources`].
+ pub fn with_ordered(mut self, ordered: bool) -> Self {
+ self.ordered = ordered;
+ self
+ }
+
+ fn map_kind(mut self, map: impl FnOnce(TableReadKind<'a>) ->
TableReadKind<'a>) -> Self {
+ self.kind = map(self.kind);
+ self
+ }
+
/// Set a filter predicate.
pub fn with_filter(self, filter: Predicate) -> Self {
- match self.0 {
- TableReadKind::Paimon(read) =>
Self(TableReadKind::Paimon(read.with_filter(filter))),
- TableReadKind::Format(read) =>
Self(TableReadKind::Format(read.with_filter(filter))),
- }
+ self.map_kind(|kind| match kind {
+ TableReadKind::Paimon(read) =>
TableReadKind::Paimon(read.with_filter(filter)),
+ TableReadKind::Format(read) =>
TableReadKind::Format(read.with_filter(filter)),
+ })
}
/// Ordinary partition conjuncts are evaluated against split metadata,
/// since partition columns may be absent from physical data files. The
/// predicate uses partition-row indices, not table-field indices.
pub(crate) fn with_partition_predicate(self, predicate: Option<Predicate>)
-> Self {
- match self.0 {
- TableReadKind::Paimon(read) => Self(TableReadKind::Paimon(
- read.with_partition_predicate(predicate),
- )),
- TableReadKind::Format(_) => self,
- }
+ self.map_kind(|kind| match kind {
+ TableReadKind::Paimon(read) => {
+ TableReadKind::Paimon(read.with_partition_predicate(predicate))
+ }
+ other => other,
+ })
}
/// Set the maximum number of concurrent BLOB range reads for this read.
@@ -185,26 +235,26 @@ impl<'a> TableRead<'a> {
source: None,
});
}
- Ok(match self.0 {
- TableReadKind::Paimon(read) => Self(TableReadKind::Paimon(
- read.with_blob_parallelism(blob_parallelism),
- )),
- TableReadKind::Format(read) => Self(TableReadKind::Format(
- read.with_blob_parallelism(blob_parallelism),
- )),
- })
+ Ok(self.map_kind(|kind| match kind {
+ TableReadKind::Paimon(read) => {
+
TableReadKind::Paimon(read.with_blob_parallelism(blob_parallelism))
+ }
+ TableReadKind::Format(read) => {
+
TableReadKind::Format(read.with_blob_parallelism(blob_parallelism))
+ }
+ }))
}
/// Pass the limit to Data Evolution and primary-key readers. These paths
/// apply it after merging/filtering and before output BLOB resolution.
pub(crate) fn with_limit(self, limit: Option<usize>) -> Self {
- match self.0 {
+ self.map_kind(|kind| match kind {
TableReadKind::Paimon(mut read) => {
read.limit = limit;
- Self(TableReadKind::Paimon(read))
+ TableReadKind::Paimon(read)
}
- TableReadKind::Format(read) => Self(TableReadKind::Format(read)),
- }
+ other => other,
+ })
}
/// Attach an engine-specific Parquet decoder-filter factory.
@@ -213,36 +263,36 @@ impl<'a> TableRead<'a> {
/// enforce the expression after the scan because an individual file may
not
/// be able to build a decoder filter. A `query-auth.enabled` read refuses
it.
pub fn with_row_filter_factory(self, factory: Arc<dyn
crate::arrow::RowFilterFactory>) -> Self {
- match self.0 {
+ self.map_kind(|kind| match kind {
TableReadKind::Paimon(read) => {
-
Self(TableReadKind::Paimon(read.with_row_filter_factory(factory)))
+ TableReadKind::Paimon(read.with_row_filter_factory(factory))
}
TableReadKind::Format(read) => {
-
Self(TableReadKind::Format(read.with_row_filter_factory(factory)))
+ TableReadKind::Format(read.with_row_filter_factory(factory))
}
- }
+ })
}
/// Override the Parquet resource budget shared by this read.
#[doc(hidden)]
pub fn with_parquet_read_budget(self, budget: Arc<ReadBudget>) -> Self {
- match self.0 {
+ self.map_kind(|kind| match kind {
TableReadKind::Paimon(read) => {
-
Self(TableReadKind::Paimon(read.with_parquet_read_budget(budget)))
+ TableReadKind::Paimon(read.with_parquet_read_budget(budget))
}
TableReadKind::Format(read) => {
-
Self(TableReadKind::Format(read.with_parquet_read_budget(budget)))
+ TableReadKind::Format(read.with_parquet_read_budget(budget))
}
- }
+ })
}
pub(crate) fn with_data_file_read_timing(self, timing:
Arc<DataFileReadTiming>) -> Self {
- match self.0 {
- TableReadKind::Paimon(read) => Self(TableReadKind::Paimon(
- read.with_data_file_read_timing(timing),
- )),
- TableReadKind::Format(read) => Self(TableReadKind::Format(read)),
- }
+ self.map_kind(|kind| match kind {
+ TableReadKind::Paimon(read) => {
+ TableReadKind::Paimon(read.with_data_file_read_timing(timing))
+ }
+ other => other,
+ })
}
/// Restore catalog authorization for transported or externally planned
splits.
@@ -266,7 +316,7 @@ impl<'a> TableRead<'a> {
.map(|field| field.name().into())
.collect();
select.extend(leaf_names(self.data_predicates()));
- if let TableReadKind::Paimon(read) = &self.0 {
+ if let TableReadKind::Paimon(read) = &self.kind {
if let Some(predicate) = &read.partition_predicate {
select.extend(leaf_names(std::slice::from_ref(predicate)));
}
@@ -321,7 +371,7 @@ impl<'a> TableRead<'a> {
};
restored.push(split.clone().planned(Some(grant)));
}
- let TableReadKind::Paimon(read) = &self.0 else {
+ let TableReadKind::Paimon(read) = &self.kind else {
return Err(unsupported(
"only Paimon data splits carry query authorization",
));
@@ -357,8 +407,10 @@ impl<'a> TableRead<'a> {
/// A `query-auth.enabled` table reads only splits from this handle's plan;
/// re-plan after a permission change.
pub fn to_arrow(&self, data_splits: &[DataSplit]) ->
crate::Result<ArrowRecordBatchStream> {
- match &self.0 {
- TableReadKind::Paimon(read) => read.to_arrow(data_splits),
+ match &self.kind {
+ TableReadKind::Paimon(read) => {
+ read.execute_splits(data_splits, self.parallelism,
self.ordered, false)
+ }
TableReadKind::Format(read) => read.to_arrow(data_splits),
}
}
@@ -373,8 +425,10 @@ impl<'a> TableRead<'a> {
data_splits: &[DataSplit],
) -> crate::Result<ArrowRecordBatchStream> {
// Decided from the splits by the `to_arrow` each branch ends in.
- match &self.0 {
- TableReadKind::Paimon(read) =>
read.to_arrow_with_row_kind(data_splits),
+ match &self.kind {
+ TableReadKind::Paimon(read) => {
+ read.execute_splits(data_splits, self.parallelism,
self.ordered, true)
+ }
TableReadKind::Format(read) => {
let schema = audit_schema_for_read_type(read.read_type(),
false)?;
prepend_insert_row_kind_stream(read.to_arrow(data_splits)?,
schema)
@@ -392,7 +446,7 @@ impl<'a> TableRead<'a> {
) -> crate::Result<ArrowRecordBatchStream> {
self.ensure_query_auth_allowed(plan)?;
plan.validate()?;
- match &self.0 {
+ match &self.kind {
TableReadKind::Paimon(read) => read.to_incremental_arrow(plan),
TableReadKind::Format(_) => Err(crate::Error::Unsupported {
message: "Format tables do not support incremental batch
read".to_string(),
@@ -412,7 +466,7 @@ impl<'a> TableRead<'a> {
) -> crate::Result<ArrowRecordBatchStream> {
self.ensure_query_auth_allowed(plan)?;
plan.validate()?;
- match &self.0 {
+ match &self.kind {
TableReadKind::Paimon(read) => read.to_audit_log_arrow(plan),
TableReadKind::Format(_) => Err(crate::Error::Unsupported {
message: "Format tables do not support audit log batch
read".to_string(),
@@ -456,6 +510,24 @@ struct PaimonTableRead<'a> {
update_read: bool,
}
+fn execution_parallelism(
+ parallelism: usize,
+ limit: Option<usize>,
+ unfiltered: bool,
+ resolves_blob: bool,
+) -> usize {
+ let parallelism = if resolves_blob {
+ parallelism.min(MAX_TOTAL_BLOB_READS)
+ } else {
+ parallelism
+ };
+ match limit {
+ Some(_) if resolves_blob => 1,
+ Some(limit) if unfiltered => parallelism.min(limit.max(1)),
+ _ => parallelism,
+ }
+}
+
impl<'a> PaimonTableRead<'a> {
/// Create a new TableRead with a specific read type (projected fields).
pub fn new(
@@ -478,6 +550,125 @@ impl<'a> PaimonTableRead<'a> {
}
}
+ fn execute_splits(
+ &self,
+ splits: &[DataSplit],
+ parallelism: usize,
+ ordered: bool,
+ include_row_kind: bool,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ let options = self.table.schema().core_options();
+ // Validate the complete grant set before splitting the work.
Per-worker
+ // validation alone could accept incompatible restricted policies.
+ let grant = self.ensure_authorized_by_splits(&options, splits)?;
+ let splits = self.matching_splits(splits)?;
+ let budget = self.parquet_read_budget()?;
+ // Preserve eager schema/configuration errors (including restricted
+ // extraction validation). Constructing the streams performs no file
I/O.
+ // Reuse the existing readers' checks instead of duplicating their
rules.
+ let mut validation = self.clone();
+ validation.parquet_read_budget = Some(Arc::clone(&budget));
+ validation.resources = None;
+ if include_row_kind {
+ drop(validation.to_arrow_with_row_kind(&splits)?);
+ } else {
+ drop(validation.to_arrow(&splits)?);
+ }
+ let mut blob_fields = options.blob_fields();
+ blob_fields.extend(
+ self.table
+ .schema()
+ .fields()
+ .iter()
+ .filter(|field| {
+
super::managed_blob_writer::managed_blob_kind(field.data_type()).is_some()
+ })
+ .map(|field| field.name().to_string()),
+ );
+ let views = options.blob_view_fields();
+ blob_fields.retain(|name| !views.contains(name));
+ let operands = super::query_auth::leaf_names(&self.data_predicates);
+ let uses = |fields: &std::collections::HashSet<String>| {
+ (grant.is_some() && !fields.is_empty())
+ || self
+ .read_type
+ .iter()
+ .any(|field| fields.contains(field.name()))
+ || !fields.is_disjoint(&operands)
+ };
+ // Captured masks and row policies may introduce hidden BLOB operands.
+ // View descriptors still require REST resolution in descriptor mode.
+ let resolves_blob = (!options.blob_as_descriptor() &&
uses(&blob_fields))
+ || (self.table.rest_env().is_some()
+ && options.blob_view_resolve_enabled()
+ && uses(&views));
+ let parallelism = execution_parallelism(
+ parallelism,
+ self.limit,
+ self.data_predicates.is_empty() &&
self.partition_predicate.is_none(),
+ resolves_blob,
+ );
+ let workers = parallelism.min(splits.len()).max(1);
+ // A later ordered group may retain its ready batch while the current
+ // group is still reading. Background prefetch must not keep a shared
+ // scheduling permit in that retained decoder and block the current
+ // group forever. Foreground reads retain shared memory admission.
+ let budget = if ordered && workers > 1 {
+ Arc::new(budget.without_prefetch())
+ } else {
+ budget
+ };
+ let blob_parallelism = self
+ .blob_parallelism
+ .min((MAX_TOTAL_BLOB_READS / workers).max(1));
+ let factory = self.execution_factory(include_row_kind, budget,
blob_parallelism);
+ Ok(super::read_execution::execute(
+ splits.into_owned(),
+ parallelism,
+ ordered,
+ self.limit,
+ factory,
+ ))
+ }
+
+ fn execution_factory(
+ &self,
+ include_row_kind: bool,
+ budget: Arc<ReadBudget>,
+ blob_parallelism: usize,
+ ) -> super::read_execution::SplitReader {
+ // The returned streams own their table/configuration; neither this
+ // TableRead nor its caller needs to outlive asynchronous execution.
+ let table = self.table.clone();
+ let read_type = self.read_type.clone();
+ let data_predicates = self.data_predicates.clone();
+ let partition_predicate = self.partition_predicate.clone();
+ let row_filter_factory = self.row_filter_factory.clone();
+ let data_file_read_timing = self.data_file_read_timing.clone();
+ let update_read = self.update_read;
+ Arc::new(move |splits, limit| {
+ let read = PaimonTableRead {
+ table: &table,
+ read_type: read_type.clone(),
+ data_predicates: data_predicates.clone(),
+ partition_predicate: partition_predicate.clone(),
+ row_filter_factory: row_filter_factory.clone(),
+ parquet_read_budget: Some(Arc::clone(&budget)),
+ // The shared budget already carries this read's resources.
+ resources: None,
+ data_file_read_timing: data_file_read_timing.clone(),
+ blob_parallelism,
+ limit,
+ update_read,
+ };
+ if include_row_kind {
+ read.to_arrow_with_row_kind(splits)
+ } else {
+ read.to_arrow(splits)
+ }
+ })
+ }
+
/// Schema (fields) that this read will produce.
pub fn read_type(&self) -> &[DataField] {
&self.read_type
@@ -2762,6 +2953,21 @@ fn pk_split_needs_merge(split: &DataSplit, dv_enabled:
bool) -> bool {
#[cfg(test)]
mod tests {
use super::*;
+
+ #[test]
+ fn split_execution_worker_policy_limits_speculation() {
+ assert_eq!(execution_parallelism(8, None, true, false), 8);
+ assert_eq!(execution_parallelism(8, Some(0), true, false), 1);
+ assert_eq!(execution_parallelism(8, Some(2), true, false), 2);
+ assert_eq!(execution_parallelism(8, Some(2), false, false), 8);
+ assert_eq!(execution_parallelism(8, Some(100), false, true), 1);
+ assert_eq!(execution_parallelism(8, None, false, true), 8);
+ assert_eq!(execution_parallelism(64, None, false, true), 64);
+ assert_eq!(execution_parallelism(128, None, false, true), 64);
+ assert_eq!(execution_parallelism(128, None, false, false), 128);
+ assert_eq!(execution_parallelism(128, Some(100), false, true), 1);
+ }
+
use crate::catalog::Identifier;
use crate::common::Options;
use crate::file_index::file_indexer_factory::{FileIndexerFactory,
BITMAP_INDEX};
diff --git a/crates/paimon/tests/parallel_table_read_test.rs
b/crates/paimon/tests/parallel_table_read_test.rs
new file mode 100644
index 00000000..ef2d7890
--- /dev/null
+++ b/crates/paimon/tests/parallel_table_read_test.rs
@@ -0,0 +1,589 @@
+// 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 arrow_array::{ArrayRef, Int32Array, LargeBinaryArray, RecordBatch,
StringArray};
+use bytes::Bytes;
+use common::incremental_helpers::{memory_table, persist_table_schema,
setup_dirs};
+use futures::TryStreamExt;
+use paimon::arrow::ReadBudget;
+use paimon::io::FileBlockCache;
+use paimon::spec::{
+ BinaryRow, BlobType, DataType, Datum, IntType, PredicateBuilder, Schema,
TableSchema,
+};
+use paimon::table::{DataSplitBuilder, Table};
+use parquet::arrow::ArrowWriter;
+use parquet::basic::Compression;
+use parquet::file::properties::WriterProperties;
+use std::ops::Range;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::sync::Notify;
+
+async fn table(mode: &str) -> Table {
+ let mut schema = Schema::builder()
+ .column("p", DataType::Int(IntType::new()))
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .partition_keys(["p"])
+ .option("read.batch-size", "1");
+ if mode == "pk" {
+ schema = schema.primary_key(["p", "id"]).option("bucket", "1");
+ } else if mode == "de" {
+ schema = schema
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true");
+ }
+ let (io, table) = memory_table(
+ &format!("memory:/parallel_table_read_{mode}"),
+ TableSchema::new(0, &schema.build().unwrap()),
+ );
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ for generation in 0..if mode == "pk" { 2 } else { 1 } {
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ for partition in 0..4 {
+ let batch = RecordBatch::try_from_iter([
+ (
+ "p",
+ Arc::new(Int32Array::from(vec![partition; 3])) as ArrayRef,
+ ),
+ ("id", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef),
+ (
+ "value",
+ Arc::new(Int32Array::from(vec![generation * 100 +
partition; 3])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ writer.write_arrow_batch(&batch).await.unwrap();
+ }
+ builder
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ }
+ table
+}
+
+fn rows(batches: &[RecordBatch]) -> Vec<(i32, i32, i32)> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ let values = ["p", "id", "value"].map(|name| {
+ batch
+ .column_by_name(name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ });
+ (0..batch.num_rows()).map(move |row| {
+ (
+ values[0].value(row),
+ values[1].value(row),
+ values[2].value(row),
+ )
+ })
+ })
+ .collect()
+}
+
+#[tokio::test]
+async fn
parallel_reads_preserve_requested_order_and_enforce_one_global_limit() {
+ for mode in ["append", "de", "pk"] {
+ let table = table(mode).await;
+ let mut splits = table
+ .new_read_builder()
+ .new_scan()
+ .plan()
+ .await
+ .unwrap()
+ .splits()
+ .to_vec();
+ assert_eq!(splits.len(), 4, "{mode}");
+ splits.reverse();
+ let predicate = PredicateBuilder::new(table.schema().fields())
+ .greater_than("id", Datum::Int(1))
+ .unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_filter(predicate);
+ let mut expected = Vec::new();
+ for split in &splits {
+ let batches = builder
+ .new_read()
+ .unwrap()
+ .to_arrow(std::slice::from_ref(split))
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ expected.extend(rows(&batches));
+ }
+ assert_eq!(expected.len(), 8);
+ if mode == "pk" {
+ assert!(expected.iter().all(|row| row.2 >= 100));
+ }
+ for parallelism in [1, 2, 8] {
+ for ordered in [false, true] {
+ for limit in [None, Some(0), Some(3), Some(100)] {
+ let mut limited = builder.clone();
+ if let Some(limit) = limit {
+ limited.with_limit(limit);
+ }
+ let read = limited
+ .new_read()
+ .unwrap()
+ .with_parallelism(parallelism)
+ .unwrap()
+ .with_ordered(ordered)
+ // One shared slot must not deadlock while workers
wait.
+ .with_parquet_read_budget(Arc::new(
+ ReadBudget::new(1, 1024 * 1024).unwrap(),
+ ));
+ let batches = tokio::time::timeout(
+ std::time::Duration::from_secs(15),
+
read.to_arrow(&splits).unwrap().try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("parallel read stalled on the shared budget")
+ .unwrap();
+ let result = rows(&batches);
+ assert_eq!(
+ result.len(),
+ limit.unwrap_or(8).min(8),
+ "{mode}, p={parallelism}, limit={limit:?}"
+ );
+ if ordered || parallelism == 1 {
+ assert_eq!(result, expected[..result.len()]);
+ } else {
+ assert!(result.iter().all(|row|
expected.contains(row)));
+ let mut unique = result.clone();
+ unique.sort_unstable();
+ unique.dedup();
+ assert_eq!(unique.len(), result.len());
+ }
+ }
+ }
+ }
+ let batches = builder
+ .new_read()
+ .unwrap()
+ .with_parallelism(3)
+ .unwrap()
+ .with_ordered(true)
+ .to_arrow_with_row_kind(&splits)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(rows(&batches), expected);
+ assert!(batches.iter().all(|batch| batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .iter()
+ .all(|kind| kind == Some("+I"))));
+ let mut zero = builder.clone();
+ zero.with_projection(&[]).unwrap().with_limit(5);
+ let batches = zero
+ .new_read()
+ .unwrap()
+ .with_parallelism(3)
+ .unwrap()
+ .to_arrow(&splits)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
5);
+ assert!(batches.iter().all(|batch| batch.num_columns() == 0));
+ }
+}
+
+#[tokio::test]
+async fn data_evolution_streaming_splits_preserve_order_and_limit() {
+ let table = table("de").await;
+ let mut planned = table
+ .new_read_builder()
+ .new_scan()
+ .plan()
+ .await
+ .unwrap()
+ .splits()
+ .to_vec();
+ assert_eq!(planned.len(), 4);
+ planned.reverse();
+
+ for mixed in [false, true] {
+ let splits = planned
+ .iter()
+ .enumerate()
+ .map(|(index, split)| {
+ let mut builder = DataSplitBuilder::new()
+ .with_snapshot(split.snapshot_id())
+ .with_partition(split.partition().clone())
+ .with_bucket(split.bucket())
+ .with_bucket_path(split.bucket_path().to_string())
+ .with_total_buckets(split.total_buckets())
+ .with_data_files(split.data_files().to_vec())
+ .with_raw_convertible(split.raw_convertible())
+ // Alternating flags keep any group of two splits mixed.
+ .with_streaming(!mixed || index % 2 == 0);
+ if let Some(files) = split.data_deletion_files() {
+ builder = builder.with_data_deletion_files(files.to_vec());
+ }
+ if let Some(ranges) = split.row_ranges() {
+ builder = builder.with_row_ranges(ranges.to_vec());
+ }
+ builder.build().unwrap()
+ })
+ .collect::<Vec<_>>();
+ let mut expected = Vec::new();
+ for split in &splits {
+ let batches = table
+ .new_read_builder()
+ .new_read()
+ .unwrap()
+ .to_arrow(std::slice::from_ref(split))
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ expected.extend(rows(&batches));
+ }
+ assert_eq!(expected.len(), 12);
+
+ for include_row_kind in [false, true] {
+ for parallelism in [1, 2] {
+ for ordered in [false, true] {
+ for limit in [None, Some(0), Some(1), Some(4), Some(7)] {
+ let label = format!(
+ "mixed={mixed}, row_kind={include_row_kind},
p={parallelism}, ordered={ordered}, limit={limit:?}"
+ );
+ let mut builder = table.new_read_builder();
+ if let Some(limit) = limit {
+ builder.with_limit(limit);
+ }
+ let read = builder
+ .new_read()
+ .unwrap()
+ .with_parallelism(parallelism)
+ .unwrap()
+ .with_ordered(ordered);
+ let stream = if include_row_kind {
+ read.to_arrow_with_row_kind(&splits)
+ } else {
+ read.to_arrow(&splits)
+ }
+ .unwrap();
+ let batches = tokio::time::timeout(
+ std::time::Duration::from_secs(15),
+ stream.try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("data-evolution split read stalled")
+ .unwrap();
+ let result = rows(&batches);
+ assert_eq!(result.len(), limit.unwrap_or(12).min(12),
"{label}");
+ if ordered || parallelism == 1 {
+ assert_eq!(result, expected[..result.len()],
"{label}");
+ } else {
+ assert!(result.iter().all(|row|
expected.contains(row)), "{label}");
+ let mut unique = result.clone();
+ unique.sort_unstable();
+ unique.dedup();
+ assert_eq!(unique.len(), result.len(), "{label}");
+ }
+ if include_row_kind {
+ assert!(
+ batches.iter().all(|batch| batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .iter()
+ .all(|kind| kind == Some("+I"))),
+ "{label}"
+ );
+ }
+ }
+ }
+ }
+ }
+ }
+}
+
+#[tokio::test]
+async fn zero_parallelism_fails_before_any_splits_are_read() {
+ let table = table("append").await;
+ assert!(table
+ .new_read_builder()
+ .new_read()
+ .unwrap()
+ .with_parallelism(0)
+ .unwrap_err()
+ .to_string()
+ .contains("parallelism"));
+}
+
+#[tokio::test]
+async fn finite_blob_reads_do_not_open_payloads_in_later_splits() {
+ for primary_key in [false, true] {
+ let mut schema = Schema::builder()
+ .column("p", DataType::Int(IntType::new()))
+ .column("id", DataType::Int(IntType::new()))
+ .column("payload", DataType::Blob(BlobType::new()))
+ .partition_keys(["p"])
+ .option("read.batch-size", "1");
+ if primary_key {
+ schema = schema.primary_key(["p", "id"]).option("bucket", "1");
+ } else {
+ schema = schema
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true");
+ }
+ let (io, table) = memory_table(
+ &format!("memory:/parallel_blob_limit_{primary_key}"),
+ TableSchema::new(0, &schema.build().unwrap()),
+ );
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ let batch = RecordBatch::try_from_iter([
+ ("p", Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef),
+ (
+ "id",
+ Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef,
+ ),
+ (
+ "payload",
+ Arc::new(LargeBinaryArray::from(vec![b"a".as_slice(), b"b",
b"c"])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ writer.write_arrow_batch(&batch).await.unwrap();
+ builder
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let plan = table.new_read_builder().new_scan().plan().await.unwrap();
+ let splits = plan.splits();
+ assert_eq!(splits.len(), 3);
+ let mut deleted = 0;
+ for split in &splits[1..] {
+ for status in
io.list_status_recursive(split.bucket_path()).await.unwrap() {
+ if status.path.ends_with(".blob") {
+ io.delete_file(&status.path).await.unwrap();
+ deleted += 1;
+ }
+ }
+ }
+ assert!(deleted > 0);
+ let mut limited = table.new_read_builder();
+ limited.with_limit(1).with_filter(
+ PredicateBuilder::new(table.schema().fields())
+ .greater_than("id", Datum::Int(0))
+ .unwrap(),
+ );
+ for ordered in [false, true] {
+ let batches = limited
+ .new_read()
+ .unwrap()
+ .with_parallelism(8)
+ .unwrap()
+ .with_ordered(ordered)
+ .to_arrow(splits)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
+ assert!(batches.iter().all(|batch| batch
+ .column_by_name("payload")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<LargeBinaryArray>()
+ .unwrap()
+ .value(0)
+ .len()
+ == 1));
+ }
+ assert!(
+ table
+ .new_read_builder()
+ .new_read()
+ .unwrap()
+ .with_parallelism(8)
+ .unwrap()
+ .to_arrow(splits)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .is_err(),
+ "the later payloads really are missing"
+ );
+ }
+}
+
+/// Delay the first file's footer until a later file starts reading column
data.
+/// A cache miss hook supplies this ordering without changing the storage
backend.
+#[derive(Debug, Default)]
+struct DelayedFirstFile {
+ first_waited: AtomicBool,
+ later_data_started: AtomicBool,
+ release_first: Notify,
+}
+
+#[async_trait::async_trait]
+impl FileBlockCache for DelayedFirstFile {
+ async fn get(&self, path: &str, range: Range<u64>) -> Option<Bytes> {
+ if path.ends_with("data-later.parquet") && range.start == 0 {
+ self.later_data_started.store(true, Ordering::SeqCst);
+ self.release_first.notify_waiters();
+ }
+ if path.ends_with("data-first.parquet") {
+ self.first_waited.store(true, Ordering::SeqCst);
+ let released = self.release_first.notified();
+ if !self.later_data_started.load(Ordering::SeqCst) {
+ released.await;
+ }
+ }
+ None
+ }
+
+ async fn put(&self, _path: &str, _offset: u64, _data: Bytes) {}
+ async fn invalidate_path(&self, _path: &str) {}
+ async fn invalidate_prefix(&self, _prefix: &str) {}
+}
+
+#[tokio::test]
+async fn ordered_split_reads_do_not_wait_on_a_later_row_groups_budget_permit()
{
+ const GROUP_ROWS: usize = 150_000;
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("read.batch-size", "1")
+ .build()
+ .unwrap();
+ let (io, table) = memory_table(
+ "memory:/ordered_shared_row_group_budget",
+ TableSchema::new(0, &schema),
+ );
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ let batch = RecordBatch::try_from_iter([(
+ "id",
+ Arc::new(Int32Array::from_iter_values(0..(2 * GROUP_ROWS) as i32)) as
ArrayRef,
+ )])
+ .unwrap();
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(Some(GROUP_ROWS))
+ .set_compression(Compression::UNCOMPRESSED)
+ .set_dictionary_enabled(false)
+ .build();
+ let mut parquet = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut parquet, batch.schema(),
Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ let metadata = writer.close().unwrap();
+ assert_eq!(metadata.num_row_groups(), 2);
+ // Parquet prefetches the final 512 KiB for metadata. Keeping the first
+ // data block outside that suffix makes the hook identify column I/O.
+ assert!(parquet.len() > 1024 * 1024);
+
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer.write_arrow_batch(&batch.slice(0, 1)).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ let template = &messages[0].new_files[0];
+ let bucket_path = format!("{}/bucket-0", table.location());
+ io.mkdirs(&bucket_path).await.unwrap();
+ let mut splits = Vec::new();
+ for name in ["data-first.parquet", "data-later.parquet"] {
+ io.new_output(&format!("{bucket_path}/{name}"))
+ .unwrap()
+ .write(Bytes::copy_from_slice(&parquet))
+ .await
+ .unwrap();
+ let mut file = template.clone();
+ file.file_name = name.to_string();
+ file.file_size = parquet.len() as i64;
+ file.row_count = (2 * GROUP_ROWS) as i64;
+ file.external_path = None;
+ splits.push(
+ DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path.clone())
+ .with_total_buckets(1)
+ .with_data_files(vec![file])
+ .build()
+ .unwrap(),
+ );
+ }
+ let order = Arc::new(DelayedFirstFile::default());
+ let io = io
+ .with_file_block_cache(order.clone(), 64 * 1024, "data")
+ .unwrap();
+ let table = Table::new(
+ io,
+ table.identifier().clone(),
+ table.location().into(),
+ table.schema().clone(),
+ None,
+ );
+ let mut builder = table.new_read_builder();
+ builder.with_limit(3);
+ let output = builder
+ .new_read()
+ .unwrap()
+ .with_parallelism(2)
+ .unwrap()
+ .with_ordered(true)
+ // Two row-group workers enable Parquet prefetch; the 1 MiB byte
+ // window admits only one projected 600 KiB row group at a time.
+ .with_parquet_read_budget(Arc::new(ReadBudget::new(2, 1024 *
1024).unwrap()))
+ .to_arrow(&splits)
+ .unwrap();
+ let batches = tokio::time::timeout(
+ std::time::Duration::from_secs(10),
+ output.try_collect::<Vec<_>>(),
+ )
+ .await
+ .expect("ordered splits deadlocked behind a later row-group permit")
+ .unwrap();
+ assert!(order.first_waited.load(Ordering::SeqCst));
+ assert!(order.later_data_started.load(Ordering::SeqCst));
+ let ids = batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(ids, [0, 1, 2]);
+}
diff --git a/crates/paimon/tests/rest_catalog_test.rs
b/crates/paimon/tests/rest_catalog_test.rs
index cc4d3f2e..164f76b2 100644
--- a/crates/paimon/tests/rest_catalog_test.rs
+++ b/crates/paimon/tests/rest_catalog_test.rs
@@ -3808,6 +3808,109 @@ async fn
test_query_auth_blob_view_predicate_limit_groups_lookups() {
}
}
+#[cfg(not(windows))]
+#[tokio::test]
+async fn test_query_auth_blob_view_limit_skips_rejected_split() {
+ let payload = b"admitted-payload".to_vec();
+ let source = written(
+ "rejected_split_source",
+ blob_schema,
+ vec![blob_batch(vec![1], vec!["Source"], vec![payload.clone()])],
+ )
+ .await;
+ let picture_id = source.table.schema().fields()[2].id();
+ let reference = |row_id| {
+ BlobViewStruct::new(source.identifier.clone(), picture_id, row_id)
+ .serialize()
+ .unwrap()
+ };
+ let target = written(
+ "rejected_split_target",
+ |options| {
+ blob_schema(
+ &[
+ options,
+ &[
+ ("blob-view-field", "picture"),
+ ("source.split.target-size", "1b"),
+ ],
+ ]
+ .concat(),
+ )
+ },
+ vec![
+ blob_batch(vec![0], vec!["Rejected"], vec![reference(99)]),
+ blob_batch(vec![1], vec!["Admitted"], vec![reference(0)]),
+ ],
+ )
+ .await;
+ let server = &target.ctx.server;
+ server.add_table_with_schema(
+ "default",
+ "rejected_split_source",
+ blob_schema(&[]),
+ source.table.location(),
+ );
+ target.set_rules(&[int_leaf(0, "id", "GREATER_THAN", 0)], &[]);
+ let mut builder = target.table.new_read_builder();
+ builder.with_limit(1);
+ let plan = builder.new_scan().plan().await.unwrap();
+ assert_eq!(plan.splits().len(), 2);
+
+ // Probe only ids so locating the rejected split cannot resolve its bad
view.
+ let mut projected = target.table.new_read_builder();
+ projected.with_projection(&["id"]).unwrap();
+ let probe = projected.new_read().unwrap();
+ let mut rejected = None;
+ let mut admitted = None;
+ for split in plan.splits() {
+ let batches = probe
+ .to_arrow(std::slice::from_ref(split))
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let returned = ids(&batches);
+ if returned.is_empty() {
+ rejected = Some(split.clone());
+ } else {
+ assert_eq!(returned, vec![1]);
+ admitted = Some(split.clone());
+ }
+ }
+ let splits = vec![rejected.unwrap(), admitted.unwrap()];
+
+ for (include_row_kind, ordered) in [(false, false), (true, true)] {
+ let label = format!("row_kind={include_row_kind}, ordered={ordered}");
+ let read = builder.new_read().unwrap().with_ordered(ordered);
+ let stream = if include_row_kind {
+ read.to_arrow_with_row_kind(&splits)
+ } else {
+ read.to_arrow(&splits)
+ }
+ .unwrap();
+ let batches = stream.try_collect::<Vec<_>>().await.unwrap();
+ let payload_batches: Vec<_> = batches
+ .iter()
+ .map(|batch| {
+ if include_row_kind {
+ let kinds = column::<StringArray>(batch, "rowkind");
+ assert_eq!(kinds.null_count(), 0, "{label}");
+ assert!(kinds.iter().all(|kind| kind == Some("+I")),
"{label}");
+ batch.project(&[1, 2, 3]).unwrap()
+ } else {
+ batch.clone()
+ }
+ })
+ .collect();
+ assert_eq!(
+ collect_blob_rows(&payload_batches),
+ vec![(1, "Admitted".to_string(), Some(payload.clone()))],
+ "{label}"
+ );
+ }
+}
+
fn schema_of(columns: &[&str], options: &[(&str, &str)]) -> Schema {
let mut builder = Schema::builder();
for name in columns {
