feat(pyamber): read Arrow rows across chunk boundaries

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
 

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