Precharge in EvWrite

Author: azevaykin

ydb/core/kqp/ut/query/kqp_query_ut.cpp

@@ -1,3 +1,4 @@
+#include <fmt/format.h>
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 
 #include <ydb/core/tx/datashard/datashard_failpoints.h>
@@ -2634,6 +2635,35 @@ Y_UNIT_TEST_SUITE(KqpQuery) {
         }
     }
 
+    Y_UNIT_TEST(ExecuteWriteQuery) {
+        using namespace fmt::literals;
+
+        TKikimrRunner kikimr;
+        auto client = kikimr.GetQueryClient();
+
+        {   // Just generate table
+            const auto sql = fmt::format(R"(
+                CREATE TABLE test_table (
+                    PRIMARY KEY (id)
+                ) AS SELECT
+                    ROW_NUMBER() OVER w AS id, data
+                FROM
+                    AS_TABLE(ListReplicate(<|data: '{data}'|>, 500000))
+                WINDOW
+                    w AS (ORDER BY data))",
+                "data"_a = std::string(137, 'a')
+            );
+            const auto result = client.ExecuteQuery(sql, NYdb::NQuery::TTxControl::NoTx()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+
+        Cerr << TInstant::Now() << " --------------- Start update ---------------\n";
+
+        const auto hangingResult = client.ExecuteQuery(R"(
+            UPDATE test_table SET data = "a"
+        )", NYdb::NQuery::TTxControl::NoTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(hangingResult.GetStatus(), EStatus::SUCCESS, hangingResult.GetIssues().ToString());
+    }    
 }
 
 } // namespace NKqp

ydb/core/tx/datashard/datashard_user_db.cpp

@@ -178,6 +178,16 @@ void TDataShardUserDb::EraseRow(
     Counters.EraseRowBytes += keyBytes + 8;
 }
 
+bool TDataShardUserDb::PrechargeRow(
+    const TTableId& tableId,
+    const TArrayRef<const TRawTypeValue> key)
+{
+    auto localTableId = Self.GetLocalTableId(tableId);
+    Y_ENSURE(localTableId != 0, "Unexpected PrechargeRow for an unknown table");
+
+    return Db.Precharge(localTableId, key, key, {}, 0, Max<ui64>(), Max<ui64>());     
+}
+
 void TDataShardUserDb::IncreaseUpdateCounters(
     const TArrayRef<const TRawTypeValue> key, 
     const TArrayRef<const NIceDb::TUpdateOp> ops) 

ydb/core/tx/datashard/datashard_user_db.h

@@ -176,6 +176,10 @@ class TDataShardUserDb final
     NMiniKQL::TEngineHostCounters& GetCounters();
     const NMiniKQL::TEngineHostCounters& GetCounters() const;
 
+    bool PrechargeRow(
+        const TTableId& tableId,
+        const TArrayRef<const TRawTypeValue> key);    
+
 private:
     static TSmallVec<TCell> ConvertTableKeys(const TArrayRef<const TRawTypeValue> key);
 

ydb/core/tx/datashard/execute_write_unit.cpp

@@ -131,8 +131,7 @@ class TExecuteWriteUnit : public TExecutionUnit {
             }
         };
 
-        for (ui32 rowIdx = 0; rowIdx < matrix.GetRowCount(); ++rowIdx)
-        {
+        auto fillKey = [&](ui32 rowIdx, TSmallVec<TRawTypeValue>& key) {
             key.clear();
             key.reserve(userTable.KeyColumnIds.size());
             for (ui16 keyColIdx = 0; keyColIdx < userTable.KeyColumnIds.size(); ++keyColIdx) {
@@ -145,6 +144,28 @@ class TExecuteWriteUnit : public TExecutionUnit {
                     key.emplace_back(cell.Data(), cell.Size(), vtypeId);
                 }
             }
+        };
+
+        // Precharge
+
+        switch (operationType) {
+            case NKikimrDataEvents::TEvWrite::TOperation::OPERATION_INSERT:
+            case NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPDATE: {
+                for (ui32 rowIdx = 0; rowIdx < matrix.GetRowCount(); ++rowIdx) {
+                    fillKey(rowIdx, key);
+                    userDb.PrechargeRow(fullTableId, key);
+                } 
+                break;
+            }
+            default:
+                break;
+        }
+
+        // Main update cycle
+
+        for (ui32 rowIdx = 0; rowIdx < matrix.GetRowCount(); ++rowIdx)
+        {
+            fillKey(rowIdx, key);
 
             switch (operationType) {
                 case NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT: {
@@ -177,6 +198,8 @@ class TExecuteWriteUnit : public TExecutionUnit {
             }
         }
 
+        // Counters
+
         switch (operationType) {
             case NKikimrDataEvents::TEvWrite::TOperation::OPERATION_UPSERT:
             case NKikimrDataEvents::TEvWrite::TOperation::OPERATION_REPLACE: