Skip to content

Commit d0170d0

Browse files
committed
Return a single element while iterating over segments
This makes the usage of the segment cursor more natural: ``` cur = conn.cursor('segment') cur.execute(sql) spooled_result = cur.fetchall() for spooled_data in spooled_result: row_mapper = RowMapperFactory().create(columns=cur._query.columns, legacy_primitive_types=False) rows = list(SegmentIterator(spooled_data, row_mapper)) ```
1 parent 3c745f7 commit d0170d0

File tree

3 files changed

+14
-7
lines changed

3 files changed

+14
-7
lines changed

tests/integration/test_dbapi_integration.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1876,12 +1876,18 @@ def test_segments_cursor(trino_connection):
18761876
start => 1,
18771877
stop => 5,
18781878
step => 1)) n""")
1879-
rows = cur.fetchall()
1880-
assert len(rows) > 0
1881-
for spooled_data, spooled_segment in rows:
1879+
segments = cur.fetchall()
1880+
assert len(segments) > 0
1881+
row_mapper = trino.mapper.RowMapperFactory().create(columns=cur._query.columns, legacy_primitive_types=False)
1882+
total = 0
1883+
for spooled_data in segments:
1884+
assert len(spooled_data.segments) == 1, "Expected SpooledData to contain a single segment"
1885+
segment = spooled_data.segments[0]
18821886
assert spooled_data.encoding == trino_connection._client_session.encoding
1883-
assert isinstance(spooled_segment.uri, str), f"Expected string for uri, got {spooled_segment.uri}"
1884-
assert isinstance(spooled_segment.ack_uri, str), f"Expected string for ack_uri, got {spooled_segment.ack_uri}"
1887+
assert isinstance(segment[0].uri, str), f"Expected string for uri, got {segment[0].uri}"
1888+
assert isinstance(segment[0].ack_uri, str), f"Expected string for ack_uri, got {segment[0].ack_uri}"
1889+
total += len(list(trino.SegmentIterator(spooled_data, row_mapper)))
1890+
assert total == 300875, f"Expected total rows 300875, got {total}"
18851891

18861892

18871893
def get_cursor(legacy_prepared_statements, run_trino):

trino/client.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1187,7 +1187,8 @@ def __iter__(self) -> Iterator[Tuple["SpooledData", "Segment"]]:
11871187
return self
11881188

11891189
def __next__(self) -> Tuple["SpooledData", "Segment"]:
1190-
return self, next(self._segments_iterator)
1190+
segment = next(self._segments_iterator)
1191+
return SpooledData(self._encoding, [segment]), segment
11911192

11921193
def __repr__(self):
11931194
return (f"SpooledData(encoding={self._encoding}, segments={list(self._segments)})")

trino/dbapi.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -742,7 +742,7 @@ def execute(self, operation, params=None):
742742
self._query = trino.client.TrinoQuery(self._request, query=operation,
743743
legacy_primitive_types=self._legacy_primitive_types,
744744
fetch_mode="segments")
745-
self._iterator = iter(self._query.execute())
745+
self._iterator = map(lambda tuple: tuple[0], iter(self._query.execute()))
746746
return self
747747

748748

0 commit comments

Comments
 (0)