Implement indexlookupjoin for non-point selection (ydb-platform#2298)

ssmike · Jun 6, 2024 · eca0972 · eca0972
1 parent ec1169f
commit eca0972
Show file tree

Hide file tree

Showing 18 changed files with 967 additions and 255 deletions.
diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
@@ -506,6 +506,8 @@ void ApplyServiceConfig(TKikimrConfiguration& kqpConfig, const TTableServiceConf
     kqpConfig.IndexAutoChooserMode = serviceConfig.GetIndexAutoChooseMode();
     kqpConfig.EnablePgConstsToParams = serviceConfig.GetEnablePgConstsToParams();
     kqpConfig.ExtractPredicateRangesLimit = serviceConfig.GetExtractPredicateRangesLimit();
+    kqpConfig.IdxLookupJoinsPrefixPointLimit = serviceConfig.GetIdxLookupJoinPointsLimit();
+    kqpConfig.OldLookupJoinBehaviour = serviceConfig.GetOldLookupJoinBehaviour();
 
     if (const auto limit = serviceConfig.GetResourceManager().GetMkqlHeavyProgramMemoryLimit()) {
         kqpConfig._KqpYqlCombinerMemoryLimit = std::max(1_GB, limit - (limit >> 2U));

diff --git a/ydb/core/kqp/compile_service/kqp_compile_service.cpp b/ydb/core/kqp/compile_service/kqp_compile_service.cpp
@@ -474,6 +474,8 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
         auto indexAutoChooser = TableServiceConfig.GetIndexAutoChooseMode();
 
         ui64 rangesLimit = TableServiceConfig.GetExtractPredicateRangesLimit();
+        ui64 idxLookupPointsLimit = TableServiceConfig.GetIdxLookupJoinPointsLimit();
+        bool oldLookupJoinBehaviour = TableServiceConfig.GetOldLookupJoinBehaviour();
 
         bool enableSequences = TableServiceConfig.GetEnableSequences();
         bool enableColumnsWithDefault = TableServiceConfig.GetEnableColumnsWithDefault();
@@ -503,8 +505,10 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
             TableServiceConfig.GetEnableSequences() != enableSequences ||
             TableServiceConfig.GetEnableColumnsWithDefault() != enableColumnsWithDefault ||
             TableServiceConfig.GetEnableOlapSink() != enableOlapSink ||
+            TableServiceConfig.GetOldLookupJoinBehaviour() != oldLookupJoinBehaviour ||
             TableServiceConfig.GetExtractPredicateRangesLimit() != rangesLimit ||
-            TableServiceConfig.GetResourceManager().GetMkqlHeavyProgramMemoryLimit() != mkqlHeavyLimit) {
+            TableServiceConfig.GetResourceManager().GetMkqlHeavyProgramMemoryLimit() != mkqlHeavyLimit ||
+            TableServiceConfig.GetIdxLookupJoinPointsLimit() != idxLookupPointsLimit) {
 
             QueryCache.Clear();
 

diff --git a/ydb/core/kqp/expr_nodes/kqp_expr_nodes.json b/ydb/core/kqp/expr_nodes/kqp_expr_nodes.json
@@ -104,7 +104,9 @@
             "Base": "TKqlReadTableRangesBase",
             "Match": {"Type": "Callable", "Name": "KqlReadTableRanges"},
             "Children": [
-                {"Index": 5, "Name": "PrefixPointsExpr", "Type": "TExprBase", "Optional": true}
+                {"Index": 5, "Name": "PrefixPointsExpr", "Type": "TExprBase", "Optional": true},
+                {"Index": 6, "Name": "PredicateExpr", "Type": "TCoLambda", "Optional": true},
+                {"Index": 7, "Name": "PredicateUsedColumns", "Type": "TCoAtomList", "Optional": true}
             ]
         },
         {
@@ -131,7 +133,9 @@
             "Match": {"Type": "Callable", "Name": "TKqlReadTableIndexRanges"},
             "Children": [
                 {"Index": 5, "Name": "Index", "Type": "TCoAtom"},
-                {"Index": 6, "Name": "PrefixPointsExpr", "Type": "TExprBase", "Optional": true}
+                {"Index": 6, "Name": "PrefixPointsExpr", "Type": "TExprBase", "Optional": true},
+                {"Index": 7, "Name": "PredicateExpr", "Type": "TCoLambda", "Optional": true},
+                {"Index": 8, "Name": "PredicateUsedColumns", "Type": "TCoAtomList", "Optional": true}
             ]
         },
         {

diff --git a/ydb/core/kqp/host/kqp_type_ann.cpp b/ydb/core/kqp/host/kqp_type_ann.cpp
@@ -340,7 +340,7 @@ TStatus AnnotateReadTableRanges(const TExprNode::TPtr& node, TExprContext& ctx,
     size_t argCount = (olapTable || index) ? 6 : 5;
 
     // prefix
-    if (!EnsureMinArgsCount(*node, argCount, ctx) && EnsureMaxArgsCount(*node, argCount + 1, ctx)) {
+    if (!EnsureMinArgsCount(*node, argCount, ctx) && EnsureMaxArgsCount(*node, argCount + 3, ctx)) {
         return TStatus::Error;
     }
 
@@ -375,8 +375,34 @@ TStatus AnnotateReadTableRanges(const TExprNode::TPtr& node, TExprContext& ctx,
     }
 
     if (TKqlReadTableRanges::Match(node.Get())) {
+        if (node->ChildrenSize() > TKqlReadTableRanges::idx_PredicateExpr) {
+            auto& lambda = node->ChildRef(TKqlReadTableRanges::idx_PredicateExpr);
+            auto rowType = GetReadTableRowType(ctx, tablesData, cluster, table.first, node->Pos(), withSystemColumns);
+            if (!rowType) {
+                return TStatus::Error;
+            }
+            if (!UpdateLambdaAllArgumentsTypes(lambda, {rowType}, ctx)) {
+                return IGraphTransformer::TStatus::Error;
+            }
+            if (!lambda->GetTypeAnn()) {
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
         node->SetTypeAnn(ctx.MakeType<TListExprType>(rowType));
     } else if (TKqlReadTableIndexRanges::Match(node.Get())) {
+        if (node->ChildrenSize() > TKqlReadTableIndexRanges::idx_PredicateExpr) {
+            auto& lambda = node->ChildRef(TKqlReadTableIndexRanges::idx_PredicateExpr);
+            auto rowType = GetReadTableRowType(ctx, tablesData, cluster, table.first, node->Pos(), withSystemColumns);
+            if (!rowType) {
+                return TStatus::Error;
+            }
+            if (!UpdateLambdaAllArgumentsTypes(lambda, {rowType}, ctx)) {
+                return IGraphTransformer::TStatus::Error;
+            }
+            if (!lambda->GetTypeAnn()) {
+                return IGraphTransformer::TStatus::Repeat;
+            }
+        }
         node->SetTypeAnn(ctx.MakeType<TListExprType>(rowType));
     } else if (TKqpReadTableRanges::Match(node.Get())) {
         node->SetTypeAnn(ctx.MakeType<TFlowExprType>(rowType));

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp b/ydb/core/kqp/opt/logical/kqp_opt_cbo.cpp
@@ -14,6 +14,25 @@ using namespace NYql::NNodes;
 
 namespace {
 
+TMaybeNode<TKqlKeyInc> GetRightTableKeyPrefix(const TKqlKeyRange& range) {
+    if (!range.From().Maybe<TKqlKeyInc>() || !range.To().Maybe<TKqlKeyInc>()) {
+        return {};
+    }
+    auto rangeFrom = range.From().Cast<TKqlKeyInc>();
+    auto rangeTo = range.To().Cast<TKqlKeyInc>();
+
+    if (rangeFrom.ArgCount() != rangeTo.ArgCount()) {
+        return {};
+    }
+    for (ui32 i = 0; i < rangeFrom.ArgCount(); ++i) {
+        if (rangeFrom.Arg(i).Raw() != rangeTo.Arg(i).Raw()) {
+            return {};
+        }
+    }
+
+    return rangeFrom;
+}
+
 /**
  * KQP specific rule to check if a LookupJoin is applicable
 */
@@ -163,4 +182,4 @@ double TKqpProviderContext::ComputeJoinCost(const TOptimizerStatistics& leftStat
 }
 
 
-}
+}
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -86,20 +86,20 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
         return output;
     }
 
-    TMaybeNode<TExprBase> PushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx) {
+    TMaybeNode<TExprBase> PushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx, const TGetParents& getParents) {
         if (!KqpCtx.Config->PredicateExtract20) {
             return node;
         }
-        TExprBase output = KqpPushExtractedPredicateToReadTable(node, ctx, KqpCtx, TypesCtx);
+        TExprBase output = KqpPushExtractedPredicateToReadTable(node, ctx, KqpCtx, TypesCtx, *getParents());
         DumpAppliedRule("PushExtractedPredicateToReadTable", node.Ptr(), output.Ptr(), ctx);
         return output;
     }
 
-    TMaybeNode<TExprBase> LatePushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx) {
+    TMaybeNode<TExprBase> LatePushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx, const TGetParents& getParents) {
         if (KqpCtx.Config->PredicateExtract20) {
             return node;
         }
-        TExprBase output = KqpPushExtractedPredicateToReadTable(node, ctx, KqpCtx, TypesCtx);
+        TExprBase output = KqpPushExtractedPredicateToReadTable(node, ctx, KqpCtx, TypesCtx, *getParents());
         DumpAppliedRule("PushExtractedPredicateToReadTable", node.Ptr(), output.Ptr(), ctx);
         return output;
     }

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp
@@ -121,6 +121,8 @@ TExprBase KqpApplyExtractMembersToReadTableRanges(TExprBase node, TExprContext&
             .ExplainPrompt(read.ExplainPrompt())
             .Index(index.Index().Cast())
             .PrefixPointsExpr(index.PrefixPointsExpr())
+            .PredicateExpr(index.PredicateExpr())
+            .PredicateUsedColumns(index.PredicateUsedColumns())
             .Done();
     }
 
@@ -132,6 +134,8 @@ TExprBase KqpApplyExtractMembersToReadTableRanges(TExprBase node, TExprContext&
             .Settings(read.Settings())
             .ExplainPrompt(read.ExplainPrompt())
             .PrefixPointsExpr(readRange.PrefixPointsExpr())
+            .PredicateExpr(readRange.PredicateExpr())
+            .PredicateUsedColumns(readRange.PredicateUsedColumns())
             .Done();
     }