This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new e0b3027d7c [python] Expose Arrow batch reader for multimodal scans
(#9537)
e0b3027d7c is described below
commit e0b3027d7c75f3b4db9bcc3efcebccabbc0372f9
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Sep 2 17:28:27 2026 +0800
[python] Expose Arrow batch reader for multimodal scans (#9537)
---
docs/docs/pypaimon/multimodal-api.mdx | 9 ++++
paimon-python/pypaimon/multimodal/query.py | 19 +++++++
paimon-python/pypaimon/read/table_read.py | 56 ++++++++++++++++++-
.../pypaimon/tests/multimodal_table_test.py | 63 ++++++++++++++++++++++
4 files changed, 146 insertions(+), 1 deletion(-)
diff --git a/docs/docs/pypaimon/multimodal-api.mdx
b/docs/docs/pypaimon/multimodal-api.mdx
index 5e6d352181..b3c59043a8 100644
--- a/docs/docs/pypaimon/multimodal-api.mdx
+++ b/docs/docs/pypaimon/multimodal-api.mdx
@@ -802,6 +802,15 @@ result = (
)
```
+Use `to_arrow_batch_reader()` to consume an ordinary scan without collecting
+all rows into one Arrow table:
+
+```python
+with docs.scan().where("category = 'lake'").to_arrow_batch_reader() as reader:
+ for batch in reader:
+ consume(batch)
+```
+
### Reading BLOB columns
`scan().read_blobs(column)` bulk-fetches a BLOB column's bytes for the filtered
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index 9046498f7c..d4491651d8 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -72,6 +72,19 @@ class ScanQuery:
plan = scan.plan()
return read_builder.new_read().to_arrow(plan.splits())
+ def to_arrow_batch_reader(self, *, blob_parallelism=None):
+ """Stream this scan as Arrow batches without collecting a table."""
+ if self._result_factory is not None:
+ raise TypeError(
+ "to_arrow_batch_reader is only supported on scan(), "
+ "not search queries."
+ )
+
+ read_builder = self._configured_read_builder()
+ splits = read_builder.new_scan().plan().splits()
+ return read_builder.new_read()._to_managed_arrow_batch_reader(
+ splits, blob_parallelism=blob_parallelism)
+
def _configured_read_builder(self, table=None):
read_builder = (
self._table if table is None else table
@@ -399,6 +412,12 @@ class _PreFilterQuery(ScanQuery):
def to_ray(self, *args, **kwargs):
raise TypeError("to_ray is only supported on scan(), not search
queries.")
+ def to_arrow_batch_reader(self, *args, **kwargs):
+ raise TypeError(
+ "to_arrow_batch_reader is only supported on scan(), "
+ "not search queries."
+ )
+
class VectorQuery(_PreFilterQuery):
"""Chainable query wrapper for vector global-index search."""
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 78b2d26f41..3472666ae5 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -40,6 +40,40 @@ from pypaimon.schema.data_types import DataField,
PyarrowFieldParser
from pypaimon.table.row.offset_row import OffsetRow
ROW_KIND_COLUMN = "_row_kind"
+_RECORD_BATCH_READER_FROM_STREAM = getattr(
+ pyarrow.ipc.RecordBatchReader, "from_stream", None)
+
+
+class _ClosableArrowBatchReader:
+
+ def __init__(self, reader, batch_iterator):
+ self._reader = reader
+ self._batch_iterator = batch_iterator
+
+ def __getattr__(self, name):
+ return getattr(self._reader, name)
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ return self._reader.read_next_batch()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.close()
+
+ def close(self):
+ try:
+ close = getattr(self._batch_iterator, "close", None)
+ if close is not None:
+ close()
+ finally:
+ close = getattr(self._reader, "close", None)
+ if close is not None:
+ close()
class _RemainingRows:
@@ -152,12 +186,32 @@ class TableRead:
def to_arrow_batch_reader(self, splits: List[Split],
blob_parallelism: Optional[int] = None) ->
pyarrow.ipc.RecordBatchReader:
+ reader, _ = self._new_arrow_batch_reader(splits, blob_parallelism)
+ return reader
+
+ def _to_managed_arrow_batch_reader(
+ self,
+ splits: List[Split],
+ blob_parallelism: Optional[int] = None):
+ reader, batch_iterator = self._new_arrow_batch_reader(
+ splits, blob_parallelism)
+ if (_RECORD_BATCH_READER_FROM_STREAM is not None
+ and hasattr(reader, "close")):
+ return _RECORD_BATCH_READER_FROM_STREAM(reader)
+ return _ClosableArrowBatchReader(reader, batch_iterator)
+
+ def _new_arrow_batch_reader(
+ self,
+ splits: List[Split],
+ blob_parallelism: Optional[int] = None):
effective_bp = self._resolve_blob_parallelism(blob_parallelism)
schema = PyarrowFieldParser.from_paimon_schema(self.read_type)
if self.include_row_kind:
schema = self._add_row_kind_to_schema(schema)
batch_iterator = self._arrow_batch_generator(splits, schema,
effective_bp)
- return pyarrow.ipc.RecordBatchReader.from_batches(schema,
batch_iterator)
+ reader_type = pyarrow.ipc.RecordBatchReader
+ reader = reader_type.from_batches(schema, batch_iterator)
+ return reader, batch_iterator
@staticmethod
def _add_row_kind_to_schema(schema: pyarrow.Schema) -> pyarrow.Schema:
diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py
b/paimon-python/pypaimon/tests/multimodal_table_test.py
index 85794d232c..ddf84f5696 100644
--- a/paimon-python/pypaimon/tests/multimodal_table_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_table_test.py
@@ -810,6 +810,69 @@ class MultimodalTableTest(unittest.TestCase):
self.assertEqual(1, result.num_rows)
self.assertEqual([1], result["id"].to_pylist())
+ def test_scan_to_arrow_batch_reader(self):
+ users = self.conn.create_table(
+ "batch_users",
+ data=[
+ {"id": 1, "age": 20},
+ {"id": 2, "age": 30},
+ {"id": 3, "age": 40},
+ ],
+ schema=_schema({"id": pa.int32(), "age": pa.int32()}),
+ options=_PARQUET_OPTIONS,
+ )
+
+ reader = (
+ users.scan()
+ .where("age >= 30")
+ .select("id")
+ .to_arrow_batch_reader()
+ )
+
+ self.assertEqual([{"id": 2}, {"id": 3}], reader.read_all().to_pylist())
+
+ closed = []
+
+ def batches(*args):
+ try:
+ yield pa.record_batch(
+ [pa.array([1], type=pa.int32())], names=["id"])
+ yield pa.record_batch(
+ [pa.array([2], type=pa.int32())], names=["id"])
+ finally:
+ closed.append(True)
+
+ with patch(
+ "pypaimon.read.table_read.TableRead._arrow_batch_generator",
+ new=batches), patch(
+ "pypaimon.read.table_read."
+ "_RECORD_BATCH_READER_FROM_STREAM", None):
+ reader = users.scan().select("id").to_arrow_batch_reader()
+ reader.read_next_batch()
+ reader.close()
+ self.assertEqual([True], closed)
+
+ closed.clear()
+ with patch(
+ "pypaimon.read.table_read.TableRead._arrow_batch_generator",
+ new=batches):
+ reader = users.scan().select("id").to_arrow_batch_reader()
+ reader.read_next_batch()
+ reader.close()
+ self.assertEqual([True], closed)
+
+ read_builder = users.raw_table.new_read_builder()
+ splits = read_builder.new_scan().plan().splits()
+ with patch(
+ "pypaimon.read.table_read."
+ "_RECORD_BATCH_READER_FROM_STREAM", None):
+ reader = read_builder.new_read().to_arrow_batch_reader(splits)
+ self.assertIsInstance(reader, pa.RecordBatchReader)
+ reader.close()
+
+ with self.assertRaisesRegex(TypeError, "only supported on scan"):
+ users.search("thirty", column="id").to_arrow_batch_reader()
+
def test_scan_with_row_id_returns_system_column(self):
users = self.conn.create_table(
"users",
