diff --git a/CHANGELOG.md b/CHANGELOG.md index d539cdae1..4eb1f834d 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: 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/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py index 152972031..7af2f8962 100644 --- a/tests/aio/query/test_query_session_pool.py +++ b/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;") diff --git a/tests/query/test_query_session_pool.py b/tests/query/test_query_session_pool.py index 413d85ebf..499febe6b 100644 --- a/tests/query/test_query_session_pool.py +++ b/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;") diff --git a/ydb/aio/query/pool.py b/ydb/aio/query/pool.py index 4e7529926..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 [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 87b25ee4c..fb760d279 100644 --- a/ydb/convert.py +++ b/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",) diff --git a/ydb/query/pool.py b/ydb/query/pool.py index b81b140e9..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 [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 a7423400d..73653a731 100644 --- a/ydb/query/pool_test.py +++ b/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)