fix crash when read keys without columns

2 files changed, 52 insertions, 7 deletions

diff --git a/ydb/core/tx/datashard/datashard__read_iterator.cpp b/ydb/core/tx/datashard/datashard__read_iterator.cpp
index 5700ad9fdb..3311680d9f 100644
--- a/ydb/core/tx/datashard/datashard__read_iterator.cpp
+++ b/ydb/core/tx/datashard/datashard__read_iterator.cpp
@@ -395,12 +395,17 @@ public:
             return EReadStatus::Done;
         }
 
-        // TODO: looks kind of ugly: we assume that cells in rowState are stored in array
-        TDbTupleRef value(&ColumnTypes[0], &rowState.Get(0), ColumnTypes.size());
+        if (ColumnTypes.size()) {
+            // TODO: looks kind of ugly: we assume that cells in rowState are stored in array
+            TDbTupleRef value(&ColumnTypes[0], &rowState.Get(0), ColumnTypes.size());
+
+            // note that if user requests key columns then they will be in
+            // rowValues and we don't have to add rowKey columns
+            BlockBuilder.AddRow(TDbTupleRef(), value);
+        } else {
+            BlockBuilder.AddRow(TDbTupleRef(), TDbTupleRef());
+        }
 
-        // note that if user requests key columns then they will be in
-        // rowValues and we don't have to add rowKey columns
-        BlockBuilder.AddRow(TDbTupleRef(), value);
         ++RowsRead;
 
         return EReadStatus::Done;
diff --git a/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp b/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
index 244a31ccda..e75678ea0e 100644
--- a/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
+++ b/ydb/core/tx/datashard/datashard_ut_read_iterator.cpp
@@ -913,7 +913,47 @@ Y_UNIT_TEST_SUITE(DataShardReadIterator) {
         }
     }
 
-    Y_UNIT_TEST(ShouldReadNoColumnsCellVec) {
+    Y_UNIT_TEST(ShouldReadNoColumnsKeysRequestCellVec) {
+        // KIKIMR-16897: no columns mean we want to calc row count
+        TTestHelper helper;
+
+        auto request = helper.GetBaseReadRequest("table-1", 1, NKikimrTxDataShard::CELLVEC);
+        request->Record.ClearColumns();
+        AddKeyQuery(*request, {3, 3, 3});
+        AddKeyQuery(*request, {1, 1, 1});
+        AddKeyQuery(*request, {5, 5, 5});
+
+        auto readResult = helper.SendRead("table-1", request.release());
+        UNIT_ASSERT(readResult);
+        CheckResult(helper.Tables["table-1"].UserTable, *readResult, {
+            std::vector<ui32>(),
+            std::vector<ui32>(),
+            std::vector<ui32>(),
+        });
+        UNIT_ASSERT_VALUES_EQUAL(readResult->GetRowsCount(), 3UL);
+    }
+
+    Y_UNIT_TEST(ShouldReadNoColumnsKeysRequestArrow) {
+        // KIKIMR-16897: no columns mean we want to calc row count
+        TTestHelper helper;
+
+        auto request = helper.GetBaseReadRequest("table-1", 1, NKikimrTxDataShard::ARROW);
+        request->Record.ClearColumns();
+        AddKeyQuery(*request, {3, 3, 3});
+        AddKeyQuery(*request, {1, 1, 1});
+        AddKeyQuery(*request, {5, 5, 5});
+
+        auto readResult = helper.SendRead("table-1", request.release());
+        UNIT_ASSERT(readResult);
+        UNIT_ASSERT_VALUES_EQUAL(readResult->Record.GetStatus().GetCode(), Ydb::StatusIds::SUCCESS);
+        UNIT_ASSERT_VALUES_EQUAL(readResult->GetRowsCount(), 3UL);
+        UNIT_ASSERT(readResult->GetArrowBatch());
+
+        auto batch = readResult->GetArrowBatch();
+        UNIT_ASSERT_VALUES_EQUAL(batch->num_rows(), 3UL);
+    }
+
+    Y_UNIT_TEST(ShouldReadNoColumnsRangeRequestCellVec) {
         // KIKIMR-16897: no columns mean we want to calc row count
         TTestHelper helper;
 
@@ -937,7 +977,7 @@ Y_UNIT_TEST_SUITE(DataShardReadIterator) {
         UNIT_ASSERT_VALUES_EQUAL(readResult->GetRowsCount(), 3UL);
     }
 
-    Y_UNIT_TEST(ShouldReadNoColumnsArrow) {
+    Y_UNIT_TEST(ShouldReadNoColumnsRangeRequestArrow) {
         // KIKIMR-16897: no columns mean we want to calc row count
         TTestHelper helper;