allow joins for non-point predicates

ydb-platform · Feb 28, 2024 · b252df2 · b252df2
1 parent bfd5754
commit b252df2
Show file tree

Hide file tree

Showing 11 changed files with 347 additions and 60 deletions.
diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
@@ -538,6 +538,7 @@ void ApplyServiceConfig(TKikimrConfiguration& kqpConfig, const TTableServiceConf
     kqpConfig.EnablePgConstsToParams = serviceConfig.GetEnablePgConstsToParams() && serviceConfig.GetEnableAstCache();
     kqpConfig.ExtractPredicateRangesLimit = serviceConfig.GetExtractPredicateRangesLimit();
     kqpConfig.EnablePerStatementQueryExecution = serviceConfig.GetEnablePerStatementQueryExecution();
+    kqpConfig.IdxLookupJoinsPrefixPointLimit = serviceConfig.GetIdxLookupJoinPointsLimit();
 
     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
@@ -487,6 +487,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
         auto indexAutoChooser = TableServiceConfig.GetIndexAutoChooseMode();
 
         ui64 rangesLimit = TableServiceConfig.GetExtractPredicateRangesLimit();
+        ui64 idxLookupPointsLimit = TableServiceConfig.GetIdxLookupJoinPointsLimit();
 
         bool enableSequences = TableServiceConfig.GetEnableSequences();
         bool enableColumnsWithDefault = TableServiceConfig.GetEnableColumnsWithDefault();
@@ -517,7 +518,8 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
             TableServiceConfig.GetEnableColumnsWithDefault() != enableColumnsWithDefault ||
             TableServiceConfig.GetEnableOlapSink() != enableOlapSink ||
             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,8 @@
             "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}
             ]
         },
         {
@@ -131,7 +132,8 @@
             "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}
             ]
         },
         {

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 + 2, ctx)) {
         return TStatus::Error;
     }
 
@@ -375,8 +375,26 @@ 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);
+            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);
+            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_extract.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_extract.cpp
@@ -121,6 +121,7 @@ TExprBase KqpApplyExtractMembersToReadTableRanges(TExprBase node, TExprContext&
             .ExplainPrompt(read.ExplainPrompt())
             .Index(index.Index().Cast())
             .PrefixPointsExpr(index.PrefixPointsExpr())
+            .PredicateExpr(index.PredicateExpr())
             .Done();
     }
 
@@ -132,6 +133,7 @@ TExprBase KqpApplyExtractMembersToReadTableRanges(TExprBase node, TExprContext&
             .Settings(read.Settings())
             .ExplainPrompt(read.ExplainPrompt())
             .PrefixPointsExpr(readRange.PrefixPointsExpr())
+            .PredicateExpr(readRange.PredicateExpr())
             .Done();
     }