[KQP] Add configuration for the custom partition pruning for DataExecuter and PartitionedExecuter (#19919)

Author: dahbka-lis

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

@@ -0,0 +1,42 @@
+#include "kqp_batch_operations.h"
+
+namespace NKikimr::NKqp::NBatchOperations {
+
+TSerializedTableRange MakePartitionRange(TMaybe<TKeyDesc::TPartitionRangeInfo> begin, TMaybe<TKeyDesc::TPartitionRangeInfo> end, size_t keySize) {
+    TVector<TCell> tableBegin;
+    TVector<TCell> tableEnd;
+
+    bool inclusiveTableBegin = (begin) ? begin->IsInclusive : false;
+    bool inclusiveTableEnd = (end) ? end->IsInclusive : false;
+
+    if (!begin || !begin->EndKeyPrefix) {
+        inclusiveTableBegin = true;
+        tableBegin.resize(keySize, TCell());
+    } else {
+        const auto& cells = begin->EndKeyPrefix.GetCells();
+        tableBegin.assign(cells.begin(), cells.end());
+    }
+
+    if (!end || !end->EndKeyPrefix) {
+        inclusiveTableEnd = true;
+    } else {
+        const auto& cells = end->EndKeyPrefix.GetCells();
+        tableEnd.assign(cells.begin(), cells.end());
+    }
+
+    return TSerializedTableRange{tableBegin, inclusiveTableBegin, tableEnd, inclusiveTableEnd};
+}
+
+TSettings ImportSettingsFromProto(const NKikimrConfig::TTableServiceConfig::TBatchOperationSettings& settings) {
+    TSettings 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::NBatchOperations

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

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <ydb/core/protos/table_service_config.pb.h>
+#include <ydb/core/scheme/scheme_tabledefs.h>
+
+#include <util/generic/fwd.h>
+
+namespace NKikimr::NKqp::NBatchOperations {
+
+TSerializedTableRange MakePartitionRange(TMaybe<TKeyDesc::TPartitionRangeInfo> begin, TMaybe<TKeyDesc::TPartitionRangeInfo> end, size_t keySize);
+
+struct TSettings {
+    ui64 MaxBatchSize = 10000;
+    ui64 MinBatchSize = 1;
+    ui64 MaxRetryDelayMs = 30000;
+    ui64 StartRetryDelayMs = 50;
+    ui64 PartitionExecutionLimit = 10;
+};
+
+TSettings ImportSettingsFromProto(const NKikimrConfig::TTableServiceConfig::TBatchOperationSettings& settings);
+
+} // namespace NKikimr::NKqp::NBatchOperations

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

@@ -6,6 +6,7 @@ ENDIF()
 
 SRCS(
     control.cpp
+    kqp_batch_operations.cpp
     kqp_event_ids.h
     kqp_event_impl.cpp
     kqp_lwtrace_probes.cpp
@@ -31,9 +32,10 @@ SRCS(
 PEERDIR(
     ydb/core/base
     ydb/core/engine
+    ydb/core/protos
+    ydb/core/scheme
     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
@@ -64,7 +66,6 @@ GENERATE_ENUM_SERIALIZATION(kqp_resolve.h)
 END()
 
 RECURSE(
-    batch
     compilation
     events
     simple

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

@@ -1,7 +1,6 @@
 #include "kqp_executer.h"
 #include "kqp_executer_impl.h"
 #include "kqp_locks_helper.h"
-#include "kqp_partition_helper.h"
 #include "kqp_planner.h"
 #include "kqp_table_resolver.h"
 #include "kqp_tasks_validate.h"
@@ -102,12 +101,13 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
         const TActorId& creator, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
         ui32 statementResultIndex, const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup,
         const TGUCSettings::TPtr& GUCSettings,
+        TPartitionPruner::TConfig partitionPrunerConfig,
         const TShardIdToTableInfoPtr& shardIdToTableInfo,
         const IKqpTransactionManagerPtr& txManager,
         const TActorId bufferActorId,
-        TMaybe<TBatchOperationSettings> batchOperationSettings = Nothing())
-        : TBase(std::move(request), std::move(asyncIoFactory), federatedQuerySetup, GUCSettings, database, userToken, counters, executerConfig,
-            userRequestContext, statementResultIndex, TWilsonKqp::DataExecuter,
+        TMaybe<NBatchOperations::TSettings> batchOperationSettings = Nothing())
+        : TBase(std::move(request), std::move(asyncIoFactory), federatedQuerySetup, GUCSettings, std::move(partitionPrunerConfig),
+            database, userToken, counters, executerConfig, userRequestContext, statementResultIndex, TWilsonKqp::DataExecuter,
             "DataExecuter", streamResult, bufferActorId, txManager, std::move(batchOperationSettings))
         , ShardIdToTableInfo(shardIdToTableInfo)
         , AllowOlapDataQuery(executerConfig.TableServiceConfig.GetAllowOlapDataQuery())
@@ -212,14 +212,14 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                     }
                 }
 
-                if (info.HasBatchOperationMaxKey()) {
-                    if (ResponseEv->BatchOperationMaxKeys.empty()) {
-                        for (auto keyId : info.GetBatchOperationKeyIds()) {
-                            ResponseEv->BatchOperationKeyIds.push_back(keyId);
+                if (!BatchOperationSettings.Empty() && info.HasSerializedEndRow()) {
+                    if (ResponseEv->EndRowColumnIds.empty()) {
+                        for (auto keyId : info.GetEndRowColumnIds()) {
+                            ResponseEv->EndRowColumnIds.push_back(keyId);
                         }
                     }
 
-                    ResponseEv->BatchOperationMaxKeys.emplace_back(info.GetBatchOperationMaxKey());
+                    ResponseEv->SerializedEndRows.emplace_back(info.GetSerializedEndRow());
                 }
             } else if (data.GetData().template Is<NKikimrKqp::TEvKqpOutputActorResultInfo>()) {
                 NKikimrKqp::TEvKqpOutputActorResultInfo info;
@@ -1651,7 +1651,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 case NKqpProto::TKqpPhyTableOperation::kLookup: {
                     auto columns = BuildKqpColumns(op, tableInfo);
                     bool isFullScan = false;
-                    auto partitions = PrunePartitions(op, stageInfo, HolderFactory(), TypeEnv(), isFullScan);
+                    auto partitions = PartitionPruner.Prune(op, stageInfo, isFullScan);
                     auto readSettings = ExtractReadSettings(op, stageInfo, HolderFactory(), TypeEnv());
 
                     if (!readSettings.ItemsLimit && isFullScan) {
@@ -1673,7 +1673,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 case NKqpProto::TKqpPhyTableOperation::kDeleteRows: {
                     YQL_ENSURE(stage.InputsSize() <= 1, "Effect stage with multiple inputs: " << stage.GetProgramAst());
 
-                    auto result = PruneEffectPartitions(op, stageInfo, HolderFactory(), TypeEnv());
+                    auto result = PartitionPruner.PruneEffect(op, stageInfo);
                     for (auto& [shardId, shardInfo] : result) {
                         YQL_ENSURE(!shardInfo.KeyReadRanges);
                         YQL_ENSURE(shardInfo.KeyWriteRanges);
@@ -2251,7 +2251,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                     if (stage.SourcesSize() > 0 && stage.GetSources(0).GetTypeCase() == NKqpProto::TKqpSource::kReadRangesSource) {
                         const auto& source = stage.GetSources(0).GetReadRangesSource();
                         bool isFullScan;
-                        SourceScanStageIdToParititions[stageInfo.Id] = PrunePartitions(source, stageInfo, HolderFactory(), TypeEnv(), isFullScan);
+                        SourceScanStageIdToParititions[stageInfo.Id] = PartitionPruner.Prune(source, stageInfo, isFullScan);
                         if (isFullScan && !source.HasItemsLimit()) {
                             Counters->Counters->FullScansExecuted->Inc();
                         }
@@ -3058,12 +3058,13 @@ 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,
-    TMaybe<TBatchOperationSettings> batchOperationSettings)
+    TPartitionPruner::TConfig partitionPrunerConfig, const TShardIdToTableInfoPtr& shardIdToTableInfo,
+    const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId,
+    TMaybe<NBatchOperations::TSettings> batchOperationSettings)
 {
     return new TKqpDataExecuter(std::move(request), database, userToken, counters, streamResult, executerConfig,
         std::move(asyncIoFactory), creator, userRequestContext, statementResultIndex, federatedQuerySetup, GUCSettings,
-        shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings));
+        std::move(partitionPrunerConfig), shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings));
 }
 
 } // namespace NKqp

ydb/core/kqp/common/kqp_batch_operations.cpp

@@ -1,10 +1,11 @@
 #pragma once
 
 #include <library/cpp/lwtrace/shuttle.h>
-#include <ydb/core/kqp/common/batch/batch_operation_settings.h>
+#include <ydb/core/kqp/common/kqp_batch_operations.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>
+#include <ydb/core/kqp/executer_actor/kqp_partition_helper.h>
 #include <ydb/core/kqp/executer_actor/shards_resolver/kqp_shards_resolver_events.h>
 #include <ydb/core/kqp/query_data/kqp_query_data.h>
 #include <ydb/core/kqp/gateway/kqp_gateway.h>
@@ -37,8 +38,9 @@ struct TEvKqpExecuter {
 
         THashSet<ui32> ParticipantNodes;
 
-        TVector<TSerializedCellVec> BatchOperationMaxKeys;
-        TVector<ui32> BatchOperationKeyIds;
+        // For BATCH operations only
+        TVector<TSerializedCellVec> SerializedEndRows;
+        TVector<ui32> EndRowColumnIds;
 
         enum class EExecutionType {
             Data,
@@ -151,8 +153,9 @@ 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,
-    TMaybe<TBatchOperationSettings> batchOperationSettings = Nothing());
+    TPartitionPruner::TConfig partitionPrunerConfig, const TShardIdToTableInfoPtr& shardIdToTableInfo,
+    const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId,
+    TMaybe<NBatchOperations::TSettings> batchOperationSettings = Nothing());
 
 IActor* CreateKqpSchemeExecuter(
     TKqpPhyTxHolder::TConstPtr phyTx, NKikimrKqp::EQueryType queryType, const TActorId& target,

ydb/core/kqp/common/kqp_batch_operations.h

@@ -82,16 +82,17 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
     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,
-    TMaybe<TBatchOperationSettings> batchOperationSettings)
+    TPartitionPruner::TConfig partitionPrunerConfig, const TShardIdToTableInfoPtr& shardIdToTableInfo,
+    const IKqpTransactionManagerPtr& txManager, const TActorId bufferActorId,
+    TMaybe<NBatchOperations::TSettings> batchOperationSettings)
 {
     if (request.Transactions.empty()) {
         // commit-only or rollback-only data transaction
         return CreateKqpDataExecuter(
             std::move(request), database, userToken, counters, false, executerConfig,
             std::move(asyncIoFactory), creator,
             userRequestContext, statementResultIndex,
-            federatedQuerySetup, /*GUCSettings*/nullptr,
+            federatedQuerySetup, /*GUCSettings*/nullptr, std::move(partitionPrunerConfig),
             shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings)
         );
     }
@@ -115,7 +116,7 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
                 std::move(request), database, userToken, counters, false, executerConfig,
                 std::move(asyncIoFactory), creator,
                 userRequestContext, statementResultIndex,
-                federatedQuerySetup, /*GUCSettings*/nullptr,
+                federatedQuerySetup, /*GUCSettings*/nullptr, std::move(partitionPrunerConfig),
                 shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings)
             );
 
@@ -131,7 +132,7 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
                 std::move(request), database, userToken, counters, true,
                 executerConfig, std::move(asyncIoFactory), creator,
                 userRequestContext, statementResultIndex,
-                federatedQuerySetup, GUCSettings,
+                federatedQuerySetup, GUCSettings, std::move(partitionPrunerConfig),
                 shardIdToTableInfo, txManager, bufferActorId, std::move(batchOperationSettings)
             );