Merge pull request #676 from ydb-platform/result_set_index

Add index to resultsets

Author: vgvoleg

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;")

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):

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