Merge 602c8fd into b765769

ydb-platform · Feb 2, 2024 · bbbeaa2 · bbbeaa2
2 parents b765769 + 602c8fd
commit bbbeaa2
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 16 deletions.
diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp
@@ -210,8 +210,8 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
         return output;
     }
 
-    TMaybeNode<TExprBase> DeleteOverLookup(TExprBase node, TExprContext& ctx) {
-        TExprBase output = KqpDeleteOverLookup(node, ctx, KqpCtx);
+    TMaybeNode<TExprBase> DeleteOverLookup(TExprBase node, TExprContext& ctx, const TGetParents& getParents) {
+        TExprBase output = KqpDeleteOverLookup(node, ctx, KqpCtx, *getParents());
         DumpAppliedRule("DeleteOverLookup", node.Ptr(), output.Ptr(), ctx);
         return output;
     }

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_effects.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log_effects.cpp
@@ -2,22 +2,55 @@
 
 #include <ydb/core/kqp/opt/kqp_opt_impl.h>
 #include <ydb/core/kqp/common/kqp_yql.h>
+#include <ydb/library/yql/core/yql_opt_utils.h>
+
+namespace {
+
+bool CanPushFlatMap(const NYql::NNodes::TCoFlatMapBase& flatMap, const NYql::TKikimrTableDescription& tableDesc, const NYql::TParentsMap& parentsMap, TVector<TString> & extraColumns) {
+    auto flatMapLambda = flatMap.Lambda();
+    if (!NYql::IsFilterFlatMap(flatMapLambda)) {
+        return false;
+    }
+
+    const auto & flatMapLambdaArgument = flatMapLambda.Args().Arg(0).Ref();
+    auto flatMapLambdaConditional = flatMapLambda.Body().Cast<NYql::NNodes::TCoConditionalValueBase>();
+
+    TSet<TString> lambdaSubset;
+    if (!HaveFieldsSubset(flatMapLambdaConditional.Predicate().Ptr(), flatMapLambdaArgument, lambdaSubset, parentsMap, true, true)) {
+        return false;
+    }
+
+    for (auto & lambdaColumn : lambdaSubset) {
+        auto columnIndex = tableDesc.GetKeyColumnIndex(lambdaColumn);
+        if (!columnIndex) {
+            return false;
+        }
+    }
+
+    extraColumns.insert(extraColumns.end(), lambdaSubset.begin(), lambdaSubset.end());
+    return true;
+}
+
+}
 
 namespace NKikimr::NKqp::NOpt {
 
 using namespace NYql;
 using namespace NYql::NNodes;
 
-TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TKqpOptimizeContext &kqpCtx) {
+TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TKqpOptimizeContext &kqpCtx, const NYql::TParentsMap& parentsMap) {
     if (!node.Maybe<TKqlDeleteRows>()) {
         return node;
     }
 
     auto deleteRows = node.Cast<TKqlDeleteRows>();
 
+    TMaybeNode<TCoFlatMap> filter;
+
     TMaybeNode<TKqlLookupTableBase> lookup;
     TMaybeNode<TKqlReadTable> read;
     TMaybeNode<TCoSkipNullMembers> skipNulMembers;
+    TMaybeNode<TKqlReadTableRanges> readranges;
 
     if (deleteRows.Input().Maybe<TKqlLookupTableBase>()) {
         lookup = deleteRows.Input().Cast<TKqlLookupTableBase>();
@@ -27,7 +60,15 @@ TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TK
     } else if (deleteRows.Input().Maybe<TKqlReadTable>()) {
         read = deleteRows.Input().Cast<TKqlReadTable>();
     } else {
-        return node;
+        TMaybeNode<TExprBase> input = deleteRows.Input();
+        if (input.Maybe<TCoFlatMap>()) {
+            filter = deleteRows.Input().Cast<TCoFlatMap>();
+            input = filter.Input();
+        }
+        readranges = input.Maybe<TKqlReadTableRanges>();
+        if (!readranges) {
+            return node;
+        }
     }
 
     TMaybeNode<TExprBase> deleteInput;
@@ -90,7 +131,30 @@ TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TK
                 .Add(structMembers)
                 .Build()
             .Done();
-    } 
+    } else if (readranges) {
+        if (!readranges.Cast().PrefixPointsExpr()) {
+            return node;
+        }
+
+        const auto& tableDesc = kqpCtx.Tables->ExistingTable(kqpCtx.Cluster, readranges.Cast().Table().Path().Value());
+        auto hint = TKqpReadTableExplainPrompt::Parse(readranges.Cast().ExplainPrompt());
+        if (hint.PointPrefixLen != tableDesc.Metadata->KeyColumnNames.size()) {
+            return node;
+        }
+
+        if (filter) {
+            TVector<TString> extraColumns;
+            if (!CanPushFlatMap(filter.Cast(), tableDesc, parentsMap, extraColumns)) {
+                return node;
+            }
+            deleteInput = Build<TCoFlatMap>(ctx, node.Pos())
+                .Lambda(filter.Lambda().Cast())
+                .Input(readranges.PrefixPointsExpr().Cast())
+                .Done();
+        } else {
+            deleteInput = readranges.PrefixPointsExpr();
+        }
+    }
 
     YQL_ENSURE(deleteInput);
 

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
@@ -466,10 +466,6 @@ TExprBase KqpPushExtractedPredicateToReadTable(TExprBase node, TExprContext& ctx
                         .Done();
                 }
             }
-        } else if (buildResult.PointPrefixLen == tableDesc.Metadata->KeyColumnNames.size()) {
-            YQL_ENSURE(prefixPointsExpr);
-            residualLambda = pointsExtractionResult.PrunedLambda;
-            buildLookup(prefixPointsExpr, input);
         }
     }
 

diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h b/ydb/core/kqp/opt/logical/kqp_opt_log_rules.h
@@ -53,7 +53,7 @@ NYql::NNodes::TExprBase KqpRewriteTakeOverIndexRead(const NYql::NNodes::TExprBas
     const TKqpOptimizeContext& kqpCtx, const NYql::TParentsMap& parentsMap);
 
 NYql::NNodes::TExprBase KqpDeleteOverLookup(const NYql::NNodes::TExprBase& node, NYql::TExprContext& ctx,
-    const TKqpOptimizeContext &kqpCtx);
+    const TKqpOptimizeContext &kqpCtx, const NYql::TParentsMap& parentsMap);
 
 NYql::NNodes::TExprBase KqpExcessUpsertInputColumns(const NYql::NNodes::TExprBase& node, NYql::TExprContext& ctx);
 

diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp
@@ -370,7 +370,7 @@ TExprNode::TPtr KeepColumnOrder(const TExprNode::TPtr& node, const TExprNode& sr
 }
 
 template<class TFieldsSet>
-bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap, bool allowDependsOn) {
+bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap, bool allowDependsOn, bool allowTrivial) {
     const TTypeAnnotationNode* argType = RemoveOptionalType(arg.GetTypeAnn());
     if (argType->GetKind() != ETypeAnnotationKind::Struct) {
         return false;
@@ -412,15 +412,15 @@ bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TField
         }
     }
 
-    return usedFields.size() < inputStructType->GetSize();
+    return allowTrivial || usedFields.size() < inputStructType->GetSize();
 }
 
 template bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TSet<TStringBuf>& usedFields, const TParentsMap& parentsMap,
-                            bool allowDependsOn);
+                            bool allowDependsOn, bool allowTrivial);
 template bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TSet<TString>& usedFields, const TParentsMap& parentsMap,
-                            bool allowDependsOn);
+                            bool allowDependsOn, bool allowTrivial);
 template bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, std::map<std::string_view, TExprNode::TPtr>& usedFields,
-                            const TParentsMap& parentsMap, bool allowDependsOn);
+                            const TParentsMap& parentsMap, bool allowDependsOn, bool allowTrivial);
 
 TExprNode::TPtr AddMembersUsedInside(const TExprNode::TPtr& start, const TExprNode& arg, TExprNode::TPtr&& members, const TParentsMap& parentsMap, TExprContext& ctx) {
     if (!members || !start || &arg == start.Get()) {

diff --git a/ydb/library/yql/core/yql_opt_utils.h b/ydb/library/yql/core/yql_opt_utils.h
@@ -32,7 +32,7 @@ TExprNode::TPtr KeepColumnOrder(const TExprNode::TPtr& node, const TExprNode& sr
 // returns true if usedFields contains subset of fields
 template<class TFieldsSet>
 bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap,
-                    bool allowDependsOn = true);
+                    bool allowDependsOn = true, bool allowTrivial = false);
 
 template<class TFieldsSet>
 TExprNode::TPtr FilterByFields(TPositionHandle position, const TExprNode::TPtr& input, const TFieldsSet& subsetFields,