Merge streamed result set parts in execute_with_retries

Stream parts that share a result_set_index are concatenated back into a
single ResultSet, so execute_with_retries returns one result set per
SELECT regardless of its size instead of one per stream part.

Author: vgvoleg

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Fixed `execute_with_retries` returning a large result set split across several `ResultSet` objects: stream parts that share a `result_set_index` are now concatenated back into a single result set, so a query is returned as one result set per SELECT regardless of its size
 * 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,15 @@ 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_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,14 @@ 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_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 convert.aggregate_result_sets_by_index([result_set async for result_set in it])
 
         return await retry_operation_async(wrapped_callee, retry_settings)
 

ydb/convert.py

@@ -460,6 +460,36 @@ def lazy_from_message(cls, message, table_client_settings=None, snapshot=None):
 ResultSet = _ResultSet
 
 
+def aggregate_result_sets_by_index(result_sets):
+    """Glue together stream parts that belong to the same result set.
+
+    The query service streams one logical result set as several response parts
+    that share a single ``result_set_index``. This concatenates the rows (and
+    arrow ``data``) of those parts back into a single result set, keeping the
+    schema from the first part that carries it.
+    """
+    merged = []
+    by_index = {}
+    for result_set in result_sets:
+        index = result_set.index
+        target = by_index.get(index) if index is not None else None
+        if target is None:
+            merged.append(result_set)
+            if index is not None:
+                by_index[index] = result_set
+            continue
+
+        target.rows.extend(result_set.rows)
+        if result_set.truncated:
+            target.truncated = True
+        if not target.columns and result_set.columns:
+            target.columns = result_set.columns
+        if result_set.data is not None:
+            target.data = result_set.data if target.data is None else target.data + result_set.data
+
+    return merged
+
+
 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([result_set for result_set in it])
 
         return retry_operation_sync(wrapped_callee, retry_settings)
 

ydb/query/pool_test.py

@@ -6,6 +6,7 @@
 from unittest.mock import MagicMock
 
 from ydb import issues
+from ydb.convert import _ResultSet, aggregate_result_sets_by_index
 from ydb.query.pool import QuerySessionPool
 from ydb.query.session import QuerySession
 
@@ -57,6 +58,56 @@ 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_data_is_concatenated(self):
+        merged = aggregate_result_sets_by_index([_rs(0, [], data=b"aa"), _rs(0, [], data=b"bb")])
+
+        self.assertEqual(merged[0].data, b"aabb")
+
+    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_empty_input_returns_empty_list(self):
+        self.assertEqual(aggregate_result_sets_by_index([]), [])
+
+
 class TestRetryOperationSync(unittest.TestCase):
     def test_retry_reacquires_invalidated_session_before_first_use(self):
         pool = _make_pool(size=1)