diff --git a/ydb/core/kqp/opt/logical/kqp_opt_log.cpp b/ydb/core/kqp/opt/logical/kqp_opt_log.cpp index 5119d7769d20..7ec98aa7f569 100644 --- 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 DeleteOverLookup(TExprBase node, TExprContext& ctx) { - TExprBase output = KqpDeleteOverLookup(node, ctx, KqpCtx); + TMaybeNode 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 index aa4dfb187fde..daed288e4486 100644 --- 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 #include +#include + +namespace { + +bool CanPushFlatMap(const NYql::NNodes::TCoFlatMapBase& flatMap, const NYql::TKikimrTableDescription& tableDesc, const NYql::TParentsMap& parentsMap, TVector & extraColumns) { + auto flatMapLambda = flatMap.Lambda(); + if (!NYql::IsFilterFlatMap(flatMapLambda)) { + return false; + } + + const auto & flatMapLambdaArgument = flatMapLambda.Args().Arg(0).Ref(); + auto flatMapLambdaConditional = flatMapLambda.Body().Cast(); + + TSet 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()) { return node; } auto deleteRows = node.Cast(); + TMaybeNode filter; + TMaybeNode lookup; TMaybeNode read; TMaybeNode skipNulMembers; + TMaybeNode readranges; if (deleteRows.Input().Maybe()) { lookup = deleteRows.Input().Cast(); @@ -27,7 +60,15 @@ TExprBase KqpDeleteOverLookup(const TExprBase& node, TExprContext& ctx, const TK } else if (deleteRows.Input().Maybe()) { read = deleteRows.Input().Cast(); } else { - return node; + TMaybeNode input = deleteRows.Input(); + if (input.Maybe()) { + filter = deleteRows.Input().Cast(); + input = filter.Input(); + } + readranges = input.Maybe(); + if (!readranges) { + return node; + } } TMaybeNode 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 extraColumns; + if (!CanPushFlatMap(filter.Cast(), tableDesc, parentsMap, extraColumns)) { + return node; + } + deleteInput = Build(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 index 5a8ff001bcd5..a4ed3c3c62bf 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 @@ -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 index 97b61b836c0b..9f4852eff0c9 100644 --- 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/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp b/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp index 6977d7981c5c..8c2c10f23a39 100644 --- a/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp +++ b/ydb/core/kqp/ut/opt/kqp_extract_predicate_unpack_ut.cpp @@ -359,82 +359,6 @@ Y_UNIT_TEST(ComplexRange) { 2); } -Y_UNIT_TEST(SqlIn) { - Test( - R"( - SELECT * FROM `/Root/SimpleKey` - WHERE Key IN AsList(100, 102, (100 + 3)) - ORDER BY Key; - )", - R"([ - [[100];["Value20"]];[[102];["Value22"]];[[103];["Value23"]] - ])"); -} - -Y_UNIT_TEST(BasicLookup) { - Test( - R"( - SELECT * FROM `/Root/SimpleKey` - WHERE Key = 100 or Key = 102 or Key = 103 or Key = null; - )", - R"([ - [[100];["Value20"]];[[102];["Value22"]];[[103];["Value23"]] - ])"); -} - -Y_UNIT_TEST(ComplexLookup) { - Test( - R"( - SELECT Key, Value FROM `/Root/SimpleKey` - WHERE Key = 100 or Key = 102 or Key = (100 + 3); - )", - R"([ - [[100];["Value20"]];[[102];["Value22"]];[[103];["Value23"]] - ])"); -} - -Y_UNIT_TEST(SqlInComplexKey) { - Test( - R"( - SELECT Key, Fk, Value FROM `/Root/ComplexKey` - WHERE (Key, Fk) IN AsList( - (1, 101), - (2, 102), - (2, 102 + 1), - ) - ORDER BY Key, Fk; - )", - R"([ - [[1];[101];["Value1"]];[[2];[102];["Value1"]];[[2];[103];["Value3"]] - ])"); -} - -Y_UNIT_TEST(BasicLookupComplexKey) { - Test( - R"( - SELECT Key, Fk, Value FROM `/Root/ComplexKey` - WHERE (Key = 1 and Fk = 101) OR - (2 = Key and 102 = Fk) OR - (2 = Key and 103 = Fk); - )", - R"([ - [[1];[101];["Value1"]];[[2];[102];["Value1"]];[[2];[103];["Value3"]] - ])"); -} - -Y_UNIT_TEST(ComplexLookupComplexKey) { - Test( - R"( - SELECT Key, Fk, Value FROM `/Root/ComplexKey` - WHERE (Key = 1 and Fk = 101) OR - (2 = Key and 102 = Fk) OR - (2 = Key and 102 + 1 = Fk); - )", - R"([ - [[1];[101];["Value1"]];[[2];[102];["Value1"]];[[2];[103];["Value3"]] - ])"); -} - Y_UNIT_TEST(PointJoin) { Test( R"( diff --git a/ydb/library/yql/core/yql_opt_utils.cpp b/ydb/library/yql/core/yql_opt_utils.cpp index a1599989dee7..1e672a88f7da 100644 --- 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 -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& usedFields, const TParentsMap& parentsMap, - bool allowDependsOn); + bool allowDependsOn, bool allowTrivial); template bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TSet& usedFields, const TParentsMap& parentsMap, - bool allowDependsOn); + bool allowDependsOn, bool allowTrivial); template bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, std::map& 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 index c89e705e96b6..eb83269d9267 100644 --- 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 bool HaveFieldsSubset(const TExprNode::TPtr& start, const TExprNode& arg, TFieldsSet& usedFields, const TParentsMap& parentsMap, - bool allowDependsOn = true); + bool allowDependsOn = true, bool allowTrivial = false); template TExprNode::TPtr FilterByFields(TPositionHandle position, const TExprNode::TPtr& input, const TFieldsSet& subsetFields,