[python] Limit scan

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)

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())

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__(

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>()?;