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/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

@@ -905,8 +905,39 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
         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()
+        # If rows is a list (standard execution), we can check len.
+        # If rows is an iterator (spooled), we can't check len easily without peeking.
+        # However, for standard execution, the first response usually contains no rows (just stats),
+        # so we need to fetch.
+        
+        # If it's a list and empty, or if we haven't finished and haven't cancelled, try to fetch more.
+        # The issue with spooled execution is that fetch() returns an iterator.
+        
+        # For standard execution: rows starts as list. fetch() returns list.
+        # For spooled execution: rows starts as list (empty). fetch() returns Iterator.
+        
+        # We need to detect if we have data. 
+        # For spooled, we might get an iterator that *yields* nothing if segments are empty?
+        # But we want to block until we have *something* to return or finished.
+        
+        # Modified logic:
+        # If _result.rows is empty list, we fetch.
+        # If _result.rows key becomes an iterator, we stop blocking and let the user 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 +951,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 +972,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: