Improve index chooser rules

7 files changed, 75 insertions, 44 deletions

diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
index e9137fb35fc..2bd14a8285c 100644
--- a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
+++ b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
@@ -426,12 +426,12 @@ void ApplyServiceConfig(TKikimrConfiguration& kqpConfig, const TTableServiceConf
     kqpConfig.EnableSequences = serviceConfig.GetEnableSequences();
     kqpConfig.BindingsMode = RemapBindingsMode(serviceConfig.GetBindingsMode());
     kqpConfig.PredicateExtract20 = serviceConfig.GetPredicateExtract20();
-    kqpConfig.EnableIndexAutoChooser = serviceConfig.GetEnableIndexAutoChooser();
+    kqpConfig.IndexAutoChooserMode = serviceConfig.GetIndexAutoChooseMode();
 }
 
 IActor* CreateKqpCompileActor(const TActorId& owner, const TKqpSettings::TConstPtr& kqpSettings,
-    const TTableServiceConfig& tableServiceConfig, 
-    const TQueryServiceConfig& queryServiceConfig, 
+    const TTableServiceConfig& tableServiceConfig,
+    const TQueryServiceConfig& queryServiceConfig,
     const TMetadataProviderConfig& metadataProviderConfig,
     TIntrusivePtr<TModuleResolverState> moduleResolverState, TIntrusivePtr<TKqpCounters> counters,
     const TString& uid, const TKqpQueryId& query, const TIntrusiveConstPtr<NACLib::TUserToken>& userToken,
@@ -440,7 +440,7 @@ IActor* CreateKqpCompileActor(const TActorId& owner, const TKqpSettings::TConstP
     NWilson::TTraceId traceId, TKqpTempTablesState::TConstPtr tempTablesState)
 {
     return new TKqpCompileActor(owner, kqpSettings, tableServiceConfig, queryServiceConfig, metadataProviderConfig,
-                                moduleResolverState, counters,  
+                                moduleResolverState, counters,
                                 uid, query, userToken, dbCounters,
                                 federatedQuerySetup, userRequestContext,
                                 std::move(traceId), std::move(tempTablesState));
diff --git a/ydb/core/kqp/compile_service/kqp_compile_service.cpp b/ydb/core/kqp/compile_service/kqp_compile_service.cpp
index 02bf1204422..1b74648f2a4 100644
--- a/ydb/core/kqp/compile_service/kqp_compile_service.cpp
+++ b/ydb/core/kqp/compile_service/kqp_compile_service.cpp
@@ -406,7 +406,7 @@ private:
         bool defaultSyntaxVersion = TableServiceConfig.GetSqlVersion();
         bool enableKqpImmediateEffects = TableServiceConfig.GetEnableKqpImmediateEffects();
 
-        bool indexAutoChooser = TableServiceConfig.GetEnableIndexAutoChooser();
+        auto indexAutoChooser = TableServiceConfig.GetIndexAutoChooseMode();
 
         TableServiceConfig.Swap(event.MutableConfig()->MutableTableServiceConfig());
         LOG_INFO(*TlsActivationContext, NKikimrServices::KQP_COMPILE_SERVICE, "Updated config");
@@ -425,7 +425,7 @@ private:
             TableServiceConfig.GetPredicateExtract20() != predicateExtract20 ||
             TableServiceConfig.GetEnableSequentialReads() != enableSequentialReads ||
             TableServiceConfig.GetEnableKqpImmediateEffects() != enableKqpImmediateEffects ||
-            TableServiceConfig.GetEnableIndexAutoChooser() != indexAutoChooser) {
+            TableServiceConfig.GetIndexAutoChooseMode() != indexAutoChooser) {
 
             LOG_NOTICE_S(*TlsActivationContext, NKikimrServices::KQP_COMPILE_SERVICE,
                 "Iterator read flags was changed. StreamLookup from " << enableKqpDataQueryStreamLookup <<
@@ -767,7 +767,7 @@ private:
     }
 
     void StartCompilation(TKqpCompileRequest&& request, const TActorContext& ctx) {
-        auto compileActor = CreateKqpCompileActor(ctx.SelfID, KqpSettings, TableServiceConfig, QueryServiceConfig, MetadataProviderConfig, ModuleResolverState, Counters, 
+        auto compileActor = CreateKqpCompileActor(ctx.SelfID, KqpSettings, TableServiceConfig, QueryServiceConfig, MetadataProviderConfig, ModuleResolverState, Counters,
              request.Uid, request.Query, request.UserToken, FederatedQuerySetup, request.DbCounters, request.UserRequestContext,
              request.CompileServiceSpan.GetTraceId(), std::move(request.TempTablesState));
         auto compileActorId = ctx.ExecutorThread.RegisterActor(compileActor, TMailboxType::HTSwap,
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
index 1445fb74ac2..d4f994c0b04 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_join.cpp
@@ -188,7 +188,7 @@ TMaybeNode<TKqlKeyInc> GetRightTableKeyPrefix(const TKqlKeyRange& range) {
 
 TExprBase BuildLookupIndex(TExprContext& ctx, const TPositionHandle pos,
     const TKqpTable& table, const TCoAtomList& columns,
-    const TExprBase& keysToLookup, const TVector<TCoAtom>& lookupNames, const TString& indexName,
+    const TExprBase& keysToLookup, const TVector<TCoAtom>& skipNullColumns, const TString& indexName,
     const TKqpOptimizeContext& kqpCtx)
 {
     if (kqpCtx.IsScanQuery()) {
@@ -198,7 +198,7 @@ TExprBase BuildLookupIndex(TExprContext& ctx, const TPositionHandle pos,
             .LookupKeys<TCoSkipNullMembers>()
                 .Input(keysToLookup)
                 .Members()
-                    .Add(lookupNames)
+                    .Add(skipNullColumns)
                     .Build()
                 .Build()
             .Columns(columns)
@@ -212,7 +212,7 @@ TExprBase BuildLookupIndex(TExprContext& ctx, const TPositionHandle pos,
         .LookupKeys<TCoSkipNullMembers>()
             .Input(keysToLookup)
             .Members()
-                .Add(lookupNames)
+                .Add(skipNullColumns)
                 .Build()
             .Build()
         .Columns(columns)
@@ -223,7 +223,7 @@ TExprBase BuildLookupIndex(TExprContext& ctx, const TPositionHandle pos,
 
 TExprBase BuildLookupTable(TExprContext& ctx, const TPositionHandle pos,
     const TKqpTable& table, const TCoAtomList& columns,
-    const TExprBase& keysToLookup, const TVector<TCoAtom>& lookupNames, const TKqpOptimizeContext& kqpCtx)
+    const TExprBase& keysToLookup, const TVector<TCoAtom>& skipNullColumns, const TKqpOptimizeContext& kqpCtx)
 {
     if (kqpCtx.IsScanQuery()) {
         YQL_ENSURE(kqpCtx.Config->EnableKqpScanQueryStreamIdxLookupJoin, "Stream lookup is not enabled for index lookup join");
@@ -232,7 +232,7 @@ TExprBase BuildLookupTable(TExprContext& ctx, const TPositionHandle pos,
             .LookupKeys<TCoSkipNullMembers>()
                 .Input(keysToLookup)
                 .Members()
-                    .Add(lookupNames)
+                    .Add(skipNullColumns)
                     .Build()
                 .Build()
             .Columns(columns)
@@ -245,7 +245,7 @@ TExprBase BuildLookupTable(TExprContext& ctx, const TPositionHandle pos,
             .LookupKeys<TCoSkipNullMembers>()
                 .Input(keysToLookup)
                 .Members()
-                    .Add(lookupNames)
+                    .Add(skipNullColumns)
                     .Build()
                 .Build()
             .Columns(columns)
@@ -257,7 +257,7 @@ TExprBase BuildLookupTable(TExprContext& ctx, const TPositionHandle pos,
         .LookupKeys<TCoSkipNullMembers>()
             .Input(keysToLookup)
             .Members()
-                .Add(lookupNames)
+                .Add(skipNullColumns)
                 .Build()
             .Build()
         .Columns(columns)
@@ -475,7 +475,7 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
         .Done();
 
     TVector<TExprBase> lookupMembers;
-    TVector<TCoAtom> lookupNames;
+    TVector<TCoAtom> skipNullColumns;
     ui32 fixedPrefix = 0;
     TSet<TString> deduplicateLeftColumns;
     for (auto& rightColumnName : rightTableDesc.Metadata->KeyColumnNames) {
@@ -544,7 +544,10 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
                 .Add<TCoAtom>().Build(rightColumnName)
                 .Add(member)
                 .Done());
-        lookupNames.emplace_back(ctx.NewAtom(join.Pos(), rightColumnName));
+
+        if (leftColumn) {
+            skipNullColumns.emplace_back(ctx.NewAtom(join.Pos(), rightColumnName));
+        }
     }
 
     if (lookupMembers.size() <= fixedPrefix) {
@@ -613,8 +616,8 @@ TMaybeNode<TExprBase> KqpJoinToIndexLookupImpl(const TDqJoin& join, TExprContext
         .Done();
 
     TExprBase lookup = indexName
-        ? BuildLookupIndex(ctx, join.Pos(), rightRead.Table(), rightRead.Columns(), keysToLookup, lookupNames, indexName, kqpCtx)
-        : BuildLookupTable(ctx, join.Pos(), rightRead.Table(), rightRead.Columns(), keysToLookup, lookupNames, kqpCtx);
+        ? BuildLookupIndex(ctx, join.Pos(), rightRead.Table(), rightRead.Columns(), keysToLookup, skipNullColumns, indexName, kqpCtx)
+        : BuildLookupTable(ctx, join.Pos(), rightRead.Table(), rightRead.Columns(), keysToLookup, skipNullColumns, kqpCtx);
 
     // Skip null keys in lookup part as for equijoin semantics null != null,
     // so we can't have nulls in lookup part
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
index 97f8ab6ec0d..6876641b2ad 100644
--- a/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
+++ b/ydb/core/kqp/opt/logical/kqp_opt_log_ranges_predext.cpp
@@ -256,30 +256,46 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
     bool prepareSuccess = extractor->Prepare(flatmap.Lambda().Ptr(), *mainTableDesc.SchemeNode, possibleKeys, ctx, typesCtx);
     YQL_ENSURE(prepareSuccess);
 
-    if (!indexName.IsValid() && !readSettings.ForcePrimary && kqpCtx.Config->EnableIndexAutoChooser) {
-        TVector<std::tuple<size_t, size_t, TMaybe<TString>>> indices;
-        {
-            auto buildResult = extractor->BuildComputeNode(mainTableDesc.Metadata->KeyColumnNames, ctx, typesCtx);
-            indices.push_back(std::make_tuple(
-                    buildResult.UsedPrefixLen,
-                    buildResult.PointPrefixLen,
-                    TMaybe<TString>()));
-        }
-        for (auto& index : mainTableDesc.Metadata->Indexes) {
-            if (index.Type == TIndexDescription::EType::GlobalSync) {
-                auto buildResult = extractor->BuildComputeNode(index.KeyColumns, ctx, typesCtx);
-                indices.push_back(std::make_tuple(
-                        buildResult.UsedPrefixLen,
-                        buildResult.PointPrefixLen,
-                        TMaybe<TString>(index.Name)));
+    if (!indexName.IsValid() && !readSettings.ForcePrimary && kqpCtx.Config->IndexAutoChooserMode != NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_DISABLED) {
+        using TIndexComparisonKey = std::tuple<bool, size_t, bool, size_t, bool>;
+        auto calcKey = [&](NYql::IPredicateRangeExtractor::TBuildResult buildResult, size_t descriptionKeyColumns, bool needsJoin) -> TIndexComparisonKey {
+
+            return std::make_tuple(
+                buildResult.PointPrefixLen >= descriptionKeyColumns,
+                buildResult.PointPrefixLen,
+                buildResult.UsedPrefixLen >= descriptionKeyColumns,
+                buildResult.UsedPrefixLen,
+                needsJoin);
+        };
+
+        TMaybe<TString> chosenIndex;
+        auto primaryBuildResult = extractor->BuildComputeNode(mainTableDesc.Metadata->KeyColumnNames, ctx, typesCtx);
+
+        if (primaryBuildResult.PointPrefixLen < mainTableDesc.Metadata->KeyColumnNames.size()) {
+            auto maxKey = calcKey(primaryBuildResult, mainTableDesc.Metadata->KeyColumnNames.size(), false);
+            for (auto& index : mainTableDesc.Metadata->Indexes) {
+                if (index.Type != TIndexDescription::EType::GlobalAsync) {
+                    auto& tableDesc = kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, mainTableDesc.Metadata->GetIndexMetadata(TString(index.Name)).first->Name);
+                    auto buildResult = extractor->BuildComputeNode(tableDesc.Metadata->KeyColumnNames, ctx, typesCtx);
+
+                    if (kqpCtx.Config->IndexAutoChooserMode == NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_ONLY_FULL_KEY && buildResult.PointPrefixLen < index.KeyColumns.size()) {
+                        continue;
+                    }
+                    if (kqpCtx.Config->IndexAutoChooserMode == NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_ONLY_POINTS && buildResult.PointPrefixLen == 0) {
+                        continue;
+                    }
+
+                    auto key = calcKey(buildResult, index.KeyColumns.size(), true);
+                    if (key > maxKey) {
+                        maxKey = key;
+                        chosenIndex = index.Name;
+                    }
+                }
             }
         }
-        if (!indices.empty()) {
-            auto it = std::max_element(indices.begin(), indices.end());
-            auto name = std::get<TMaybe<TString>>(*it);
-            if (name) {
-                indexName = ctx.NewAtom(read.Pos(), *name);
-            }
+
+        if (chosenIndex) {
+            indexName = ctx.NewAtom(read.Pos(), *chosenIndex);
         }
     }
 
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
index 141a0307c68..5887882481f 100644
--- a/ydb/core/kqp/provider/yql_kikimr_settings.h
+++ b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -152,7 +152,7 @@ struct TKikimrConfiguration : public TKikimrSettings, public NCommon::TSettingDi
     bool EnablePreparedDdl = false;
     bool EnableSequences = false;
     NSQLTranslation::EBindingsMode BindingsMode = NSQLTranslation::EBindingsMode::ENABLED;
-    bool EnableIndexAutoChooser = false;
+    NKikimrConfig::TTableServiceConfig::EIndexAutoChooseMode IndexAutoChooserMode = NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_DISABLED;
 };
 
 }
diff --git a/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp b/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp
index 302324cad7f..d03aa1f721b 100644
--- a/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp
+++ b/ydb/core/kqp/ut/opt/kqp_ne_ut.cpp
@@ -3873,7 +3873,12 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
     }
 
     Y_UNIT_TEST(PrimaryView) {
-        auto kikimr = DefaultKikimrRunner();
+        TKikimrSettings settings;
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetIndexAutoChooseMode(NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_MAX_USED_PREFIX);
+        settings.SetAppConfig(appConfig);
+
+        TKikimrRunner kikimr(settings);
         auto db = kikimr.GetTableClient();
         auto session = db.CreateSession().GetValueSync().GetSession();
         CreateSampleTablesWithIndex(session);
@@ -3892,7 +3897,7 @@ Y_UNIT_TEST_SUITE(KqpNewEngine) {
     Y_UNIT_TEST(AutoChooseIndex) {
         TKikimrSettings settings;
         NKikimrConfig::TAppConfig appConfig;
-        appConfig.MutableTableServiceConfig()->SetEnableIndexAutoChooser(true);
+        appConfig.MutableTableServiceConfig()->SetIndexAutoChooseMode(NKikimrConfig::TTableServiceConfig_EIndexAutoChooseMode_MAX_USED_PREFIX);
         settings.SetAppConfig(appConfig);
 
         TKikimrRunner kikimr(settings);
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index b43691634d0..3e9524c5f4e 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -1381,7 +1381,14 @@ message TTableServiceConfig {
     }
 
     optional EChannelTransportVersion ChannelTransportVersion = 46 [default = CTV_OOB_PICKLE_1_0];
-    optional bool EnableIndexAutoChooser = 49 [default = false];
+
+    enum EIndexAutoChooseMode {
+        DISABLED = 0;
+        ONLY_FULL_KEY = 1; // Accepts iff point prefix includes all key columns
+        ONLY_POINTS = 2; // Accepts iff point prefix len > 0
+        MAX_USED_PREFIX = 3;
+    }
+    optional EIndexAutoChooseMode IndexAutoChooseMode = 50 [default = DISABLED];
 };
 
 message TQueryServiceConfig {