From 25e62399b65ec80404671c658ea96d3d3b139604 Mon Sep 17 00:00:00 2001 From: carloea2 Date: Thu, 15 Oct 2026 00:34:55 -0600 Subject: [PATCH] feat(pyamber): read Arrow rows across chunk boundaries --- amber/src/main/python/core/models/tuple.py | 13 ++----------- .../src/test/python/core/models/test_tuple.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 11 deletions(-) diff --git a/amber/src/main/python/core/models/tuple.py b/amber/src/main/python/core/models/tuple.py index 1f319b6a4f4..715245255d7 100644 --- a/amber/src/main/python/core/models/tuple.py +++ b/amber/src/main/python/core/models/tuple.py @@ -60,7 +60,6 @@ def __init__(self, table: pyarrow.Table): """ self._table = table self._current_idx = 0 - self._current_chunk = 0 def __iter__(self) -> Iterator[Callable]: """ @@ -71,18 +70,10 @@ def __iter__(self) -> Iterator[Callable]: def __next__(self) -> Callable: """ Provide the field accessor of the next tuple. - If current chunk is exhausted, move to the first tuple of the next chunk. """ - if self._table.num_columns == 0: - # empty table + if self._table.num_columns == 0 or self._current_idx >= self._table.num_rows: raise StopIteration - if self._current_idx >= len(self._table.column(0).chunks[self._current_chunk]): - self._current_idx = 0 - self._current_chunk += 1 - if self._current_chunk >= self._table.column(0).num_chunks: - raise StopIteration - chunk_idx = self._current_chunk tuple_idx = self._current_idx def field_accessor(field_name: str) -> Field: @@ -91,7 +82,7 @@ def field_accessor(field_name: str) -> Field: This abstracts and hides the underlying implementation of the tuple data storage from the user. """ - value = self._table.column(field_name).chunks[chunk_idx][tuple_idx].as_py() + value = self._table.column(field_name)[tuple_idx].as_py() field_type = self._table.schema.field(field_name).type field_metadata = self._table.schema.field(field_name).metadata diff --git a/amber/src/test/python/core/models/test_tuple.py b/amber/src/test/python/core/models/test_tuple.py index aff43faa881..b48cb06e9ba 100644 --- a/amber/src/test/python/core/models/test_tuple.py +++ b/amber/src/test/python/core/models/test_tuple.py @@ -659,6 +659,25 @@ def test_tuple_with_null_large_binary(self): tuple_ = tuples[0] assert tuple_["large_binary_field"] is None + def test_arrow_columns_can_have_different_chunk_boundaries(self): + arrow_table = pyarrow.table( + { + "a": pyarrow.chunked_array([[1, 2], [3]]), + "b": pyarrow.chunked_array([[10], [20, 30]]), + } + ) + + tuples = [ + Tuple({name: accessor for name in arrow_table.column_names}) + for accessor in ArrowTableTupleProvider(arrow_table) + ] + + assert [tuple_.as_dict() for tuple_ in tuples] == [ + {"a": 1, "b": 10}, + {"a": 2, "b": 20}, + {"a": 3, "b": 30}, + ] + def test_binary_field_round_trips_through_arrow(self): # cast_to_schema pickles a non-bytes BINARY field behind a # "pickle " sentinel; the arrow field accessor must recognise