This is an automated email from the ASF dual-hosted git repository.
fresh-borzoni pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 350d9362 feat: Add RecordBatchLogReader for bounded log reading (#446)
350d9362 is described below
commit 350d9362249914daa9e60f2ff57012702b408d20
Author: Kaiqi Dong <[email protected]>
AuthorDate: Thu May 14 11:11:14 2026 +0200
feat: Add RecordBatchLogReader for bounded log reading (#446)
* Add RecordBatchLogReader for bounded log reading
* address comments
* update doc
* update doc and inline comments
* rebase and follow up after rebase of new changes, and fix a corner issue
* feedback
* run tests in thread to avoid asyncio event loop starvation
* address feedback
---
bindings/python/example/example.py | 10 +-
bindings/python/fluss/__init__.pyi | 24 +
bindings/python/src/table.rs | 293 +++++-----
bindings/python/test/test_log_table.py | 124 +++++
crates/fluss/src/client/table/mod.rs | 2 +
crates/fluss/src/client/table/reader.rs | 701 ++++++++++++++++++++++++
crates/fluss/src/client/table/scanner.rs | 138 ++++-
website/docs/user-guide/python/api-reference.md | 3 +
website/docs/user-guide/rust/api-reference.md | 44 ++
9 files changed, 1169 insertions(+), 170 deletions(-)
diff --git a/bindings/python/example/example.py
b/bindings/python/example/example.py
index 0149996c..23ccc6d1 100644
--- a/bindings/python/example/example.py
+++ b/bindings/python/example/example.py
@@ -294,8 +294,14 @@ async def main():
except Exception as e:
print(f"Could not convert to Pandas: {e}")
- # TODO: support to_arrow_batch_reader()
- # which is reserved for streaming use cases
+ # to_arrow_batch_reader() — returns a lazy PyArrow RecordBatchReader
+ batch_scanner_reader = await
table.new_scan().create_record_batch_log_scanner()
+ batch_scanner_reader.subscribe_buckets(
+ {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}
+ )
+ arrow_reader = batch_scanner_reader.to_arrow_batch_reader()
+ reader_table = pa.Table.from_batches(list(arrow_reader),
schema=arrow_reader.schema)
+ print(f"\nVia to_arrow_batch_reader(): {reader_table.num_rows} rows")
# TODO: support to_duckdb()
diff --git a/bindings/python/fluss/__init__.pyi
b/bindings/python/fluss/__init__.pyi
index 18095c01..b5bfdfab 100644
--- a/bindings/python/fluss/__init__.pyi
+++ b/bindings/python/fluss/__init__.pyi
@@ -898,6 +898,26 @@ class LogScanner:
or timeout expires.
"""
...
+ def to_arrow_batch_reader(self) -> pa.RecordBatchReader:
+ """Create a lazy Arrow RecordBatchReader that reads until latest
offsets.
+
+ Returns a ``pyarrow.RecordBatchReader`` that lazily polls batches one
at
+ a time (streaming). Prefer this when you want to process batches
without
+ holding the full result in memory at once.
+
+ Do not call ``poll_arrow`` / ``poll_record_batch`` on this scanner
while
+ iterating the reader; they share the same underlying scanner state.
+ Overlapping calls are not supported. Use one active
+ polling/consumption path at a time.
+
+ Requires a batch-based scanner (created with
``new_scan().create_record_batch_log_scanner()``).
+ You must call ``subscribe()``, ``subscribe_buckets()``,
``subscribe_partition()``,
+ or ``subscribe_partition_buckets()`` first.
+
+ Returns:
+ ``pyarrow.RecordBatchReader`` yielding ``RecordBatch`` objects.
+ """
+ ...
async def to_pandas(self) -> pd.DataFrame:
"""Convert all data to Pandas DataFrame.
@@ -910,6 +930,10 @@ class LogScanner:
async def to_arrow(self) -> pa.Table:
"""Convert all data to Arrow Table.
+ Batches are collected in Rust then combined into one table (no
per-batch
+ Python iteration). Do not interleave with ``poll_arrow`` /
``poll_record_batch``
+ for the same subscription session; overlapping use is not supported.
+
Requires a batch-based scanner (created with
new_scan().create_record_batch_log_scanner()).
Reads from currently subscribed buckets until reaching their latest
offsets.
diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs
index 9ee84d76..4133bed4 100644
--- a/bindings/python/src/table.rs
+++ b/bindings/python/src/table.rs
@@ -18,10 +18,10 @@
use crate::TOKIO_RUNTIME;
use crate::*;
use arrow::array::RecordBatch as ArrowRecordBatch;
+use arrow::record_batch::RecordBatchReader as _;
use arrow_pyarrow::{FromPyArrow, ToPyArrow};
use arrow_schema::SchemaRef;
use fluss::record::to_arrow_schema;
-use fluss::rpc::message::OffsetSpec;
use indexmap::IndexMap;
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::{PyIndexError, PyRuntimeError, PyTypeError};
@@ -2014,6 +2014,38 @@ fn get_type_name(value: &Bound<PyAny>) -> String {
.unwrap_or_else(|_| "unknown".to_string())
}
+/// Thin Python iterator over [`fcore::client::SyncRecordBatchLogReader`].
+/// Used internally as the backing iterator for
+/// ``pa.RecordBatchReader.from_batches()``; not registered on the module.
+#[pyclass]
+struct PyRecordBatchLogReader {
+ sync_reader: fcore::client::SyncRecordBatchLogReader,
+}
+
+#[pymethods]
+impl PyRecordBatchLogReader {
+ fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+ slf
+ }
+
+ fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyAny>>> {
+ let result = py.detach(|| self.sync_reader.next().transpose());
+
+ match result {
+ Ok(Some(batch)) => {
+ let py_batch = batch
+ .to_pyarrow(py)
+ .map_err(|e| FlussError::new_err(format!("Failed to
convert batch: {e}")))?;
+ Ok(Some(py_batch.unbind()))
+ }
+ Ok(None) => Ok(None),
+ Err(arrow_err) => Err(FlussError::new_err(format!(
+ "Error reading batch: {arrow_err}"
+ ))),
+ }
+ }
+}
+
/// Wraps the two scanner variants so we never have an impossible state
/// (both None or both Some).
enum ScannerKind {
@@ -2066,8 +2098,6 @@ pub struct LogScanner {
projected_schema: SchemaRef,
/// The projected row type to use for record-based scanning
projected_row_type: Arc<fcore::metadata::RowType>,
- /// Cache for partition_id -> partition_name mapping (avoids repeated
list_partition_infos calls)
- partition_name_cache: Arc<std::sync::RwLock<Option<HashMap<i64, String>>>>,
}
#[pymethods]
@@ -2307,11 +2337,75 @@ impl LogScanner {
})
}
+ /// Create a lazy Arrow RecordBatchReader that reads until latest offsets.
+ ///
+ /// This is a **blocking / synchronous** API: construction queries the
+ /// server for latest offsets (via ``block_on``), and each
+ /// ``RecordBatchReader.__next__()`` call blocks the calling thread until
+ /// the next batch is available. It is suitable for Arrow interop
+ /// (feeding into DuckDB, Polars, etc.) but should not be used
+ /// from ``asyncio`` coroutines -- see issue #545 for a planned
+ /// asyncio-native streaming alternative.
+ /// TODO(#545): Add asyncio-native streaming counterpart.
+ ///
+ /// Returns a PyArrow RecordBatchReader that lazily polls batches one at a
+ /// time. This is more memory-efficient than ``to_arrow()`` which loads all
+ /// data into a single table.
+ ///
+ /// **Concurrency:** While this reader is alive, ``subscribe*`` and
+ /// ``unsubscribe*`` calls on the scanner are rejected with an error.
+ /// You should also avoid calling ``poll_arrow`` / ``poll_record_batch``
+ /// on the same scanner — these are not blocked by the guard, but they
+ /// share the underlying fetch buffer with the reader and would
+ /// interleave batches between both consumers. Drop the reader before
+ /// resuming any of these operations.
+ ///
+ /// You must call subscribe(), subscribe_buckets(), subscribe_partition(),
+ /// or subscribe_partition_buckets() first.
+ ///
+ /// Returns:
+ /// ``pyarrow.RecordBatchReader`` yielding ``RecordBatch`` objects
+ fn to_arrow_batch_reader(&self, py: Python) -> PyResult<Py<PyAny>> {
+ let scanner = self.kind.as_batch()?;
+
+ let sync_reader = py
+ .detach(|| {
+ TOKIO_RUNTIME.block_on(async {
+ let reader =
fcore::client::RecordBatchLogReader::new_until_latest(
+ scanner.new_shared_handle(),
+ &self.admin,
+ )
+ .await?;
+ Ok::<_, fcore::error::Error>(
+
reader.to_record_batch_reader(TOKIO_RUNTIME.handle().clone()),
+ )
+ })
+ })
+ .map_err(|e| FlussError::from_core_error(&e))?;
+
+ let py_schema = sync_reader
+ .schema()
+ .to_pyarrow(py)
+ .map_err(|e| FlussError::new_err(format!("Failed to convert
schema: {e}")))?;
+
+ let py_iter = Py::new(py, PyRecordBatchLogReader { sync_reader })?;
+
+ let pyarrow = py.import("pyarrow")?;
+ let batch_reader = pyarrow
+ .getattr("RecordBatchReader")?
+ .call_method1("from_batches", (py_schema, py_iter))?;
+
+ Ok(batch_reader.into())
+ }
+
/// Convert all data to Arrow Table.
///
/// Reads from currently subscribed buckets until reaching their latest
offsets.
/// Works for both partitioned and non-partitioned tables.
///
+ /// Materializes batches in Rust
(``RecordBatchLogReader::collect_all_batches``)
+ /// then builds one PyArrow table, avoiding per-batch Python iteration.
+ ///
/// You must call subscribe(), subscribe_buckets(), subscribe_partition(),
or subscribe_partition_buckets() first.
///
/// Returns:
@@ -2319,29 +2413,29 @@ impl LogScanner {
fn to_arrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let kind = Arc::clone(&self.kind);
let admin = Arc::clone(&self.admin);
- let table_info = self.table_info.clone();
let projected_schema = self.projected_schema.clone();
- let partition_name_cache = Arc::clone(&self.partition_name_cache);
future_into_py(py, async move {
let scanner = kind.as_batch()?;
- let subscribed = scanner.get_subscribed_buckets();
- if subscribed.is_empty() {
- return Err(FlussError::new_err(
- "No buckets subscribed. Call subscribe(),
subscribe_buckets(), subscribe_partition(), or subscribe_partition_buckets()
first.",
- ));
- }
- let all_batches = Self::collect_all_batches(
- scanner,
+ let mut reader =
fcore::client::RecordBatchLogReader::new_until_latest(
+ scanner.new_shared_handle(),
&admin,
- &table_info,
- &subscribed,
- &partition_name_cache,
)
- .await?;
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+
+ let scan_batches = reader
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
- Python::attach(|py| Self::batches_to_arrow_table(py, all_batches,
&projected_schema))
+ let batches: Vec<Arc<ArrowRecordBatch>> = scan_batches
+ .into_iter()
+ .map(|sb| Arc::new(sb.into_batch()))
+ .collect();
+
+ Python::attach(|py| Self::batches_to_arrow_table(py, batches,
&projected_schema))
})
}
@@ -2357,30 +2451,30 @@ impl LogScanner {
fn to_pandas<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let kind = Arc::clone(&self.kind);
let admin = Arc::clone(&self.admin);
- let table_info = self.table_info.clone();
let projected_schema = self.projected_schema.clone();
- let partition_name_cache = Arc::clone(&self.partition_name_cache);
future_into_py(py, async move {
let scanner = kind.as_batch()?;
- let subscribed = scanner.get_subscribed_buckets();
- if subscribed.is_empty() {
- return Err(FlussError::new_err(
- "No buckets subscribed. Call subscribe(),
subscribe_buckets(), subscribe_partition(), or subscribe_partition_buckets()
first.",
- ));
- }
- let all_batches = Self::collect_all_batches(
- scanner,
+ let mut reader =
fcore::client::RecordBatchLogReader::new_until_latest(
+ scanner.new_shared_handle(),
&admin,
- &table_info,
- &subscribed,
- &partition_name_cache,
)
- .await?;
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+
+ let scan_batches = reader
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+
+ let batches: Vec<Arc<ArrowRecordBatch>> = scan_batches
+ .into_iter()
+ .map(|sb| Arc::new(sb.into_batch()))
+ .collect();
Python::attach(|py| {
- let arrow_table = Self::batches_to_arrow_table(py,
all_batches, &projected_schema)?;
+ let arrow_table = Self::batches_to_arrow_table(py, batches,
&projected_schema)?;
arrow_table.call_method0(py, "to_pandas")
})
})
@@ -2442,7 +2536,6 @@ impl LogScanner {
table_info,
projected_schema,
projected_row_type,
- partition_name_cache: Arc::new(std::sync::RwLock::new(None)),
}
}
@@ -2466,138 +2559,6 @@ impl LogScanner {
Utils::combine_batches_to_table(py, batches)
}
}
-
- /// Query stopping offsets and poll until all subscribed buckets are fully
read.
- /// Returns collected Arrow record batches.
- async fn collect_all_batches(
- scanner: &fcore::client::RecordBatchLogScanner,
- admin: &fcore::client::FlussAdmin,
- table_info: &fcore::metadata::TableInfo,
- subscribed: &[(fcore::metadata::TableBucket, i64)],
- partition_name_cache: &std::sync::RwLock<Option<HashMap<i64, String>>>,
- ) -> PyResult<Vec<Arc<ArrowRecordBatch>>> {
- let is_partitioned = scanner.is_partitioned();
- let table_path = &table_info.table_path;
- let table_id = table_info.table_id;
-
- // 1. Query latest offsets
- let mut stopping_offsets: HashMap<fcore::metadata::TableBucket, i64> =
if !is_partitioned {
- let bucket_ids: Vec<i32> = subscribed.iter().map(|(tb, _)|
tb.bucket_id()).collect();
- let offsets = admin
- .list_offsets(table_path, &bucket_ids, OffsetSpec::Latest)
- .await
- .map_err(|e| FlussError::from_core_error(&e))?;
- offsets
- .into_iter()
- .filter(|(_, offset)| *offset > 0)
- .map(|(bucket_id, offset)| {
- (
- fcore::metadata::TableBucket::new(table_id, bucket_id),
- offset,
- )
- })
- .collect()
- } else {
- let cached = partition_name_cache.read().unwrap().clone();
- let partition_id_to_name = match cached {
- Some(map) => map,
- None => {
- let infos = admin
- .list_partition_infos(table_path)
- .await
- .map_err(|e| FlussError::from_core_error(&e))?;
- let map: HashMap<i64, String> = infos
- .into_iter()
- .map(|info| (info.get_partition_id(),
info.get_partition_name()))
- .collect();
- *partition_name_cache.write().unwrap() = Some(map.clone());
- map
- }
- };
-
- let mut by_partition: HashMap<i64, Vec<i32>> = HashMap::new();
- for (tb, _) in subscribed {
- if let Some(partition_id) = tb.partition_id() {
- by_partition
- .entry(partition_id)
- .or_default()
- .push(tb.bucket_id());
- }
- }
-
- let mut result = HashMap::new();
- for (partition_id, bucket_ids) in by_partition {
- let partition_name =
partition_id_to_name.get(&partition_id).ok_or_else(|| {
- FlussError::new_err(format!("Unknown partition_id:
{partition_id}"))
- })?;
- let offsets = admin
- .list_partition_offsets(
- table_path,
- partition_name,
- &bucket_ids,
- OffsetSpec::Latest,
- )
- .await
- .map_err(|e| FlussError::from_core_error(&e))?;
- for (bucket_id, offset) in offsets {
- if offset > 0 {
- let tb =
fcore::metadata::TableBucket::new_with_partition(
- table_id,
- Some(partition_id),
- bucket_id,
- );
- result.insert(tb, offset);
- }
- }
- }
- result
- };
-
- // 2. Poll until all buckets reach their stopping offsets
- let mut all_batches = Vec::new();
- while !stopping_offsets.is_empty() {
- let scan_batches = scanner
- .poll(Duration::from_millis(500))
- .await
- .map_err(|e| FlussError::from_core_error(&e))?;
-
- if scan_batches.is_empty() {
- continue;
- }
-
- for scan_batch in scan_batches {
- let table_bucket = scan_batch.bucket().clone();
- let Some(&stop_at) = stopping_offsets.get(&table_bucket) else {
- continue;
- };
-
- let base_offset = scan_batch.base_offset();
- let last_offset = scan_batch.last_offset();
-
- if base_offset >= stop_at {
- stopping_offsets.remove(&table_bucket);
- continue;
- }
-
- let batch = if last_offset >= stop_at {
- let num_to_keep = (stop_at - base_offset) as usize;
- let b = scan_batch.into_batch();
- let limit = num_to_keep.min(b.num_rows());
- b.slice(0, limit)
- } else {
- scan_batch.into_batch()
- };
-
- all_batches.push(Arc::new(batch));
-
- if last_offset >= stop_at - 1 {
- stopping_offsets.remove(&table_bucket);
- }
- }
- }
-
- Ok(all_batches)
- }
}
#[cfg(test)]
diff --git a/bindings/python/test/test_log_table.py
b/bindings/python/test/test_log_table.py
index 2f560bcf..50b9078b 100644
--- a/bindings/python/test/test_log_table.py
+++ b/bindings/python/test/test_log_table.py
@@ -388,6 +388,130 @@ async def test_to_arrow_and_to_pandas(connection, admin):
await admin.drop_table(table_path, ignore_if_not_exists=False)
+async def test_to_arrow_batch_reader(connection, admin):
+ """Test to_arrow_batch_reader() returns a lazy PyArrow
RecordBatchReader."""
+ table_path = fluss.TablePath("fluss", "py_test_to_arrow_batch_reader")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+
+ schema = fluss.Schema(
+ pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())])
+ )
+ table_descriptor = fluss.TableDescriptor(schema)
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ writer = table.new_append().create_writer()
+
+ pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name",
pa.string())])
+ writer.write_arrow_batch(
+ pa.RecordBatch.from_arrays(
+ [pa.array([10, 20, 30], type=pa.int32()), pa.array(["x", "y",
"z"])],
+ schema=pa_schema,
+ )
+ )
+ await writer.flush()
+
+ num_buckets = (await admin.get_table_info(table_path)).num_buckets
+
+ scanner = await table.new_scan().create_record_batch_log_scanner()
+ scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in
range(num_buckets)})
+
+ # to_arrow_batch_reader() is a blocking/sync API; run in a thread to
+ # avoid starving the asyncio event loop (see docstring warning).
+ def _read_all():
+ reader = scanner.to_arrow_batch_reader()
+ assert isinstance(reader, pa.RecordBatchReader)
+ assert reader.schema == pa_schema
+
+ batches = list(reader)
+ total_rows = sum(b.num_rows for b in batches)
+ assert total_rows == 3
+
+ result_table = pa.Table.from_batches(batches, schema=pa_schema)
+ assert result_table.column("id").to_pylist() == [10, 20, 30]
+ assert result_table.column("name").to_pylist() == ["x", "y", "z"]
+
+ await asyncio.to_thread(_read_all)
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_to_arrow_batch_reader_drop_and_guard(connection, admin):
+ """Test reader-active guard and Drop cleanup on mid-iteration drop."""
+ table_path = fluss.TablePath("fluss", "py_test_batch_reader_drop_guard")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+
+ schema = fluss.Schema(
+ pa.schema([pa.field("id", pa.int32()), pa.field("name", pa.string())])
+ )
+ table_descriptor = fluss.TableDescriptor(schema)
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ writer = table.new_append().create_writer()
+
+ pa_schema = pa.schema([pa.field("id", pa.int32()), pa.field("name",
pa.string())])
+ # Write multiple separate flushes so the server stores multiple log
+ # batches per bucket. This makes it likely that the reader's first poll
+ # only drains a subset, leaving real work for the Drop cleanup loop.
+ num_flushes = 10
+ rows_per_flush = 200
+ total_rows = num_flushes * rows_per_flush
+ for f in range(num_flushes):
+ start = f * rows_per_flush
+ writer.write_arrow_batch(
+ pa.RecordBatch.from_arrays(
+ [
+ pa.array(
+ list(range(start, start + rows_per_flush)),
type=pa.int32()
+ ),
+ pa.array(
+ [f"row_{i}" for i in range(start, start +
rows_per_flush)]
+ ),
+ ],
+ schema=pa_schema,
+ )
+ )
+ await writer.flush()
+
+ num_buckets = (await admin.get_table_info(table_path)).num_buckets
+
+ scanner = await table.new_scan().create_record_batch_log_scanner()
+ scanner.subscribe_buckets({i: fluss.EARLIEST_OFFSET for i in
range(num_buckets)})
+
+ # to_arrow_batch_reader() is a blocking/sync API; run all blocking
+ # interactions in a thread to avoid starving the asyncio event loop.
+ def _test_guard_and_drop():
+ # --- Guard blocks subscribe / unsubscribe while reader is active ---
+ reader = scanner.to_arrow_batch_reader()
+ with pytest.raises(fluss.FlussError, match="RecordBatchLogReader is
active"):
+ scanner.subscribe_buckets({0: fluss.EARLIEST_OFFSET})
+ with pytest.raises(fluss.FlussError, match="RecordBatchLogReader is
active"):
+ scanner.unsubscribe(0)
+
+ # --- Drop mid-iteration: read one batch, then discard ---
+ first_batch = next(reader)
+ assert first_batch.num_rows > 0
+ del reader
+
+ # --- Drop unsubscribed leftover buckets: creating a reader without
+ # re-subscribing must fail with "No buckets subscribed" ---
+ with pytest.raises(fluss.FlussError, match="No buckets subscribed"):
+ scanner.to_arrow_batch_reader()
+
+ # --- Guard cleared after drop: scanner is reusable from a fresh
subscribe ---
+ scanner.subscribe_buckets(
+ {i: fluss.EARLIEST_OFFSET for i in range(num_buckets)}
+ )
+ reader2 = scanner.to_arrow_batch_reader()
+ batches = list(reader2)
+ assert sum(b.num_rows for b in batches) == total_rows
+
+ await asyncio.to_thread(_test_guard_and_drop)
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
async def test_partitioned_table_append_scan(connection, admin,
wait_for_table_ready):
"""Test append and scan on a partitioned log table."""
table_path = fluss.TablePath("fluss", "py_test_partitioned_log_append")
diff --git a/crates/fluss/src/client/table/mod.rs
b/crates/fluss/src/client/table/mod.rs
index ba1edd2f..e116bbb4 100644
--- a/crates/fluss/src/client/table/mod.rs
+++ b/crates/fluss/src/client/table/mod.rs
@@ -29,12 +29,14 @@ mod lookup;
mod log_fetch_buffer;
mod partition_getter;
+mod reader;
mod remote_log;
mod scanner;
mod upsert;
pub use append::{AppendWriter, TableAppend};
pub use lookup::{LookupResult, Lookuper, PrefixKeyLookuper, TableLookup,
TablePrefixLookup};
+pub use reader::{RecordBatchLogReader, SyncRecordBatchLogReader};
pub use remote_log::{
DEFAULT_REMOTE_FILE_DOWNLOAD_THREAD_NUM,
DEFAULT_SCANNER_REMOTE_LOG_PREFETCH_NUM,
};
diff --git a/crates/fluss/src/client/table/reader.rs
b/crates/fluss/src/client/table/reader.rs
new file mode 100644
index 00000000..0a08803d
--- /dev/null
+++ b/crates/fluss/src/client/table/reader.rs
@@ -0,0 +1,701 @@
+// 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 log reader that polls until stopping offsets, then terminates.
+//!
+//! Unlike [`RecordBatchLogScanner`] which is unbounded (continuous streaming),
+//! [`RecordBatchLogReader`] reads log data up to a finite set of stopping
+//! offsets and then signals completion. This enables "snapshot-style" reads
+//! from a streaming log: capture the latest offsets, then consume all data
+//! up to those offsets.
+//!
+//! The reader **takes ownership** of the scanner (move, not clone). Once the
+//! scanner is moved into a reader, the compiler prevents concurrent polls.
+//!
+//! The reader also provides a synchronous
[`arrow::record_batch::RecordBatchReader`]
+//! adapter via [`RecordBatchLogReader::to_record_batch_reader`] for Arrow
+//! ecosystem interop and FFI consumers (Python, C++).
+
+use crate::client::admin::FlussAdmin;
+use crate::client::table::RecordBatchLogScanner;
+use crate::error::{Error, Result};
+use crate::metadata::TableBucket;
+use crate::record::ScanBatch;
+use crate::rpc::message::OffsetSpec;
+use arrow::record_batch::RecordBatch;
+use arrow_schema::SchemaRef;
+use log::warn;
+use std::collections::{HashMap, VecDeque};
+use std::time::Duration;
+
+const DEFAULT_POLL_TIMEOUT: Duration = Duration::from_millis(500);
+
+/// Bounded log reader that consumes log data up to specified stopping offsets.
+///
+/// This type wraps a [`RecordBatchLogScanner`] and adds stopping semantics:
+/// it polls batches from the scanner, filters/slices them against per-bucket
+/// stopping offsets, and signals completion when all buckets are caught up.
+///
+/// The reader takes **ownership** of the scanner. Once moved in, no other code
+/// can poll the same scanner concurrently.
+///
+/// # Construction
+///
+/// Use [`RecordBatchLogReader::new_until_latest`] for the common case of
+/// reading all currently-available data, or
[`RecordBatchLogReader::new_until_offsets`]
+/// for custom stopping offsets.
+///
+/// # Async iteration
+///
+/// Call [`next_batch`](RecordBatchLogReader::next_batch) repeatedly to get
+/// [`ScanBatch`]es lazily, one at a time. Returns `None` when all buckets
+/// have reached their stopping offsets.
+///
+/// # Sync adapter
+///
+/// Call
[`to_record_batch_reader`](RecordBatchLogReader::to_record_batch_reader)
+/// to get a synchronous [`arrow::record_batch::RecordBatchReader`] suitable
+/// for Arrow FFI consumers.
+pub struct RecordBatchLogReader {
+ scanner: RecordBatchLogScanner,
+ stopping_offsets: HashMap<TableBucket, i64>,
+ buffer: VecDeque<ScanBatch>,
+ schema: SchemaRef,
+}
+
+impl RecordBatchLogReader {
+ /// Create a reader that reads until the latest offsets at the time of
creation.
+ ///
+ /// Queries the server for the current latest offset of each subscribed
+ /// bucket, then reads until those offsets are reached. Buckets whose
+ /// subscribed offset already meets or exceeds the latest offset are
+ /// excluded (nothing to read).
+ ///
+ /// Partition metadata is fetched once during construction; no caching
+ /// is needed since each reader is typically short-lived.
+ pub async fn new_until_latest(
+ scanner: RecordBatchLogScanner,
+ admin: &FlussAdmin,
+ ) -> Result<Self> {
+ // Acquire the guard first so no concurrent unsubscribe can mutate
+ // state between reading subscriptions and using them.
+ scanner.try_set_reader_active()?;
+
+ let subscribed = scanner.get_subscribed_buckets();
+ if subscribed.is_empty() {
+ scanner.clear_reader_active();
+ return Err(Error::IllegalArgument {
+ message: "No buckets subscribed. Call subscribe() before
creating a reader."
+ .to_string(),
+ });
+ }
+
+ let stopping_offsets = match query_latest_offsets(admin, &scanner,
&subscribed).await {
+ Ok(o) => o,
+ Err(e) => {
+ scanner.clear_reader_active();
+ return Err(e);
+ }
+ };
+ let schema = scanner.schema();
+
+ Ok(Self {
+ scanner,
+ stopping_offsets,
+ buffer: VecDeque::new(),
+ schema,
+ })
+ }
+
+ /// Create a reader with explicit stopping offsets per bucket.
+ ///
+ /// # NOTE: Every key in `stopping_offsets` **must** correspond to a
bucket that is
+ /// currently subscribed on the `scanner`. If a stopping offset refers to a
+ /// bucket that will never appear in polled batches, the reader will loop
+ /// indefinitely waiting for data that never arrives.
+ ///
+ /// Use [`new_until_latest`](Self::new_until_latest) for the common case;
+ /// it queries the server and builds a validated stopping-offset map
+ /// automatically.
+ pub fn new_until_offsets(
+ scanner: RecordBatchLogScanner,
+ stopping_offsets: HashMap<TableBucket, i64>,
+ ) -> Result<Self> {
+ scanner.try_set_reader_active()?;
+ let schema = scanner.schema();
+ Ok(Self {
+ scanner,
+ stopping_offsets,
+ buffer: VecDeque::new(),
+ schema,
+ })
+ }
+
+ /// Returns the Arrow schema for batches produced by this reader.
+ pub fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ /// Drain all remaining batches until stopping offsets are satisfied.
+ ///
+ /// This is a convenience for callers (e.g. bindings building a single
Arrow
+ /// table) that want to materialize the full result in Rust without
per-batch
+ /// iteration.
+ pub async fn collect_all_batches(&mut self) -> Result<Vec<ScanBatch>> {
+ let mut out = Vec::new();
+ while let Some(b) = self.next_batch().await? {
+ out.push(b);
+ }
+ Ok(out)
+ }
+
+ /// Fetch the next [`ScanBatch`], or `None` if all buckets are caught up.
+ ///
+ /// Each call may internally poll multiple batches from the scanner,
+ /// buffer them, and return one at a time. Batches that cross a stopping
+ /// offset boundary are sliced to exclude records at or beyond the stop
point.
+ ///
+ /// Completed buckets are unsubscribed from the scanner to avoid wasting
+ /// network traffic on data the reader will discard.
+ pub async fn next_batch(&mut self) -> Result<Option<ScanBatch>> {
+ loop {
+ if let Some(batch) = self.buffer.pop_front() {
+ return Ok(Some(batch));
+ }
+
+ if self.stopping_offsets.is_empty() {
+ return Ok(None);
+ }
+
+ let scan_batches = self.scanner.poll(DEFAULT_POLL_TIMEOUT).await?;
+
+ if scan_batches.is_empty() {
+ continue;
+ }
+
+ let completed =
+ filter_batches(scan_batches, &mut self.stopping_offsets, &mut
self.buffer);
+
+ // Use the `_sync` unsubscribe variants here: the active-reader
+ // guard rejects calls to the async `unsubscribe*` methods, but
+ // the reader is allowed to clean up its own completed buckets.
+ // The sync variants do the same map removal without the guard
+ // check, and the partitioned/non-partitioned mismatch they
+ // silently ignore is unreachable since the reader inherits the
+ // scanner's partition mode.
+ for tb in completed {
+ if let Some(partition_id) = tb.partition_id() {
+ self.scanner
+ .unsubscribe_partition_sync(partition_id,
tb.bucket_id());
+ } else {
+ self.scanner.unsubscribe_sync(tb.bucket_id());
+ }
+ }
+ }
+ }
+
+ /// Convert this async reader into a synchronous
[`arrow::record_batch::RecordBatchReader`].
+ ///
+ /// The returned adapter calls [`tokio::runtime::Handle::block_on`] on each
+ /// iterator step. **Do not** call this from inside a Tokio worker thread
+ /// while the same runtime is driving async work (nested `block_on` can
+ /// panic or deadlock). Prefer
[`next_batch`](RecordBatchLogReader::next_batch)
+ /// in async Rust code. This is intended for sync/FFI boundaries (C++, some
+ /// Python call paths).
+ pub fn to_record_batch_reader(
+ self,
+ handle: tokio::runtime::Handle,
+ ) -> SyncRecordBatchLogReader {
+ SyncRecordBatchLogReader {
+ reader: self,
+ handle,
+ }
+ }
+}
+
+/// Best-effort cleanup when the reader is dropped before all buckets reach
+/// their stopping offsets (early `break`, an exception in the consumer, etc.).
+///
+/// Why this matters even though we own the scanner:
+///
+/// In pure Rust, dropping the reader drops the owned `RecordBatchLogScanner`,
+/// which decrements the `Arc<LogScannerInner>` to zero and frees the inner
+/// state. Subscriptions die with it, so this `Drop` is a no-op in that path.
+///
+/// In the binding layer (Python today, C++/Elixir later), the binding holds
+/// its own `Arc<LogScannerInner>` and uses
+/// [`RecordBatchLogScanner::new_shared_handle`] to obtain a second handle for
+/// the reader. When the reader is dropped mid-iteration the inner state stays
+/// alive — and any buckets the reader hadn't yet completed remain in
+/// `LogScannerStatus.bucket_status_map`. The user's next operations on the
+/// original `LogScanner` would then see "ghost" subscriptions (extra buckets
+/// being polled, stale offsets, etc.).
+///
+/// The `next_batch` loop already calls `unsubscribe` on each completed bucket,
+/// so `stopping_offsets` accurately reflects the still-active set when `Drop`
+/// runs. We unsubscribe each remaining bucket synchronously via the
+/// `_sync` escape hatches (the underlying `LogScannerStatus` ops don't await),
+/// so this is safe to call from any context — sync, async, a Tokio worker, or
+/// a Python thread holding the GIL.
+///
+/// After cleanup, the `reader_active` guard is cleared so that the original
+/// scanner (held by the binding layer) can accept new subscriptions again.
+///
+/// Caveats:
+/// - Batches already buffered in `LogFetcher.log_fetch_buffer` for an
+/// unsubscribed bucket are not drained here. They'll either be filtered out
+/// by the next `RecordBatchLogReader` (via the "bucket not in
+/// stopping_offsets" branch) or surface to a direct `poll_arrow` caller, who
+/// was sharing scanner state in the first place.
+/// - `Drop` cannot return errors. The `_sync` variants no-op on
+/// partitioned/non-partitioned mismatch, but that mismatch is unreachable
+/// here because the reader was constructed from this scanner and inherited
+/// its partition mode.
+impl Drop for RecordBatchLogReader {
+ fn drop(&mut self) {
+ for (tb, _) in self.stopping_offsets.drain() {
+ if let Some(partition_id) = tb.partition_id() {
+ self.scanner
+ .unsubscribe_partition_sync(partition_id, tb.bucket_id());
+ } else {
+ self.scanner.unsubscribe_sync(tb.bucket_id());
+ }
+ }
+ self.scanner.clear_reader_active();
+ }
+}
+
+/// Synchronous adapter that implements
[`arrow::record_batch::RecordBatchReader`].
+///
+/// Created via [`RecordBatchLogReader::to_record_batch_reader`].
+/// Blocks the current thread on each `next()` call using the provided
+/// Tokio runtime handle.
+///
+/// The iterator yields plain [`RecordBatch`]es (bucket/offset metadata from
+/// [`ScanBatch`] is stripped to satisfy the Arrow trait contract).
+pub struct SyncRecordBatchLogReader {
+ reader: RecordBatchLogReader,
+ handle: tokio::runtime::Handle,
+}
+
+impl Iterator for SyncRecordBatchLogReader {
+ type Item = std::result::Result<RecordBatch, arrow::error::ArrowError>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ match self.handle.block_on(self.reader.next_batch()) {
+ Ok(Some(scan_batch)) => Some(Ok(scan_batch.into_batch())),
+ Ok(None) => None,
+ Err(e) =>
Some(Err(arrow::error::ArrowError::ExternalError(Box::new(e)))),
+ }
+ }
+}
+
+impl arrow::record_batch::RecordBatchReader for SyncRecordBatchLogReader {
+ fn schema(&self) -> SchemaRef {
+ self.reader.schema()
+ }
+}
+
+/// Query latest offsets for all subscribed buckets, handling both partitioned
+/// and non-partitioned tables.
+///
+/// Buckets whose subscribed offset already meets or exceeds the latest offset
+/// are excluded from the result (there is nothing to read). A `latest_offset`
+/// of `0` means the bucket is empty and is silently skipped; a negative value
+/// is unexpected from the server and is logged as a warning before being
+/// skipped.
+async fn query_latest_offsets(
+ admin: &FlussAdmin,
+ scanner: &RecordBatchLogScanner,
+ subscribed: &[(TableBucket, i64)],
+) -> Result<HashMap<TableBucket, i64>> {
+ let table_path = scanner.table_path();
+
+ if !scanner.is_partitioned() {
+ let bucket_ids: Vec<i32> = subscribed.iter().map(|(tb, _)|
tb.bucket_id()).collect();
+
+ let offsets = admin
+ .list_offsets(table_path, &bucket_ids, OffsetSpec::Latest)
+ .await?;
+
+ let subscribed_offset_by_bucket: HashMap<i32, i64> = subscribed
+ .iter()
+ .map(|(tb, off)| (tb.bucket_id(), *off))
+ .collect();
+
+ let table_id = scanner.table_id();
+ Ok(offsets
+ .into_iter()
+ .filter(|(bucket_id, latest_offset)| {
+ if *latest_offset < 0 {
+ warn!(
+ "Server returned negative latest offset
{latest_offset} for bucket {bucket_id} of table {table_id}; skipping bucket."
+ );
+ return false;
+ }
+ if *latest_offset == 0 {
+ return false;
+ }
+ let Some(&subscribed_offset) =
subscribed_offset_by_bucket.get(bucket_id)
+ else {
+ return false;
+ };
+ subscribed_offset < *latest_offset
+ })
+ .map(|(bucket_id, offset)| (TableBucket::new(table_id, bucket_id),
offset))
+ .collect())
+ } else {
+ query_partitioned_offsets(admin, scanner, subscribed).await
+ }
+}
+
+/// Query offsets for partitioned table subscriptions.
+///
+/// Partition metadata is fetched once per reader construction (not cached),
+/// since each [`RecordBatchLogReader`] is typically short-lived and consumed.
+async fn query_partitioned_offsets(
+ admin: &FlussAdmin,
+ scanner: &RecordBatchLogScanner,
+ subscribed: &[(TableBucket, i64)],
+) -> Result<HashMap<TableBucket, i64>> {
+ let table_path = scanner.table_path();
+ let table_id = scanner.table_id();
+
+ let partition_infos = admin.list_partition_infos(table_path).await?;
+ let partition_id_to_name: HashMap<i64, String> = partition_infos
+ .into_iter()
+ .map(|info| (info.get_partition_id(), info.get_partition_name()))
+ .collect();
+
+ let subscribed_offset_map: HashMap<TableBucket, i64> =
subscribed.iter().cloned().collect();
+
+ let mut by_partition: HashMap<i64, Vec<i32>> = HashMap::new();
+ for (tb, _) in subscribed {
+ if let Some(partition_id) = tb.partition_id() {
+ by_partition
+ .entry(partition_id)
+ .or_default()
+ .push(tb.bucket_id());
+ }
+ }
+
+ let mut result: HashMap<TableBucket, i64> = HashMap::new();
+
+ for (partition_id, bucket_ids) in by_partition {
+ let partition_name =
+ partition_id_to_name
+ .get(&partition_id)
+ .ok_or_else(|| Error::UnexpectedError {
+ message: format!("Unknown partition_id: {partition_id}"),
+ source: None,
+ })?;
+
+ let offsets = admin
+ .list_partition_offsets(table_path, partition_name, &bucket_ids,
OffsetSpec::Latest)
+ .await?;
+
+ for (bucket_id, latest_offset) in offsets {
+ if latest_offset < 0 {
+ warn!(
+ "Server returned negative latest offset {latest_offset}
for bucket {bucket_id} of partition {partition_id} (table {table_id}); skipping
bucket."
+ );
+ continue;
+ }
+ if latest_offset == 0 {
+ continue;
+ }
+ let tb = TableBucket::new_with_partition(table_id,
Some(partition_id), bucket_id);
+ let Some(&subscribed_offset) = subscribed_offset_map.get(&tb) else
{
+ continue;
+ };
+ if subscribed_offset < latest_offset {
+ result.insert(tb, latest_offset);
+ }
+ }
+ }
+
+ Ok(result)
+}
+
+/// Filter and slice scan batches against per-bucket stopping offsets.
+///
+/// For each batch:
+/// - If the batch's bucket is not in `stopping_offsets`, skip it.
+/// - If `base_offset >= stop_at`, the bucket is exhausted; remove from map.
+/// - If `last_offset >= stop_at`, slice to keep only records before stop_at.
+/// - Otherwise, keep the full batch.
+///
+/// Accepted batches with at least one row are pushed to `buffer`; empty
+/// batches (e.g. a server-emitted batch containing no rows, or a slice that
+/// reduces to zero rows) are dropped so consumers never observe an empty
+/// `ScanBatch`. Returns the list of buckets that completed (were removed
+/// from `stopping_offsets`).
+fn filter_batches(
+ scan_batches: Vec<ScanBatch>,
+ stopping_offsets: &mut HashMap<TableBucket, i64>,
+ buffer: &mut VecDeque<ScanBatch>,
+) -> Vec<TableBucket> {
+ let mut completed = Vec::new();
+
+ for scan_batch in scan_batches {
+ let bucket = scan_batch.bucket().clone();
+ let Some(&stop_at) = stopping_offsets.get(&bucket) else {
+ continue;
+ };
+
+ let base_offset = scan_batch.base_offset();
+ let last_offset = scan_batch.last_offset();
+
+ if base_offset >= stop_at {
+ stopping_offsets.remove(&bucket);
+ completed.push(bucket);
+ continue;
+ }
+
+ let kept_batch = if last_offset >= stop_at {
+ let num_to_keep = (stop_at - base_offset) as usize;
+ let b = scan_batch.into_batch();
+ let limit = num_to_keep.min(b.num_rows());
+ ScanBatch::new(bucket.clone(), b.slice(0, limit), base_offset)
+ } else {
+ scan_batch
+ };
+
+ if kept_batch.batch().num_rows() > 0 {
+ buffer.push_back(kept_batch);
+ }
+
+ if last_offset >= stop_at - 1 {
+ stopping_offsets.remove(&bucket);
+ completed.push(bucket);
+ }
+ }
+
+ completed
+}
+
+// TODO: Add Rust-level end-to-end tests with `FlussTestingCluster` (feature
+// `integration_tests`) covering `new_until_latest`, partitioned tables,
+// and `new_until_offsets` stopping semantics. Drop cleanup and the
+// reader-active guard are covered by the Python integration test
+// `test_to_arrow_batch_reader_drop_and_guard`.
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::Int32Array;
+ use arrow_schema::{DataType, Field, Schema};
+ use std::sync::Arc;
+
+ fn test_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]))
+ }
+
+ fn make_batch(values: &[i32]) -> RecordBatch {
+ RecordBatch::try_new(
+ test_schema(),
+ vec![Arc::new(Int32Array::from(values.to_vec()))],
+ )
+ .unwrap()
+ }
+
+ fn make_scan_batch(bucket: TableBucket, base_offset: i64, values: &[i32])
-> ScanBatch {
+ ScanBatch::new(bucket, make_batch(values), base_offset)
+ }
+
+ fn bucket(id: i32) -> TableBucket {
+ TableBucket::new(1, id)
+ }
+
+ #[test]
+ fn filter_batch_entirely_before_stop() {
+ let mut offsets = HashMap::from([(bucket(0), 100)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(0), 10, &[1, 2, 3])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 1);
+ assert_eq!(buffer[0].batch().num_rows(), 3);
+ assert!(offsets.contains_key(&bucket(0)));
+ assert!(completed.is_empty());
+ }
+
+ #[test]
+ fn filter_batch_crossing_stop_offset_is_sliced() {
+ let mut offsets = HashMap::from([(bucket(0), 12)]);
+ let mut buffer = VecDeque::new();
+
+ // base_offset=10, 5 rows -> offsets 10,11,12,13,14; stop_at=12 ->
keep 2
+ let batches = vec![make_scan_batch(bucket(0), 10, &[1, 2, 3, 4, 5])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 1);
+ assert_eq!(buffer[0].batch().num_rows(), 2);
+ assert!(!offsets.contains_key(&bucket(0)));
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_batch_at_or_after_stop_offset_is_skipped() {
+ let mut offsets = HashMap::from([(bucket(0), 10)]);
+ let mut buffer = VecDeque::new();
+
+ // base_offset=10, stop_at=10 -> base >= stop, skip entirely
+ let batches = vec![make_scan_batch(bucket(0), 10, &[1, 2, 3])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert!(buffer.is_empty());
+ assert!(!offsets.contains_key(&bucket(0)));
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_batch_ending_exactly_at_stop_minus_one() {
+ let mut offsets = HashMap::from([(bucket(0), 13)]);
+ let mut buffer = VecDeque::new();
+
+ // base_offset=10, 3 rows -> offsets 10,11,12; last_offset=12,
stop_at=13
+ // last_offset (12) >= stop_at - 1 (12) => bucket done
+ let batches = vec![make_scan_batch(bucket(0), 10, &[1, 2, 3])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 1);
+ assert_eq!(buffer[0].batch().num_rows(), 3);
+ assert!(!offsets.contains_key(&bucket(0)));
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_unknown_bucket_is_ignored() {
+ let mut offsets = HashMap::from([(bucket(0), 100)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(99), 0, &[1, 2])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert!(buffer.is_empty());
+ assert!(offsets.contains_key(&bucket(0)));
+ assert!(completed.is_empty());
+ }
+
+ #[test]
+ fn filter_multiple_buckets_independent_tracking() {
+ let mut offsets = HashMap::from([(bucket(0), 12), (bucket(1), 5)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![
+ make_scan_batch(bucket(0), 10, &[1, 2, 3]), // last=12, stop=12 ->
keep 2, done
+ make_scan_batch(bucket(1), 0, &[10, 20, 30]), // last=2, stop=5 ->
keep all, not done
+ ];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 2);
+ assert_eq!(buffer[0].batch().num_rows(), 2); // bucket 0: sliced
+ assert_eq!(buffer[1].batch().num_rows(), 3); // bucket 1: full
+ assert!(!offsets.contains_key(&bucket(0))); // bucket 0: done
+ assert!(offsets.contains_key(&bucket(1))); // bucket 1: still tracking
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_empty_batch_at_stop() {
+ let mut offsets = HashMap::from([(bucket(0), 5)]);
+ let mut buffer = VecDeque::new();
+
+ // empty batch: base_offset=5, 0 rows -> last_offset = base-1 = 4
+ // base_offset (5) >= stop_at (5) -> skip, remove
+ let batches = vec![make_scan_batch(bucket(0), 5, &[])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert!(buffer.is_empty());
+ assert!(!offsets.contains_key(&bucket(0)));
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_drops_empty_batch_before_stop() {
+ // Empty batch well below the stop offset: base=5, 0 rows -> last=4,
stop=100.
+ // base_offset (5) < stop_at (100) and last_offset (4) < stop_at (100),
+ // so it falls into the "keep full batch" branch but must not surface
to
+ // the consumer because it has zero rows.
+ let mut offsets = HashMap::from([(bucket(0), 100)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(0), 5, &[])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert!(buffer.is_empty());
+ assert!(offsets.contains_key(&bucket(0)));
+ assert!(completed.is_empty());
+ }
+
+ #[test]
+ fn filter_single_row_batch_before_stop() {
+ let mut offsets = HashMap::from([(bucket(0), 10)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(0), 5, &[42])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 1);
+ assert_eq!(buffer[0].batch().num_rows(), 1);
+ assert!(offsets.contains_key(&bucket(0)));
+ assert!(completed.is_empty());
+ }
+
+ #[test]
+ fn filter_single_row_batch_at_stop_boundary() {
+ let mut offsets = HashMap::from([(bucket(0), 5)]);
+ let mut buffer = VecDeque::new();
+
+ // base_offset=4, 1 row -> last_offset=4, stop=5
+ // last < stop -> keep all; last (4) >= stop-1 (4) -> done
+ let batches = vec![make_scan_batch(bucket(0), 4, &[42])];
+ let completed = filter_batches(batches, &mut offsets, &mut buffer);
+
+ assert_eq!(buffer.len(), 1);
+ assert_eq!(buffer[0].batch().num_rows(), 1);
+ assert!(!offsets.contains_key(&bucket(0)));
+ assert_eq!(completed, vec![bucket(0)]);
+ }
+
+ #[test]
+ fn filter_preserves_scan_batch_metadata() {
+ let mut offsets = HashMap::from([(bucket(3), 100)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(3), 42, &[1, 2])];
+ filter_batches(batches, &mut offsets, &mut buffer);
+
+ let sb = &buffer[0];
+ assert_eq!(*sb.bucket(), bucket(3));
+ assert_eq!(sb.base_offset(), 42);
+ }
+
+ #[test]
+ fn filter_sliced_batch_preserves_base_offset() {
+ let mut offsets = HashMap::from([(bucket(0), 12)]);
+ let mut buffer = VecDeque::new();
+
+ let batches = vec![make_scan_batch(bucket(0), 10, &[1, 2, 3, 4, 5])];
+ filter_batches(batches, &mut offsets, &mut buffer);
+
+ let sb = &buffer[0];
+ assert_eq!(sb.base_offset(), 10);
+ assert_eq!(*sb.bucket(), bucket(0));
+ }
+}
diff --git a/crates/fluss/src/client/table/scanner.rs
b/crates/fluss/src/client/table/scanner.rs
index c6228e59..86870991 100644
--- a/crates/fluss/src/client/table/scanner.rs
+++ b/crates/fluss/src/client/table/scanner.rs
@@ -257,6 +257,10 @@ pub struct LogScanner {
///
/// More efficient than [`LogScanner`] for batch-level analytics where
per-record
/// metadata (offsets, timestamps) is not needed.
+///
+/// This type is intentionally **not** `Clone`. To perform a bounded read, move
+/// the scanner into a [`crate::client::RecordBatchLogReader`] — the compiler
+/// then prevents concurrent polls by construction.
pub struct RecordBatchLogScanner {
inner: Arc<LogScannerInner>,
}
@@ -269,6 +273,10 @@ struct LogScannerInner {
log_scanner_status: Arc<LogScannerStatus>,
log_fetcher: LogFetcher,
is_partitioned_table: bool,
+ arrow_schema: SchemaRef,
+ /// Guards against subscription changes while a
+ /// [`crate::client::RecordBatchLogReader`] is iterating.
+ reader_active: std::sync::atomic::AtomicBool,
}
impl LogScannerInner {
@@ -280,6 +288,20 @@ impl LogScannerInner {
projected_fields: Option<Vec<usize>>,
) -> Result<Self> {
let log_scanner_status = Arc::new(LogScannerStatus::new());
+
+ let full_row_type = table_info.get_row_type();
+ let arrow_schema = match &projected_fields {
+ Some(indices) => {
+ let projected_fields_vec: Vec<_> = indices
+ .iter()
+ .map(|&i| full_row_type.fields()[i].clone())
+ .collect();
+ let projected_row_type =
crate::metadata::RowType::new(projected_fields_vec);
+ to_arrow_schema(&projected_row_type)?
+ }
+ None => to_arrow_schema(full_row_type)?,
+ };
+
Ok(Self {
table_path: table_info.table_path.clone(),
table_id: table_info.table_id,
@@ -288,15 +310,31 @@ impl LogScannerInner {
log_scanner_status: log_scanner_status.clone(),
log_fetcher: LogFetcher::new(
table_info.clone(),
- connections.clone(),
- metadata.clone(),
+ connections,
+ metadata,
log_scanner_status.clone(),
config,
projected_fields,
)?,
+ arrow_schema,
+ reader_active: std::sync::atomic::AtomicBool::new(false),
})
}
+ fn check_no_active_reader(&self) -> Result<()> {
+ if self
+ .reader_active
+ .load(std::sync::atomic::Ordering::Acquire)
+ {
+ return Err(Error::IllegalArgument {
+ message: "Cannot modify subscriptions while a
RecordBatchLogReader is active. \
+ Drop the reader first."
+ .to_string(),
+ });
+ }
+ Ok(())
+ }
+
async fn poll_records(&self, timeout: Duration) -> Result<ScanRecords> {
let start = Instant::now();
let deadline = start + timeout;
@@ -337,6 +375,7 @@ impl LogScannerInner {
}
async fn subscribe(&self, bucket: i32, offset: i64) -> Result<()> {
+ self.check_no_active_reader()?;
if self.is_partitioned_table {
return Err(Error::UnsupportedOperation {
message: "The table is a partitioned table, please use
\"subscribe_partition\" to \
@@ -354,6 +393,7 @@ impl LogScannerInner {
}
async fn subscribe_buckets(&self, bucket_offsets: &HashMap<i32, i64>) ->
Result<()> {
+ self.check_no_active_reader()?;
if self.is_partitioned_table {
return Err(Error::UnsupportedOperation {
message:
@@ -376,6 +416,7 @@ impl LogScannerInner {
bucket: i32,
offset: i64,
) -> Result<()> {
+ self.check_no_active_reader()?;
if !self.is_partitioned_table {
return Err(Error::UnsupportedOperation {
message: "The table is not a partitioned table, please use
\"subscribe\" to \
@@ -397,6 +438,7 @@ impl LogScannerInner {
&self,
partition_bucket_offsets: &HashMap<(PartitionId, i32), i64>,
) -> Result<()> {
+ self.check_no_active_reader()?;
if !self.is_partitioned_table {
return Err(UnsupportedOperation {
message: "The table is not a partitioned table, please use
\"subscribe_buckets\" \
@@ -431,6 +473,7 @@ impl LogScannerInner {
}
async fn unsubscribe(&self, bucket: i32) -> Result<()> {
+ self.check_no_active_reader()?;
if self.is_partitioned_table {
return Err(Error::UnsupportedOperation {
message:
@@ -446,6 +489,7 @@ impl LogScannerInner {
}
async fn unsubscribe_partition(&self, partition_id: PartitionId, bucket:
i32) -> Result<()> {
+ self.check_no_active_reader()?;
if !self.is_partitioned_table {
return Err(Error::UnsupportedOperation {
message: "Can't unsubscribe a partition for a non-partitioned
table.".to_string(),
@@ -615,6 +659,95 @@ impl RecordBatchLogScanner {
) -> Result<()> {
self.inner.unsubscribe_partition(partition_id, bucket).await
}
+
+ /// Returns the Arrow schema for batches produced by this scanner.
+ pub fn schema(&self) -> SchemaRef {
+ self.inner.arrow_schema.clone()
+ }
+
+ pub fn table_path(&self) -> &TablePath {
+ &self.inner.table_path
+ }
+
+ pub fn table_id(&self) -> TableId {
+ self.inner.table_id
+ }
+
+ /// Creates a new handle to the same underlying scanner state.
+ ///
+ /// Binding layers (Python, C++) that hold the scanner behind shared
+ /// ownership (`Arc`) cannot move it into a
[`crate::client::RecordBatchLogReader`].
+ /// This method produces a second handle so the reader can take ownership
+ /// while the binding retains its reference for subscription management.
+ ///
+ /// **Not intended for general use** — prefer moving the scanner directly.
+ #[doc(hidden)]
+ pub fn new_shared_handle(&self) -> Self {
+ RecordBatchLogScanner {
+ inner: Arc::clone(&self.inner),
+ }
+ }
+
+ /// Atomically marks the scanner as having an active reader.
+ ///
+ /// Returns `Err(IllegalArgument)` if another reader is already active on
+ /// this scanner — only one [`crate::client::RecordBatchLogReader`] may
+ /// iterate per scanner at a time. This mirrors Java's
+ /// `LogScannerImpl.acquire()` single-consumer guard.
+ pub(crate) fn try_set_reader_active(&self) -> Result<()> {
+ use std::sync::atomic::Ordering;
+ self.inner
+ .reader_active
+ .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+ .map(|_| ())
+ .map_err(|_| Error::IllegalArgument {
+ message: "Another RecordBatchLogReader is already active on
this scanner. \
+ Drop the existing reader first."
+ .to_string(),
+ })
+ }
+
+ /// Clears the active-reader guard, re-enabling subscription changes.
+ pub(crate) fn clear_reader_active(&self) {
+ self.inner
+ .reader_active
+ .store(false, std::sync::atomic::Ordering::Release);
+ }
+
+ /// Synchronous, infallible counterpart to
[`unsubscribe`](Self::unsubscribe).
+ ///
+ /// Exists so [`crate::client::RecordBatchLogReader`]'s `Drop` impl can
+ /// release lingering subscriptions without `.await`. The async version is
+ /// also synchronous under the hood (it only acquires a lock and removes
+ /// from a map — no IO), so this exposes the same work without the
+ /// async wrapper. Silently no-ops on partitioned/non-partitioned mismatch
+ /// because `Drop` cannot return errors; callers must pick the correct
+ /// variant.
+ ///
+ /// **Not intended for general use** — prefer the async [`unsubscribe`].
+ pub(crate) fn unsubscribe_sync(&self, bucket: i32) {
+ if self.inner.is_partitioned_table {
+ return;
+ }
+ let table_bucket = TableBucket::new(self.inner.table_id, bucket);
+ self.inner
+ .log_scanner_status
+ .unassign_scan_buckets(from_ref(&table_bucket));
+ }
+
+ /// Synchronous, infallible counterpart to
+ /// [`unsubscribe_partition`](Self::unsubscribe_partition). See
+ /// [`unsubscribe_sync`](Self::unsubscribe_sync) for rationale.
+ pub(crate) fn unsubscribe_partition_sync(&self, partition_id: PartitionId,
bucket: i32) {
+ if !self.inner.is_partitioned_table {
+ return;
+ }
+ let table_bucket =
+ TableBucket::new_with_partition(self.inner.table_id,
Some(partition_id), bucket);
+ self.inner
+ .log_scanner_status
+ .unassign_scan_buckets(from_ref(&table_bucket));
+ }
}
struct LogFetcher {
@@ -2009,6 +2142,7 @@ mod tests {
let result = validate_scan_support(&table_path, &table_info);
assert!(result.is_ok());
}
+
#[tokio::test]
async fn prepare_fetch_log_requests_uses_configured_fetch_params() ->
Result<()> {
let table_path = TablePath::new("db".to_string(), "tbl".to_string());
diff --git a/website/docs/user-guide/python/api-reference.md
b/website/docs/user-guide/python/api-reference.md
index dc252b68..9bf0b690 100644
--- a/website/docs/user-guide/python/api-reference.md
+++ b/website/docs/user-guide/python/api-reference.md
@@ -181,9 +181,12 @@ Builder for creating a `PrefixLookuper`. Obtain via
`TableLookup.lookup_by(colum
| `await .poll(timeout_ms) -> ScanRecords` | Poll
individual records (record scanner only) |
| `await .poll_arrow(timeout_ms) -> pa.Table` | Poll as
Arrow Table (batch scanner only) |
| `await .poll_record_batch(timeout_ms) -> list[RecordBatch]` | Poll batches
with metadata (batch scanner only) |
+| `.to_arrow_batch_reader() -> pa.RecordBatchReader` | Lazy Arrow
RecordBatchReader reading until latest offsets (batch scanner only) |
| `await .to_arrow() -> pa.Table` | Read all
subscribed data as Arrow Table (batch scanner only) |
| `await .to_pandas() -> pd.DataFrame` | Read all
subscribed data as DataFrame (batch scanner only) |
+> **Note:** Overlapping `poll_*` / `to_arrow*` / `to_arrow_batch_reader` calls
on the same underlying scanner are not supported. Use only one active
polling/consumption path at a time.
+
## `ScanRecords`
Returned by `LogScanner.poll()`. Records are grouped by bucket.
diff --git a/website/docs/user-guide/rust/api-reference.md
b/website/docs/user-guide/rust/api-reference.md
index 03054f0f..5d3068b5 100644
--- a/website/docs/user-guide/rust/api-reference.md
+++ b/website/docs/user-guide/rust/api-reference.md
@@ -151,6 +151,8 @@ Complete API reference for the Fluss Rust client.
## `RecordBatchLogScanner`
+Overlapping `poll` calls on clones that share state, or `poll` concurrent with
`RecordBatchLogReader::next_batch`, are not supported. Use one active
polling/consumption call at a time per underlying scanner state.
+
| Method
| Description
|
|-----------------------------------------------------------------------------------------------------------|----------------------------------------------------------|
| `async fn subscribe(&self, bucket_id: i32, start_offset: i64) -> Result<()>`
| Subscribe to a bucket
|
@@ -162,6 +164,48 @@ Complete API reference for the Fluss Rust client.
| `async fn poll(&self, timeout: Duration) -> Result<Vec<ScanBatch>>`
| Poll for Arrow record batches
|
| `fn is_partitioned(&self) -> bool`
| Check if the table is partitioned
|
| `fn get_subscribed_buckets(&self) -> Vec<(TableBucket, i64)>`
| Get all current subscriptions as (bucket,
offset) pairs |
+| `fn schema(&self) -> SchemaRef`
| Get the Arrow schema for batches produced by
this scanner|
+| `fn table_path(&self) -> &TablePath`
| Get the table path
|
+| `fn table_id(&self) -> TableId`
| Get the table ID
|
+
+## `RecordBatchLogReader`
+
+Bounded log reader that consumes data up to specified stopping offsets, then
terminates.
+Unlike `RecordBatchLogScanner` which polls indefinitely, this reader stops
automatically.
+
+| Method
| Description
|
+|-------------------------------------------------------------------------------------------------------------|----------------------------------------------------------|
+| `async fn new_until_latest(scanner: RecordBatchLogScanner, admin:
&FlussAdmin) -> Result<Self>` | Read until the latest offsets at
time of creation |
+| `fn new_until_offsets(scanner: RecordBatchLogScanner, stopping_offsets:
HashMap<TableBucket, i64>) -> Result<Self>` | Read until custom stopping
offsets per bucket |
+| `async fn next_batch(&mut self) -> Result<Option<ScanBatch>>`
| Get the next batch with bucket/offset
metadata, or `None` when all buckets caught up |
+| `async fn collect_all_batches(&mut self) -> Result<Vec<ScanBatch>>`
| Drain all batches (with metadata) until
stopping offsets are satisfied |
+| `fn schema(&self) -> SchemaRef`
| Arrow schema for produced batches
|
+| `fn to_record_batch_reader(self, handle: tokio::runtime::Handle) ->
SyncRecordBatchLogReader` | Sync adapter implementing
`arrow::RecordBatchReader` (see below) |
+
+## `SyncRecordBatchLogReader`
+
+Synchronous adapter for `RecordBatchLogReader`. Created via
+`RecordBatchLogReader::to_record_batch_reader(handle)`.
+
+Implements both [`Iterator`] and [`arrow::record_batch::RecordBatchReader`],
so it
+plugs into the wider Arrow ecosystem — FFI, PyArrow's
+`pa.RecordBatchReader.from_batches`, the C++ Arrow `RecordBatchReader`
interface,
+DataFusion sources, etc.
+
+Each `next()` call drives the underlying async reader via
+`tokio::runtime::Handle::block_on`. **Do not call from inside a Tokio worker
+thread that belongs to the same runtime** — nested `block_on` panics. Prefer
+`RecordBatchLogReader::next_batch` in async Rust code; use this adapter only at
+sync/FFI boundaries.
+
+Bucket and offset metadata carried by `ScanBatch` is **dropped** here, because
+the Arrow trait contract yields plain `RecordBatch`. If you need offsets or
+bucket identity per batch, use `next_batch` instead.
+
+| Method |
Description |
+|-----------------------------------------------------------------|--------------------------------------------------|
+| `fn next(&mut self) -> Option<Result<RecordBatch, ArrowError>>` | Iterator:
next batch, or `None` when caught up |
+| `fn schema(&self) -> SchemaRef` | Arrow
schema for produced batches |
## `ScanRecord`
