summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--ydb/core/kqp/compute_actor/kqp_compute_actor.cpp5
-rw-r--r--ydb/core/kqp/compute_actor/kqp_compute_actor.h4
-rw-r--r--ydb/core/kqp/proxy_service/kqp_proxy_service.cpp8
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp25
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_actor.h4
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_factory.cpp10
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_factory.h4
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp189
-rw-r--r--ydb/core/kqp/runtime/kqp_stream_lookup_worker.h4
-rw-r--r--ydb/core/kqp/runtime/kqp_vector_index_levels_cache.cpp158
-rw-r--r--ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h216
-rw-r--r--ydb/core/kqp/runtime/ya.make4
-rw-r--r--ydb/core/kqp/ut/indexes/vector/kqp_indexes_vector_ut.cpp238
-rw-r--r--ydb/core/protos/table_service_config.proto3
14 files changed, 796 insertions, 76 deletions
diff --git a/ydb/core/kqp/compute_actor/kqp_compute_actor.cpp b/ydb/core/kqp/compute_actor/kqp_compute_actor.cpp
index 277e1d29633..0accc33a3eb 100644
--- a/ydb/core/kqp/compute_actor/kqp_compute_actor.cpp
+++ b/ydb/core/kqp/compute_actor/kqp_compute_actor.cpp
@@ -165,10 +165,11 @@ IDqOutputConsumer::TPtr TKqpTaskRunnerExecutionContext::CreateOutputConsumer(con
NYql::NDq::IDqAsyncIoFactory::TPtr CreateKqpAsyncIoFactory(
TIntrusivePtr<TKqpCounters> counters,
std::optional<TKqpFederatedQuerySetup> federatedQuerySetup,
- std::shared_ptr<NYql::NDq::IS3ActorsFactory> s3ActorsFactory
+ std::shared_ptr<NYql::NDq::IS3ActorsFactory> s3ActorsFactory,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache
) {
auto factory = MakeIntrusive<NYql::NDq::TDqAsyncIoFactory>();
- RegisterStreamLookupActorFactory(*factory, counters);
+ RegisterStreamLookupActorFactory(*factory, counters, std::move(vectorIndexLevelsCache));
RegisterKqpReadActor(*factory, counters);
RegisterKqpWriteActor(*factory, counters);
RegisterSequencerActorFactory(*factory, counters);
diff --git a/ydb/core/kqp/compute_actor/kqp_compute_actor.h b/ydb/core/kqp/compute_actor/kqp_compute_actor.h
index 18946047797..28cf272388d 100644
--- a/ydb/core/kqp/compute_actor/kqp_compute_actor.h
+++ b/ydb/core/kqp/compute_actor/kqp_compute_actor.h
@@ -3,6 +3,7 @@
#include <ydb/core/kqp/compute_actor/kqp_compute_events.h>
#include <ydb/core/kqp/counters/kqp_counters.h>
#include <ydb/core/kqp/federated_query/kqp_federated_query_helpers.h>
+#include <ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h>
#include <ydb/core/kqp/runtime/scheduler/kqp_schedulable_actor.h>
#include <ydb/core/scheme/scheme_tabledefs.h>
#include <ydb/library/yql/dq/actors/compute/dq_compute_actor.h>
@@ -89,7 +90,8 @@ IActor* CreateKqpScanFetcher(const NKikimrKqp::TKqpSnapshot& snapshot, std::vect
NYql::NDq::IDqAsyncIoFactory::TPtr CreateKqpAsyncIoFactory(
TIntrusivePtr<TKqpCounters> counters,
std::optional<TKqpFederatedQuerySetup> federatedQuerySetup,
- std::shared_ptr<NYql::NDq::IS3ActorsFactory> s3ActorsFactory);
+ std::shared_ptr<NYql::NDq::IS3ActorsFactory> s3ActorsFactory,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache);
} // namespace NKqp
} // namespace NKikimr
diff --git a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
index 4c953d49bdc..9695af9ea3b 100644
--- a/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
+++ b/ydb/core/kqp/proxy_service/kqp_proxy_service.cpp
@@ -4,6 +4,7 @@
#include <ydb/core/actorlib_impl/long_timer.h>
#include <ydb/core/base/appdata.h>
+#include <ydb/core/base/counters.h>
#include <ydb/core/base/feature_flags.h>
#include <ydb/core/base/location.h>
#include <ydb/core/base/path.h>
@@ -224,11 +225,15 @@ public:
void Bootstrap(const TActorContext &ctx) {
NLwTraceMonPage::ProbeRegistry().AddProbesList(LWTRACE_GET_PROBES(KQP_PROVIDER));
Counters = MakeIntrusive<TKqpCounters>(AppData()->Counters, &TlsActivationContext->AsActorContext());
+ VectorIndexLevelsCache = MakeIntrusive<TVectorIndexLevelsCache>(
+ GetServiceCounters(AppData()->Counters, "kqp"));
+ ctx.Register(CreateVectorIndexLevelsCacheMaintainer(
+ VectorIndexLevelsCache, GetKqpResourceManager(), TableServiceConfig.GetResourceManager()));
FeatureFlags = AppData()->FeatureFlags;
WorkloadManagerConfig = AppData()->WorkloadManagerConfig;
// NOTE: some important actors are constructed within next call
FederatedQuerySetup = FederatedQuerySetupFactory->Make(ctx.ActorSystem());
- AsyncIoFactory = CreateKqpAsyncIoFactory(Counters, FederatedQuerySetup, S3ActorsFactory);
+ AsyncIoFactory = CreateKqpAsyncIoFactory(Counters, FederatedQuerySetup, S3ActorsFactory, VectorIndexLevelsCache);
ModuleResolverState = MakeIntrusive<TModuleResolverState>();
LocalSessions = std::make_unique<TLocalSessionsRegistry>(AppData()->RandomProvider);
@@ -2022,6 +2027,7 @@ private:
TIntrusivePtr<TExecuterMutableConfig> ExecuterConfig;
TIntrusivePtr<TKqpCounters> Counters;
+ TIntrusivePtr<TVectorIndexLevelsCache> VectorIndexLevelsCache;
std::unique_ptr<TLocalSessionsRegistry> LocalSessions;
std::shared_ptr<TKqpProxySharedResources> KqpProxySharedResources;
std::shared_ptr<NYql::NDq::IS3ActorsFactory> S3ActorsFactory;
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp
index 64630ae6d83..356d8117a73 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp
@@ -65,7 +65,7 @@ TKqpStreamLockSettings BuildStreamLockSettings(
class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLookupActor>, public NYql::NDq::IDqComputeActorAsyncInput {
public:
TKqpStreamLookupActor(NYql::NDq::IDqAsyncIoFactory::TInputTransformArguments&& args, NKikimrKqp::TKqpStreamLookupSettings&& settings,
- TIntrusivePtr<TKqpCounters> counters)
+ TIntrusivePtr<TKqpCounters> counters, TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache)
: LogPrefix(TStringBuilder() << "StreamLookupActor, inputIndex: " << args.InputIndex << ", CA Id " << args.ComputeActorId)
, InputIndex(args.InputIndex)
, Input(args.TransformInput)
@@ -97,12 +97,13 @@ public:
args.TaskId,
args.TypeEnv,
args.HolderFactory,
- args.InputDesc))
+ args.InputDesc, vectorIndexLevelsCache))
, MaxTotalBytesQuota(MaxTotalBytesQuotaStreamLookup())
, MaxRowsProcessing(MaxRowsProcessingStreamLookup())
, MaxInFlightReads(MaxInFlightReadsStreamLookup())
, MaxBytesPerFetch(MaxBytesPerFetchStreamLookup())
, Counters(counters)
+ , VectorIndexLevelsCache(std::move(vectorIndexLevelsCache))
, LookupActorSpan(TWilsonKqp::LookupActor, std::move(args.TraceId), "LookupActor")
{
IngressStats.Level = args.StatsLevel;
@@ -435,7 +436,7 @@ private:
void PassAway() final {
Counters->StreamLookupActorsCount->Dec();
-
+
if (!LockSendTime.empty()) {
TInstant now = AppData()->TimeProvider->Now();
TDuration maxInFlightTime = TDuration::Zero();
@@ -447,7 +448,7 @@ private:
}
Counters->MaxInFlightLockTimeOnExit->Collect(maxInFlightTime.MilliSeconds());
}
-
+
{
auto alloc = BindAllocator();
Input.Clear();
@@ -947,7 +948,7 @@ private:
auto lockState = TLockState(requestId, shardId);
lockState.State = TLockState::EState::Running;
Reads.insertLock(requestId, std::move(lockState));
-
+
LockSendTime[requestId] = AppData()->TimeProvider->Now();
}
@@ -958,7 +959,7 @@ private:
<< ", status: " << record.GetStatus());
ui64 requestId = record.GetRequestId();
-
+
if (auto it = LockSendTime.find(requestId); it != LockSendTime.end()) {
Counters->LockLatencyHistogram->Collect((AppData()->TimeProvider->Now() - it->second).MilliSeconds());
LockSendTime.erase(it);
@@ -1059,7 +1060,7 @@ private:
if (hasModifiedRows) {
ProcessInputRows();
}
-
+
if (hasUnmodifiedRows) {
Send(ComputeActorId, new TEvNewAsyncInputDataArrived(InputIndex));
}
@@ -1302,6 +1303,7 @@ private:
const ui64 QuerySpanId;
TReads Reads;
bool SentResultsAvailable = false;
+ THashMap<ui64, TString> CacheKeys;
NUdf::EFetchStatus LastFetchStatus = NUdf::EFetchStatus::Yield;
TPartitioning::TCPtr Partitioning;
const TDuration SchemeCacheRequestTimeout;
@@ -1333,17 +1335,20 @@ private:
size_t MaxRowsDefaultQuota = 0;
TIntrusivePtr<TKqpCounters> Counters;
+ TIntrusivePtr<TVectorIndexLevelsCache> VectorIndexLevelsCache;
+
NWilson::TSpan LookupActorSpan;
NWilson::TSpan LookupActorStateSpan;
-
+
THashMap<ui64, TInstant> LockSendTime;
};
} // namespace
std::pair<NYql::NDq::IDqComputeActorAsyncInput*, NActors::IActor*> CreateStreamLookupActor(NYql::NDq::IDqAsyncIoFactory::TInputTransformArguments&& args,
- NKikimrKqp::TKqpStreamLookupSettings&& settings, TIntrusivePtr<TKqpCounters> counters) {
- auto actor = new TKqpStreamLookupActor(std::move(args), std::move(settings), counters);
+ NKikimrKqp::TKqpStreamLookupSettings&& settings, TIntrusivePtr<TKqpCounters> counters,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache) {
+ auto actor = new TKqpStreamLookupActor(std::move(args), std::move(settings), counters, std::move(vectorIndexLevelsCache));
return {actor, actor};
}
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_actor.h b/ydb/core/kqp/runtime/kqp_stream_lookup_actor.h
index 78839d1f2d9..06a5ba0973a 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_actor.h
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_actor.h
@@ -1,6 +1,7 @@
#pragma once
#include <ydb/core/kqp/counters/kqp_counters.h>
+#include <ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h>
#include <ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io.h>
#include <ydb/core/protos/kqp.pb.h>
@@ -8,7 +9,8 @@ namespace NKikimr {
namespace NKqp {
std::pair<NYql::NDq::IDqComputeActorAsyncInput*, NActors::IActor*> CreateStreamLookupActor(NYql::NDq::IDqAsyncIoFactory::TInputTransformArguments&& args,
- NKikimrKqp::TKqpStreamLookupSettings&& settings, TIntrusivePtr<TKqpCounters>);
+ NKikimrKqp::TKqpStreamLookupSettings&& settings, TIntrusivePtr<TKqpCounters>,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache);
void InterceptStreamLookupActorPipeCache(NActors::TActorId);
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_factory.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_factory.cpp
index d244303fd95..d9b552c248e 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_factory.cpp
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_factory.cpp
@@ -4,10 +4,12 @@
namespace NKikimr {
namespace NKqp {
-void RegisterStreamLookupActorFactory(NYql::NDq::TDqAsyncIoFactory& factory, TIntrusivePtr<TKqpCounters> counters) {
- factory.RegisterInputTransform<NKikimrKqp::TKqpStreamLookupSettings>("StreamLookupInputTransformer", [counters](NKikimrKqp::TKqpStreamLookupSettings&& settings,
- NYql::NDq::TDqAsyncIoFactory::TInputTransformArguments&& args) {
- return CreateStreamLookupActor(std::move(args), std::move(settings), counters);
+void RegisterStreamLookupActorFactory(NYql::NDq::TDqAsyncIoFactory& factory, TIntrusivePtr<TKqpCounters> counters,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache) {
+ factory.RegisterInputTransform<NKikimrKqp::TKqpStreamLookupSettings>("StreamLookupInputTransformer",
+ [counters, vectorIndexLevelsCache](NKikimrKqp::TKqpStreamLookupSettings&& settings,
+ NYql::NDq::TDqAsyncIoFactory::TInputTransformArguments&& args) {
+ return CreateStreamLookupActor(std::move(args), std::move(settings), counters, vectorIndexLevelsCache);
});
}
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_factory.h b/ydb/core/kqp/runtime/kqp_stream_lookup_factory.h
index db49e005fa9..70e6fe019aa 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_factory.h
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_factory.h
@@ -1,12 +1,14 @@
#pragma once
#include <ydb/core/kqp/counters/kqp_counters.h>
+#include <ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h>
#include <ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io_factory.h>
namespace NKikimr {
namespace NKqp {
-void RegisterStreamLookupActorFactory(NYql::NDq::TDqAsyncIoFactory& factory, TIntrusivePtr<NKqp::TKqpCounters>);
+void RegisterStreamLookupActorFactory(NYql::NDq::TDqAsyncIoFactory& factory, TIntrusivePtr<NKqp::TKqpCounters>,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache);
} // namespace NKqp
} // namespace NKikimr
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
index c44471b33be..39b08617ed8 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.cpp
@@ -1,5 +1,6 @@
#include "kqp_stream_lookup_worker.h"
#include "kqp_stream_lookup_join_helpers.h"
+#include "kqp_vector_index_levels_cache.h"
#include <ydb/core/kqp/common/kqp_resolve.h>
#include <ydb/core/kqp/common/kqp_types.h>
@@ -42,9 +43,19 @@ TStreamLookupShardReadResult::~TStreamLookupShardReadResult() {
namespace {
+// Check if this is a vector index level table
+bool IsVectorIndexLevelTable(const TString& tablePath) {
+ return tablePath.EndsWith(NTableIndex::NKMeans::LevelTable);
+}
+
+struct TCacheData {
+ ui64 PendingReads = 0;
+ TOwnedCellVecBatch Batch;
+};
struct TReadState {
std::vector<TOwnedTableRange> PendingKeys;
+ std::optional<TString> CacheKey;
TMaybe<TOwnedCellVec> LastProcessedKey;
ui32 FirstUnprocessedQuery = 0;
@@ -106,8 +117,12 @@ TTableId TKqpStreamLookupWorker::GetTableId() const {
class TKqpLookupRows : public TKqpStreamLookupWorker {
public:
TKqpLookupRows(TLookupSettings&& settings, const NMiniKQL::TTypeEnvironment& typeEnv,
- const NMiniKQL::THolderFactory& holderFactory)
+ const NMiniKQL::THolderFactory& holderFactory, TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache)
: TKqpStreamLookupWorker(std::move(settings), typeEnv, holderFactory)
+ , VectorIndexLevelsCache(
+ IsVectorIndexLevelTable(Settings.TablePath) && vectorIndexLevelsCache && vectorIndexLevelsCache->MaxBytes() > 0
+ ? std::move(vectorIndexLevelsCache)
+ : nullptr)
{
}
@@ -145,7 +160,7 @@ public:
const i32 lookupKeySize = std::min(Settings.KeyColumns.size(), Settings.InputColumns.size());
NMiniKQL::TStringProviderBackend backend;
std::vector<TCell> keyCells(lookupKeySize);
-
+
for (size_t colId = 0; colId < Settings.InputColumns.size(); ++colId) {
const auto& lookupKeyColumn = Settings.InputColumns[colId];
if (0 <= lookupKeyColumn.KeyOrder) {
@@ -160,6 +175,16 @@ public:
}
virtual void AddInputRowImpl(std::vector<TCell> keyCells) {
+ std::optional<TString> cacheKey;
+ if (VectorIndexLevelsCache) {
+ cacheKey = TSerializedCellVec::Serialize(keyCells);
+ auto result = VectorIndexLevelsCache->Get(Settings.TableId.PathId, cacheKey.value());
+ if (result && !result->BatchRows.empty()) {
+ PendingBatches.push_back(result);
+ return;
+ }
+ }
+
if (keyCells.size() < Settings.KeyColumns.size()) {
// build prefix range [[key_prefix, NULL, ..., NULL], [key_prefix, +inf, ..., +inf])
std::vector<TCell> fromCells(Settings.KeyColumns.size());
@@ -175,6 +200,10 @@ public:
// full pk, build point
UnprocessedKeys.emplace_back(std::move(keyCells));
}
+
+ if (cacheKey.has_value()) {
+ CacheKeys.emplace_back(std::move(cacheKey.value()));
+ }
}
void RebuildRequest(const ui64 shardId, const ui64& prevReadId, ui64& newReadId) final {
@@ -259,6 +288,28 @@ public:
UnprocessedKeys.pop_front();
auto partitions = partitioning->GetIntersectionWithRange(GetKeyColumnTypes(), range);
+ if (!CacheKeys.empty()) {
+ ui64 pendingReads = 0;
+ for(auto [shardId, range] : partitions) {
+ THolder<TEvDataShard::TEvRead> request(new TEvDataShard::TEvRead());
+ std::vector<TOwnedTableRange> ranges = {range};
+ ui64 nextReadId = ++readId;
+ FillReadRequest(nextReadId, request, ranges);
+ ScheduledReads.emplace_back(shardId, std::move(request));
+ ++pendingReads;
+ YQL_ENSURE(ReadStateByReadId.emplace(
+ nextReadId,
+ TReadState{
+ .PendingKeys = std::move(ranges),
+ .CacheKey = CacheKeys.front(),
+ }).second);
+ }
+
+ CacheKeysMap[CacheKeys.front()] = {.PendingReads = pendingReads};
+ CacheKeys.pop_front();
+ continue;
+ }
+
for (auto [shardId, range] : partitions) {
if (range.Point) {
pointsPerShard[shardId].push_back(std::move(range));
@@ -292,7 +343,7 @@ public:
}
bool HasPendingResults() final {
- return !ReadResults.empty();
+ return !ReadResults.empty() || !PendingBatches.empty();
}
void AddResult(TStreamLookupShardReadResult result) final {
@@ -309,10 +360,66 @@ public:
ReadResults.emplace_back(std::move(result));
}
+ void ProcessResultRow(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, TConstArrayRef<TCell> resultRow, TReadResultStats& resultStats, ui64 shardId, bool hasReadStats) {
+ YQL_ENSURE(resultRow.size() <= Settings.Columns.size(), "Result columns mismatch");
+
+ if (Settings.VectorTopK && Settings.VectorTopK->DistinctColumnsSize()) {
+ TVector<TCell> uniqueKey;
+
+ for (auto& colIdx: Settings.VectorTopK->GetDistinctColumns()) {
+ YQL_ENSURE(colIdx < resultRow.size(), "Unique column index is too large");
+ uniqueKey.push_back(resultRow.at(colIdx));
+ }
+
+ TString serializedKey = TSerializedCellVec::Serialize(uniqueKey);
+
+ if (UniqueKeys.contains(serializedKey)) {
+ return;
+ }
+ UniqueKeys.insert(serializedKey);
+ }
+
+ NUdf::TUnboxedValue* rowItems = nullptr;
+ auto row = HolderFactory.CreateDirectArrayHolder(Settings.Columns.size(), rowItems);
+ i64 storageRowSize = 0;
+ for (size_t colIndex = 0, resultColIndex = 0; colIndex < Settings.Columns.size(); ++colIndex) {
+ const auto& column = Settings.Columns[colIndex];
+ if (IsSystemColumn(column.Name)) {
+ NMiniKQL::FillSystemColumn(rowItems[colIndex], shardId, column.Id, column.PType);
+ } else {
+ YQL_ENSURE(resultColIndex < resultRow.size());
+ storageRowSize += resultRow[resultColIndex].Size();
+ rowItems[colIndex] = NMiniKQL::GetCellValue(resultRow[resultColIndex], column.PType);
+ ++resultColIndex;
+ }
+ }
+
+ batch.push_back(std::move(row));
+ storageRowSize = std::max(storageRowSize, (i64)8);
+
+
+ if (!Settings.VectorTopK || !hasReadStats) {
+ resultStats.ReadRowsCount += 1;
+ resultStats.ReadBytesCount += storageRowSize;
+ }
+
+ resultStats.ResultRowsCount += 1;
+ resultStats.ResultBytesCount += storageRowSize;
+ }
+
TReadResultStats ReplyResult(NKikimr::NMiniKQL::TUnboxedValueBatch& batch, i64 ) final {
TReadResultStats resultStats;
batch.clear();
+ while(!PendingBatches.empty()) {
+ auto& frontline = PendingBatches.front();
+ for(const auto& row: frontline->BatchRows) {
+ ProcessResultRow(batch, row, resultStats, 0, false);
+ }
+
+ PendingBatches.pop_front();
+ }
+
while (!ReadResults.empty() && !resultStats.SizeLimitExceeded) {
auto& result = ReadResults.front();
if (!result.UnprocessedResultRow && Settings.VectorTopK &&
@@ -322,52 +429,30 @@ public:
}
for (; result.UnprocessedResultRow < result.ReadResult->Get()->GetRowsCount(); ++result.UnprocessedResultRow) {
const auto& resultRow = result.ReadResult->Get()->GetCells(result.UnprocessedResultRow);
- YQL_ENSURE(resultRow.size() <= Settings.Columns.size(), "Result columns mismatch");
-
- if (Settings.VectorTopK && Settings.VectorTopK->DistinctColumnsSize()) {
- TVector<TCell> uniqueKey;
- for (auto& colIdx: Settings.VectorTopK->GetDistinctColumns()) {
- YQL_ENSURE(colIdx < resultRow.size(), "Unique column index is too large");
- uniqueKey.push_back(resultRow.at(colIdx));
- }
- TString serializedKey = TSerializedCellVec::Serialize(uniqueKey);
- if (UniqueKeys.contains(serializedKey)) {
- continue;
- }
- UniqueKeys.insert(serializedKey);
- }
-
- NUdf::TUnboxedValue* rowItems = nullptr;
- auto row = HolderFactory.CreateDirectArrayHolder(Settings.Columns.size(), rowItems);
-
- i64 storageRowSize = 0;
- for (size_t colIndex = 0, resultColIndex = 0; colIndex < Settings.Columns.size(); ++colIndex) {
- const auto& column = Settings.Columns[colIndex];
- if (IsSystemColumn(column.Name)) {
- NMiniKQL::FillSystemColumn(rowItems[colIndex], result.ShardId, column.Id, column.PType);
- } else {
- YQL_ENSURE(resultColIndex < resultRow.size());
- storageRowSize += resultRow[resultColIndex].Size();
- rowItems[colIndex] = NMiniKQL::GetCellValue(resultRow[resultColIndex], column.PType);
- ++resultColIndex;
- }
- }
-
- batch.push_back(std::move(row));
- storageRowSize = std::max(storageRowSize, (i64)8);
-
- if (!Settings.VectorTopK || !result.ReadResult->Get()->Record.HasStats()) {
- resultStats.ReadRowsCount += 1;
- resultStats.ReadBytesCount += storageRowSize;
- }
- resultStats.ResultRowsCount += 1;
- resultStats.ResultBytesCount += storageRowSize;
+ ProcessResultRow(batch, resultRow, resultStats, result.ShardId, result.ReadResult->Get()->Record.HasStats());
}
if (result.UnprocessedResultRow == result.ReadResult->Get()->GetRowsCount()) {
if (result.ReadResult->Get()->Record.GetFinished()) {
// delete finished read
auto it = ReadStateByReadId.find(result.ReadResult->Get()->Record.GetReadId());
+ if (it->second.CacheKey.has_value()) {
+ auto entryIt = CacheKeysMap.find(it->second.CacheKey.value());
+ YQL_ENSURE(entryIt != CacheKeysMap.end());
+ auto& entry = entryIt->second;
+ entry.PendingReads--;
+ for(ui32 idx = 0; idx < result.ReadResult->Get()->GetRowsCount(); idx++) {
+ entry.Batch.Append(result.ReadResult->Get()->GetCells(idx));
+ }
+
+ if (entry.PendingReads == 0) {
+ auto ptr = MakeIntrusive<TCachedLevelTableData>();
+ ptr->BatchRows = std::move(entry.Batch);
+ VectorIndexLevelsCache->Put(Settings.TableId.PathId, it->second.CacheKey.value(), ptr);
+ CacheKeysMap.erase(entryIt);
+ }
+ }
+
ReadStateByReadId.erase(it);
}
@@ -417,7 +502,8 @@ public:
return UnprocessedKeys.empty()
&& ReadStateByReadId.empty()
&& ReadResults.empty()
- && ScheduledReads.empty();
+ && ScheduledReads.empty()
+ && PendingBatches.empty();
}
void ResetRowsProcessing(ui64 readId) final {
@@ -463,7 +549,7 @@ private:
}
}
- if (Settings.VectorTopK) {
+ if (Settings.VectorTopK && !VectorIndexLevelsCache) {
*record.MutableVectorTopK() = *Settings.VectorTopK;
}
@@ -495,11 +581,15 @@ private:
}
private:
+ std::unordered_map<TString, TCacheData> CacheKeysMap;
+ std::deque<TString> CacheKeys;
std::deque<TOwnedTableRange> UnprocessedKeys;
std::deque<std::pair<ui64, THolder<TEvDataShard::TEvRead>>> ScheduledReads;
std::unordered_map<ui64, TReadState> ReadStateByReadId;
std::deque<TStreamLookupShardReadResult> ReadResults;
std::unordered_set<TString> UniqueKeys;
+ TIntrusivePtr<TVectorIndexLevelsCache> VectorIndexLevelsCache;
+ std::deque<TCachedLevelTableDataPtr> PendingBatches;
};
class TKqpJoinRows : public TKqpStreamLookupWorker {
@@ -1231,7 +1321,8 @@ private:
std::unique_ptr<TKqpStreamLookupWorker> CreateStreamLookupWorker(NKikimrKqp::TKqpStreamLookupSettings&& settings,
ui64 taskId,
const NMiniKQL::TTypeEnvironment& typeEnv, const NMiniKQL::THolderFactory& holderFactory,
- const NYql::NDqProto::TTaskInput& inputDesc) {
+ const NYql::NDqProto::TTaskInput& inputDesc,
+ TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache) {
TLookupSettings preparedSettings;
preparedSettings.TablePath = std::move(settings.GetTable().GetPath());
@@ -1333,7 +1424,7 @@ std::unique_ptr<TKqpStreamLookupWorker> CreateStreamLookupWorker(NKikimrKqp::TKq
case NKqpProto::EStreamLookupStrategy::LOOKUP:
case NKqpProto::EStreamLookupStrategy::UNIQUE:
case NKqpProto::EStreamLookupStrategy::LOCK_AND_LOOKUP:
- return std::make_unique<TKqpLookupRows>(std::move(preparedSettings), typeEnv, holderFactory);
+ return std::make_unique<TKqpLookupRows>(std::move(preparedSettings), typeEnv, holderFactory, vectorIndexLevelsCache);
case NKqpProto::EStreamLookupStrategy::JOIN:
case NKqpProto::EStreamLookupStrategy::SEMI_JOIN:
return std::make_unique<TKqpJoinRows>(std::move(preparedSettings), taskId, typeEnv, holderFactory, inputDesc);
@@ -1349,8 +1440,8 @@ std::unique_ptr<TKqpStreamLookupWorker> CreateLookupWorker(TLookupSettings&& set
AFL_ENSURE(!settings.KeepRowsOrder);
AFL_ENSURE(!settings.AllowNullKeysPrefixSize);
AFL_ENSURE(settings.InputColumns.size() <= settings.KeyColumns.size());
-
- return std::make_unique<TKqpLookupRows>(std::move(settings), typeEnv, holderFactory);
+
+ return std::make_unique<TKqpLookupRows>(std::move(settings), typeEnv, holderFactory, nullptr);
}
} // namespace NKqp
diff --git a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h
index 4667ff0cdff..348093c84f5 100644
--- a/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h
+++ b/ydb/core/kqp/runtime/kqp_stream_lookup_worker.h
@@ -1,5 +1,7 @@
#pragma once
+#include "kqp_vector_index_levels_cache.h"
+
#include <ydb/core/protos/kqp.pb.h>
#include <yql/essentials/minikql/mkql_node.h>
#include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
@@ -120,7 +122,7 @@ protected:
std::unique_ptr<TKqpStreamLookupWorker> CreateStreamLookupWorker(NKikimrKqp::TKqpStreamLookupSettings&& settings,
ui64 taskId,
const NMiniKQL::TTypeEnvironment& typeEnv, const NMiniKQL::THolderFactory& holderFactory,
- const NYql::NDqProto::TTaskInput& inputDesc);
+ const NYql::NDqProto::TTaskInput& inputDesc, TIntrusivePtr<TVectorIndexLevelsCache> vectorIndexLevelsCache);
std::unique_ptr<TKqpStreamLookupWorker> CreateLookupWorker(TLookupSettings&& settings,
const NMiniKQL::TTypeEnvironment& typeEnv, const NMiniKQL::THolderFactory& holderFactory);
diff --git a/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.cpp b/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.cpp
new file mode 100644
index 00000000000..6fd13f807a8
--- /dev/null
+++ b/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.cpp
@@ -0,0 +1,158 @@
+#include "kqp_vector_index_levels_cache.h"
+#include "kqp_read_iterator_common.h"
+
+#include <ydb/core/base/appdata.h>
+#include <ydb/core/base/domain.h>
+#include <ydb/core/base/tablet_pipecache.h>
+#include <ydb/core/mon/mon.h>
+#include <ydb/core/scheme/scheme_tablecell.h>
+#include <ydb/core/tx/scheme_board/events.h>
+#include <ydb/core/tx/scheme_board/subscriber.h>
+#include <ydb/core/resource_pools/resource_pool_settings.h>
+#include <ydb/core/cms/console/configs_dispatcher.h>
+#include <ydb/core/cms/console/console.h>
+
+#include <ydb/library/actors/core/actor_bootstrapped.h>
+#include <ydb/library/actors/core/hfunc.h>
+#include <ydb/library/actors/core/log.h>
+#include <ydb/library/services/services.pb.h>
+
+#include <library/cpp/monlib/service/pages/templates.h>
+
+#include <util/generic/algorithm.h>
+#include <util/generic/hash.h>
+
+namespace NKikimr::NKqp {
+
+namespace {
+
+constexpr ui64 LevelCacheTxId = std::numeric_limits<ui64>::max() - 1;
+
+#define LOG_E(stream) LOG_ERROR_S(*TlsActivationContext, NKikimrServices::KQP_COMPUTE, "VectorIndexLevelsCacheMaintainer: " << stream)
+#define LOG_N(stream) LOG_NOTICE_S(*TlsActivationContext, NKikimrServices::KQP_COMPUTE, "VectorIndexLevelsCacheMaintainer: " << stream)
+
+
+class TVectorIndexLevelsCacheMaintainer
+ : public TActorBootstrapped<TVectorIndexLevelsCacheMaintainer>
+{
+ using TBase = TActorBootstrapped<TVectorIndexLevelsCacheMaintainer>;
+
+public:
+ explicit TVectorIndexLevelsCacheMaintainer(
+ TIntrusivePtr<TVectorIndexLevelsCache> cache,
+ std::shared_ptr<NRm::IKqpResourceManager> rm,
+ const NKikimrConfig::TTableServiceConfig::TResourceManager& initialConfig)
+ : Cache_(std::move(cache))
+ , ResourceManager(rm)
+ , Tx(MakeIntrusive<NRm::TTxState>(ResourceManager, LevelCacheTxId, TInstant::Now(), NKikimr::NResourcePool::DEFAULT_POOL_ID, 100.0, AppData()->TenantName, false))
+ , RmConfig(initialConfig)
+ {}
+
+ struct TEvPrivate {
+ enum EEv {
+ EvIncreaseCacheSize = EventSpaceBegin(TEvents::ES_PRIVATE),
+ };
+
+ struct TEvIncreaseCacheSize: public TEventLocal<TEvIncreaseCacheSize, EEv::EvIncreaseCacheSize> {};
+ };
+
+ void Bootstrap() {
+ Become(&TThis::StateWork);
+ UpdateCacheSize();
+
+ ui32 tableServiceConfigKind = (ui32) NKikimrConsole::TConfigItem::TableServiceConfigItem;
+ Send(NConsole::MakeConfigsDispatcherID(SelfId().NodeId()),
+ new NConsole::TEvConfigsDispatcher::TEvSetConfigSubscriptionRequest({tableServiceConfigKind}),
+ IEventHandle::FlagTrackDelivery);
+ }
+
+ STATEFN(StateWork) {
+ switch (ev->GetTypeRewrite()) {
+ IgnoreFunc(NConsole::TEvConfigsDispatcher::TEvSetConfigSubscriptionResponse);
+ hFunc(NConsole::TEvConsole::TEvConfigNotificationRequest, Handle);
+ hFunc(TEvents::TEvUndelivered, Handle);
+ cFunc(TEvents::TSystem::Poison, PassAway);
+ cFunc(TEvPrivate::EvIncreaseCacheSize, UpdateCacheSize);
+ }
+ }
+
+ void PassAway() override {
+ TBase::PassAway();
+ }
+
+ void Handle(TEvents::TEvUndelivered::TPtr& ev) {
+ switch (ev->Get()->SourceType) {
+ case NConsole::TEvConfigsDispatcher::EvSetConfigSubscriptionRequest:
+ LOG_E("Failed to deliver subscription request to config dispatcher");
+ break;
+ case NConsole::TEvConsole::EvConfigNotificationResponse:
+ LOG_E("Failed to deliver config notification response");
+ break;
+ default:
+ LOG_E("Undelivered event with unexpected source type: " << ev->Get()->SourceType);
+ break;
+ }
+ }
+
+ void Handle(NConsole::TEvConsole::TEvConfigNotificationRequest::TPtr& ev) {
+ auto &event = ev->Get()->Record;
+
+ auto newConfig = event.GetConfig().GetTableServiceConfig();
+
+ RmConfig.Swap(newConfig.MutableResourceManager());
+ UpdateCacheSize();
+ }
+
+ void UpdateCacheSize() {
+ i64 maxCurrentSizeBytes = Cache_->MaxBytes();
+ i64 leftBytes = static_cast<i64>(maxCurrentSizeBytes) - static_cast<i64>(Cache_->Bytes());
+ ui64 increaseBatchSize = RmConfig.GetKqpLevelCacheIncreaseBatchSizeBytes();
+
+ ui64 maxAllowedSizeBytes = RmConfig.GetKqpLevelCacheMaxSizeBytes();
+
+ if ((maxCurrentSizeBytes == 0 || leftBytes < static_cast<i64>(increaseBatchSize))
+ && static_cast<ui64>(maxCurrentSizeBytes) + increaseBatchSize <= maxAllowedSizeBytes) {
+
+ auto res = ResourceManager->AllocateResources(*Tx, 1, NRm::TKqpResourcesRequest{.Memory=increaseBatchSize});
+ if (res) {
+ Cache_->SetMaxBytes(maxCurrentSizeBytes + increaseBatchSize);
+ LOG_N("Altered max bytes to " << HumanReadableSize(maxCurrentSizeBytes + increaseBatchSize, ESizeFormat::SF_BYTES)
+ << ", prev size " << HumanReadableSize(maxCurrentSizeBytes, ESizeFormat::SF_BYTES));
+ }
+
+ } else if (maxAllowedSizeBytes < static_cast<ui64>(maxCurrentSizeBytes)) {
+ ui64 deficit = static_cast<ui64>(maxCurrentSizeBytes) - maxAllowedSizeBytes;
+ ui64 change = std::min(increaseBatchSize, deficit);
+ ResourceManager->FreeResources(*Tx, 1, NRm::TKqpResourcesRequest{.Memory=change});
+ i64 newSize = maxCurrentSizeBytes - static_cast<i64>(change);
+ Cache_->SetMaxBytes(newSize);
+ LOG_N("Altered max bytes to " << HumanReadableSize(newSize, ESizeFormat::SF_BYTES)
+ << ", prev size " << HumanReadableSize(maxCurrentSizeBytes, ESizeFormat::SF_BYTES));
+ }
+
+ Schedule(TDuration::Seconds(1), new TEvPrivate::TEvIncreaseCacheSize);
+ }
+
+private:
+ TIntrusivePtr<TVectorIndexLevelsCache> Cache_;
+
+ std::shared_ptr<NRm::IKqpResourceManager> ResourceManager;
+ TIntrusivePtr<NRm::TTxState> Tx;
+ NKikimrConfig::TTableServiceConfig::TResourceManager RmConfig;
+};
+
+#undef LOG_E
+#undef LOG_N
+
+} // anonymous namespace
+
+IActor* CreateVectorIndexLevelsCacheMaintainer(
+ TIntrusivePtr<TVectorIndexLevelsCache> cache,
+ std::shared_ptr<NRm::IKqpResourceManager> rm,
+ const NKikimrConfig::TTableServiceConfig::TResourceManager& initialConfig)
+{
+ return new TVectorIndexLevelsCacheMaintainer(std::move(cache), std::move(rm), initialConfig);
+}
+
+
+} // namespace NKikimr::NKqp
diff --git a/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h b/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h
new file mode 100644
index 00000000000..3761d15f46d
--- /dev/null
+++ b/ydb/core/kqp/runtime/kqp_vector_index_levels_cache.h
@@ -0,0 +1,216 @@
+#pragma once
+
+#include <ydb/core/base/events.h>
+#include <ydb/core/scheme/scheme_pathid.h>
+#include <ydb/core/tx/datashard/datashard.h>
+
+#include <ydb/library/actors/core/actorid.h>
+#include <ydb/library/actors/core/event_local.h>
+#include <ydb/library/actors/core/events.h>
+
+#include <ydb/core/scheme/scheme_tablecell.h>
+#include <ydb/core/kqp/rm_service/kqp_rm_service.h>
+#include <library/cpp/cache/cache.h>
+#include <library/cpp/monlib/dynamic_counters/counters.h>
+
+#include <util/datetime/base.h>
+#include <util/generic/hash.h>
+#include <util/generic/ptr.h>
+#include <util/generic/string.h>
+#include <util/generic/vector.h>
+#include <util/system/rwlock.h>
+
+#include <atomic>
+
+namespace NKikimr::NKqp {
+
+// Cached level table data - stores rows in a format ready for zero-copy access
+// Uses shared_ptr for safe concurrent access without copying
+struct TCachedLevelTableData : public TAtomicRefCount<TCachedLevelTableData> {
+ // Store serialized cell vectors directly - no deserialization needed
+ TOwnedCellVecBatch BatchRows;
+
+ TCachedLevelTableData() = default;
+};
+
+using TCachedLevelTableDataPtr = TIntrusivePtr<TCachedLevelTableData>;
+
+struct TLevelTableCacheKey {
+ TPathId TableId;
+ TString SerializedKeys;
+
+ bool operator==(const TLevelTableCacheKey& other) const {
+ return TableId == other.TableId && SerializedKeys == other.SerializedKeys;
+ }
+
+ size_t Hash() const {
+ return CombineHashes(THash<TPathId>()(TableId), THash<TString>()(SerializedKeys));
+ }
+};
+
+} // namespace NKikimr::NKqp
+
+template <>
+struct THash<NKikimr::NKqp::TLevelTableCacheKey> {
+ size_t operator()(const NKikimr::NKqp::TLevelTableCacheKey& key) const {
+ return key.Hash();
+ }
+};
+
+namespace NKikimr::NKqp {
+
+// LRU cache entry wrapper
+struct TLevelTableCacheEntry {
+ TCachedLevelTableDataPtr Data;
+ size_t ByteSize = 0; // approximate memory footprint, computed once at insertion
+
+ TLevelTableCacheEntry() = default;
+ TLevelTableCacheEntry(TCachedLevelTableDataPtr data, size_t byteSize)
+ : Data(std::move(data))
+ , ByteSize(byteSize)
+ {}
+};
+
+// Approximate memory footprint of cached data. Used both to drive the bytes
+// gauge and to give operators a sense of cache pressure.
+inline size_t EstimateCachedDataBytes(const TCachedLevelTableData& data) {
+ return data.BatchRows.DataSizeEstimate();
+}
+
+class TVectorIndexLevelsCache : public TThrRefBase {
+public:
+ using TPtr = TIntrusivePtr<TVectorIndexLevelsCache>;
+
+public:
+
+ // Try to get cached data - returns shared pointer (zero-copy)
+ TCachedLevelTableDataPtr Get(const TPathId& tableId, const TString& serializedKeys) {
+ TWriteGuard guard(Lock_);
+
+ TLevelTableCacheKey key{tableId, serializedKeys};
+ auto it = Cache_.Find(key);
+ if (it != Cache_.End()) {
+ Hits_->Inc();
+ return it.Value().Data;
+ }
+ Misses_->Inc();
+ return nullptr;
+ }
+
+ // Store data in cache
+ void Put(const TPathId& tableId, const TString& serializedKeys, TCachedLevelTableDataPtr data) {
+ const size_t bytes = data ? EstimateCachedDataBytes(*data) : 0;
+
+ TWriteGuard guard(Lock_);
+
+ TLevelTableCacheKey key{tableId, serializedKeys};
+
+ if (auto existing = Cache_.FindWithoutPromote(key); existing != Cache_.End()) {
+ return;
+ } else {
+ Evict(bytes);
+ Cache_.Insert(key, TLevelTableCacheEntry(std::move(data), bytes));
+ BytesTotal_ += bytes;
+ Inserts_->Inc();
+ RefreshGauges();
+ }
+ }
+
+ void Evict(ui64 bytesToAdd = 0) {
+ // Byte-cap eviction. The TLRUCache only caps by entry count; with
+ // multi-KB rows (embedding columns) one entry can be megabytes, so
+ // the entry-count cap alone can let total bytes balloon.
+ auto maxBytes = MaxCacheBytes_.load();
+ while (BytesTotal_ + bytesToAdd > maxBytes && Cache_.Size() >= 1) {
+ auto oldest = Cache_.FindOldest();
+ Y_ENSURE(oldest != Cache_.End());
+ BytesTotal_ -= oldest.Value().ByteSize;
+ Cache_.Erase(oldest);
+ Invalidations_->Inc();
+ }
+ }
+
+ // Number of cached entries. Cheap; primarily useful for tests.
+ size_t Size() const {
+ TReadGuard guard(Lock_);
+ return Cache_.Size();
+ }
+
+ size_t MaxBytes() const { return MaxCacheBytes_.load(); }
+
+ void SetMaxBytes(ui64 maxBytes) {
+ MaxCacheBytes_.store(maxBytes);
+
+ TWriteGuard guard(Lock_);
+
+ Evict();
+
+ RefreshGauges();
+ }
+
+ // Counter getters for monitoring (cumulative since process start).
+ ui64 HitsTotal() const { return Hits_ ? Hits_->Val() : 0; }
+ ui64 MissesTotal() const { return Misses_ ? Misses_->Val() : 0; }
+ ui64 InsertsTotal() const { return Inserts_ ? Inserts_->Val() : 0; }
+ ui64 InvalidationsTotal() const { return Invalidations_ ? Invalidations_->Val() : 0; }
+
+ // Approximate bytes currently held in the cache. O(1).
+ size_t Bytes() const {
+ TReadGuard guard(Lock_);
+ return BytesTotal_;
+ }
+
+ // Counters argument: a parent group under which the cache exposes its
+ // metrics. Pass nullptr in tests / contexts without a counters tree — the
+ // cache attaches to a private group and metric updates become no-op-ish.
+ // maxCacheBytes caps total cached bytes (rows + metadata)
+ explicit TVectorIndexLevelsCache(::NMonitoring::TDynamicCounterPtr counters = nullptr)
+ : MaxCacheBytes_(0)
+ , Cache_(MaxCacheEntries_)
+ {
+ auto group = counters
+ ? counters->GetSubgroup("subsystem", "VectorIndexLevelsCache")
+ : MakeIntrusive<::NMonitoring::TDynamicCounters>();
+
+ MaxSizeGauge_ = group->GetCounter("MaxSize", false);
+ BytesGauge_ = group->GetCounter("Bytes", false);
+ EntriesGauge_ = group->GetCounter("Entries", false);
+ Hits_ = group->GetCounter("Hits", true);
+ Misses_ = group->GetCounter("Misses", true);
+ Inserts_ = group->GetCounter("Inserts", true);
+ Invalidations_ = group->GetCounter("Invalidations", true);
+ }
+
+private:
+ // Level tables are immutable, so cache can be long-lived. Entry-count
+ // cap is a coarse upper bound; the byte cap below is the active limit.
+ static constexpr size_t MaxCacheEntries_ = 10000;
+
+ // Must be called with the write lock held.
+ void RefreshGauges() {
+ BytesGauge_->Set(BytesTotal_);
+ EntriesGauge_->Set(Cache_.Size());
+ MaxSizeGauge_->Set(MaxCacheBytes_.load());
+ }
+
+ std::atomic<ui64> MaxCacheBytes_;
+ size_t BytesTotal_ = 0;
+
+ mutable TRWMutex Lock_;
+ TLRUCache<TLevelTableCacheKey, TLevelTableCacheEntry> Cache_;
+
+ ::NMonitoring::TDynamicCounters::TCounterPtr MaxSizeGauge_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr BytesGauge_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr EntriesGauge_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr Hits_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr Misses_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr Inserts_;
+ ::NMonitoring::TDynamicCounters::TCounterPtr Invalidations_;
+};
+
+IActor* CreateVectorIndexLevelsCacheMaintainer(
+ TIntrusivePtr<TVectorIndexLevelsCache> cache,
+ std::shared_ptr<NRm::IKqpResourceManager> rm,
+ const NKikimrConfig::TTableServiceConfig::TResourceManager& initialConfig);
+
+} // namespace NKikimr::NKqp \ No newline at end of file
diff --git a/ydb/core/kqp/runtime/ya.make b/ydb/core/kqp/runtime/ya.make
index 2871362863a..c98dd18cb71 100644
--- a/ydb/core/kqp/runtime/ya.make
+++ b/ydb/core/kqp/runtime/ya.make
@@ -27,6 +27,8 @@ SRCS(
kqp_stream_lock_worker.cpp
kqp_stream_lock_worker.h
kqp_transport.cpp
+ kqp_vector_index_levels_cache.cpp
+ kqp_vector_index_levels_cache.h
kqp_vector_actor.cpp
kqp_write_actor_settings.cpp
kqp_write_actor.cpp
@@ -51,9 +53,11 @@ PEERDIR(
ydb/core/formats
ydb/core/kqp/common
ydb/core/kqp/common/buffer
+ ydb/core/mon
ydb/core/persqueue/events
ydb/core/protos
ydb/core/scheme
+ ydb/core/tx/scheme_board
ydb/core/ydb_convert
ydb/library/aclib
ydb/library/yql/dq/actors
diff --git a/ydb/core/kqp/ut/indexes/vector/kqp_indexes_vector_ut.cpp b/ydb/core/kqp/ut/indexes/vector/kqp_indexes_vector_ut.cpp
index 6e6582a5596..5bc89a334aa 100644
--- a/ydb/core/kqp/ut/indexes/vector/kqp_indexes_vector_ut.cpp
+++ b/ydb/core/kqp/ut/indexes/vector/kqp_indexes_vector_ut.cpp
@@ -1361,7 +1361,10 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
Y_UNIT_TEST_QUAD(VectorSearchPushdown, Covered, Followers) {
auto setting = NKikimrKqp::TKqpSetting();
- auto serverSettings = TKikimrSettings()
+ NKikimrConfig::TAppConfig appConfig;
+ appConfig.MutableTableServiceConfig()->MutableResourceManager()->SetKqpLevelCacheMaxSizeBytes(300_MB);
+
+ auto serverSettings = TKikimrSettings(appConfig)
// SetUseRealThreads(false) is required to capture events (!) but then you have to do kikimr.RunCall() for everything
.SetUseRealThreads(false)
.SetEnableForceFollowers(Followers)
@@ -1418,7 +1421,10 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
// Check that level & posting are read from followers
UNIT_ASSERT(isFollower == (Followers && (shardType == levelType || shardType == postingType)));
auto & read = ev->Get<TEvDataShard::TEvRead>()->Record;
- if (shardType == (Covered ? mainType : postingType)) {
+ if (shardType == levelType) {
+ // Level table reads do full scan for caching (VectorTopK cleared by cache layer)
+ UNIT_ASSERT(!read.HasVectorTopK());
+ } else if (shardType == (Covered ? mainType : postingType)) {
// Non-covering index does topK on main table, covering does it on posting
UNIT_ASSERT(!read.HasVectorTopK());
} else {
@@ -1427,10 +1433,7 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
auto & topK = read.GetVectorTopK();
// Check that target and limit are pushed down
UNIT_ASSERT(topK.GetTargetVector() == "\x67\x71\x02");
- if (shardType == levelType) {
- // Equal to pragma
- UNIT_ASSERT(topK.GetLimit() == 2);
- } else if (shardType == (Covered ? postingType : mainType)) {
+ if (shardType == (Covered ? postingType : mainType)) {
// Equal to LIMIT
UNIT_ASSERT(topK.GetLimit() == 3);
}
@@ -1549,6 +1552,229 @@ Y_UNIT_TEST_SUITE(KqpVectorIndexes) {
WITH (parallel=2);
)", EnableIndexStreamWrite);
}
+
+ Y_UNIT_TEST(VectorIndexLevelTableCache) {
+ NKikimrConfig::TAppConfig appConfig;
+ appConfig.MutableTableServiceConfig()->MutableResourceManager()->SetKqpLevelCacheMaxSizeBytes(300_MB);
+
+ auto setting = NKikimrKqp::TKqpSetting();
+ auto serverSettings = TKikimrSettings(appConfig)
+ // SetUseRealThreads(false) is required to capture events
+ .SetUseRealThreads(false)
+ .SetKqpSettings({setting});
+
+ TKikimrRunner kikimr(serverSettings);
+ auto runtime = kikimr.GetTestServer().GetRuntime();
+
+ const int flags = F_NULLABLE;
+ auto db = kikimr.RunCall([&] { return kikimr.GetTableClient(); });
+ auto session = kikimr.RunCall([&] { return DoCreateTableAndVectorIndex(db, flags); });
+
+ // Resolve level table shards to track reads
+ constexpr ui32 levelType = 1, postingType = 2, mainType = 3;
+ THashMap<TActorId, ui32> actorTypes;
+ auto resolveActors = [&](const char* tableName, ui32 type) {
+ auto shards = GetTableShards(&kikimr.GetTestServer(), runtime->AllocateEdgeActor(), tableName);
+ for (auto shardId: shards) {
+ auto actorId = ResolveTablet(*runtime, shardId);
+ actorTypes[actorId] = type;
+ }
+ };
+ resolveActors("/Root/TestTable/index1/indexImplLevelTable", levelType);
+ resolveActors("/Root/TestTable/index1/indexImplPostingTable", postingType);
+ resolveActors("/Root/TestTable", mainType);
+
+ // Counter for level table reads
+ std::atomic<int> levelTableReadCount{0};
+ std::atomic<int> postingTableReadCount{0};
+
+ auto captureEvents = [&](TTestActorRuntimeBase&, TAutoPtr<IEventHandle>& ev) {
+ if (ev->GetTypeRewrite() == TEvDataShard::TEvRead::EventType) {
+ ui32 shardType = actorTypes[ev->GetRecipientRewrite()];
+ if (shardType == levelType) {
+ levelTableReadCount++;
+ } else if (shardType == postingType) {
+ postingTableReadCount++;
+ }
+ }
+ return false;
+ };
+ runtime->SetEventFilter(captureEvents);
+
+ const TString query1(Q_(R"(
+ pragma ydb.KMeansTreeSearchTopSize = "2";
+ $TargetEmbedding = String::HexDecode("677102");
+ SELECT * FROM `/Root/TestTable`
+ VIEW index1 ORDER BY Knn::CosineDistance(emb, $TargetEmbedding)
+ LIMIT 3
+ )"));
+
+ // First query - should read from level table
+ {
+ levelTableReadCount = 0;
+ postingTableReadCount = 0;
+
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query1, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+ // First query should read from level table
+ UNIT_ASSERT_C(levelTableReadCount.load() > 0,
+ "First query should read from level table, but read count is " << levelTableReadCount.load());
+ Cerr << "First query: level table reads = " << levelTableReadCount.load()
+ << ", posting table reads = " << postingTableReadCount.load() << Endl;
+ }
+
+ int firstLevelReads = levelTableReadCount.load();
+
+ // Second query with the same target - should NOT read from level table (cached)
+ {
+ levelTableReadCount = 0;
+ postingTableReadCount = 0;
+
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query1, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+ Cerr << "Second query: level table reads = " << levelTableReadCount.load()
+ << ", posting table reads = " << postingTableReadCount.load() << Endl;
+
+ // Second query should NOT read from level table (cached)
+ UNIT_ASSERT_C(levelTableReadCount.load() == 0,
+ "Second query should use cached level table data, but level table read count is "
+ << levelTableReadCount.load() << " (first query had " << firstLevelReads << ")");
+ }
+
+ // Third query - should still use cache
+ {
+ levelTableReadCount = 0;
+ postingTableReadCount = 0;
+
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query1, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+
+ Cerr << "Third query: level table reads = " << levelTableReadCount.load()
+ << ", posting table reads = " << postingTableReadCount.load() << Endl;
+
+ // Third query should NOT read from level table (still cached)
+ UNIT_ASSERT_C(levelTableReadCount.load() == 0,
+ "Third query should use cached level table data, but level table read count is "
+ << levelTableReadCount.load());
+ }
+ }
+
+ Y_UNIT_TEST(VectorIndexLevelTableCacheAcrossSnapshots) {
+ // Test that level table cache works across different snapshots
+ // Level tables are immutable after index creation, so cache should hit
+ // regardless of snapshot changes
+
+ NKikimrConfig::TAppConfig appConfig;
+ appConfig.MutableTableServiceConfig()->MutableResourceManager()->SetKqpLevelCacheMaxSizeBytes(300_MB);
+ auto setting = NKikimrKqp::TKqpSetting();
+ auto serverSettings = TKikimrSettings(appConfig)
+ .SetUseRealThreads(false)
+ .SetKqpSettings({setting});
+
+ TKikimrRunner kikimr(serverSettings);
+ auto runtime = kikimr.GetTestServer().GetRuntime();
+
+ const int flags = F_NULLABLE;
+ auto db = kikimr.RunCall([&] { return kikimr.GetTableClient(); });
+ auto session = kikimr.RunCall([&] { return DoCreateTableAndVectorIndex(db, flags); });
+
+ // Resolve level table shards to track reads
+ constexpr ui32 levelType = 1;
+ THashMap<TActorId, ui32> actorTypes;
+ auto shards = GetTableShards(&kikimr.GetTestServer(), runtime->AllocateEdgeActor(),
+ "/Root/TestTable/index1/indexImplLevelTable");
+ for (auto shardId: shards) {
+ auto actorId = ResolveTablet(*runtime, shardId);
+ actorTypes[actorId] = levelType;
+ }
+
+ std::atomic<int> levelTableReadCount{0};
+ auto captureEvents = [&](TTestActorRuntimeBase&, TAutoPtr<IEventHandle>& ev) {
+ if (ev->GetTypeRewrite() == TEvDataShard::TEvRead::EventType) {
+ if (actorTypes[ev->GetRecipientRewrite()] == levelType) {
+ levelTableReadCount++;
+ }
+ }
+ return false;
+ };
+ runtime->SetEventFilter(captureEvents);
+
+ const TString query(Q_(R"(
+ pragma ydb.KMeansTreeSearchTopSize = "2";
+ $TargetEmbedding = String::HexDecode("677102");
+ SELECT * FROM `/Root/TestTable`
+ VIEW index1 ORDER BY Knn::CosineDistance(emb, $TargetEmbedding)
+ LIMIT 3
+ )"));
+
+ // First query - should read from level table (cache miss)
+ {
+ levelTableReadCount = 0;
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+ UNIT_ASSERT_C(levelTableReadCount.load() > 0,
+ "First query should read from level table, got " << levelTableReadCount.load());
+ Cerr << "Query 1 (first session tx): level reads = " << levelTableReadCount.load() << Endl;
+ }
+
+ // Second query in NEW TRANSACTION (different snapshot!) - should still hit cache
+ // because level tables are immutable
+ {
+ levelTableReadCount = 0;
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+ Cerr << "Query 2 (new tx, different snapshot): level reads = " << levelTableReadCount.load() << Endl;
+ UNIT_ASSERT_C(levelTableReadCount.load() == 0,
+ "Second query should use cache despite different snapshot, got " << levelTableReadCount.load());
+ }
+
+ // Third query in yet another transaction - should still use cache
+ {
+ levelTableReadCount = 0;
+ auto result = kikimr.RunCall([&] {
+ return session.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+ Cerr << "Query 3 (another new tx): level reads = " << levelTableReadCount.load() << Endl;
+ UNIT_ASSERT_C(levelTableReadCount.load() == 0,
+ "Third query should use cache, got " << levelTableReadCount.load());
+ }
+
+ // Create a NEW session (completely different connection) - should still hit cache
+ auto session2 = kikimr.RunCall([&] {
+ return db.CreateSession().GetValueSync().GetSession();
+ });
+ {
+ levelTableReadCount = 0;
+ auto result = kikimr.RunCall([&] {
+ return session2.ExecuteDataQuery(query, TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx())
+ .ExtractValueSync();
+ });
+ UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+ Cerr << "Query 4 (new session): level reads = " << levelTableReadCount.load() << Endl;
+ UNIT_ASSERT_C(levelTableReadCount.load() == 0,
+ "Query from new session should use cache, got " << levelTableReadCount.load());
+ }
+ }
+
}
}
diff --git a/ydb/core/protos/table_service_config.proto b/ydb/core/protos/table_service_config.proto
index a13664141e8..14312e4dbc7 100644
--- a/ydb/core/protos/table_service_config.proto
+++ b/ydb/core/protos/table_service_config.proto
@@ -59,6 +59,9 @@ message TTableServiceConfig {
optional uint64 MaxNonParallelDataQueryTasksLimit = 28 [default = 1000];
optional bool VerboseMemoryLimitException = 29 [default = false];
+
+ optional uint64 KqpLevelCacheMaxSizeBytes = 31 [default = 0];
+ optional uint64 KqpLevelCacheIncreaseBatchSizeBytes = 32 [default = 33554432];
}
message TSpillingServiceConfig {
generated by cgit v1.3 (git 2.53.0) at 2026-07-21 03:46:08 +0000