From ec0e7a28066e9255bd155e1262a10cfe76327b77 Mon Sep 17 00:00:00 2001
From: Nikolay Shestakov <tesseract@ydb.tech>
Date: Wed, 18 Sep 2024 18:07:35 +0500
Subject: [PATCH] Optimize size of PQTabletConfig (#9375) (#9426)

---
 ydb/core/protos/pqconfig.proto                |  1 +
 .../schemeshard__operation_common.h           | 29 +++++++++++++++----
 .../tx/schemeshard/schemeshard_info_types.h   |  4 +++
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/ydb/core/protos/pqconfig.proto b/ydb/core/protos/pqconfig.proto
index b3d5b1431521..71d32d7c4eb7 100644
--- a/ydb/core/protos/pqconfig.proto
+++ b/ydb/core/protos/pqconfig.proto
@@ -421,6 +421,7 @@ message TPQTabletConfig {
     }
     optional TPartitionStrategy PartitionStrategy = 35;
 
+    // The field is filled in only for the PQ tablet. Contains information about linked partitions for constructing a partial PartitionGraph.
     repeated TPartition AllPartitions = 36; // filled by schemeshard
 
     optional TOffloadConfig OffloadConfig = 38;
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_common.h b/ydb/core/tx/schemeshard/schemeshard__operation_common.h
index 064d8bf250e0..33813a3f9b10 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation_common.h
+++ b/ydb/core/tx/schemeshard/schemeshard__operation_common.h
@@ -928,20 +928,37 @@ class TConfigureParts: public TSubOperationState {
             config.SetVersion(pqGroup.AlterData->AlterVersion);
         }
 
+        THashSet<ui32> linkedPartitions;
+
         for(const auto& pq : pqShard.Partitions) {
             config.AddPartitionIds(pq->PqId);
 
             auto& partition = *config.AddPartitions();
             FillPartition(partition, pq.Get(), 0);
+
+            linkedPartitions.insert(pq->PqId);
+            linkedPartitions.insert(pq->ParentPartitionIds.begin(), pq->ParentPartitionIds.end());
+            linkedPartitions.insert(pq->ChildPartitionIds.begin(), pq->ChildPartitionIds.end());
+            for (auto c : pq->ChildPartitionIds) {
+                auto it = pqGroup.Partitions.find(c);
+                if (it == pqGroup.Partitions.end()) {
+                    continue;
+                }
+                linkedPartitions.insert(it->second->ParentPartitionIds.begin(), it->second->ParentPartitionIds.end());
+            }
         }
 
-        for(const auto& p : pqGroup.Shards) {
-            const auto& pqShard = p.second;
-            const auto& tabletId = context.SS->ShardInfos[p.first].TabletID;
-            for (const auto& pq : pqShard->Partitions) {
-                auto& partition = *config.AddAllPartitions();
-                FillPartition(partition, pq.Get(), ui64(tabletId));
+        for(auto lp : linkedPartitions) {
+            auto it = pqGroup.Partitions.find(lp);
+            if (it == pqGroup.Partitions.end()) {
+                continue;
             }
+
+            auto* partitionInfo = it->second;
+            const auto& tabletId = context.SS->ShardInfos[partitionInfo->ShardIdx].TabletID;
+
+            auto& partition = *config.AddAllPartitions();
+            FillPartition(partition, partitionInfo, ui64(tabletId));
         }
     }
 
diff --git a/ydb/core/tx/schemeshard/schemeshard_info_types.h b/ydb/core/tx/schemeshard/schemeshard_info_types.h
index c227382a424c..890807a1e899 100644
--- a/ydb/core/tx/schemeshard/schemeshard_info_types.h
+++ b/ydb/core/tx/schemeshard/schemeshard_info_types.h
@@ -940,6 +940,8 @@ struct TTopicTabletInfo : TSimpleRefCount<TTopicTabletInfo> {
         THashSet<ui32> ParentPartitionIds;
         THashSet<ui32> ChildPartitionIds;
 
+        TShardIdx ShardIdx;
+
         void SetStatus(const TActorContext& ctx, ui32 value) {
             if (value >= NKikimrPQ::ETopicPartitionStatus::Active &&
                 value <= NKikimrPQ::ETopicPartitionStatus::Deleted) {
@@ -1133,6 +1135,8 @@ struct TTopicInfo : TSimpleRefCount<TTopicInfo> {
     TTopicStats Stats;
 
     void AddPartition(TShardIdx shardIdx, TTopicTabletInfo::TTopicPartitionInfo* partition) {
+        partition->ShardIdx = shardIdx;
+
         TTopicTabletInfo::TPtr& pqShard = Shards[shardIdx];
         if (!pqShard) {
             pqShard.Reset(new TTopicTabletInfo());