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 d071c906 [python] Limit scan (#613)
d071c906 is described below
commit d071c9066a22fe25009259fa4dbaaca9ee7324e3
Author: Anton Borisov <[email protected]>
AuthorDate: Wed Jun 10 23:31:36 2026 +0100
[python] Limit scan (#613)
---
bindings/python/example/log_table.py | 42 ++++
bindings/python/example/pk_table.py | 15 ++
bindings/python/fluss/__init__.pyi | 82 +++++++
bindings/python/src/lib.rs | 1 +
bindings/python/src/table.rs | 252 +++++++++++++++++----
bindings/python/test/test_batch_scanner.py | 221 ++++++++++++++++++
crates/fluss/src/client/table/batch_scanner.rs | 151 +++++++-----
crates/fluss/src/test_utils.rs | 22 +-
crates/fluss/tests/integration/batch_scanner.rs | 87 ++++++-
website/docs/user-guide/python/api-reference.md | 14 ++
.../docs/user-guide/python/example/log-tables.md | 14 ++
.../python/example/primary-key-tables.md | 12 +
12 files changed, 803 insertions(+), 110 deletions(-)
diff --git a/bindings/python/example/log_table.py
b/bindings/python/example/log_table.py
index 37da1da9..018d056b 100644
--- a/bindings/python/example/log_table.py
+++ b/bindings/python/example/log_table.py
@@ -94,6 +94,9 @@ async def _run(conn):
await admin.create_table(table_path, table_descriptor,
ignore_if_exists=True)
print(f"Created table: {table_path}")
+ # A fresh table briefly reports "not leader" until bucket leaders are
elected.
+ await _await_bucket_leader(admin, table_path)
+
table_info = await admin.get_table_info(table_path)
print(f"Table info: {table_info}")
print(f"Table ID: {table_info.table_id}")
@@ -242,12 +245,30 @@ async def _run(conn):
await _scan_batch(table, num_buckets)
await _scan_records(table, num_buckets)
await _projection(table, num_buckets)
+ await _limit_scan(table, num_buckets)
await _context_manager_demo(conn, table_path)
await admin.drop_table(table_path, ignore_if_not_exists=True)
print(f"\nDropped table: {table_path}")
+async def _await_bucket_leader(admin, table_path, *, attempts=60, delay_s=0.5):
+ """Poll until the bucket leader is elected, so bucket-level requests on a
+ just-created table don't fail with "not leader or follower"."""
+ for _ in range(attempts):
+ try:
+ await admin.list_offsets(
+ table_path, bucket_ids=[0],
offset_spec=fluss.OffsetSpec.earliest()
+ )
+ return
+ except fluss.FlussError:
+ await asyncio.sleep(delay_s)
+ # Final attempt (outside the guard) surfaces the real error, not a timeout.
+ await admin.list_offsets(
+ table_path, bucket_ids=[0], offset_spec=fluss.OffsetSpec.earliest()
+ )
+
+
async def _scan_batch(table, num_buckets):
print("\n--- Batch scanner: to_arrow() / to_pandas() ---")
scanner = await table.new_scan().create_record_batch_log_scanner()
@@ -363,6 +384,27 @@ async def _projection(table, num_buckets):
print(f"Projected columns: {list(df_named.columns)}")
+async def _limit_scan(table, num_buckets):
+ print("\n--- Limit scan: one-shot bounded BatchScanner (per bucket) ---")
+ table_id = table.get_table_info().table_id
+ total = 0
+ for bucket_id in range(num_buckets):
+ bucket = fluss.TableBucket(table_id, bucket_id)
+ scanner = (
+
table.new_scan().limit(EXPECTED_ROWS).create_bucket_batch_scanner(bucket)
+ )
+ batch = await scanner.next_batch()
+ if batch is not None:
+ assert batch.bucket == bucket
+ total += batch.batch.num_rows
+ # One-shot: the scanner is spent after the first batch.
+ assert await scanner.next_batch() is None
+ assert total == EXPECTED_ROWS, (
+ f"Limit scan across buckets returned {total} rows, expected
{EXPECTED_ROWS}"
+ )
+ print(f"Limit scan across {num_buckets} bucket(s) returned {total} rows")
+
+
async def _context_manager_demo(conn, table_path):
print("\n--- Async context manager (auto-flush on exit) ---")
table = await conn.get_table(table_path)
diff --git a/bindings/python/example/pk_table.py
b/bindings/python/example/pk_table.py
index 68a7c9ea..19df7741 100644
--- a/bindings/python/example/pk_table.py
+++ b/bindings/python/example/pk_table.py
@@ -86,6 +86,7 @@ async def _run(conn):
await _lookup(table)
await _delete(table)
await _partial_update(table)
+ await _limit_scan(table)
await admin.drop_table(table_path, ignore_if_not_exists=True)
print(f"\nDropped PK table: {table_path}")
@@ -221,5 +222,19 @@ async def _partial_update(table):
)
+async def _limit_scan(table):
+ print("\n--- Limit scan: bounded BatchScanner over current rows (per
bucket) ---")
+ table_info = table.get_table_info()
+ total = 0
+ for bucket_id in range(table_info.num_buckets):
+ bucket = fluss.TableBucket(table_info.table_id, bucket_id)
+ scanner =
table.new_scan().limit(100).create_bucket_batch_scanner(bucket)
+ arrow_table = await scanner.to_arrow()
+ total += arrow_table.num_rows
+ # Users 1 and 2 remain (user 3 was deleted; user 1 was updated in place).
+ assert total == 2, f"Limit scan returned {total} current rows, expected 2"
+ print(f"Limit scan across {table_info.num_buckets} bucket(s) returned
{total} rows")
+
+
if __name__ == "__main__":
asyncio.run(main())
diff --git a/bindings/python/fluss/__init__.pyi
b/bindings/python/fluss/__init__.pyi
index 7d5bfa73..7b808dc9 100644
--- a/bindings/python/fluss/__init__.pyi
+++ b/bindings/python/fluss/__init__.pyi
@@ -503,6 +503,34 @@ class TableScan:
Self for method chaining.
"""
...
+ def limit(self, n: int) -> "TableScan":
+ """Set a positive row limit for the scan.
+
+ A limit enables ``create_bucket_batch_scanner()`` for a one-shot
+ bounded scan. The log scanners do not support limit pushdown and reject
+ a configured limit.
+
+ Args:
+ n: The maximum number of rows to scan. Must be positive.
+
+ Returns:
+ Self for method chaining.
+ """
+ ...
+ def create_bucket_batch_scanner(self, bucket: TableBucket) -> BatchScanner:
+ """Create a one-shot bounded scanner over a single bucket.
+
+ Requires a limit configured via ``limit()``. Creation is cheap; the
+ scan RPC runs lazily on the first ``next_batch()``.
+
+ Args:
+ bucket: The bucket to scan. Its ``table_id`` must match this table
+ and its ``bucket_id`` must be in range.
+
+ Returns:
+ BatchScanner for a single bounded scan of ``bucket``.
+ """
+ ...
async def create_log_scanner(self) -> LogScanner:
"""Create a record-based log scanner.
@@ -976,6 +1004,60 @@ class LogScanner:
def __repr__(self) -> str: ...
def __aiter__(self) -> AsyncIterator[Union[ScanRecord, RecordBatch]]: ...
+@final
+class BatchScanner:
+ """One-shot bounded scanner over a single bucket.
+
+ Obtain via
``table.new_scan().limit(n).create_bucket_batch_scanner(bucket)``.
+ The scan runs lazily on the first ``next_batch()`` (or
``collect_all_batches()``
+ / ``to_arrow()`` / ``to_pandas()``), yields its single batch once, then is
+ spent. Honors the configured limit and any projection.
+
+ Example:
+ ```python
+ table_id = table.get_table_info().table_id
+ scanner = table.new_scan().limit(100).create_bucket_batch_scanner(
+ fluss.TableBucket(table_id, 0)
+ )
+ table_data = await scanner.to_arrow()
+ ```
+ """
+
+ @property
+ def bucket(self) -> TableBucket:
+ """The bucket scanned by this batch scanner."""
+ ...
+ async def next_batch(self) -> Optional[RecordBatch]:
+ """Run the scan and return its batch, or ``None`` once the scanner is
spent.
+
+ The scan RPC runs on the first call; subsequent calls return ``None``.
+ The scan is not retried — an error leaves the scanner spent, so create
a
+ new one to retry.
+
+ Returns:
+ A RecordBatch on the first call, then ``None``.
+ """
+ ...
+ async def collect_all_batches(self) -> List[RecordBatch]:
+ """Drain the scanner into all of its batches.
+
+ Returns:
+ List of RecordBatch objects (a single element for a limit scan).
+ """
+ ...
+ async def to_arrow(self) -> pa.Table:
+ """Drain the scanner into a single PyArrow Table.
+
+ Returns:
+ PyArrow Table with the scanned rows, or an empty table with the
+ projected schema when the scan yields nothing.
+ """
+ ...
+ async def to_pandas(self) -> pd.DataFrame:
+ """Drain the scanner into a Pandas DataFrame."""
+ ...
+ def __repr__(self) -> str: ...
+
@final
class Schema:
def __new__(
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index 2d71491a..45b5092b 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -120,6 +120,7 @@ fn _fluss(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PrefixLookuper>()?;
m.add_class::<Schema>()?;
m.add_class::<LogScanner>()?;
+ m.add_class::<BatchScanner>()?;
m.add_class::<LakeSnapshot>()?;
m.add_class::<TableBucket>()?;
m.add_class::<ChangeType>()?;
diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs
index e18c74d4..cbdcbc01 100644
--- a/bindings/python/src/table.rs
+++ b/bindings/python/src/table.rs
@@ -21,6 +21,7 @@ use arrow::array::RecordBatch as ArrowRecordBatch;
use arrow::record_batch::RecordBatchReader as _;
use arrow_pyarrow::{FromPyArrow, ToPyArrow};
use arrow_schema::SchemaRef;
+use fcore::client::LimitBatchScanner;
use fcore::metadata::{DataField, DataType, MapType, RowType};
use fcore::row::binary_array::{FlussArray, FlussArrayWriter};
use fcore::row::binary_map::{FlussMap, FlussMapWriter};
@@ -39,6 +40,7 @@ use pyo3_async_runtimes::tokio::future_into_py;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
+use tokio::sync::Mutex;
// Time conversion constants
const MILLIS_PER_SECOND: i64 = 1_000;
@@ -427,6 +429,7 @@ pub struct TableScan {
metadata: Arc<fcore::client::Metadata>,
table_info: fcore::metadata::TableInfo,
projection: Option<ProjectionType>,
+ limit: Option<i32>,
}
/// Scanner type for internal use
@@ -461,6 +464,57 @@ impl TableScan {
slf
}
+ /// Set a positive row limit, enabling `create_bucket_batch_scanner()`.
+ ///
+ /// Args:
+ /// n: Maximum number of rows to scan. Must be positive.
+ ///
+ /// Returns:
+ /// Self for method chaining.
+ pub fn limit(mut slf: PyRefMut<'_, Self>, n: i32) -> PyResult<PyRefMut<'_,
Self>> {
+ if n <= 0 {
+ return Err(FlussError::new_err(format!(
+ "Scan limit must be positive, got {n}"
+ )));
+ }
+ slf.limit = Some(n);
+ Ok(slf)
+ }
+
+ /// Create a one-shot bounded scanner over a single bucket.
+ ///
+ /// Requires a limit set via `limit()`; the scan runs on the first
+ /// `next_batch()`.
+ ///
+ /// Args:
+ /// bucket: Bucket to scan; must belong to this table.
+ ///
+ /// Returns:
+ /// A BatchScanner for `bucket`.
+ pub fn create_bucket_batch_scanner(&self, bucket: &TableBucket) ->
PyResult<BatchScanner> {
+ let limit = self.limit.ok_or_else(|| {
+ FlussError::new_err("create_bucket_batch_scanner requires a limit
set via .limit(n)")
+ })?;
+
+ let conn = self.connection.clone();
+ let _guard = TOKIO_RUNTIME.enter();
+ let table =
+ fcore::client::FlussTable::new(&conn, self.metadata.clone(),
self.table_info.clone());
+
+ let projection = self.projection.clone();
+ let projection_indices = resolve_projection_indices(&projection,
&self.table_info)?;
+ let scan = apply_projection(table.new_scan(), projection)?
+ .limit(limit)
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ let scanner = scan
+ .create_bucket_batch_scanner(bucket.to_core())
+ .map_err(|e| FlussError::from_core_error(&e))?;
+
+ let (projected_schema, _) =
+ calculate_projected_types(&self.table_info, projection_indices)?;
+ Ok(BatchScanner::new(scanner, bucket.clone(), projected_schema))
+ }
+
/// Create a record-based log scanner.
///
/// Use this scanner with `poll()` to get individual records with metadata
@@ -501,6 +555,13 @@ impl TableScan {
py: Python<'py>,
scanner_type: ScannerType,
) -> PyResult<Bound<'py, PyAny>> {
+ if let Some(limit) = self.limit {
+ return Err(FlussError::new_err(format!(
+ "Log scanners don't support limit pushdown (requested limit:
{limit}). \
+ Use create_bucket_batch_scanner() for a bounded scan."
+ )));
+ }
+
let conn = self.connection.clone();
let metadata = self.metadata.clone();
let table_info = self.table_info.clone();
@@ -638,6 +699,7 @@ impl FlussTable {
metadata: self.metadata.clone(),
table_info: self.table_info.clone(),
projection: None,
+ limit: None,
}
}
@@ -2630,32 +2692,14 @@ impl LogScanner {
/// Returns:
/// PyArrow Table containing all data from subscribed buckets
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 projected_schema = self.projected_schema.clone();
-
- future_into_py(py, async move {
- let scanner = kind.as_batch()?;
-
- let mut reader =
fcore::client::RecordBatchLogReader::new_until_latest(
- scanner.new_shared_handle(),
- &admin,
- )
- .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| Self::batches_to_arrow_table(py, batches,
&projected_schema))
- })
+ future_into_py(
+ py,
+ Self::scan_to_arrow_table(
+ Arc::clone(&self.kind),
+ Arc::clone(&self.admin),
+ self.projected_schema.clone(),
+ ),
+ )
}
/// Convert all data to Pandas DataFrame.
@@ -2671,31 +2715,9 @@ impl LogScanner {
let kind = Arc::clone(&self.kind);
let admin = Arc::clone(&self.admin);
let projected_schema = self.projected_schema.clone();
-
future_into_py(py, async move {
- let scanner = kind.as_batch()?;
-
- let mut reader =
fcore::client::RecordBatchLogReader::new_until_latest(
- scanner.new_shared_handle(),
- &admin,
- )
- .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, batches,
&projected_schema)?;
- arrow_table.call_method0(py, "to_pandas")
- })
+ let table = Self::scan_to_arrow_table(kind, admin,
projected_schema).await?;
+ Python::attach(|py| table.call_method0(py, "to_pandas"))
})
}
@@ -2758,6 +2780,29 @@ impl LogScanner {
}
}
+ /// Read until the latest offsets and build one PyArrow Table.
+ async fn scan_to_arrow_table(
+ kind: Arc<ScannerKind>,
+ admin: Arc<fcore::client::FlussAdmin>,
+ projected_schema: SchemaRef,
+ ) -> PyResult<Py<PyAny>> {
+ let scanner = kind.as_batch()?;
+ let mut reader = fcore::client::RecordBatchLogReader::new_until_latest(
+ scanner.new_shared_handle(),
+ &admin,
+ )
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ let batches: Vec<Arc<ArrowRecordBatch>> = reader
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?
+ .into_iter()
+ .map(|sb| Arc::new(sb.into_batch()))
+ .collect();
+ Python::attach(|py| Self::batches_to_arrow_table(py, batches,
&projected_schema))
+ }
+
/// Convert Arrow record batches to a PyArrow Table (or empty table if no
batches).
fn batches_to_arrow_table(
py: Python<'_>,
@@ -2780,6 +2825,113 @@ impl LogScanner {
}
}
+/// One-shot bounded scanner over a single bucket.
+///
+/// Obtained via
`table.new_scan().limit(n).create_bucket_batch_scanner(bucket)`.
+/// The scan runs on the first `next_batch()` and yields its single batch once,
+/// then is spent. Honors the configured limit and any projection.
+#[pyclass]
+pub struct BatchScanner {
+ inner: Arc<Mutex<LimitBatchScanner>>,
+ bucket: TableBucket,
+ projected_schema: SchemaRef,
+}
+
+#[pymethods]
+impl BatchScanner {
+ /// The bucket scanned by this batch scanner.
+ #[getter]
+ fn bucket(&self) -> TableBucket {
+ self.bucket.clone()
+ }
+
+ /// Run the scan and return its batch, or `None` once the scanner is spent.
+ ///
+ /// The scan runs on the first call and is not retried; on error, create a
+ /// new scanner.
+ fn next_batch<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ future_into_py(py, async move {
+ let mut scanner = inner.lock().await;
+ let batch = scanner
+ .next_batch()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ Python::attach(|py| match batch {
+ Some(sb) => Ok(Some(Py::new(py,
RecordBatch::from_scan_batch(sb))?)),
+ None => Ok(None),
+ })
+ })
+ }
+
+ /// Drain the scanner into all of its batches.
+ fn collect_all_batches<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py,
PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ future_into_py(py, async move {
+ let mut scanner = inner.lock().await;
+ let batches = scanner
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?;
+ Python::attach(|py| {
+ batches
+ .into_iter()
+ .map(|sb| Py::new(py, RecordBatch::from_scan_batch(sb)))
+ .collect::<PyResult<Vec<_>>>()
+ })
+ })
+ }
+
+ /// Drain the scanner into a PyArrow Table (empty, with the projected
schema,
+ /// when the scan yields nothing).
+ fn to_arrow<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+ future_into_py(
+ py,
+ Self::scan_to_arrow_table(Arc::clone(&self.inner),
self.projected_schema.clone()),
+ )
+ }
+
+ /// Drain the scanner into a Pandas DataFrame.
+ fn to_pandas<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+ let inner = Arc::clone(&self.inner);
+ let projected_schema = self.projected_schema.clone();
+ future_into_py(py, async move {
+ let table = Self::scan_to_arrow_table(inner,
projected_schema).await?;
+ Python::attach(|py| table.call_method0(py, "to_pandas"))
+ })
+ }
+
+ fn __repr__(&self) -> String {
+ format!("BatchScanner(bucket={})", self.bucket.__str__())
+ }
+}
+
+impl BatchScanner {
+ fn new(scanner: LimitBatchScanner, bucket: TableBucket, projected_schema:
SchemaRef) -> Self {
+ Self {
+ inner: Arc::new(Mutex::new(scanner)),
+ bucket,
+ projected_schema,
+ }
+ }
+
+ /// Drain the scanner into one PyArrow Table.
+ async fn scan_to_arrow_table(
+ inner: Arc<Mutex<LimitBatchScanner>>,
+ projected_schema: SchemaRef,
+ ) -> PyResult<Py<PyAny>> {
+ let mut scanner = inner.lock().await;
+ let batches = scanner
+ .collect_all_batches()
+ .await
+ .map_err(|e| FlussError::from_core_error(&e))?
+ .into_iter()
+ .map(|sb| Arc::new(sb.into_batch()))
+ .collect();
+ Python::attach(|py| LogScanner::batches_to_arrow_table(py, batches,
&projected_schema))
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/bindings/python/test/test_batch_scanner.py
b/bindings/python/test/test_batch_scanner.py
new file mode 100644
index 00000000..6e30fff6
--- /dev/null
+++ b/bindings/python/test/test_batch_scanner.py
@@ -0,0 +1,221 @@
+# 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.
+
+"""Integration tests for the one-shot limit-based BatchScanner.
+
+Mirrors crates/fluss/tests/integration/batch_scanner.rs.
+"""
+
+import pyarrow as pa
+import pytest
+
+import fluss
+
+
+async def test_returns_appended_rows_then_none(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_log")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+ schema = pa.schema([pa.field("c1", pa.int32()), pa.field("c2",
pa.string())])
+ table_descriptor = fluss.TableDescriptor(
+ fluss.Schema(schema), bucket_count=1, bucket_keys=["c1"]
+ )
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ append_writer = table.new_append().create_writer()
+ append_writer.write_arrow_batch(
+ pa.RecordBatch.from_arrays(
+ [
+ pa.array([1, 2, 3, 4, 5], pa.int32()),
+ pa.array(["a", "b", "c", "d", "e"]),
+ ],
+ schema=schema,
+ )
+ )
+ await append_writer.flush()
+
+ bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+ scanner = table.new_scan().limit(3).create_bucket_batch_scanner(bucket)
+ assert scanner.bucket == bucket
+
+ first = await scanner.next_batch()
+ assert first is not None
+ assert first.bucket == bucket
+ # The server may return fewer rows than the limit, but never more.
+ assert 0 < first.batch.num_rows <= 3
+ assert await scanner.next_batch() is None
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_reads_primary_key_table(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_pk")
+ 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())]),
+ primary_keys=["id"],
+ )
+ table_descriptor = fluss.TableDescriptor(schema, bucket_count=1)
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ upsert_writer = table.new_upsert().create_writer()
+ expected = {1: "a", 2: "b", 3: "c", 4: "d", 5: "e"}
+ for id_, name in expected.items():
+ upsert_writer.upsert({"id": id_, "name": name})
+ await upsert_writer.flush()
+
+ bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+ scanner = table.new_scan().limit(3).create_bucket_batch_scanner(bucket)
+ first = await scanner.next_batch()
+ assert first is not None
+
+ rows = first.batch.to_pydict()
+ assert 0 < len(rows["id"]) <= 3
+ assert all(expected[i] == name for i, name in zip(rows["id"],
rows["name"]))
+ assert await scanner.next_batch() is None
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_to_arrow_and_collect(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_to_arrow")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+ schema = pa.schema([pa.field("c1", pa.int32()), pa.field("c2",
pa.string())])
+ table_descriptor = fluss.TableDescriptor(
+ fluss.Schema(schema), bucket_count=1, bucket_keys=["c1"]
+ )
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ append_writer = table.new_append().create_writer()
+ append_writer.write_arrow_batch(
+ pa.RecordBatch.from_arrays(
+ [pa.array([10, 20], pa.int32()), pa.array(["x", "y"])],
schema=schema
+ )
+ )
+ await append_writer.flush()
+ table_id = table.get_table_info().table_id
+
+ batches = (
+ await table.new_scan()
+ .limit(10)
+ .create_bucket_batch_scanner(fluss.TableBucket(table_id, 0))
+ .collect_all_batches()
+ )
+ assert [b.batch.num_rows for b in batches] == [2]
+
+ arrow = (
+ await table.new_scan()
+ .limit(10)
+ .create_bucket_batch_scanner(fluss.TableBucket(table_id, 0))
+ .to_arrow()
+ )
+ assert arrow.to_pydict() == {"c1": [10, 20], "c2": ["x", "y"]}
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_projection_skips_middle_column(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_projection")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+ schema = pa.schema(
+ [
+ pa.field("c1", pa.int32()),
+ pa.field("c2", pa.string()),
+ pa.field("c3", pa.int64()),
+ ]
+ )
+ table_descriptor = fluss.TableDescriptor(
+ fluss.Schema(schema), bucket_count=1, bucket_keys=["c1"]
+ )
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ append_writer = table.new_append().create_writer()
+ append_writer.write_arrow_batch(
+ pa.RecordBatch.from_arrays(
+ [
+ pa.array([1, 2], pa.int32()),
+ pa.array(["a", "b"]),
+ pa.array([100, 200], pa.int64()),
+ ],
+ schema=schema,
+ )
+ )
+ await append_writer.flush()
+
+ bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+ arrow = (
+ await table.new_scan()
+ .project_by_name(["c1", "c3"])
+ .limit(10)
+ .create_bucket_batch_scanner(bucket)
+ .to_arrow()
+ )
+ assert arrow.to_pydict() == {"c1": [1, 2], "c3": [100, 200]}
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_construction_errors(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_errors")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+ schema = fluss.Schema(pa.schema([pa.field("c1", pa.int32())]))
+ table_descriptor = fluss.TableDescriptor(schema, bucket_count=1,
bucket_keys=["c1"])
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ table_id = table.get_table_info().table_id
+
+ for bad in (0, -5):
+ with pytest.raises(fluss.FlussError):
+ table.new_scan().limit(bad)
+
+ with pytest.raises(fluss.FlussError):
+
table.new_scan().create_bucket_batch_scanner(fluss.TableBucket(table_id, 0))
+
+ for bad_bucket in (
+ fluss.TableBucket(table_id + 9999, 0),
+ fluss.TableBucket(table_id, 99),
+ ):
+ with pytest.raises(fluss.FlussError):
+ table.new_scan().limit(1).create_bucket_batch_scanner(bad_bucket)
+
+ with pytest.raises(fluss.FlussError):
+ await table.new_scan().limit(5).create_log_scanner()
+ with pytest.raises(fluss.FlussError):
+ await table.new_scan().limit(5).create_record_batch_log_scanner()
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
+
+
+async def test_rejects_non_arrow_log_format(connection, admin):
+ table_path = fluss.TablePath("fluss", "py_test_bs_indexed")
+ await admin.drop_table(table_path, ignore_if_not_exists=True)
+ schema = fluss.Schema(pa.schema([pa.field("c1", pa.int32())]))
+ table_descriptor = fluss.TableDescriptor(
+ schema, bucket_count=1, bucket_keys=["c1"], log_format="INDEXED"
+ )
+ await admin.create_table(table_path, table_descriptor,
ignore_if_exists=False)
+
+ table = await connection.get_table(table_path)
+ bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+ with pytest.raises(fluss.FlussError):
+ table.new_scan().limit(1).create_bucket_batch_scanner(bucket)
+
+ await admin.drop_table(table_path, ignore_if_not_exists=False)
diff --git a/crates/fluss/src/client/table/batch_scanner.rs
b/crates/fluss/src/client/table/batch_scanner.rs
index cc0585f3..5d0cf0c6 100644
--- a/crates/fluss/src/client/table/batch_scanner.rs
+++ b/crates/fluss/src/client/table/batch_scanner.rs
@@ -37,7 +37,6 @@ use crate::rpc::RpcClient;
use crate::rpc::message::LimitScanRequest;
use arrow::array::RecordBatch;
use arrow::compute::concat_batches;
-use arrow_schema::SchemaRef;
use byteorder::{ByteOrder, LittleEndian};
use bytes::Bytes;
use std::collections::HashMap;
@@ -177,25 +176,17 @@ fn decode_log_batch(
) -> Result<(RecordBatch, i64)> {
let row_type = Arc::new(table_info.get_row_type().clone());
let full_schema = to_arrow_schema(table_info.get_row_type())?;
- let read_context = match projected_fields {
- None => ArrowReadContext::new(full_schema.clone(), row_type.clone(),
false),
- Some(fields) => ArrowReadContext::with_projection_pushdown(
- full_schema.clone(),
- row_type.clone(),
- fields.to_vec(),
- false,
- )?,
- };
-
- let target_schema: SchemaRef = match projected_fields {
- None => full_schema,
- Some(fields) => {
-
ArrowReadContext::project_schema(to_arrow_schema(table_info.get_row_type())?,
fields)?
- }
- };
+ // A limit scan returns every column (never projected server-side); decode
+ // the full batch and project after, like the KV path. Pushdown here would
+ // misparse the full-column body and corrupt the buffers.
+ let read_context = ArrowReadContext::new(full_schema.clone(),
row_type.clone(), false);
if raw.is_empty() {
- return Ok((RecordBatch::new_empty(target_schema), 0));
+ let empty = RecordBatch::new_empty(full_schema);
+ return Ok((
+ project_batch(empty, table_info.get_row_type(), projected_fields)?,
+ 0,
+ ));
}
let mut batches: Vec<RecordBatch> = Vec::new();
@@ -211,17 +202,21 @@ fn decode_log_batch(
let base_offset = base_offset.unwrap_or(0);
let merged = if batches.is_empty() {
- RecordBatch::new_empty(target_schema)
+ RecordBatch::new_empty(full_schema)
} else if batches.len() == 1 {
batches.into_iter().next().unwrap()
} else {
- concat_batches(&target_schema, batches.iter()).map_err(|e|
Error::UnexpectedError {
+ concat_batches(&full_schema, batches.iter()).map_err(|e|
Error::UnexpectedError {
message: format!("Failed to concatenate log record batches: {e}"),
source: None,
})?
};
- Ok(take_last_rows(merged, base_offset, limit))
+ let (trimmed, base_offset) = take_last_rows(merged, base_offset, limit);
+ Ok((
+ project_batch(trimmed, table_info.get_row_type(), projected_fields)?,
+ base_offset,
+ ))
}
/// Decode a KV limit-scan [`ValueRecordBatch`] into a single Arrow
@@ -408,51 +403,36 @@ mod tests {
DEFAULT_NON_ZSTD_COMPRESSION_LEVEL,
};
use crate::metadata::{
- Column, DataField, DataType, DataTypes, PhysicalTablePath, Schema,
TableDescriptor,
- TableInfo, TablePath,
+ Column, DataField, DataType, DataTypes, PhysicalTablePath, Schema,
TableInfo, TablePath,
};
use crate::record::MemoryLogRecordsArrowBuilder;
use crate::row::GenericRow;
use crate::row::binary::BinaryWriter;
use crate::row::compacted::CompactedRowWriter;
+ use crate::test_utils::build_table_info_with_columns;
use arrow::array::{Array, Int32Array, Int64Array};
fn build_two_col_table_info() -> TableInfo {
- let row_type = DataTypes::row(vec![
- DataField::new("id", DataTypes::int(), None),
- DataField::new("name", DataTypes::string(), None),
- ]);
- let schema = Schema::builder()
- .with_row_type(&row_type)
- .build()
- .expect("schema build");
- let descriptor = TableDescriptor::builder()
- .schema(schema)
- .distributed_by(Some(1), vec![])
- .build()
- .expect("descriptor build");
- TableInfo::of(
+ build_table_info_with_columns(
TablePath::new("db".to_string(), "tbl".to_string()),
42,
1,
- descriptor,
- 0,
- 0,
+ vec![
+ DataField::new("id", DataTypes::int(), None),
+ DataField::new("name", DataTypes::string(), None),
+ ],
)
}
- fn build_log_records(
- table_info: &TableInfo,
- base_offset: i64,
- rows: &[(i32, &str)],
- ) -> Vec<u8> {
- let row_type = table_info.get_row_type();
- let table_path = table_info.table_path.clone();
+ /// Encode `rows` (built against `table_info`'s row type) as one Arrow log
batch.
+ fn build_log_batch(table_info: &TableInfo, rows: &[GenericRow]) -> Vec<u8>
{
let table_info_arc = Arc::new(table_info.clone());
- let physical = Arc::new(PhysicalTablePath::of(Arc::new(table_path)));
+ let physical = Arc::new(PhysicalTablePath::of(Arc::new(
+ table_info.table_path.clone(),
+ )));
let mut builder = MemoryLogRecordsArrowBuilder::new(
1,
- row_type,
+ table_info.get_row_type(),
false,
ArrowCompressionInfo {
compression_type: ArrowCompressionType::None,
@@ -462,20 +442,33 @@ mod tests {
Arc::new(ArrowCompressionRatioEstimator::default()),
)
.expect("builder");
-
- for (i, (id, name)) in rows.iter().enumerate() {
- let mut row = GenericRow::new(2);
- row.set_field(0, *id);
- row.set_field(1, *name);
+ for (i, row) in rows.iter().enumerate() {
let record = WriteRecord::for_append(
Arc::clone(&table_info_arc),
physical.clone(),
(i + 1) as i32,
- &row,
+ row,
);
builder.append(&record).expect("append");
}
- let mut data = builder.build().expect("build log batch");
+ builder.build().expect("build log batch")
+ }
+
+ fn build_log_records(
+ table_info: &TableInfo,
+ base_offset: i64,
+ rows: &[(i32, &str)],
+ ) -> Vec<u8> {
+ let rows: Vec<GenericRow> = rows
+ .iter()
+ .map(|(id, name)| {
+ let mut row = GenericRow::new(2);
+ row.set_field(0, *id);
+ row.set_field(1, *name);
+ row
+ })
+ .collect();
+ let mut data = build_log_batch(table_info, &rows);
// Builder always writes base_log_offset=0; patch it so tests can
verify
// BatchScanner faithfully propagates whatever offset the server
returned.
let bytes = base_offset.to_le_bytes();
@@ -531,6 +524,52 @@ mod tests {
assert_eq!(batch.schema().field(0).name(), "id");
}
+ /// Projection skipping a middle variable-length column — catches a
+ /// full-column body being misparsed as the projected schema.
+ #[test]
+ fn decode_log_batch_projection_skips_middle_variable_length_column() {
+ let table_info = build_table_info_with_columns(
+ TablePath::new("db".to_string(), "tbl".to_string()),
+ 43,
+ 1,
+ vec![
+ DataField::new("c1", DataTypes::int(), None),
+ DataField::new("c2", DataTypes::string(), None),
+ DataField::new("c3", DataTypes::bigint(), None),
+ ],
+ );
+ let rows: Vec<GenericRow> = [(1, "alice", 100i64), (2, "bob", 200i64)]
+ .iter()
+ .map(|(c1, c2, c3)| {
+ let mut row = GenericRow::new(3);
+ row.set_field(0, *c1);
+ row.set_field(1, *c2);
+ row.set_field(2, *c3);
+ row
+ })
+ .collect();
+ let raw = build_log_batch(&table_info, &rows);
+
+ let (batch, _) = decode_log_batch(&table_info, Some(&[0usize,
2usize]), raw, usize::MAX)
+ .expect("decode projected");
+ assert_eq!(batch.num_columns(), 2);
+ assert_eq!(batch.num_rows(), 2);
+ assert_eq!(batch.schema().field(0).name(), "c1");
+ assert_eq!(batch.schema().field(1).name(), "c3");
+ let c1 = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let c3 = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!((c1.value(0), c1.value(1)), (1, 2));
+ assert_eq!((c3.value(0), c3.value(1)), (100, 200));
+ }
+
#[test]
fn decode_log_batch_truncates_to_last_limit_rows() {
let table_info = build_two_col_table_info();
diff --git a/crates/fluss/src/test_utils.rs b/crates/fluss/src/test_utils.rs
index ec192a50..e1f31bf9 100644
--- a/crates/fluss/src/test_utils.rs
+++ b/crates/fluss/src/test_utils.rs
@@ -25,9 +25,25 @@ use std::collections::HashMap;
use std::sync::Arc;
pub(crate) fn build_table_info(table_path: TablePath, table_id: i64, buckets:
i32) -> TableInfo {
- let row_type = DataTypes::row(vec![DataField::new("id", DataTypes::int(),
None)]);
- let schema_builder = Schema::builder().with_row_type(&row_type);
- let schema = schema_builder.build().expect("schema build");
+ build_table_info_with_columns(
+ table_path,
+ table_id,
+ buckets,
+ vec![DataField::new("id", DataTypes::int(), None)],
+ )
+}
+
+pub(crate) fn build_table_info_with_columns(
+ table_path: TablePath,
+ table_id: i64,
+ buckets: i32,
+ columns: Vec<DataField>,
+) -> TableInfo {
+ let row_type = DataTypes::row(columns);
+ let schema = Schema::builder()
+ .with_row_type(&row_type)
+ .build()
+ .expect("schema build");
let table_descriptor = TableDescriptor::builder()
.schema(schema)
.distributed_by(Some(buckets), vec![])
diff --git a/crates/fluss/tests/integration/batch_scanner.rs
b/crates/fluss/tests/integration/batch_scanner.rs
index 0b484a8c..443d0518 100644
--- a/crates/fluss/tests/integration/batch_scanner.rs
+++ b/crates/fluss/tests/integration/batch_scanner.rs
@@ -19,7 +19,7 @@
#[cfg(test)]
mod batch_scanner_test {
use crate::integration::utils::{create_table, get_shared_cluster};
- use arrow::array::{Int32Array, StringArray, record_batch};
+ use arrow::array::{Int32Array, Int64Array, StringArray, record_batch};
use fluss::metadata::{DataTypes, LogFormat, Schema, TableBucket,
TableDescriptor, TablePath};
use fluss::row::GenericRow;
use std::collections::HashMap;
@@ -93,6 +93,91 @@ mod batch_scanner_test {
);
}
+ /// End-to-end projection skipping the middle `c2` string column.
+ #[tokio::test]
+ async fn batch_scanner_projects_non_contiguous_columns() {
+ let cluster = get_shared_cluster();
+ let connection = cluster.get_fluss_connection().await;
+ let admin = connection.get_admin().expect("admin");
+
+ let table_path = TablePath::new("fluss",
"test_batch_scanner_projection");
+ let descriptor = TableDescriptor::builder()
+ .schema(
+ Schema::builder()
+ .column("c1", DataTypes::int())
+ .column("c2", DataTypes::string())
+ .column("c3", DataTypes::bigint())
+ .build()
+ .expect("schema"),
+ )
+ // Single bucket so a single BatchScanner sees every row.
+ .distributed_by(Some(1), vec!["c1".to_string()])
+ .build()
+ .expect("descriptor");
+ create_table(&admin, &table_path, &descriptor).await;
+
+ let table = connection.get_table(&table_path).await.expect("table");
+ let writer = table
+ .new_append()
+ .expect("append")
+ .create_writer()
+ .expect("writer");
+
+ let batch = record_batch!(
+ ("c1", Int32, [1, 2, 3]),
+ ("c2", Utf8, ["a", "b", "c"]),
+ ("c3", Int64, [100, 200, 300])
+ )
+ .unwrap();
+ writer.append_arrow_batch(batch).expect("append batch");
+ writer.flush().await.expect("flush");
+
+ let table_info = table.get_table_info();
+ let bucket = TableBucket::new(table_info.table_id, 0);
+
+ let mut scanner = table
+ .new_scan()
+ .project(&[0, 2])
+ .expect("project")
+ .limit(10)
+ .expect("limit")
+ .create_bucket_batch_scanner(bucket.clone())
+ .expect("create batch scanner");
+
+ let first = scanner
+ .next_batch()
+ .await
+ .expect("poll")
+ .expect("first batch should be Some");
+
+ let rows = first.batch();
+ assert_eq!(rows.num_columns(), 2, "projected to c1 + c3");
+ assert_eq!(rows.schema().field(0).name(), "c1");
+ assert_eq!(rows.schema().field(1).name(), "c3");
+
+ let c1 = rows
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .expect("c1 Int32");
+ let c3 = rows
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .expect("c3 Int64");
+ // Every (c1, c3) pair must match what we appended (c2 is dropped).
+ let expected: HashMap<i32, i64> = [(1, 100), (2, 200), (3,
300)].into();
+ for i in 0..rows.num_rows() {
+ assert_eq!(
+ expected.get(&c1.value(i)),
+ Some(&c3.value(i)),
+ "projected row ({}, {}) does not match appended data",
+ c1.value(i),
+ c3.value(i)
+ );
+ }
+ }
+
/// Limit scan on a primary-key table: decodes the value-record batch and
/// honors the limit. Exercises the KV wire path (distinct from the log
one).
#[tokio::test]
diff --git a/website/docs/user-guide/python/api-reference.md
b/website/docs/user-guide/python/api-reference.md
index 9bf0b690..341919c0 100644
--- a/website/docs/user-guide/python/api-reference.md
+++ b/website/docs/user-guide/python/api-reference.md
@@ -94,8 +94,10 @@ Supports `async with` statement (async context manager).
|----------------------------------------------------------|---------------------------------------------------------------------|
| `.project(indices) -> TableScan` | Project columns
by index |
| `.project_by_name(names) -> TableScan` | Project columns
by name |
+| `.limit(n) -> TableScan` | Set a positive
row limit (enables `create_bucket_batch_scanner`; rejected by log scanners) |
| `await .create_log_scanner() -> LogScanner` | Create
record-based scanner (for `poll()`) |
| `await .create_record_batch_log_scanner() -> LogScanner` | Create
batch-based scanner (for `poll_arrow()`, `to_arrow()`, etc.) |
+| `.create_bucket_batch_scanner(bucket) -> BatchScanner` | Bounded scan of
one bucket (requires `limit`; runs on first `next_batch()`) |
## `TableAppend`
@@ -187,6 +189,18 @@ Builder for creating a `PrefixLookuper`. Obtain via
`TableLookup.lookup_by(colum
> **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.
+## `BatchScanner`
+
+One-shot bounded scan of a single bucket. Obtain via
`table.new_scan().limit(n).create_bucket_batch_scanner(bucket)`. The scan runs
on the first call below, yields its single batch once, then is spent (create a
new scanner to scan again).
+
+| Method | Description
|
+|-----------------------------------------------------|-----------------------------------------------------|
+| `.bucket -> TableBucket` | The bucket being
scanned (property) |
+| `await .next_batch() -> RecordBatch \| None` | Run the scan; returns
the batch once, then `None` |
+| `await .collect_all_batches() -> list[RecordBatch]` | Drain into a list of
batches |
+| `await .to_arrow() -> pa.Table` | Drain into a single
Arrow Table |
+| `await .to_pandas() -> pd.DataFrame` | Drain into a Pandas
DataFrame |
+
## `ScanRecords`
Returned by `LogScanner.poll()`. Records are grouped by bucket.
diff --git a/website/docs/user-guide/python/example/log-tables.md
b/website/docs/user-guide/python/example/log-tables.md
index 4dbe2567..9f456394 100644
--- a/website/docs/user-guide/python/example/log-tables.md
+++ b/website/docs/user-guide/python/example/log-tables.md
@@ -127,3 +127,17 @@ scanner = await table.new_scan().project([0,
2]).create_record_batch_log_scanner
# or by name
scanner = await table.new_scan().project_by_name(["id",
"score"]).create_record_batch_log_scanner()
```
+
+## Limit Scan
+
+For a bounded read of up to `n` rows from a single bucket, use a batch scanner
instead of subscribing. It issues one request; poll it with `next_batch()`
until it returns `None`.
+
+```python
+bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+scanner = table.new_scan().limit(10).create_bucket_batch_scanner(bucket)
+
+while (batch := await scanner.next_batch()) is not None:
+ print(f"rows: {batch.batch.num_rows}")
+```
+
+`to_arrow()`, `to_pandas()`, and `collect_all_batches()` drain the scan in one
call instead. Limit applies per bucket; scan each bucket to cover a
multi-bucket table.
diff --git a/website/docs/user-guide/python/example/primary-key-tables.md
b/website/docs/user-guide/python/example/primary-key-tables.md
index cd61e508..3cbb8d57 100644
--- a/website/docs/user-guide/python/example/primary-key-tables.md
+++ b/website/docs/user-guide/python/example/primary-key-tables.md
@@ -59,3 +59,15 @@ partial_writer =
table.new_upsert().partial_update_by_name(["id", "age"]).create
partial_writer.upsert({"id": 1, "age": 27}) # only updates age
await partial_writer.flush()
```
+
+## Limit Scan
+
+To read up to `n` rows of a bucket's current state without supplying keys, use
a batch scanner. The server returns the deduplicated current rows as Arrow
batches — convenient for previews or DataFusion sources.
+
+```python
+bucket = fluss.TableBucket(table.get_table_info().table_id, 0)
+scanner = table.new_scan().limit(10).create_bucket_batch_scanner(bucket)
+arrow_table = await scanner.to_arrow()
+```
+
+Limit applies per bucket; scan each bucket to cover a multi-bucket table.
