Reduce memory usage of query service result sets

Reading large result sets via the query service kept far more memory
alive than the decoded data itself required.

The converted result set referenced the source protobuf message through
`_ResultSet.columns` and through `_Row._columns` on every row. With the
upb protobuf backend, any reference into a message pins its entire arena,
so the raw `message.rows` payload stayed alive for the whole lifetime of
the result set — the decoded Python objects and the full source protobuf
were held simultaneously.

Detach the result set from the source arena: copy the schema into a
standalone message and store a shared tuple of column names on rows
instead of the proto columns, so the source arena is released right
after conversion. Also copy the Arrow `data` payload out of the arena.

Additionally give the row/struct dict subclasses `__slots__` to drop the
redundant per-row instance `__dict__`, and make `_DotDict.__getattr__`
raise `AttributeError` instead of leaking `KeyError` (which had broken
`copy.copy`/`copy.deepcopy` of rows).

Peak RSS for a streaming read of 20x20k rows (results kept) drops from
~953 MB to ~298 MB.

Author: vgvoleg

ydb/convert.py

@@ -24,11 +24,16 @@ def _initialize():
 
 
 class _DotDict(dict):
+    __slots__ = ()
+
     def __init__(self, *args, **kwargs):
         super(_DotDict, self).__init__(*args, **kwargs)
 
     def __getattr__(self, item):
-        return self[item]
+        try:
+            return self[item]
+        except KeyError:
+            raise AttributeError(item)
 
 
 def _is_decimal_signed(hi_value):
@@ -83,7 +88,7 @@ def _pb_to_dict(type_pb, value_pb, table_client_settings):
 
 
 class _Struct(_DotDict):
-    pass
+    __slots__ = ()
 
 
 def _pb_to_struct(type_pb, value_pb, table_client_settings):
@@ -351,6 +356,16 @@ def _unwrap_optionality(column):
     return _to_native_map.get(current_type), c_type
 
 
+def _detach_columns(columns):
+    # The columns container references the source protobuf message, which keeps
+    # the whole arena (including already-parsed message.rows) alive for as long
+    # as the result set is held. Copy the schema into a standalone message so the
+    # source arena can be released right after conversion.
+    holder = _apis.ydb_value.ResultSet()
+    holder.columns.extend(columns)
+    return holder.columns
+
+
 class _ResultSet(object):
     __slots__ = ("columns", "rows", "truncated", "snapshot", "index", "format", "arrow_format_meta", "data")
 
@@ -375,20 +390,25 @@ def from_message(cls, message, table_client_settings=None, snapshot=None, index=
             for column in message.columns:
                 column_parsers.append(_unwrap_optionality(column))
 
+        # Names are the only per-row metadata we need. Storing this shared tuple
+        # (instead of the proto columns) keeps rows detached from the source
+        # protobuf arena, so it can be freed once conversion is done.
+        column_names = tuple(column.name for column in message.columns)
+
         for row_proto in message.rows:
-            row = _Row(message.columns)
-            for column, value, column_info in zip(message.columns, row_proto.items, column_parsers):
+            row = _Row(column_names)
+            for name, value, column_info in zip(column_names, row_proto.items, column_parsers):
                 v_type = value.WhichOneof("value")
                 if v_type == "null_flag_value":
-                    row[column.name] = None
+                    row[name] = None
                     continue
 
                 while v_type == "nested_value":
                     value = value.nested_value
                     v_type = value.WhichOneof("value")
 
                 column_parser, unwrapped_type = column_info
-                row[column.name] = column_parser(unwrapped_type, value, table_client_settings)
+                row[name] = column_parser(unwrapped_type, value, table_client_settings)
             rows.append(row)
 
         from ydb.query import QueryResultSetFormat, ArrowFormatMeta
@@ -399,9 +419,11 @@ def from_message(cls, message, table_client_settings=None, snapshot=None, index=
         if message.HasField("arrow_format_meta"):
             arrow_meta = ArrowFormatMeta.from_proto(message.arrow_format_meta)
 
-        data = message.data if message.data else None
+        data = bytes(message.data) if message.data else None
 
-        return cls(message.columns, rows, message.truncated, snapshot, index, result_format, arrow_meta, data)
+        return cls(
+            _detach_columns(message.columns), rows, message.truncated, snapshot, index, result_format, arrow_meta, data
+        )
 
     @classmethod
     def lazy_from_message(cls, message, table_client_settings=None, snapshot=None):
@@ -422,15 +444,17 @@ def lazy_from_message(cls, message, table_client_settings=None, snapshot=None):
 
 
 class _Row(_DotDict):
+    __slots__ = ("_columns",)
+
     def __init__(self, columns):
         super(_Row, self).__init__()
         self._columns = columns
 
     def __getitem__(self, key):
         if isinstance(key, int):
-            return self[self._columns[key].name]
+            return self[self._columns[key]]
         elif isinstance(key, slice):
-            return tuple(map(lambda x: self[x.name], self._columns[key]))
+            return tuple(self[name] for name in self._columns[key])
         else:
             return super(_Row, self).__getitem__(key)
 
@@ -455,6 +479,8 @@ def get(self):
 
 
 class _LazyRow(_DotDict):
+    __slots__ = ("_columns", "_table_client_settings")
+
     def __init__(self, columns, proto_row, table_client_settings, parsers):
         super(_LazyRow, self).__init__()
         self._columns = columns

ydb/table_test.py

@@ -1,5 +1,7 @@
+import copy
+
 from unittest import mock
-from . import issues
+from . import issues, convert, types, _apis
 
 from .retries import (
     retry_operation_impl,
@@ -9,6 +11,74 @@
 )
 
 
+def _build_int_result_set(n_rows, n_cols):
+    result_set = _apis.ydb_value.ResultSet()
+    for col_idx in range(n_cols):
+        column = result_set.columns.add()
+        column.name = "column_%d" % col_idx
+        column.type.type_id = types.PrimitiveType.Int64._idn_
+    for row_idx in range(n_rows):
+        row = result_set.rows.add()
+        for col_idx in range(n_cols):
+            row.items.add().int64_value = row_idx * 1000 + col_idx
+    return result_set
+
+
+def test_result_set_row_access():
+    message = _build_int_result_set(n_rows=1, n_cols=3)
+    row = convert.ResultSet.from_message(message).rows[0]
+
+    assert row["column_1"] == 1
+    assert row.column_1 == 1
+    assert row[1] == 1
+    assert row[0:2] == (0, 1)
+    assert dict(row) == {"column_0": 0, "column_1": 1, "column_2": 2}
+
+
+def test_result_set_row_has_no_instance_dict():
+    # Rows must not carry a per-instance __dict__: it is pure memory overhead
+    # multiplied by every row in a large result set.
+    message = _build_int_result_set(n_rows=1, n_cols=3)
+    row = convert.ResultSet.from_message(message).rows[0]
+    assert not hasattr(row, "__dict__")
+
+
+def test_result_set_row_missing_attribute_raises_attribute_error():
+    message = _build_int_result_set(n_rows=1, n_cols=1)
+    row = convert.ResultSet.from_message(message).rows[0]
+
+    assert not hasattr(row, "definitely_missing")
+    try:
+        row.definitely_missing
+    except AttributeError:
+        pass
+    else:
+        raise AssertionError("expected AttributeError for missing attribute")
+
+
+def test_result_set_row_is_copyable():
+    message = _build_int_result_set(n_rows=1, n_cols=3)
+    row = convert.ResultSet.from_message(message).rows[0]
+
+    assert dict(copy.copy(row)) == dict(row)
+    assert dict(copy.deepcopy(row)) == dict(row)
+
+
+def test_result_set_detached_from_source_message():
+    # The converted result set must not hold a reference into the source
+    # protobuf: otherwise the whole arena (raw rows included) stays alive for
+    # the lifetime of the result, doubling memory usage on large reads.
+    message = _build_int_result_set(n_rows=2, n_cols=2)
+    result = convert.ResultSet.from_message(message)
+
+    message.Clear()  # emulate the source proto being dropped/reused by the stream
+
+    assert [column.name for column in result.columns] == ["column_0", "column_1"]
+    assert result.rows[0]["column_1"] == 1
+    assert result.rows[1][0] == 1000
+    assert result.rows[1][0:2] == (1000, 1001)
+
+
 def test_retry_operation_impl(monkeypatch):
     monkeypatch.setattr("random.random", lambda: 0.5)
     monkeypatch.setattr(