25-1: Fixes for BATCH UPDATE/DELETE operations (#18135)

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/kqp_resolve.cpp

@@ -7,14 +7,20 @@
 #include <ydb/core/tx/datashard/range_ops.h>
 #endif
 
+#include <yql/essentials/minikql/mkql_node_builder.h>
+
 namespace NKikimr {
 namespace NKqp {
 
 using namespace NMiniKQL;
 using namespace NYql;
 using namespace NYql::NNodes;
 
-NUdf::TUnboxedValue MakeDefaultValueByType(const NKikimr::NMiniKQL::TType* type) {
+NUdf::TUnboxedValue MakeDefaultValueByType(NKikimr::NMiniKQL::TType* type) {
+    bool isOptional;
+    type = UnpackOptional(type, isOptional);
+    Y_ABORT_UNLESS(type->IsData(), "%s", type->GetKindAsStr());
+
     auto dataType = static_cast<const NKikimr::NMiniKQL::TDataType*>(type);
     switch (dataType->GetSchemeType()) {
         case NUdf::TDataType<bool>::Id:

ydb/core/kqp/common/kqp_resolve.h

@@ -233,7 +233,7 @@ class TKqpTableKeys {
     THashMap<TTableId, TTable> TablesById;
 };
 
-NUdf::TUnboxedValue MakeDefaultValueByType(const NKikimr::NMiniKQL::TType* type);
+NUdf::TUnboxedValue MakeDefaultValueByType(NKikimr::NMiniKQL::TType* type);
 
 TVector<TCell> MakeKeyCells(const NKikimr::NUdf::TUnboxedValue& value, const TVector<NScheme::TTypeInfo>& keyColumnTypes,
     const TVector<ui32>& keyColumnIndices, const NMiniKQL::TTypeEnvironment& typeEnv, bool copyValues);

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

@@ -891,6 +891,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

@@ -482,6 +482,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 {