Merge streamed result set parts in execute_with_retries

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Fixed `execute_with_retries` returning a large result set split across several `ResultSet` objects: row-format stream parts that share a `result_set_index` are now concatenated back into a single result set, so a large SELECT is returned as one result set instead of one per stream chunk
 * Introduce `ydb.observability` — vendor-neutral tracing entrypoint with a `TracingProvider` interface; `enable_tracing(provider)` accepts any implementation (custom or OpenTelemetry) and replaces the previously installed one. The SDK core no longer imports `opentelemetry`, so tracing can be enabled without the OpenTelemetry packages by supplying a custom provider
 * When tracing is enabled the SDK appends a `ydb-sdk-tracing/0.1.0` token to the `x-ydb-sdk-build-info` header, so the server can distinguish requests from tracing-enabled clients
 * Add client-side metrics through the same vendor-neutral `ydb.observability` layer: `enable_metrics(provider)` / `disable_metrics()`, with `ydb.opentelemetry.enable_metrics(meter_provider=None)` as the OpenTelemetry convenience. Instruments cover client operation duration/failures, retry duration/attempts, and query session pool state. `QuerySessionPool` gains an optional `name` argument for the pool metric label. When metrics are enabled the SDK appends a `ydb-sdk-metrics/0.1.0` token to the `x-ydb-sdk-build-info` header

tests/aio/query/test_query_session_pool.py

@@ -23,6 +23,38 @@ async def test_oneshot_query_normal(self, pool: QuerySessionPool):
         res = await pool.execute_with_retries("select 1;")
         assert len(res) == 1
 
+    @pytest.mark.asyncio
+    async def test_oneshot_query_merges_large_result_set_parts(self, pool: QuerySessionPool):
+        row_count = 100_000
+        query = f"SELECT * FROM AS_TABLE(ListMap(ListFromRange(0ul, {row_count}ul), ($x) -> (<|id: $x|>)))"
+        res = await pool.execute_with_retries(query)
+        assert len(res) == 1
+        assert res[0].index == 0
+        assert len(res[0].rows) == row_count
+
+    @pytest.mark.asyncio
+    async def test_oneshot_query_arrow_parts_are_intact(self, pool: QuerySessionPool):
+        row_count = 100_000
+        query = f"SELECT * FROM AS_TABLE(ListMap(ListFromRange(0ul, {row_count}ul), ($x) -> (<|id: $x|>)))"
+        res = await pool.execute_with_retries(query, result_set_format=ydb.QueryResultSetFormat.ARROW)
+
+        # Arrow parts are passed through untouched (never merged): each result
+        # set is one self-describing record batch for result set index 0.
+        assert len(res) >= 1
+        for result_set in res:
+            assert result_set.index == 0
+            assert result_set.data
+            assert result_set.arrow_format_meta.schema
+            assert not result_set.rows
+
+        pa = pytest.importorskip("pyarrow")
+        total = 0
+        for result_set in res:
+            schema = pa.ipc.read_schema(pa.py_buffer(result_set.arrow_format_meta.schema))
+            batch = pa.ipc.read_record_batch(pa.py_buffer(result_set.data), schema)
+            total += batch.num_rows
+        assert total == row_count
+
     @pytest.mark.asyncio
     async def test_oneshot_ddl_query(self, pool: QuerySessionPool):
         await pool.execute_with_retries("drop table if exists Queen;")

tests/query/test_query_session_pool.py

@@ -28,6 +28,36 @@ def test_oneshot_query_result_set_index(self, pool: QuerySessionPool):
         indexes = [result_set.index for result_set in res]
         assert indexes == [0, 1, 2]
 
+    def test_oneshot_query_merges_large_result_set_parts(self, pool: QuerySessionPool):
+        row_count = 100_000
+        query = f"SELECT * FROM AS_TABLE(ListMap(ListFromRange(0ul, {row_count}ul), ($x) -> (<|id: $x|>)))"
+        res = pool.execute_with_retries(query)
+        assert len(res) == 1
+        assert res[0].index == 0
+        assert len(res[0].rows) == row_count
+
+    def test_oneshot_query_arrow_parts_are_intact(self, pool: QuerySessionPool):
+        row_count = 100_000
+        query = f"SELECT * FROM AS_TABLE(ListMap(ListFromRange(0ul, {row_count}ul), ($x) -> (<|id: $x|>)))"
+        res = pool.execute_with_retries(query, result_set_format=ydb.QueryResultSetFormat.ARROW)
+
+        # Arrow parts are passed through untouched (never merged): each result
+        # set is one self-describing record batch for result set index 0.
+        assert len(res) >= 1
+        for result_set in res:
+            assert result_set.index == 0
+            assert result_set.data
+            assert result_set.arrow_format_meta.schema
+            assert not result_set.rows
+
+        pa = pytest.importorskip("pyarrow")
+        total = 0
+        for result_set in res:
+            schema = pa.ipc.read_schema(pa.py_buffer(result_set.arrow_format_meta.schema))
+            batch = pa.ipc.read_record_batch(pa.py_buffer(result_set.data), schema)
+            total += batch.num_rows
+        assert total == row_count
+
     def test_oneshot_ddl_query(self, pool: QuerySessionPool):
         pool.execute_with_retries("create table Queen(key UInt64, PRIMARY KEY (key));")
         pool.execute_with_retries("drop table Queen;")

ydb/aio/query/pool.py

@@ -237,7 +237,7 @@ async def execute_with_retries(
         async def wrapped_callee():
             async with self.checkout(timeout=retry_settings.max_session_acquire_timeout) as session:
                 it = await session.execute(query, parameters, *args, **kwargs)
-                return [result_set async for result_set in it]
+                return await convert.aggregate_result_sets_by_index_async(it)
 
         return await retry_operation_async(wrapped_callee, retry_settings)
 

ydb/convert.py

@@ -395,6 +395,20 @@ def __init__(
         self.arrow_format_meta = arrow_format_meta
         self.data = data
 
+    def _extend(self, other):
+        """Merge another stream part of the same result set into this one.
+
+        The query service streams one logical VALUE result set as several
+        response parts sharing a ``result_set_index``; this concatenates
+        ``other``'s rows onto ``self`` and carries the schema over from the
+        first part that provides it.
+        """
+        self.rows.extend(other.rows)
+        if other.truncated:
+            self.truncated = True
+        if not self.columns and other.columns:
+            self.columns = other.columns
+
     @classmethod
     def from_message(cls, message, table_client_settings=None, snapshot=None, index=None):
         rows = []
@@ -460,6 +474,53 @@ def lazy_from_message(cls, message, table_client_settings=None, snapshot=None):
 ResultSet = _ResultSet
 
 
+class _ResultSetsAccumulator:
+    """Reassembles streamed result-set parts in a single pass.
+
+    The query service streams one logical result set as several response parts
+    that share a ``result_set_index``. Parts are fed one at a time via
+    :meth:`add`: VALUE parts sharing an index are concatenated into a single
+    result set, while ARROW parts — each already carrying its own schema and
+    record-batch ``data`` — are kept as separate, independently decodable
+    result sets.
+    """
+
+    __slots__ = ("result_sets", "_by_index")
+
+    def __init__(self):
+        self.result_sets = []
+        self._by_index = {}
+
+    def add(self, result_set):
+        index = result_set.index
+        if index is None or result_set.data is not None:
+            self.result_sets.append(result_set)
+            return
+
+        target = self._by_index.get(index)
+        if target is None:
+            self._by_index[index] = result_set
+            self.result_sets.append(result_set)
+        else:
+            target._extend(result_set)
+
+
+def aggregate_result_sets_by_index(result_sets):
+    """Merge a sync stream of result-set parts into one result set per index."""
+    accumulator = _ResultSetsAccumulator()
+    for result_set in result_sets:
+        accumulator.add(result_set)
+    return accumulator.result_sets
+
+
+async def aggregate_result_sets_by_index_async(result_sets):
+    """Merge an async stream of result-set parts into one result set per index."""
+    accumulator = _ResultSetsAccumulator()
+    async for result_set in result_sets:
+        accumulator.add(result_set)
+    return accumulator.result_sets
+
+
 class _Row(_DotDict):
     __slots__ = ("_columns",)
 

ydb/query/pool.py

@@ -264,7 +264,7 @@ def execute_with_retries(
         def wrapped_callee():
             with self.checkout(timeout=retry_settings.max_session_acquire_timeout) as session:
                 it = session.execute(query, parameters, *args, **kwargs)
-                return [result_set for result_set in it]
+                return convert.aggregate_result_sets_by_index(it)
 
         return retry_operation_sync(wrapped_callee, retry_settings)
 

ydb/query/pool_test.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import asyncio
 import threading
 import time
 import unittest
 from unittest.mock import MagicMock
 
 from ydb import issues
+from ydb.convert import _ResultSet, aggregate_result_sets_by_index, aggregate_result_sets_by_index_async
 from ydb.query.pool import QuerySessionPool
 from ydb.query.session import QuerySession
 
@@ -57,6 +59,77 @@ def release_after_delay():
             t.join()
 
 
+def _rs(index, rows, columns=None, truncated=False, data=None):
+    return _ResultSet(
+        columns=["id"] if columns is None else columns,
+        rows=list(rows),
+        truncated=truncated,
+        index=index,
+        data=data,
+    )
+
+
+class TestAggregateResultSetsByIndex(unittest.TestCase):
+    def test_merges_parts_with_same_index_into_one_result_set(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1, 2]), _rs(0, [3, 4]), _rs(0, [5])])
+
+        self.assertEqual(len(merged), 1)
+        self.assertEqual(merged[0].index, 0)
+        self.assertEqual(merged[0].rows, [1, 2, 3, 4, 5])
+
+    def test_keeps_distinct_indexes_separate_and_ordered(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1]), _rs(0, [2]), _rs(1, [3]), _rs(2, [4]), _rs(2, [5])])
+
+        self.assertEqual([rs.index for rs in merged], [0, 1, 2])
+        self.assertEqual([rs.rows for rs in merged], [[1, 2], [3], [4, 5]])
+
+    def test_schema_kept_from_first_part_when_later_parts_omit_it(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1], columns=["id", "name"]), _rs(0, [2], columns=[])])
+
+        self.assertEqual(merged[0].columns, ["id", "name"])
+        self.assertEqual(merged[0].rows, [1, 2])
+
+    def test_truncated_flag_is_propagated_from_any_part(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1], truncated=False), _rs(0, [2], truncated=True)])
+
+        self.assertTrue(merged[0].truncated)
+
+    def test_arrow_parts_are_not_merged(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [], data=b"aa"), _rs(0, [], data=b"bb")])
+
+        self.assertEqual([rs.data for rs in merged], [b"aa", b"bb"])
+
+    def test_interleaved_parts_are_merged_by_index(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1]), _rs(1, [2]), _rs(0, [3]), _rs(1, [4])])
+
+        self.assertEqual([rs.index for rs in merged], [0, 1])
+        self.assertEqual([rs.rows for rs in merged], [[1, 3], [2, 4]])
+
+    def test_schema_filled_from_later_part_when_first_omits_it(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [1], columns=[]), _rs(0, [2], columns=["id", "name"])])
+
+        self.assertEqual(merged[0].columns, ["id", "name"])
+        self.assertEqual(merged[0].rows, [1, 2])
+
+    def test_parts_without_index_are_not_merged(self):
+        merged = aggregate_result_sets_by_index([_rs(None, [1]), _rs(None, [2])])
+
+        self.assertEqual([rs.rows for rs in merged], [[1], [2]])
+
+    def test_empty_input_returns_empty_list(self):
+        self.assertEqual(aggregate_result_sets_by_index([]), [])
+
+    def test_async_stream_is_merged_in_one_pass(self):
+        async def stream():
+            for part in [_rs(0, [1]), _rs(0, [2]), _rs(1, [3]), _rs(1, [4])]:
+                yield part
+
+        merged = asyncio.run(aggregate_result_sets_by_index_async(stream()))
+
+        self.assertEqual([rs.index for rs in merged], [0, 1])
+        self.assertEqual([rs.rows for rs in merged], [[1, 2], [3, 4]])
+
+
 class TestRetryOperationSync(unittest.TestCase):
     def test_retry_reacquires_invalidated_session_before_first_use(self):
         pool = _make_pool(size=1)