Lazily load spooled result set segments

When using the spooling protocol, TrinoQuery.fetch() materializes all
segments into a single list, causing out-of-memory errors on large
result sets. Return SegmentIterator directly so rows are decoded on
demand, keeping memory usage constant regardless of result set size.

Author: wendigo

tests/development_server.py

@@ -51,7 +51,7 @@ def start_development_server(port=None, trino_version=TRINO_VERSION):
         network = Network().create()
         supports_spooling_protocol = TRINO_VERSION == "latest" or int(TRINO_VERSION) >= 466
         if supports_spooling_protocol:
-            localstack = LocalStackContainer(image="localstack/localstack:latest", region_name="us-east-1") \
+            localstack = LocalStackContainer(image="localstack/localstack:4.14.0", region_name="us-east-1") \
                 .with_name("localstack") \
                 .with_network(network) \
                 .with_bind_ports(4566, 4566) \

tests/integration/test_dbapi_integration.py

@@ -1900,6 +1900,79 @@ def test_segments_cursor(trino_connection):
     assert total == 300875, f"Expected total rows 300875, got {total}"
 
 
+@pytest.mark.skipif(
+    trino_version() <= 466,
+    reason="spooling protocol was introduced in version 466"
+)
+def test_spooled_segments_lazy_fetchone(trino_connection):
+    """Verify that spooled results can be consumed row-by-row via fetchone()
+    without materializing the entire result set in memory."""
+    if trino_connection._client_session.encoding is None:
+        pytest.skip("spooling requires an encoding")
+
+    cur = trino_connection.cursor()
+    cur.execute("""SELECT l.*
+    FROM tpch.tiny.lineitem l, TABLE(sequence(
+        start => 1,
+        stop => 5,
+        step => 1)) n""")
+
+    # The underlying result rows should be an iterator, not a list
+    result_rows = cur._query._result._rows
+    assert not isinstance(result_rows, list), (
+        f"Expected lazy iterator for spooled results, got {type(result_rows)}"
+    )
+
+    # Consume rows one by one and count them
+    count = 0
+    while cur.fetchone() is not None:
+        count += 1
+    assert count == 300875, f"Expected 300875 rows, got {count}"
+
+
+@pytest.mark.skipif(
+    trino_version() <= 466,
+    reason="spooling protocol was introduced in version 466"
+)
+def test_spooled_segments_fetchmany(trino_connection):
+    """Verify that fetchmany() works correctly with lazily loaded spooled segments."""
+    if trino_connection._client_session.encoding is None:
+        pytest.skip("spooling requires an encoding")
+
+    cur = trino_connection.cursor()
+    cur.execute("SELECT * FROM tpch.tiny.lineitem")
+
+    batch = cur.fetchmany(100)
+    assert len(batch) == 100
+
+    total = len(batch)
+    while True:
+        batch = cur.fetchmany(1000)
+        if not batch:
+            break
+        total += len(batch)
+    assert total == 60175, f"Expected 60175 rows, got {total}"
+
+
+@pytest.mark.skipif(
+    trino_version() <= 466,
+    reason="spooling protocol was introduced in version 466"
+)
+def test_spooled_segments_iterator_protocol(trino_connection):
+    """Verify that cursor iteration works correctly with spooled segments."""
+    if trino_connection._client_session.encoding is None:
+        pytest.skip("spooling requires an encoding")
+
+    cur = trino_connection.cursor()
+    cur.execute("SELECT * FROM tpch.tiny.lineitem")
+
+    count = 0
+    for row in cur:
+        count += 1
+        assert isinstance(row, list)
+    assert count == 60175, f"Expected 60175 rows, got {count}"
+
+
 def get_cursor(legacy_prepared_statements, run_trino):
     host, port = run_trino
 

trino/client.py

@@ -904,9 +904,30 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
         rows = self._row_mapper.map(status.rows) if self._row_mapper else status.rows
         self._result = TrinoResult(self, rows)
 
-        # Execute should block until at least one row is received or query is finished or cancelled
-        while not self.finished and not self.cancelled and len(self._result.rows) == 0:
-            self._result.rows += self.fetch()
+        """
+        Execute should block until at least one row is received or query is finished or cancelled.
+
+        In the direct protocol, rows is a list that we can check for length. Usually the first response contains just
+        stats so we need to continue fetching next uri until we get some rows or query is finished or cancelled.
+
+        In the spooling protocol, rows start as empty list and eventually fetch returns the rows as lazy iterator,
+        we can't check length of an iterator easily without peeking.
+
+        So, if we get rows as non empty list or iterator, we stop blocking and return it to the caller to consume it.
+        """
+        while not self.finished and not self.cancelled:
+            if isinstance(self._result.rows, list) and len(self._result.rows) == 0:
+                new_rows = self.fetch()
+                if isinstance(new_rows, list):
+                    # Direct protocol - append rows to a list
+                    self._result.rows += new_rows
+                else:
+                    # Spooling protocol - replace rows with a lazy iterator
+                    self._result.rows = new_rows
+                    break
+            else:
+                break
+
         return self._result
 
     def _update_state(self, status):
@@ -920,7 +941,7 @@ def _update_state(self, status):
         if status.columns:
             self._columns = status.columns
 
-    def fetch(self) -> List[Union[List[Any]], Any]:
+    def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]:
         """Continue fetching data for the current query_id"""
         try:
             response = self._request.get(self._request.next_uri)
@@ -941,7 +962,8 @@ def fetch(self) -> List[Union[List[Any]], Any]:
             spooled = self._to_segments(rows)
             if self._fetch_mode == "segments":
                 return spooled
-            return list(SegmentIterator(spooled, self._row_mapper))
+            # Return iterator directly, do NOT materialize with list()
+            return SegmentIterator(spooled, self._row_mapper)
         elif isinstance(status.rows, list):
             return self._row_mapper.map(rows)
         else: