From a2ca99eba175bf6d9275932c33bfbefb81c028cd Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 17 Jul 2026 17:03:27 +0300 Subject: [PATCH 1/7] 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. --- CHANGELOG.md | 1 + tests/aio/query/test_query_session_pool.py | 9 ++++ tests/query/test_query_session_pool.py | 8 ++++ ydb/aio/query/pool.py | 2 +- ydb/convert.py | 30 +++++++++++++ ydb/query/pool.py | 2 +- ydb/query/pool_test.py | 51 ++++++++++++++++++++++ 7 files changed, 101 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d539cdae1..03a6471bf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py index 152972031..78e4e4e55 100644 --- a/tests/aio/query/test_query_session_pool.py +++ b/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;") diff --git a/tests/query/test_query_session_pool.py b/tests/query/test_query_session_pool.py index 413d85ebf..736f1c211 100644 --- a/tests/query/test_query_session_pool.py +++ b/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;") diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index 4e7529926..e73563f2d 100644 --- a/ydb/aio/query/pool.py +++ b/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) diff --git a/ydb/convert.py b/ydb/convert.py index 87b25ee4c..33e0c783b 100644 --- a/ydb/convert.py +++ b/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",) diff --git a/ydb/query/pool.py b/ydb/query/pool.py index b81b140e9..858eca67e 100644 --- a/ydb/query/pool.py +++ b/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) diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index a7423400d..c34ae39d3 100644 --- a/ydb/query/pool_test.py +++ b/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) From bf1d392bf25675fe7eb8305798b38f1ba51c5198 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 17 Jul 2026 17:11:35 +0300 Subject: [PATCH 2/7] Cover remaining branches of aggregate_result_sets_by_index --- ydb/query/pool_test.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index c34ae39d3..658318379 100644 --- a/ydb/query/pool_test.py +++ b/ydb/query/pool_test.py @@ -104,6 +104,22 @@ def test_interleaved_parts_are_merged_by_index(self): 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_arrow_data_filled_from_later_part_when_first_has_none(self): + merged = aggregate_result_sets_by_index([_rs(0, [], data=None), _rs(0, [], data=b"bb")]) + + self.assertEqual(merged[0].data, b"bb") + + 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([]), []) From d864078d03ae32afe7a511bfb69fe9d8566abcec Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 17 Jul 2026 17:30:11 +0300 Subject: [PATCH 3/7] Join arrow data chunks once to avoid quadratic concatenation --- ydb/convert.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/ydb/convert.py b/ydb/convert.py index 33e0c783b..b53731a8f 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -470,6 +470,7 @@ def aggregate_result_sets_by_index(result_sets): """ merged = [] by_index = {} + data_chunks = {} for result_set in result_sets: index = result_set.index target = by_index.get(index) if index is not None else None @@ -477,6 +478,8 @@ def aggregate_result_sets_by_index(result_sets): merged.append(result_set) if index is not None: by_index[index] = result_set + if result_set.data is not None: + data_chunks[index] = [result_set.data] continue target.rows.extend(result_set.rows) @@ -485,7 +488,14 @@ def aggregate_result_sets_by_index(result_sets): 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 + data_chunks.setdefault(index, []).append(result_set.data) + + # Join arrow ``data`` chunks once per result set: the parts are schema-less, + # EOS-less record batches, so concatenating them under the shared schema + # rebuilds a valid IPC stream. Doing it in one pass keeps this linear — + # concatenating bytes on every part would be O(total_size^2). + for index, chunks in data_chunks.items(): + by_index[index].data = chunks[0] if len(chunks) == 1 else b"".join(chunks) return merged From 28bed353f4abf1e75dfe4726c640a58a2d9cf951 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 17 Jul 2026 17:35:21 +0300 Subject: [PATCH 4/7] Move stream-part merging onto _ResultSet._extend/_seal --- ydb/convert.py | 62 ++++++++++++++++++++++++++++-------------- ydb/query/pool_test.py | 6 ++-- 2 files changed, 46 insertions(+), 22 deletions(-) diff --git a/ydb/convert.py b/ydb/convert.py index b53731a8f..9ad324169 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -381,7 +381,17 @@ def _detach_columns(columns): class _ResultSet(object): - __slots__ = ("columns", "rows", "truncated", "snapshot", "index", "format", "arrow_format_meta", "data") + __slots__ = ( + "columns", + "rows", + "truncated", + "snapshot", + "index", + "format", + "arrow_format_meta", + "data", + "_data_chunks", + ) def __init__( self, columns, rows, truncated, snapshot=None, index=None, format=None, arrow_format_meta=None, data=None @@ -394,6 +404,32 @@ def __init__( self.format = format self.arrow_format_meta = arrow_format_meta self.data = data + self._data_chunks = None + + def _extend(self, other): + """Merge another stream part of the same result set into this one. + + The query service streams one logical result set as several response + parts sharing a ``result_set_index``; this appends ``other``'s payload + to ``self``. Arrow ``data`` chunks are accumulated in a list and joined + once by :meth:`_seal`, so merging many parts stays linear instead of + re-concatenating bytes on every part. + """ + self.rows.extend(other.rows) + if other.truncated: + self.truncated = True + if not self.columns and other.columns: + self.columns = other.columns + if other.data is not None: + if self._data_chunks is None: + self._data_chunks = [self.data] if self.data is not None else [] + self._data_chunks.append(other.data) + + def _seal(self): + """Collapse accumulated arrow ``data`` chunks into a single ``bytes``.""" + if self._data_chunks is not None: + self.data = b"".join(self._data_chunks) + self._data_chunks = None @classmethod def from_message(cls, message, table_client_settings=None, snapshot=None, index=None): @@ -470,7 +506,6 @@ def aggregate_result_sets_by_index(result_sets): """ merged = [] by_index = {} - data_chunks = {} for result_set in result_sets: index = result_set.index target = by_index.get(index) if index is not None else None @@ -478,24 +513,11 @@ def aggregate_result_sets_by_index(result_sets): merged.append(result_set) if index is not None: by_index[index] = result_set - if result_set.data is not None: - data_chunks[index] = [result_set.data] - 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: - data_chunks.setdefault(index, []).append(result_set.data) - - # Join arrow ``data`` chunks once per result set: the parts are schema-less, - # EOS-less record batches, so concatenating them under the shared schema - # rebuilds a valid IPC stream. Doing it in one pass keeps this linear — - # concatenating bytes on every part would be O(total_size^2). - for index, chunks in data_chunks.items(): - by_index[index].data = chunks[0] if len(chunks) == 1 else b"".join(chunks) + else: + target._extend(result_set) + + for result_set in merged: + result_set._seal() return merged diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index 658318379..2744e2ab5 100644 --- a/ydb/query/pool_test.py +++ b/ydb/query/pool_test.py @@ -94,9 +94,11 @@ def test_truncated_flag_is_propagated_from_any_part(self): 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")]) + merged = aggregate_result_sets_by_index( + [_rs(0, [], data=b"aa"), _rs(0, [], data=b"bb"), _rs(0, [], data=b"cc")] + ) - self.assertEqual(merged[0].data, b"aabb") + self.assertEqual(merged[0].data, b"aabbcc") 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])]) From 3df06664b00c22e2bc329f4b5f295af236adcdc2 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Fri, 17 Jul 2026 17:54:01 +0300 Subject: [PATCH 5/7] Merge only VALUE result-set parts, leave Arrow parts untouched --- CHANGELOG.md | 2 +- ydb/convert.py | 51 +++++++++++++----------------------------- ydb/query/pool_test.py | 13 +++-------- 3 files changed, 19 insertions(+), 47 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 03a6471bf..4eb1f834d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +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 +* 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 diff --git a/ydb/convert.py b/ydb/convert.py index 9ad324169..271e7cc60 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -381,17 +381,7 @@ def _detach_columns(columns): class _ResultSet(object): - __slots__ = ( - "columns", - "rows", - "truncated", - "snapshot", - "index", - "format", - "arrow_format_meta", - "data", - "_data_chunks", - ) + __slots__ = ("columns", "rows", "truncated", "snapshot", "index", "format", "arrow_format_meta", "data") def __init__( self, columns, rows, truncated, snapshot=None, index=None, format=None, arrow_format_meta=None, data=None @@ -404,32 +394,20 @@ def __init__( self.format = format self.arrow_format_meta = arrow_format_meta self.data = data - self._data_chunks = None def _extend(self, other): """Merge another stream part of the same result set into this one. - The query service streams one logical result set as several response - parts sharing a ``result_set_index``; this appends ``other``'s payload - to ``self``. Arrow ``data`` chunks are accumulated in a list and joined - once by :meth:`_seal`, so merging many parts stays linear instead of - re-concatenating bytes on every part. + 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 - if other.data is not None: - if self._data_chunks is None: - self._data_chunks = [self.data] if self.data is not None else [] - self._data_chunks.append(other.data) - - def _seal(self): - """Collapse accumulated arrow ``data`` chunks into a single ``bytes``.""" - if self._data_chunks is not None: - self.data = b"".join(self._data_chunks) - self._data_chunks = None @classmethod def from_message(cls, message, table_client_settings=None, snapshot=None, index=None): @@ -500,25 +478,26 @@ 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. + that share a single ``result_set_index``. VALUE parts are concatenated back + into one result set. ARROW parts are left untouched: each already carries + its own schema and record-batch ``data`` and is independently decodable, so + they stay as separate result sets instead of splicing opaque bytes together. """ 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 index is None or result_set.data is not None: + merged.append(result_set) + continue + + target = by_index.get(index) if target is None: + by_index[index] = result_set merged.append(result_set) - if index is not None: - by_index[index] = result_set else: target._extend(result_set) - for result_set in merged: - result_set._seal() - return merged diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index 2744e2ab5..e42e6cc39 100644 --- a/ydb/query/pool_test.py +++ b/ydb/query/pool_test.py @@ -93,12 +93,10 @@ def test_truncated_flag_is_propagated_from_any_part(self): 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"), _rs(0, [], data=b"cc")] - ) + 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(merged[0].data, b"aabbcc") + 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])]) @@ -112,11 +110,6 @@ def test_schema_filled_from_later_part_when_first_omits_it(self): self.assertEqual(merged[0].columns, ["id", "name"]) self.assertEqual(merged[0].rows, [1, 2]) - def test_arrow_data_filled_from_later_part_when_first_has_none(self): - merged = aggregate_result_sets_by_index([_rs(0, [], data=None), _rs(0, [], data=b"bb")]) - - self.assertEqual(merged[0].data, b"bb") - def test_parts_without_index_are_not_merged(self): merged = aggregate_result_sets_by_index([_rs(None, [1]), _rs(None, [2])]) From 21c0d102f9f7aff324ab72297982b23ae5e1095b Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 20 Jul 2026 11:31:35 +0300 Subject: [PATCH 6/7] Aggregate result-set stream in one pass without an intermediate list --- ydb/aio/query/pool.py | 2 +- ydb/convert.py | 50 +++++++++++++++++++++++++++++------------- ydb/query/pool.py | 2 +- ydb/query/pool_test.py | 13 ++++++++++- 4 files changed, 49 insertions(+), 18 deletions(-) diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index e73563f2d..7a71c5f73 100644 --- a/ydb/aio/query/pool.py +++ b/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 convert.aggregate_result_sets_by_index([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) diff --git a/ydb/convert.py b/ydb/convert.py index 271e7cc60..fb760d279 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -474,31 +474,51 @@ 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. +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 single ``result_set_index``. VALUE parts are concatenated back - into one result set. ARROW parts are left untouched: each already carries - its own schema and record-batch ``data`` and is independently decodable, so - they stay as separate result sets instead of splicing opaque bytes together. + 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. """ - merged = [] - by_index = {} - for result_set in 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: - merged.append(result_set) - continue + self.result_sets.append(result_set) + return - target = by_index.get(index) + target = self._by_index.get(index) if target is None: - by_index[index] = result_set - merged.append(result_set) + self._by_index[index] = result_set + self.result_sets.append(result_set) else: target._extend(result_set) - return merged + +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): diff --git a/ydb/query/pool.py b/ydb/query/pool.py index 858eca67e..869884102 100644 --- a/ydb/query/pool.py +++ b/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 convert.aggregate_result_sets_by_index([result_set for result_set in it]) + return convert.aggregate_result_sets_by_index(it) return retry_operation_sync(wrapped_callee, retry_settings) diff --git a/ydb/query/pool_test.py b/ydb/query/pool_test.py index e42e6cc39..73653a731 100644 --- a/ydb/query/pool_test.py +++ b/ydb/query/pool_test.py @@ -1,12 +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 +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 @@ -118,6 +119,16 @@ def test_parts_without_index_are_not_merged(self): 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): From bace7024182d1920263a7ca5ce3920bb66289f96 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 20 Jul 2026 12:47:56 +0300 Subject: [PATCH 7/7] Add integration test for Arrow-format execute_with_retries --- tests/aio/query/test_query_session_pool.py | 23 ++++++++++++++++++++++ tests/query/test_query_session_pool.py | 22 +++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py index 78e4e4e55..7af2f8962 100644 --- a/tests/aio/query/test_query_session_pool.py +++ b/tests/aio/query/test_query_session_pool.py @@ -32,6 +32,29 @@ async def test_oneshot_query_merges_large_result_set_parts(self, pool: QuerySess 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;") diff --git a/tests/query/test_query_session_pool.py b/tests/query/test_query_session_pool.py index 736f1c211..499febe6b 100644 --- a/tests/query/test_query_session_pool.py +++ b/tests/query/test_query_session_pool.py @@ -36,6 +36,28 @@ def test_oneshot_query_merges_large_result_set_parts(self, pool: QuerySessionPoo 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;")