From 2894330e69a7a409d089a3dbe8b5423d11c50b97 Mon Sep 17 00:00:00 2001
From: carloea2 <carloea2@uci.edu>
Date: Tue, 3 Nov 2026 21:39:00 -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 | 32 +++++++++++++++++++
 2 files changed, 34 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..1ce50f2bc26 100644
--- a/amber/src/test/python/core/models/test_tuple.py
+++ b/amber/src/test/python/core/models/test_tuple.py
@@ -659,6 +659,38 @@ def test_tuple_with_null_large_binary(self):
         tuple_ = tuples[0]
         assert tuple_["large_binary_field"] is None
 
+    @pytest.mark.parametrize(
+        ("a_chunks", "b_chunks"),
+        [
+            ([[1, 2], [3]], [[10], [20, 30]]),
+            ([[1], [2, 3]], [[10, 20], [30]]),
+        ],
+        ids=["second-column-shorter", "second-column-longer"],
+    )
+    def test_arrow_columns_can_have_different_chunk_boundaries(
+        self, a_chunks, b_chunks
+    ):
+        arrow_table = pyarrow.table(
+            {
+                "a": pyarrow.chunked_array(a_chunks),
+                "b": pyarrow.chunked_array(b_chunks),
+            }
+        )
+
+        tuples = [
+            Tuple({name: accessor for name in arrow_table.column_names})
+            for accessor in ArrowTableTupleProvider(arrow_table)
+        ]
+
+        expected = [
+            {"a": 1, "b": 10},
+            {"a": 2, "b": 20},
+            {"a": 3, "b": 30},
+        ]
+        assert len(tuples) == len(expected)
+        for tuple_, row in zip(tuples, expected):
+            assert tuple_.as_dict() == row
+
     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