Fix OOM on large spooled result sets

Previously, TrinoQuery.fetch() eagerness caused all segments to load into memory at once when using fault-tolerant execution. This led to OOM errors on large datasets.

Changes:

- Enable lazy loading by returning SegmentIterator directly in fetch().

- Update execute() to handle result rows as iterators instead of requiring lists.

- Add unit test to verify lazy fetching implementation.

Author: gopinathnelluri

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
 

tests/unit/test_spooling.py

@@ -0,0 +1,84 @@
+
+import unittest
+from unittest.mock import MagicMock, patch
+from trino.client import TrinoQuery, TrinoRequest, ClientSession, TrinoResult
+from trino.client import SegmentIterator
+
+class TestTrinoQueryLazy(unittest.TestCase):
+    def setUp(self):
+        self.mock_request = MagicMock(spec=TrinoRequest)
+        self.client_session = ClientSession("user")
+        self.mock_request.client_session = self.client_session
+
+    def test_fetch_returns_iterator_for_spooled_segments(self):
+        # Mock the initial POST response
+        post_response = MagicMock()
+        post_response.id = "query_1"
+        post_response.stats = {}
+        post_response.info_uri = "info"
+        post_response.next_uri = "next_1"
+        post_response.rows = [] # No rows initially
+        
+        self.mock_request.process.return_value = post_response
+        self.mock_request.post.return_value = MagicMock()
+        
+        query = TrinoQuery(self.mock_request, "SELECT 1")
+        
+        # Execute should return empty result initially but try to fetch
+        # We need to mock fetch behavior too since execute calls it if rows are empty
+        
+        # Mock the GET response for fetch()
+        get_response_status = MagicMock()
+        get_response_status.next_uri = None # Finished
+        get_response_status.stats = {}
+        # Status rows as dict indicates spooling protocol
+        get_response_status.rows = {
+            "encoding": "json",
+            "segments": [
+                {"type": "spooled", "uri": "u1", "ackUri": "a1", "metadata": {"segmentSize": "10", "uncompressedSize": "10"}}
+            ],
+            "metadata": {}
+        }
+        
+        # When execute calls fetch(), it calls request.get -> process -> returns get_response_status
+        self.mock_request.process.side_effect = [post_response, get_response_status]
+        self.mock_request.get.return_value = MagicMock()
+        
+        # Mock _to_segments to return a list of decodable segments
+        # We can just verify that fetch returns a SegmentIterator
+        # But _to_segments is internal.
+        
+        # We need to patch SegmentIterator or check the return type
+        
+        result = query.execute()
+        
+        # Verify result.rows is a SegmentIterator, NOT a list
+        self.assertIsInstance(result.rows, SegmentIterator)
+        self.assertNotIsInstance(result.rows, list)
+        
+    def test_fetch_returns_list_for_normal_segments(self):
+         # Mock the initial POST response
+        post_response = MagicMock()
+        post_response.id = "query_1"
+        post_response.stats = {}
+        post_response.info_uri = "info"
+        post_response.next_uri = "next_1"
+        post_response.rows = [] 
+        
+        # Mock the GET response for fetch()
+        get_response_status = MagicMock()
+        get_response_status.next_uri = None
+        get_response_status.stats = {}
+        get_response_status.rows = [[1], [2]] # Normal list rows
+        
+        self.mock_request.process.side_effect = [post_response, get_response_status]
+        
+        query = TrinoQuery(self.mock_request, "SELECT 1")
+        result = query.execute()
+        
+        # Verify result.rows is a list (appended)
+        self.assertIsInstance(result.rows, list)
+        self.assertEqual(result.rows, [[1], [2]])
+
+if __name__ == '__main__':
+    unittest.main()

trino/client.py

@@ -904,9 +904,32 @@ 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
+        
+        For Standard Execution, rows is a list, we can check len. the first response usually contains no rows (just stats),
+        so we need to continue fetching until we get some rows or query is finished or cancelled.
+        
+        For Spooled Execution, rows start as empty list and eventually fetch returns the rows as iterator, 
+        we can't check len 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):
+                     self._result.rows += new_rows
+                 else:
+                     # It's an iterator (spooled segments), replace rows with it
+                     self._result.rows = new_rows
+                     # We have an iterator now, so we can return result to user
+                     break
+            else:
+                 # We have data (list with items or an iterator), so return
+                 break
+                 
         return self._result
 
     def _update_state(self, status):
@@ -920,7 +943,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 +964,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: