Add return_not_null_data_as_optional parameter to read_table

CHANGELOG.md

@@ -1,4 +1,5 @@
 * Fixed `OSError: [Errno 22] Invalid argument` when reading native `Datetime64` value before 1970 (negative Unix timestamps) on Windows: conversion now uses epoch arithmetic instead of `datetime.utcfromtimestamp`
+* Add `return_not_null_data_as_optional` parameter to `read_table` (sync and async) — allows reading `NOT NULL` columns as non-optional types, matching the existing gRPC field and the C++ SDK's `ReturnNotNullAsOptional`. By default `read_table` still wraps `NOT NULL` columns in `Optional` for backward compatibility
 
 ## 3.31.4 ##
 * Fixed async `QuerySessionPool` permanently losing a pool slot when `acquire()` was cancelled while a new session was being created: `asyncio.CancelledError` no longer leaks the pool size counter, so a pool under deadline-driven cancellations can no longer end up exhausted and blocking forever. A cancelled or interrupted session attach now also closes the session instead of orphaning it server-side

ydb/_session_impl.py

@@ -353,6 +353,7 @@ def read_table_request_factory(
     ordered=False,
     row_limit=None,
     use_snapshot=None,
+    return_not_null_data_as_optional=None,
 ):
     request = _apis.ydb_table.ReadTableRequest()
     request.path = path
@@ -383,6 +384,14 @@ def read_table_request_factory(
                 request.use_snapshot = _apis.FeatureFlag.DISABLED
         else:
             request.use_snapshot = use_snapshot
+    if return_not_null_data_as_optional is not None:
+        if isinstance(return_not_null_data_as_optional, bool):
+            if return_not_null_data_as_optional:
+                request.return_not_null_data_as_optional = _apis.FeatureFlag.ENABLED
+            else:
+                request.return_not_null_data_as_optional = _apis.FeatureFlag.DISABLED
+        else:
+            request.return_not_null_data_as_optional = return_not_null_data_as_optional
     return session_state.attach_request(request)
 
 

ydb/aio/table.py

@@ -44,6 +44,7 @@ async def read_table(
         row_limit=None,
         settings=None,
         use_snapshot=None,
+        return_not_null_data_as_optional=None,
     ):  # pylint: disable=W0236
         request = _session_impl.read_table_request_factory(
             self._state,
@@ -53,6 +54,7 @@ async def read_table(
             ordered,
             row_limit,
             use_snapshot=use_snapshot,
+            return_not_null_data_as_optional=return_not_null_data_as_optional,
         )
         stream_it = await self._driver(
             request,

ydb/table.py

@@ -1053,6 +1053,7 @@ def read_table(
         row_limit=None,
         settings=None,
         use_snapshot=None,
+        return_not_null_data_as_optional=None,
     ):
         """
         Perform an read table request.
@@ -1772,6 +1773,7 @@ def read_table(
         row_limit=None,
         settings=None,
         use_snapshot=None,
+        return_not_null_data_as_optional=None,
     ):
         """
         Perform an read table request.
@@ -1794,6 +1796,7 @@ def read_table(
             ordered,
             row_limit,
             use_snapshot=use_snapshot,
+            return_not_null_data_as_optional=return_not_null_data_as_optional,
         )
         stream_it = self._driver(
             request,
@@ -2032,6 +2035,7 @@ def async_read_table(
         row_limit=None,
         settings=None,
         use_snapshot=None,
+        return_not_null_data_as_optional=None,
     ):
         """
         Perform an read table request.
@@ -2056,6 +2060,7 @@ def async_read_table(
             ordered,
             row_limit,
             use_snapshot=use_snapshot,
+            return_not_null_data_as_optional=return_not_null_data_as_optional,
         )
         stream_it = self._driver(
             request,

ydb/table_test.py

@@ -4,7 +4,7 @@
 
 from unittest import mock
 from . import issues, convert, types, _apis, scheme, _session_impl
-from .table import SystemViewSchemeEntry, TableClient
+from .table import SystemViewSchemeEntry, TableClient, TableClientSettings
 
 from .retries import (
     retry_operation_impl,
@@ -336,3 +336,38 @@ def future(self, request, stub, method, wrap_fn, settings, wrap_args, *rest):
     assert driver.request.path == "/local/.sys/partition_stats"
     assert isinstance(entry, SystemViewSchemeEntry)
     assert entry.sys_view_name == "partition_stats"
+
+
+def _read_table_session_state():
+    state = _session_impl.SessionState(TableClientSettings())
+    state.set_id("test-session-id")
+    return state
+
+
+def test_read_table_request_not_null_as_optional_enabled():
+    request = _session_impl.read_table_request_factory(
+        _read_table_session_state(), "/local/table", return_not_null_data_as_optional=True
+    )
+    assert request.return_not_null_data_as_optional == _apis.FeatureFlag.ENABLED
+
+
+def test_read_table_request_not_null_as_optional_disabled():
+    request = _session_impl.read_table_request_factory(
+        _read_table_session_state(), "/local/table", return_not_null_data_as_optional=False
+    )
+    assert request.return_not_null_data_as_optional == _apis.FeatureFlag.DISABLED
+
+
+def test_read_table_request_not_null_as_optional_unset_by_default():
+    request = _session_impl.read_table_request_factory(_read_table_session_state(), "/local/table")
+    assert request.return_not_null_data_as_optional == _apis.FeatureFlag.STATUS_UNSPECIFIED
+
+
+def test_read_table_request_not_null_as_optional_raw_status():
+    # A raw FeatureFlag.Status (not a bool) is passed through unchanged.
+    request = _session_impl.read_table_request_factory(
+        _read_table_session_state(),
+        "/local/table",
+        return_not_null_data_as_optional=_apis.FeatureFlag.DISABLED,
+    )
+    assert request.return_not_null_data_as_optional == _apis.FeatureFlag.DISABLED