Fixes + new features for BATCH UPDATE/DELETE (#14994)

Author: dahbka-lis

ydb/core/kqp/common/batch/batch_operation_settings.cpp

@@ -0,0 +1,17 @@
+#include "batch_operation_settings.h"
+
+namespace NKikimr::NKqp {
+
+TBatchOperationSettings SetBatchOperationSettings(const NKikimrConfig::TTableServiceConfig::TBatchOperationSettings& settings) {
+    TBatchOperationSettings res;
+
+    res.MaxBatchSize = settings.GetMaxBatchSize();
+    res.MinBatchSize = settings.GetMinBatchSize();
+    res.MaxRetryDelayMs = settings.GetMaxRetryDelayMs();
+    res.StartRetryDelayMs = settings.GetStartRetryDelayMs();
+    res.PartitionExecutionLimit = settings.GetPartitionExecutionLimit();
+
+    return res;
+}
+
+} // namespace NKikimr::NKqp

ydb/core/kqp/common/batch/batch_operation_settings.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <ydb/core/protos/table_service_config.pb.h>
+#include <util/generic/fwd.h>
+
+namespace NKikimr::NKqp {
+
+struct TBatchOperationSettings {
+    ui64 MaxBatchSize = 10000;
+    ui64 MinBatchSize = 1;
+    ui64 MaxRetryDelayMs = 30000;
+    ui64 StartRetryDelayMs = 50;
+    ui64 PartitionExecutionLimit = 10;
+};
+
+TBatchOperationSettings SetBatchOperationSettings(const NKikimrConfig::TTableServiceConfig::TBatchOperationSettings& settings);
+
+} // namespace NKikimr::NKqp

ydb/core/kqp/common/batch/params.h

@@ -4,12 +4,12 @@
 
 namespace NKikimr::NKqp::NBatchParams {
 
-const TString Header = "$_kqp_batch_";
-const TString IsFirstQuery = Header + "is_first_query";
-const TString IsLastQuery = Header + "is_last_query";
+const TString Header = "%kqp%batch_";
 const TString IsInclusiveLeft = Header + "is_inclusive_left";
 const TString IsInclusiveRight = Header + "is_inclusive_right";
 const TString Begin = Header + "begin_"; // begin_N
 const TString End = Header + "end_";     // end_N
+const TString BeginPrefixSize = Begin + "prefix_size";
+const TString EndPrefixSize = End + "prefix_size";
 
-} // namespace NKikimr::NKqp::NPartitionedExecuter
+} // namespace NKikimr::NKqp::NBatchParams

ydb/core/kqp/common/batch/ya.make

@@ -0,0 +1,13 @@
+LIBRARY()
+
+SRCS(
+    batch_operation_settings.cpp
+)
+
+PEERDIR(
+    ydb/core/protos
+)
+
+YQL_LAST_ABI_VERSION()
+
+END()

ydb/core/kqp/common/simple/kqp_event_ids.h

@@ -65,7 +65,8 @@ struct TKqpExecuterEvents {
         EvProgress,
         EvStreamDataAck,
         EvTableResolveStatus,
-        EvShardsResolveStatus
+        EvShardsResolveStatus,
+        EvDelayedExecution
     };
 };
 

ydb/core/kqp/common/ya.make

@@ -29,6 +29,7 @@ PEERDIR(
     ydb/core/engine
     ydb/core/kqp/expr_nodes
     ydb/core/kqp/common/simple
+    ydb/core/kqp/common/batch
     ydb/core/kqp/common/compilation
     ydb/core/kqp/common/events
     ydb/core/kqp/common/shutdown
@@ -58,6 +59,7 @@ GENERATE_ENUM_SERIALIZATION(kqp_yql.h)
 END()
 
 RECURSE(
+    batch
     compilation
     events
     simple

ydb/core/kqp/counters/kqp_counters.cpp

@@ -888,6 +888,14 @@ TKqpCounters::TKqpCounters(const ::NMonitoring::TDynamicCounterPtr& counters, co
     QueryStatMemFinishBytes = KqpGroup->GetCounter("Query/Stat/MemFinishBytes", true);
     QueryStatMemConvertBytes = KqpGroup->GetCounter("Query/Stat/MemConvertBytes", true);
 
+    /* Statistics batch operations */
+    BatchOperationUpdateRows = KqpGroup->GetCounter("BatchOperation/Update/Rows", true);
+    BatchOperationUpdateBytes = KqpGroup->GetCounter("BatchOperation/Update/Bytes", true);
+
+    BatchOperationDeleteRows = KqpGroup->GetCounter("BatchOperation/Delete/Rows", true);
+    BatchOperationDeleteBytes = KqpGroup->GetCounter("BatchOperation/Delete/Bytes", true);
+
+    BatchOperationRetries = KqpGroup->GetCounter("BatchOperation/Retries", true);
 }
 
 ::NMonitoring::TDynamicCounterPtr TKqpCounters::GetKqpCounters() const {

ydb/core/kqp/counters/kqp_counters.h

@@ -479,6 +479,12 @@ class TKqpCounters : public TKqpCountersBase, public NYql::NDq::TSpillingCounter
     ::NMonitoring::TDynamicCounters::TCounterPtr QueryStatMemFinishBytes;
     ::NMonitoring::TDynamicCounters::TCounterPtr QueryStatMemConvertBytes;
 
+    // Statistics batch operations
+    ::NMonitoring::TDynamicCounters::TCounterPtr BatchOperationUpdateRows;
+    ::NMonitoring::TDynamicCounters::TCounterPtr BatchOperationUpdateBytes;
+    ::NMonitoring::TDynamicCounters::TCounterPtr BatchOperationDeleteRows;
+    ::NMonitoring::TDynamicCounters::TCounterPtr BatchOperationDeleteBytes;
+    ::NMonitoring::TDynamicCounters::TCounterPtr BatchOperationRetries;
 };
 
 struct TKqpRequestCounters : public TThrRefBase {

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -104,10 +104,11 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
         const TGUCSettings::TPtr& GUCSettings,
         const TShardIdToTableInfoPtr& shardIdToTableInfo,
         const IKqpTransactionManagerPtr& txManager,
-        const TActorId bufferActorId)
+        const TActorId bufferActorId,
+        TMaybe<TBatchOperationSettings> batchOperationSettings = Nothing())
         : TBase(std::move(request), std::move(asyncIoFactory), federatedQuerySetup, GUCSettings, database, userToken, counters, tableServiceConfig,
             userRequestContext, statementResultIndex, TWilsonKqp::DataExecuter,
-            "DataExecuter", streamResult, bufferActorId, txManager)
+            "DataExecuter", streamResult, bufferActorId, txManager, std::move(batchOperationSettings))
         , ShardIdToTableInfo(shardIdToTableInfo)
         , AllowOlapDataQuery(tableServiceConfig.GetAllowOlapDataQuery())
         , WaitCAStatsTimeout(TDuration::MilliSeconds(tableServiceConfig.GetQueryLimits().GetWaitCAStatsTimeoutMs()))
@@ -213,6 +214,16 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                         TxManager->AddLock(lock.GetDataShard(), lock);
                     }
                 }
+
+                if (info.HasBatchOperationMaxKey()) {
+                    if (ResponseEv->BatchOperationMaxKeys.empty()) {
+                        for (auto keyId : info.GetBatchOperationKeyIds()) {
+                            ResponseEv->BatchOperationKeyIds.push_back(keyId);
+                        }
+                    }
+
+                    ResponseEv->BatchOperationMaxKeys.emplace_back(info.GetBatchOperationMaxKey());
+                }
             } else if (data.GetData().template Is<NKikimrKqp::TEvKqpOutputActorResultInfo>()) {
                 NKikimrKqp::TEvKqpOutputActorResultInfo info;
                 YQL_ENSURE(data.GetData().UnpackTo(&info), "Failed to unpack settings");
@@ -3055,11 +3066,12 @@ IActor* CreateKqpDataExecuter(IKqpGateway::TExecPhysicalRequest&& request, const
     NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, const TActorId& creator,
     const TIntrusivePtr<TUserRequestContext>& userRequestContext, ui32 statementResultIndex,
     const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup, const TGUCSettings::TPtr& GUCSettings,
-    const TShardIdToTableInfoPtr& shardIdToTableInfo, const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId)
+    const TShardIdToTableInfoPtr& shardIdToTableInfo, const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId,
+    TMaybe<TBatchOperationSettings> batchOperationSettings)
 {
     return new TKqpDataExecuter(std::move(request), database, userToken, counters, streamResult, tableServiceConfig,
         std::move(asyncIoFactory), creator, userRequestContext, statementResultIndex, federatedQuerySetup, GUCSettings,
-        shardIdToTableInfo, txManager, bufferActorId);
+        shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings));
 }
 
 } // namespace NKqp

ydb/core/kqp/executer_actor/kqp_executer.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <library/cpp/lwtrace/shuttle.h>
+#include <ydb/core/kqp/common/batch/batch_operation_settings.h>
 #include <ydb/core/kqp/common/kqp_tx.h>
 #include <ydb/core/kqp/common/kqp_event_ids.h>
 #include <ydb/core/kqp/common/kqp_user_request_context.h>
@@ -36,6 +37,9 @@ struct TEvKqpExecuter {
 
         THashSet<ui32> ParticipantNodes;
 
+        TVector<TSerializedCellVec> BatchOperationMaxKeys;
+        TVector<ui32> BatchOperationKeyIds;
+
         enum class EExecutionType {
             Data,
             Scan,
@@ -105,6 +109,16 @@ struct TEvKqpExecuter {
         NYql::TIssues Issues;
         TDuration CpuTime;
     };
+
+    struct TEvTxDelayedExecution : public TEventLocal<TEvTxDelayedExecution,
+        TKqpExecuterEvents::EvDelayedExecution>
+    {
+        TEvTxDelayedExecution(size_t partitionIdx)
+            : PartitionIdx(partitionIdx)
+        {}
+
+        size_t PartitionIdx;
+    };
 };
 
 struct TKqpFederatedQuerySetup;
@@ -115,7 +129,8 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
     NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, TPreparedQueryHolder::TConstPtr preparedQuery, const TActorId& creator,
     const TIntrusivePtr<TUserRequestContext>& userRequestContext, ui32 statementResultIndex,
     const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup, const TGUCSettings::TPtr& GUCSettings,
-    const TShardIdToTableInfoPtr& shardIdToTableInfo, const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId);
+    const TShardIdToTableInfoPtr& shardIdToTableInfo, const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId,
+    TMaybe<TBatchOperationSettings> batchOperationSettings = Nothing());
 
 IActor* CreateKqpSchemeExecuter(
     TKqpPhyTxHolder::TConstPtr phyTx, NKikimrKqp::EQueryType queryType, const TActorId& target,