From f948e45530f0107422cbc317a92776b8d1e47fc1 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Sat, 14 Jun 2025 20:20:02 +0300 Subject: [PATCH] Add index to resultsets --- tests/query/test_query_session_pool.py | 6 ++++++ ydb/convert.py | 9 +++++---- ydb/query/base.py | 6 +++++- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/query/test_query_session_pool.py b/tests/query/test_query_session_pool.py index 4c88ae77..6b5a7bf4 100644 --- a/tests/query/test_query_session_pool.py +++ b/tests/query/test_query_session_pool.py @@ -17,6 +17,12 @@ def test_oneshot_query_normal(self, pool: QuerySessionPool): res = pool.execute_with_retries("select 1;") assert len(res) == 1 + def test_oneshot_query_result_set_index(self, pool: QuerySessionPool): + res = pool.execute_with_retries("select 1; select 2; select 3") + assert len(res) == 3 + indexes = [result_set.index for result_set in res] + assert indexes == [0, 1, 2] + 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/convert.py b/ydb/convert.py index f48e024a..485b5897 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -352,16 +352,17 @@ def _unwrap_optionality(column): class _ResultSet(object): - __slots__ = ("columns", "rows", "truncated", "snapshot") + __slots__ = ("columns", "rows", "truncated", "snapshot", "index") - def __init__(self, columns, rows, truncated, snapshot=None): + def __init__(self, columns, rows, truncated, snapshot=None, index=None): self.columns = columns self.rows = rows self.truncated = truncated self.snapshot = snapshot + self.index = index @classmethod - def from_message(cls, message, table_client_settings=None, snapshot=None): + def from_message(cls, message, table_client_settings=None, snapshot=None, index=None): rows = [] # prepare column parsers before actuall parsing column_parsers = [] @@ -384,7 +385,7 @@ def from_message(cls, message, table_client_settings=None, snapshot=None): column_parser, unwrapped_type = column_info row[column.name] = column_parser(unwrapped_type, value, table_client_settings) rows.append(row) - return cls(message.columns, rows, message.truncated, snapshot) + return cls(message.columns, rows, message.truncated, snapshot, index) @classmethod def lazy_from_message(cls, message, table_client_settings=None, snapshot=None): diff --git a/ydb/query/base.py b/ydb/query/base.py index 52a6312e..2c16716c 100644 --- a/ydb/query/base.py +++ b/ydb/query/base.py @@ -208,7 +208,11 @@ def wrap_execute_query_response( session._last_query_stats = response_pb.exec_stats if response_pb.HasField("result_set"): - return convert.ResultSet.from_message(response_pb.result_set, settings) + return convert.ResultSet.from_message( + response_pb.result_set, + settings, + index=response_pb.result_set_index, + ) return None