This is an automated email from the ASF dual-hosted git repository.
XiaoHongbo-Hope 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 ed26b66a9c [python] Avoid Arrow offset overflow when merging read
batches (#9209)
ed26b66a9c is described below
commit ed26b66a9c0e98a11dbc482c707f623142fc3d2c
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Aug 13 23:20:53 2026 +0800
[python] Avoid Arrow offset overflow when merging read batches (#9209)
### Purpose
Fix Data Evolution reads failing when ordinary variable-length columns
exceed the Arrow 2 GiB offset limit across input batches:
```text
pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays,
consider casting input from `string` to `large_string` first.
```
`MergeAllBatchReader` previously collected every supplier batch and
forced each column into one Arrow array. This could overflow the 32-bit
offsets used by `string`, `binary`, and similar types.
### Changes
- Stream supplier batches instead of materializing the complete range.
- Coalesce small inputs into bounded output batches to avoid excessive
one-row batches.
- Recursively track 32-bit offsets: payload bytes for STRING/BINARY and
child
element counts for ARRAY/MAP, including nested values.
- Flush before concatenation reaches any Arrow 32-bit offset limit.
- Preserve buffered remainders in `DataEvolutionMergeReader` without
unbounded concatenation.
- Keep the original Arrow schema; columns are not changed to
`large_string` or `large_binary`.
### Tests
- Verify the original data and batch slicing behavior.
- Verify concatenation flushes before the Arrow offset limit.
- Verify two single-row ARRAY batches with 1.1 billion child elements
each do
not overflow when coalesced.
- Verify nested ARRAY and MAP offsets are checked recursively.
- Verify 10,000 misaligned one-row inputs produce 10 bounded output
batches rather than 10,000 batches.
- Related Data Evolution, deletion vector, format, and BLOB tests: 184
passed, 49 subtests passed.
- flake8 and `git diff --check` passed.
---
.../pypaimon/read/reader/concat_batch_reader.py | 187 ++++++++++--------
.../pypaimon/tests/concat_batch_reader_test.py | 218 +++++++++++++++++++++
2 files changed, 328 insertions(+), 77 deletions(-)
diff --git a/paimon-python/pypaimon/read/reader/concat_batch_reader.py
b/paimon-python/pypaimon/read/reader/concat_batch_reader.py
index 7a0c89f18b..831e41012a 100644
--- a/paimon-python/pypaimon/read/reader/concat_batch_reader.py
+++ b/paimon-python/pypaimon/read/reader/concat_batch_reader.py
@@ -16,10 +16,10 @@
# under the License.
import collections
+import struct
from typing import Callable, Dict, List, Optional, Tuple
import pyarrow as pa
-import pyarrow.dataset as ds
from pyarrow import RecordBatch
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
@@ -29,7 +29,65 @@ from pypaimon.schema.data_types import DataField,
PyarrowFieldParser
from pypaimon.table.row.blob import Blob
from pypaimon.utils.range import Range
-_MIN_BATCH_SIZE_TO_REFILL = 1024
+_MAX_ARROW_OFFSET = (1 << 31) - 1
+
+
+def _offset_bounds(array, width=4):
+ offsets = array.buffers()[1]
+ if offsets is None:
+ return 0, 0
+ value_format = "=i" if width == 4 else "=q"
+ start = struct.unpack_from(
+ value_format, offsets, array.offset * width)[0]
+ end = struct.unpack_from(
+ value_format, offsets, (array.offset + len(array)) * width)[0]
+ return start, end
+
+
+def _collect_offset_usage(array, path, usage):
+ data_type = array.type
+ if pa.types.is_string(data_type) or pa.types.is_binary(data_type):
+ start, end = _offset_bounds(array)
+ usage[path] = end - start
+ return
+
+ if pa.types.is_list(data_type) or pa.types.is_map(data_type):
+ start, end = _offset_bounds(array)
+ usage[path] = end - start
+ if pa.types.is_map(data_type):
+ _collect_offset_usage(
+ array.keys.slice(start, end - start), path + ("key",), usage)
+ _collect_offset_usage(
+ array.items.slice(start, end - start), path + ("item",), usage)
+ else:
+ _collect_offset_usage(
+ array.values.slice(start, end - start), path + ("item",),
usage)
+ return
+
+ if pa.types.is_large_list(data_type):
+ start, end = _offset_bounds(array, 8)
+ _collect_offset_usage(
+ array.values.slice(start, end - start), path + ("item",), usage)
+ return
+
+ if pa.types.is_fixed_size_list(data_type):
+ size = data_type.list_size
+ start = array.offset * size
+ _collect_offset_usage(
+ array.values.slice(start, len(array) * size), path + ("item",),
usage)
+ return
+
+ if pa.types.is_struct(data_type):
+ for index in range(data_type.num_fields):
+ _collect_offset_usage(
+ array.field(index), path + ("field", index), usage)
+
+
+def _batch_offset_usage(batch):
+ usage = {}
+ for index in range(batch.num_columns):
+ _collect_offset_usage(batch.column(index), (index,), usage)
+ return usage
class _BlobFileState:
@@ -83,76 +141,65 @@ class ConcatBatchReader(RecordBatchReader):
class MergeAllBatchReader(RecordBatchReader):
"""
- A reader that accepts multiple reader suppliers and concatenates all their
arrow batches
- into one big batch. This is useful when you want to merge all data from
multiple sources
- into a single batch for processing.
+ Read multiple suppliers as one bounded stream.
+
+ Small input batches are combined up to ``batch_size`` without allowing one
+ output batch to exceed Arrow's 32-bit offset limit.
"""
def __init__(self, reader_suppliers: List[Callable], batch_size: int =
1024):
- self.reader_suppliers = reader_suppliers
- self.merged_batch: Optional[RecordBatch] = None
- self.reader = None
- self._batch_size = batch_size
+ self._reader = ConcatBatchReader(reader_suppliers)
+ self._remainder: Optional[RecordBatch] = None
+ self._batch_size = max(1, batch_size)
def read_arrow_batch(self) -> Optional[RecordBatch]:
- if self.reader:
- try:
- return self.reader.read_next_batch()
- except StopIteration:
- return None
+ batches = []
+ num_rows = 0
+ offset_usage = {}
+ while True:
+ batch = self._remainder
+ self._remainder = None
+ if batch is None:
+ batch = self._reader.read_arrow_batch()
+ if batch is None:
+ break
+ if batch.num_rows == 0:
+ continue
- all_batches = []
+ take = min(batch.num_rows, self._batch_size - num_rows)
+ piece = batch.slice(0, take)
+ piece_usage = _batch_offset_usage(piece)
+ if batches and any(
+ offset_usage.get(path, 0) + value > _MAX_ARROW_OFFSET
+ for path, value in piece_usage.items()):
+ self._remainder = batch
+ break
- # Read all batches from all reader suppliers
- for supplier in self.reader_suppliers:
- reader = supplier()
- if reader is None:
- continue
- try:
- while True:
- batch = reader.read_arrow_batch()
- if batch is None:
- break
- all_batches.append(batch)
- finally:
- reader.close()
-
- # Concatenate all batches into one big batch
- if all_batches:
- # For PyArrow < 17.0.0, use Table.concat_tables approach
- # Convert batches to tables and concatenate
- tables = [pa.Table.from_batches([batch]) for batch in all_batches]
- if len(tables) == 1:
- # Single table, just get the first batch
- self.merged_batch = tables[0].to_batches()[0]
- else:
- # Multiple tables, concatenate them
- concatenated_table = pa.concat_tables(tables)
- # Convert back to a single batch by taking all batches and
combining
- all_concatenated_batches = concatenated_table.to_batches()
- if len(all_concatenated_batches) == 1:
- self.merged_batch = all_concatenated_batches[0]
- else:
- # If still multiple batches, we need to manually combine
them
- # This shouldn't happen with concat_tables, but just in
case
- combined_arrays = []
- for i in range(len(all_concatenated_batches[0].columns)):
- column_arrays = [batch.column(i) for batch in
all_concatenated_batches]
- combined_arrays.append(pa.concat_arrays(column_arrays))
- self.merged_batch = pa.RecordBatch.from_arrays(
- combined_arrays,
- schema=all_concatenated_batches[0].schema
- )
- else:
- self.merged_batch = None
+ batches.append(piece)
+ num_rows += take
+ for path, value in piece_usage.items():
+ offset_usage[path] = offset_usage.get(path, 0) + value
+ if take < batch.num_rows:
+ self._remainder = batch.slice(take)
+ if num_rows == self._batch_size or any(
+ value >= _MAX_ARROW_OFFSET
+ for value in offset_usage.values()):
+ break
+
+ if not batches:
return None
- dataset = ds.InMemoryDataset(self.merged_batch)
- self.reader = dataset.scanner(batch_size=self._batch_size).to_reader()
- return self.reader.read_next_batch()
+ if len(batches) == 1:
+ return batches[0]
+
+ columns = [
+ pa.concat_arrays([batch.column(i) for batch in batches])
+ for i in range(batches[0].num_columns)
+ ]
+ return pa.RecordBatch.from_arrays(columns, schema=batches[0].schema)
def close(self) -> None:
- self.merged_batch = None
- self.reader = None
+ self._remainder = None
+ self._reader.close()
class DataEvolutionMergeReader(RecordBatchReader):
@@ -198,22 +245,8 @@ class DataEvolutionMergeReader(RecordBatchReader):
for i, reader in enumerate(self.readers):
if reader is not None:
if self._buffers[i] is not None:
- remainder = self._buffers[i]
+ batches[i] = self._buffers[i]
self._buffers[i] = None
- if remainder.num_rows >= _MIN_BATCH_SIZE_TO_REFILL:
- batches[i] = remainder
- else:
- new_batch = reader.read_arrow_batch()
- if new_batch is not None and new_batch.num_rows > 0:
- combined_arrays = [
- pa.concat_arrays([remainder.column(j),
new_batch.column(j)])
- for j in range(remainder.num_columns)
- ]
- batches[i] = pa.RecordBatch.from_arrays(
- combined_arrays, schema=remainder.schema
- )
- else:
- batches[i] = remainder
else:
batch = reader.read_arrow_batch()
if batch is None:
diff --git a/paimon-python/pypaimon/tests/concat_batch_reader_test.py
b/paimon-python/pypaimon/tests/concat_batch_reader_test.py
new file mode 100644
index 0000000000..1504d55530
--- /dev/null
+++ b/paimon-python/pypaimon/tests/concat_batch_reader_test.py
@@ -0,0 +1,218 @@
+# 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.
+
+import unittest
+from unittest.mock import patch
+
+import pyarrow as pa
+
+from pypaimon.read.reader.concat_batch_reader import (
+ DataEvolutionMergeReader,
+ MergeAllBatchReader,
+)
+from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
+
+
+class _BatchReader(RecordBatchReader):
+
+ def __init__(self, batches):
+ self._batches = iter(batches)
+
+ def read_arrow_batch(self):
+ return next(self._batches, None)
+
+ def close(self):
+ pass
+
+
+class ConcatBatchReaderTest(unittest.TestCase):
+
+ def test_merge_all_coalesces_small_batches(self):
+ batches = [
+ pa.record_batch([pa.array(["a", "b"])], names=["value"]),
+ pa.record_batch([pa.array(["c", "d", "e"])], names=["value"]),
+ ]
+ reader = MergeAllBatchReader(
+ [lambda batch=batch: _BatchReader([batch]) for batch in batches],
+ batch_size=2,
+ )
+
+ actual = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ actual.append(batch.column(0).to_pylist())
+
+ self.assertEqual(actual, [["a", "b"], ["c", "d"], ["e"]])
+
+ def test_merge_all_flushes_before_arrow_offset_limit(self):
+ batches = [
+ pa.record_batch([pa.array(["aaaa"])], names=["value"]),
+ pa.record_batch([pa.array(["bbbb"])], names=["value"]),
+ ]
+ reader = MergeAllBatchReader(
+ [lambda batch=batch: _BatchReader([batch]) for batch in batches],
+ batch_size=2,
+ )
+
+ with patch(
+ "pypaimon.read.reader.concat_batch_reader._MAX_ARROW_OFFSET",
+ 4):
+ actual = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ actual.append(batch.column(0).to_pylist())
+
+ self.assertEqual(actual, [["aaaa"], ["bbbb"]])
+
+ def test_merge_all_flushes_before_list_offset_limit(self):
+ child_count = 1100000000
+
+ def batch():
+ values = pa.nulls(child_count)
+ offsets = pa.array([0, child_count], type=pa.int32())
+ return pa.record_batch(
+ [pa.ListArray.from_arrays(offsets, values)], names=["value"])
+
+ batches = [batch(), batch()]
+ self.assertTrue(all(item.nbytes < child_count for item in batches))
+ reader = MergeAllBatchReader(
+ [lambda item=item: _BatchReader([item]) for item in batches],
+ batch_size=2,
+ )
+
+ actual = []
+ while True:
+ item = reader.read_arrow_batch()
+ if item is None:
+ break
+ actual.append(item.column(0).value_lengths()[0].as_py())
+
+ self.assertEqual(actual, [child_count, child_count])
+
+ def test_merge_all_checks_nested_and_map_offsets(self):
+ arrays = [
+ pa.array([["aaa"]], type=pa.list_(pa.string())),
+ pa.array(
+ [[("a", 1), ("b", 2)]],
+ type=pa.map_(pa.string(), pa.int32()),
+ ),
+ ]
+ for array in arrays:
+ with self.subTest(data_type=array.type):
+ batches = [
+ pa.record_batch([array], names=["value"]),
+ pa.record_batch([array], names=["value"]),
+ ]
+ reader = MergeAllBatchReader(
+ [
+ lambda item=item: _BatchReader([item])
+ for item in batches
+ ],
+ batch_size=2,
+ )
+ with patch(
+ "pypaimon.read.reader.concat_batch_reader."
+ "_MAX_ARROW_OFFSET",
+ 3):
+ sizes = []
+ while True:
+ item = reader.read_arrow_batch()
+ if item is None:
+ break
+ sizes.append(item.num_rows)
+
+ self.assertEqual(sizes, [1, 1])
+
+ def test_misaligned_small_files_keep_bounded_batch_count(self):
+ row_count = 10000
+ left = MergeAllBatchReader([
+ lambda value=value: _BatchReader([
+ pa.record_batch([pa.array([value])], names=["left"])
+ ])
+ for value in range(row_count)
+ ])
+ right = _BatchReader([
+ pa.record_batch(
+ [pa.array(range(start, min(start + 1024, row_count)))],
+ names=["right"],
+ )
+ for start in range(0, row_count, 1024)
+ ])
+ reader = DataEvolutionMergeReader(
+ row_offsets=[0, 1],
+ field_offsets=[0, 0],
+ readers=[left, right],
+ schema=pa.schema([("left", pa.int64()), ("right", pa.int64())]),
+ )
+
+ batch_sizes = []
+ values = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ batch_sizes.append(batch.num_rows)
+ values.extend(batch.column(0).to_pylist())
+
+ self.assertEqual(batch_sizes, [1024] * 9 + [784])
+ self.assertEqual(values, list(range(row_count)))
+
+ def test_merge_reader_does_not_join_buffered_remainders(self):
+ left = _BatchReader([
+ pa.record_batch([pa.array([0])], names=["left"]),
+ pa.record_batch([pa.array([1, 2])], names=["left"]),
+ ])
+ right = _BatchReader([
+ pa.record_batch([pa.array([10, 11])], names=["right"]),
+ pa.record_batch([pa.array([12])], names=["right"]),
+ ])
+ reader = DataEvolutionMergeReader(
+ row_offsets=[0, 1],
+ field_offsets=[0, 0],
+ readers=[left, right],
+ schema=pa.schema([("left", pa.int64()), ("right", pa.int64())]),
+ )
+
+ with patch.object(
+ pa, "concat_arrays",
+ side_effect=AssertionError("must preserve Arrow chunks")):
+ actual = []
+ while True:
+ batch = reader.read_arrow_batch()
+ if batch is None:
+ break
+ actual.extend(
+ {"left": left, "right": right}
+ for left, right in zip(
+ batch.column(0).to_pylist(),
+ batch.column(1).to_pylist(),
+ )
+ )
+
+ self.assertEqual(actual, [
+ {"left": 0, "right": 10},
+ {"left": 1, "right": 11},
+ {"left": 2, "right": 12},
+ ])
+
+
+if __name__ == "__main__":
+ unittest.main()
