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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; 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, table_info: fcore::metadata::TableInfo, projection: Option, + limit: Option, } /// 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> { + 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 { + 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> { + 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> { - 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> = 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> = 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, + admin: Arc, + projected_schema: SchemaRef, + ) -> PyResult> { + 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> = 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>, + 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> { + 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> { + 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::>>() + }) + }) + } + + /// 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> { + 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> { + 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>, + projected_schema: SchemaRef, + ) -> PyResult> { + 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 = 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 { - 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 { 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 { + let rows: Vec = 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 = [(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::() + .unwrap(); + let c3 = batch + .column(1) + .as_any() + .downcast_ref::() + .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, +) -> 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::() + .expect("c1 Int32"); + let c3 = rows + .column(1) + .as_any() + .downcast_ref::() + .expect("c3 Int64"); + // Every (c1, c3) pair must match what we appended (c2 is dropped). + let expected: HashMap = [(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.