change memory planner for CS scan (#7372)

Author: ivanmorozov333

ydb/core/base/events.h

@@ -180,6 +180,7 @@ struct TKikimrEvents : TEvents {
         ES_SS_BG_TASKS = 4257, 
         ES_LIMITER = 4258,
         ES_MEMORY = 4259,
+        ES_GROUPED_ALLOCATIONS_MANAGER = 4260,
     };
 };
 

ydb/core/driver_lib/run/kikimr_services_initializers.cpp

@@ -187,6 +187,9 @@
 #include <ydb/core/tx/limiter/usage/config.h>
 #include <ydb/core/tx/limiter/usage/service.h>
 
+#include <ydb/core/tx/limiter/grouped_memory/usage/config.h>
+#include <ydb/core/tx/limiter/grouped_memory/usage/service.h>
+
 #include <ydb/core/backup/controller/tablet.h>
 
 #include <ydb/services/ext_index/common/config.h>
@@ -2216,6 +2219,26 @@ void TKqpServiceInitializer::InitializeServices(NActors::TActorSystemSetup* setu
     }
 }
 
+TGroupedMemoryLimiterInitializer::TGroupedMemoryLimiterInitializer(const TKikimrRunConfig& runConfig)
+    : IKikimrServicesInitializer(runConfig) {
+}
+
+void TGroupedMemoryLimiterInitializer::InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) {
+    NOlap::NGroupedMemoryManager::TConfig serviceConfig;
+    Y_ABORT_UNLESS(serviceConfig.DeserializeFromProto(Config.GetGroupedMemoryLimiterConfig()));
+
+    if (serviceConfig.IsEnabled()) {
+        TIntrusivePtr<::NMonitoring::TDynamicCounters> tabletGroup = GetServiceCounters(appData->Counters, "tablets");
+        TIntrusivePtr<::NMonitoring::TDynamicCounters> countersGroup = tabletGroup->GetSubgroup("type", "TX_GROUPED_MEMORY_LIMITER");
+
+        auto service = NOlap::NGroupedMemoryManager::TScanMemoryLimiterOperator::CreateService(serviceConfig, countersGroup);
+
+        setup->LocalServices.push_back(std::make_pair(
+            NOlap::NGroupedMemoryManager::TScanMemoryLimiterOperator::MakeServiceId(NodeId),
+            TActorSetupCmd(service, TMailboxType::HTSwap, appData->UserPoolId)));
+    }
+}
+
 TCompDiskLimiterInitializer::TCompDiskLimiterInitializer(const TKikimrRunConfig& runConfig)
     : IKikimrServicesInitializer(runConfig) {
 }

ydb/core/driver_lib/run/kikimr_services_initializers.h

@@ -404,6 +404,12 @@ class TCompDiskLimiterInitializer: public IKikimrServicesInitializer {
     void InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) override;
 };
 
+class TGroupedMemoryLimiterInitializer: public IKikimrServicesInitializer {
+public:
+    TGroupedMemoryLimiterInitializer(const TKikimrRunConfig& runConfig);
+    void InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) override;
+};
+
 class TCompConveyorInitializer: public IKikimrServicesInitializer {
 public:
     TCompConveyorInitializer(const TKikimrRunConfig& runConfig);

ydb/core/driver_lib/run/run.cpp

@@ -1556,6 +1556,10 @@ TIntrusivePtr<TServiceInitializersList> TKikimrRunner::CreateServiceInitializers
         sil->AddServiceInitializer(new TCompDiskLimiterInitializer(runConfig));
     }
 
+    if (serviceMask.EnableGroupedMemoryLimiter) {
+        sil->AddServiceInitializer(new TGroupedMemoryLimiterInitializer(runConfig));
+    }
+
     if (serviceMask.EnableScanConveyor) {
         sil->AddServiceInitializer(new TScanConveyorInitializer(runConfig));
     }

ydb/core/driver_lib/run/service_mask.h

@@ -78,6 +78,7 @@ union TBasicKikimrServicesMask {
         bool EnableDatabaseMetadataCache:1;
         bool EnableGraphService:1;
         bool EnableCompDiskLimiter:1;
+        bool EnableGroupedMemoryLimiter:1;
     };
 
     struct {

ydb/core/driver_lib/run/ya.make

@@ -101,7 +101,7 @@ PEERDIR(
     ydb/core/security
     ydb/core/security/ldap_auth_provider
     ydb/core/statistics/aggregator
-    ydb/core/statistics/service    
+    ydb/core/statistics/service
     ydb/core/sys_view/processor
     ydb/core/sys_view/service
     ydb/core/tablet
@@ -113,6 +113,7 @@ PEERDIR(
     ydb/core/tx/coordinator
     ydb/core/tx/conveyor/service
     ydb/core/tx/limiter/service
+    ydb/core/tx/limiter/grouped_memory/usage
     ydb/core/tx/datashard
     ydb/core/tx/long_tx_service
     ydb/core/tx/long_tx_service/public

ydb/core/formats/arrow/arrow_helpers.cpp

@@ -875,6 +875,18 @@ std::shared_ptr<arrow::RecordBatch> ReallocateBatch(std::shared_ptr<arrow::Recor
     return DeserializeBatch(SerializeBatch(original, arrow::ipc::IpcWriteOptions::Defaults()), original->schema());
 }
 
+std::shared_ptr<arrow::Table> ReallocateBatch(const std::shared_ptr<arrow::Table>& original) {
+    if (!original) {
+        return original;
+    }
+    auto batches = NArrow::SliceToRecordBatches(original);
+    for (auto&& i : batches) {
+        i = NArrow::TStatusValidator::GetValid(
+            NArrow::NSerialization::TNativeSerializer().Deserialize(NArrow::NSerialization::TNativeSerializer().SerializeFull(i)));
+    }
+    return NArrow::TStatusValidator::GetValid(arrow::Table::FromRecordBatches(batches));
+}
+
 std::shared_ptr<arrow::RecordBatch> MergeColumns(const std::vector<std::shared_ptr<arrow::RecordBatch>>& batches) {
     std::vector<std::shared_ptr<arrow::Array>> columns;
     std::vector<std::shared_ptr<arrow::Field>> fields;

ydb/core/formats/arrow/arrow_helpers.h

@@ -103,6 +103,7 @@ std::partial_ordering ColumnsCompare(const std::vector<std::shared_ptr<arrow::Ar
 bool ScalarLess(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr<arrow::Scalar>& y);
 bool ScalarLess(const arrow::Scalar& x, const arrow::Scalar& y);
 std::shared_ptr<arrow::RecordBatch> ReallocateBatch(std::shared_ptr<arrow::RecordBatch> original);
+std::shared_ptr<arrow::Table> ReallocateBatch(const std::shared_ptr<arrow::Table>& original);
 
 bool HasNulls(const std::shared_ptr<arrow::Array>& column);
 

ydb/core/kqp/ut/olap/datatime64_ut.cpp

@@ -170,6 +170,7 @@ Y_UNIT_TEST_SUITE(KqpDatetime64ColumnShard) {
         runnerSettings.WithSampleTables = false;
 
         TTestHelper testHelper(runnerSettings);
+        Tests::NCommon::TLoggerInit(testHelper.GetKikimr()).SetComponents({ NKikimrServices::GROUPED_MEMORY_LIMITER }, "CS").Initialize();
 
         TVector<TTestHelper::TColumnSchema> schema = {
             TTestHelper::TColumnSchema().SetName("id").SetType(NScheme::NTypeIds::Int64).SetNullable(false),

ydb/core/kqp/ut/olap/helpers/aggregation.cpp

@@ -13,6 +13,7 @@ void TestAggregationsBase(const std::vector<TAggregationTestCase>& cases) {
 
     TLocalHelper(kikimr).CreateTestOlapTable();
     auto tableClient = kikimr.GetTableClient();
+    Tests::NCommon::TLoggerInit(kikimr).SetComponents({ NKikimrServices::GROUPED_MEMORY_LIMITER, NKikimrServices::TX_COLUMNSHARD_SCAN }, "CS").Initialize();
 
     {
         WriteTestData(kikimr, "/Root/olapStore/olapTable", 10000, 3000000, 1000);
@@ -49,10 +50,11 @@ void TestAggregationsInternal(const std::vector<TAggregationTestCase>& cases) {
     Tests::TServer::TPtr server = new Tests::TServer(settings);
 
     auto runtime = server->GetRuntime();
+    Tests::NCommon::TLoggerInit(runtime).Initialize();
+    Tests::NCommon::TLoggerInit(runtime).SetComponents({ NKikimrServices::GROUPED_MEMORY_LIMITER }, "CS").Initialize();
     auto sender = runtime->AllocateEdgeActor();
 
     InitRoot(server, sender);
-    Tests::NCommon::TLoggerInit(runtime).Initialize();
 
     ui32 numShards = 1;
     ui32 numIterations = 10;