allow reading from followers for vector index tables (#18175)

gridnevvvit · web-flow · commit 8703fb9c5089 · 2025-05-12T11:35:58.000+03:00
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -967,7 +967,7 @@ class TKqpExecuterBase : public TActor<TDerived> {
             } else {
                 settings = *stageInfo.Meta.ResolvedSinkSettings;
             }
-        
+
             auto& lockTxId = TasksGraph.GetMeta().LockTxId;
             if (lockTxId) {
                 settings.SetLockTxId(*lockTxId);
@@ -1178,6 +1178,8 @@ class TKqpExecuterBase : public TActor<TDerived> {
             NKikimrTxDataShard::TKqpReadRangesSourceSettings* settings = input.Meta.SourceSettings;
             FillTableMeta(stageInfo, settings->MutableTable());
 
+            settings->SetIsTableImmutable(source.GetIsTableImmutable());
+
             for (auto& keyColumn : keyTypes) {
                 auto columnType = NScheme::ProtoColumnTypeFromTypeInfoMod(keyColumn, "");
                 *settings->AddKeyColumnTypeInfos() = columnType.TypeInfo ?
diff --git a/ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp b/ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp
@@ -367,6 +367,11 @@ void BuildStreamLookupChannels(TKqpTasksGraph& graph, const TStageInfo& stageInf
     streamLookupTransform.InputType = streamLookup.GetLookupKeysType();
     streamLookupTransform.OutputType = streamLookup.GetResultType();
 
+    if (streamLookup.GetIsTableImmutable()) {
+        settings->SetAllowUseFollowers(true);
+        settings->SetIsTableImmutable(true);
+    }
+
     for (ui32 taskId = 0; taskId < inputStageInfo.Tasks.size(); ++taskId) {
         auto& originTask = graph.GetTask(inputStageInfo.Tasks[taskId]);
         auto& targetTask = graph.GetTask(stageInfo.Tasks[taskId]);
@@ -1159,18 +1164,26 @@ void FillInputDesc(const TKqpTasksGraph& tasksGraph, NYql::NDqProto::TTaskInput&
             inputDesc.MutableSource()->SetWatermarksMode(input.WatermarksMode);
             if (Y_LIKELY(input.Meta.SourceSettings)) {
                 enableMetering = true;
-                if (snapshot.IsValid()) {
+                YQL_ENSURE(input.Meta.SourceSettings->HasTable());
+                bool isTableImmutable = input.Meta.SourceSettings->GetIsTableImmutable();
+
+                if (snapshot.IsValid() && !isTableImmutable) {
                     input.Meta.SourceSettings->MutableSnapshot()->SetStep(snapshot.Step);
                     input.Meta.SourceSettings->MutableSnapshot()->SetTxId(snapshot.TxId);
                 }
 
-                if (tasksGraph.GetMeta().UseFollowers) {
-                    input.Meta.SourceSettings->SetUseFollowers(tasksGraph.GetMeta().UseFollowers);
+                if (tasksGraph.GetMeta().UseFollowers || isTableImmutable) {
+                    input.Meta.SourceSettings->SetUseFollowers(tasksGraph.GetMeta().UseFollowers || isTableImmutable);
                 }
 
                 if (serializeAsyncIoSettings) {
                     inputDesc.MutableSource()->MutableSettings()->PackFrom(*input.Meta.SourceSettings);
                 }
+
+                if (isTableImmutable) {
+                    input.Meta.SourceSettings->SetAllowInconsistentReads(true);
+                }
+
             } else {
                 YQL_ENSURE(input.SourceSettings);
                 inputDesc.MutableSource()->MutableSettings()->CopyFrom(*input.SourceSettings);
@@ -1209,21 +1222,25 @@ void FillInputDesc(const TKqpTasksGraph& tasksGraph, NYql::NDqProto::TTaskInput&
         if (input.Meta.StreamLookupSettings) {
             enableMetering = true;
             YQL_ENSURE(input.Meta.StreamLookupSettings);
-            if (snapshot.IsValid()) {
+            bool isTableImmutable = input.Meta.StreamLookupSettings->GetIsTableImmutable();
+
+            if (snapshot.IsValid() && !isTableImmutable) {
                 input.Meta.StreamLookupSettings->MutableSnapshot()->SetStep(snapshot.Step);
                 input.Meta.StreamLookupSettings->MutableSnapshot()->SetTxId(snapshot.TxId);
             } else {
-                YQL_ENSURE(tasksGraph.GetMeta().AllowInconsistentReads, "Expected valid snapshot or enabled inconsistent read mode");
+                YQL_ENSURE(tasksGraph.GetMeta().AllowInconsistentReads || isTableImmutable, "Expected valid snapshot or enabled inconsistent read mode");
                 input.Meta.StreamLookupSettings->SetAllowInconsistentReads(true);
             }
 
-            if (lockTxId) {
+            if (lockTxId && !isTableImmutable) {
                 input.Meta.StreamLookupSettings->SetLockTxId(*lockTxId);
                 input.Meta.StreamLookupSettings->SetLockNodeId(tasksGraph.GetMeta().LockNodeId);
             }
-            if (tasksGraph.GetMeta().LockMode) {
+
+            if (tasksGraph.GetMeta().LockMode && !isTableImmutable) {
                 input.Meta.StreamLookupSettings->SetLockMode(*tasksGraph.GetMeta().LockMode);
             }
+
             transformProto->MutableSettings()->PackFrom(*input.Meta.StreamLookupSettings);
         } else if (input.Meta.SequencerSettings) {
             transformProto->MutableSettings()->PackFrom(*input.Meta.SequencerSettings);
diff --git a/ydb/core/kqp/provider/yql_kikimr_datasource.cpp b/ydb/core/kqp/provider/yql_kikimr_datasource.cpp
@@ -383,6 +383,7 @@ class TKiSourceLoadTableMetadataTransformer : public TGraphTransformerBase {
                     do {
                         auto nextImplTable = implTable->Next;
                         auto& desc = SessionCtx->Tables().GetOrAddTable(implTable->Cluster, SessionCtx->GetDatabase(), implTable->Name);
+                        SessionCtx->Tables().AddIndexImplTableToMainTableMapping(tablePath, implTable->Name);
                         desc.Metadata = std::move(implTable);
                         desc.Load(ctx, sysColumnsEnabled);
                         implTable = std::move(nextImplTable);
diff --git a/ydb/core/kqp/provider/yql_kikimr_provider.cpp b/ydb/core/kqp/provider/yql_kikimr_provider.cpp
@@ -204,6 +204,8 @@ TKikimrTableDescription& TKikimrTablesData::GetOrAddTable(const TString& cluster
     return Tables[std::make_pair(cluster, tablePath)];
 }
 
+
+
 TKikimrTableDescription& TKikimrTablesData::GetTable(const TString& cluster, const TString& table) {
     auto tablePath = table;
     if (TempTablesState) {
@@ -220,6 +222,26 @@ TKikimrTableDescription& TKikimrTablesData::GetTable(const TString& cluster, con
     return *desc;
 }
 
+bool TKikimrTablesData::IsTableImmutable(const TStringBuf& cluster, const TStringBuf& path) {
+    auto mainTableImpl = GetMainTableIfTableIsImplTableOfIndex(cluster, path);
+    if (mainTableImpl) {
+
+        for(const auto& index: mainTableImpl->Metadata->Indexes) {
+            if (index.Type != TIndexDescription::EType::GlobalSyncVectorKMeansTree) {
+                continue;
+            }
+
+            for(const auto& implTable: index.GetImplTables()) {
+                TString implTablePath = TStringBuilder() << mainTableImpl->Metadata->Name << "/" << index.Name << "/" << implTable;
+                if (path == implTablePath)
+                    return true;
+            }
+        }
+    }
+
+    return false;
+}
+
 const TKikimrTableDescription& TKikimrTablesData::ExistingTable(const TStringBuf& cluster,
     const TStringBuf& table) const
 {
diff --git a/ydb/core/kqp/provider/yql_kikimr_provider.h b/ydb/core/kqp/provider/yql_kikimr_provider.h
@@ -207,6 +207,23 @@ class TKikimrTablesData : public TThrRefBase {
         return Tables;
     }
 
+    void AddIndexImplTableToMainTableMapping(const TString& mainTable, const TString& indexTable) {
+        auto [it, success] = IndexTableToMainTable.emplace(indexTable, mainTable);
+        if (!success) {
+            YQL_ENSURE(it->second == mainTable);
+        }
+    }
+
+    const TKikimrTableDescription* GetMainTableIfTableIsImplTableOfIndex(const TStringBuf& cluster, const TStringBuf& id) {
+        auto it = IndexTableToMainTable.find(id);
+        if (it == IndexTableToMainTable.end()) {
+            return nullptr;
+        }
+        return &ExistingTable(cluster, it->second);
+    }
+
+    bool IsTableImmutable(const TStringBuf& cluster, const TStringBuf& path);
+
     std::optional<TString> GetTempTablePath(const TStringBuf& table) const;
 
     void Reset() {
@@ -220,6 +237,8 @@ class TKikimrTablesData : public TThrRefBase {
 private:
     THashMap<std::pair<TString, TString>, TKikimrTableDescription> Tables;
     NKikimr::NKqp::TKqpTempTablesState::TConstPtr TempTablesState;
+
+    THashMap<TString, TString> IndexTableToMainTable;
 };
 
 enum class TYdbOperation : ui64 {
diff --git a/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp b/ydb/core/kqp/query_compiler/kqp_query_compiler.cpp
@@ -1029,6 +1029,7 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             auto tableMeta = TablesData->ExistingTable(Cluster, settings.Table().Cast().Path()).Metadata;
             YQL_ENSURE(tableMeta);
 
+            readProto.SetIsTableImmutable(TablesData->IsTableImmutable(Cluster, settings.Table().Cast().Path()));
             {
 
                 THashMap<TString, const TExprNode*> columnsMap;
@@ -1473,6 +1474,8 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             auto tableMeta = TablesData->ExistingTable(Cluster, streamLookup.Table().Path()).Metadata;
             YQL_ENSURE(tableMeta);
 
+            streamLookupProto.SetIsTableImmutable(TablesData->IsTableImmutable(Cluster, streamLookup.Table().Path()));
+
             FillTablesMap(streamLookup.Table(), streamLookup.Columns(), tablesMap);
             FillTableId(streamLookup.Table(), *streamLookupProto.MutableTable());
 
diff --git a/ydb/core/kqp/runtime/kqp_read_actor.cpp b/ydb/core/kqp/runtime/kqp_read_actor.cpp
@@ -1044,6 +1044,10 @@ class TKqpReadActor : public TActorBootstrapped<TKqpReadActor>, public NYql::NDq
             BrokenLocks.push_back(lock);
         }
 
+        if (UseFollowers) {
+            YQL_ENSURE(Locks.empty());
+        }
+
         CA_LOG_D("Taken " << Locks.size() << " locks");
         Reads[id].SerializedContinuationToken = record.GetContinuationToken();
 
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp
@@ -24,6 +24,7 @@ namespace {
 
 static constexpr TDuration SCHEME_CACHE_REQUEST_TIMEOUT = TDuration::Seconds(10);
 NActors::TActorId MainPipeCacheId = NKikimr::MakePipePerNodeCacheID(false);
+NActors::TActorId FollowersPipeCacheId = NKikimr::MakePipePerNodeCacheID(true);
 
 class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLookupActor>, public NYql::NDq::IDqComputeActorAsyncInput {
 public:
@@ -37,6 +38,8 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
         , Alloc(args.Alloc)
         , Snapshot(settings.GetSnapshot().GetStep(), settings.GetSnapshot().GetTxId())
         , AllowInconsistentReads(settings.GetAllowInconsistentReads())
+        , UseFollowers(settings.GetAllowUseFollowers())
+        , PipeCacheId(UseFollowers ? FollowersPipeCacheId : MainPipeCacheId)
         , LockTxId(settings.HasLockTxId() ? settings.GetLockTxId() : TMaybe<ui64>())
         , NodeLockId(settings.HasLockNodeId() ? settings.GetLockNodeId() : TMaybe<ui32>())
         , LockMode(settings.HasLockMode() ? settings.GetLockMode() : TMaybe<NKikimrDataEvents::ELockMode>())
@@ -207,11 +210,11 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
                 Counters->SentIteratorCancels->Inc();
                 auto cancel = MakeHolder<TEvDataShard::TEvReadCancel>();
                 cancel->Record.SetReadId(id);
-                Send(MainPipeCacheId, new TEvPipeCache::TEvForward(cancel.Release(), state.ShardId, false));
+                Send(PipeCacheId, new TEvPipeCache::TEvForward(cancel.Release(), state.ShardId, false));
             }
         }
 
-        Send(MainPipeCacheId, new TEvPipeCache::TEvUnlink(0));
+        Send(PipeCacheId, new TEvPipeCache::TEvUnlink(0));
         TActorBootstrapped<TKqpStreamLookupActor>::PassAway();
 
         LookupActorSpan.End();
@@ -344,6 +347,14 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
             Locks.push_back(lock);
         }
 
+        if (UseFollowers) {
+            YQL_ENSURE(Locks.empty());
+            if (!record.GetFinished()) {
+                RuntimeError("read from follower returned partial data.", NYql::NDqProto::StatusIds::INTERNAL_ERROR);
+                return;
+            }
+        }
+
         if (!Snapshot.IsValid()) {
             Snapshot = IKqpGateway::TKqpSnapshot(record.GetSnapshot().GetStep(), record.GetSnapshot().GetTxId());
         }
@@ -429,7 +440,7 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
             request->Record.SetMaxRows(defaultSettings.GetMaxRows());
             request->Record.SetMaxBytes(defaultSettings.GetMaxBytes());
 
-            Send(MainPipeCacheId, new TEvPipeCache::TEvForward(request.Release(), read.ShardId, true),
+            Send(PipeCacheId, new TEvPipeCache::TEvForward(request.Release(), read.ShardId, true),
                 IEventHandle::FlagTrackDelivery);
 
             CA_LOG_D("TEvReadAck was sent to shard: " << read.ShardId);
@@ -562,7 +573,7 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
             << ", lockTxId=" << record.GetLockTxId()
             << ", lockNodeId=" << record.GetLockNodeId());
 
-        Send(MainPipeCacheId, new TEvPipeCache::TEvForward(request.Release(), shardId, true),
+        Send(PipeCacheId, new TEvPipeCache::TEvForward(request.Release(), shardId, true),
             IEventHandle::FlagTrackDelivery, 0, LookupActorSpan.GetTraceId());
 
         read.State = EReadState::Running;
@@ -699,6 +710,8 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
     std::shared_ptr<NKikimr::NMiniKQL::TScopedAlloc> Alloc;
     IKqpGateway::TKqpSnapshot Snapshot;
     const bool AllowInconsistentReads;
+    const bool UseFollowers;
+    const TActorId PipeCacheId;
     const TMaybe<ui64> LockTxId;
     const TMaybe<ui32> NodeLockId;
     const TMaybe<NKikimrDataEvents::ELockMode> LockMode;
diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto
@@ -823,6 +823,8 @@ message TKqpStreamLookupSettings {
     optional bool AllowNullKeys = 12;
     optional uint32 AllowNullKeysPrefixSize = 13;
     optional NKikimrDataEvents.ELockMode LockMode = 14;
+    optional bool AllowUseFollowers = 15;
+    optional bool IsTableImmutable = 16;
 }
 
 message TKqpSequencerSettings {
diff --git a/ydb/core/protos/kqp_physical.proto b/ydb/core/protos/kqp_physical.proto
@@ -304,6 +304,7 @@ message TKqpPhyCnStreamLookup {
     bool KeepRowsOrder = 7;
     bool AllowNullKeys = 8;
     uint32 AllowNullKeysPrefixSize = 9;
+    bool IsTableImmutable = 10;
 }
 
 message TKqpPhyCnSequencer {
@@ -348,6 +349,7 @@ message TKqpReadRangesSource {
 
     repeated string SkipNullKeys = 8;
     uint64 SequentialInFlightShards = 9;
+    bool IsTableImmutable = 10;
 }
 
 message TKqpExternalSource {
diff --git a/ydb/core/protos/tx_datashard.proto b/ydb/core/protos/tx_datashard.proto
@@ -269,6 +269,8 @@ message TKqpReadRangesSourceSettings {
     optional NKikimrDataEvents.ELockMode LockMode = 20;
     optional uint32 OptionalSorting = 21;
     optional bool IsBatch = 22 [default = false];
+
+    optional bool IsTableImmutable = 23 [default = false];
 }
 
 message TKqpTaskInfo {
@@ -1549,7 +1551,7 @@ message TEvLocalKMeansRequest {
     optional uint64 Seed = 9;
 
     optional EKMeansState Upload = 11;
-    
+
     optional uint32 K = 10;
     optional uint32 NeedsRounds = 14;
 

Original file line number	Diff line number	Diff line change
`@@ -1044,6 +1044,10 @@ class TKqpReadActor : public TActorBootstrapped<TKqpReadActor>, public NYql::NDq`
`1044`	`1044`	`BrokenLocks.push_back(lock);`
`1045`	`1045`	`}`
`1046`	`1046`
	`1047`	`+ if (UseFollowers) {`
	`1048`	`+ YQL_ENSURE(Locks.empty());`
	`1049`	`+ }`
	`1050`	`+`
`1047`	`1051`	`CA_LOG_D("Taken " << Locks.size() << " locks");`
`1048`	`1052`	`Reads[id].SerializedContinuationToken = record.GetContinuationToken();`
`1049`	`1053`
Original file line number	Diff line number	Diff line change
`@@ -823,6 +823,8 @@ message TKqpStreamLookupSettings {`
`823`	`823`	`optional bool AllowNullKeys = 12;`
`824`	`824`	`optional uint32 AllowNullKeysPrefixSize = 13;`
`825`	`825`	`optional NKikimrDataEvents.ELockMode LockMode = 14;`
	`826`	`+ optional bool AllowUseFollowers = 15;`
	`827`	`+ optional bool IsTableImmutable = 16;`
`826`	`828`	`}`
`827`	`829`
`828`	`830`	`message TKqpSequencerSettings {`
Original file line number	Diff line number	Diff line change
`@@ -304,6 +304,7 @@ message TKqpPhyCnStreamLookup {`
`304`	`304`	`bool KeepRowsOrder = 7;`
`305`	`305`	`bool AllowNullKeys = 8;`
`306`	`306`	`uint32 AllowNullKeysPrefixSize = 9;`
	`307`	`+ bool IsTableImmutable = 10;`
`307`	`308`	`}`
`308`	`309`
`309`	`310`	`message TKqpPhyCnSequencer {`
`@@ -348,6 +349,7 @@ message TKqpReadRangesSource {`
`348`	`349`
`349`	`350`	`repeated string SkipNullKeys = 8;`
`350`	`351`	`uint64 SequentialInFlightShards = 9;`
	`352`	`+ bool IsTableImmutable = 10;`
`351`	`353`	`}`
`352`	`354`
`353`	`355`	`message TKqpExternalSource {`