diff options
| author | aneporada <[email protected]> | 2025-08-07 14:15:45 +0300 |
|---|---|---|
| committer | aneporada <[email protected]> | 2025-08-07 14:32:20 +0300 |
| commit | cd8fda2d174e97c10e8f59e0286039aa77a5df89 (patch) | |
| tree | 5b1356e34c9537d604fe7bf9607ffd05b809501c | |
| parent | acba301ee4565c2e4e621d0d250c5e4c56c91f70 (diff) | |
Push non-member equality filter expressions over join
commit_hash:6bc558d09291d35872e4f2a9da0ada2f6fe5e4f3
10 files changed, 372 insertions, 4 deletions
diff --git a/yql/essentials/core/common_opt/yql_co_flow2.cpp b/yql/essentials/core/common_opt/yql_co_flow2.cpp index 8fec283e275..222c177f5ed 100644 --- a/yql/essentials/core/common_opt/yql_co_flow2.cpp +++ b/yql/essentials/core/common_opt/yql_co_flow2.cpp @@ -1090,7 +1090,7 @@ TExprNode::TPtr PullUpFlatMapOverEquiJoin(const TExprNode::TPtr& node, TExprCont return node; } - static const TStringBuf canaryBaseName = "_yql_canary_"; + static const TStringBuf canaryBaseName = YqlCanaryColumnName; THashMap<TStringBuf, THashSet<TStringBuf>> joinKeysByLabel = CollectEquiJoinKeyColumnsByLabel(*joinTree); const auto renames = LoadJoinRenameMap(*settings); @@ -1109,7 +1109,7 @@ TExprNode::TPtr PullUpFlatMapOverEquiJoin(const TExprNode::TPtr& node, TExprCont const TTypeAnnotationNode* itemType = input.List().Ref().GetTypeAnn()->Cast<TListExprType>()->GetItemType(); auto structType = itemType->Cast<TStructExprType>(); for (auto& si : structType->GetItems()) { - if (si->GetName().find(canaryBaseName, 0) == 0) { + if (IsNoPullColumn(si->GetName())) { // EquiJoin already processed return node; } diff --git a/yql/essentials/core/common_opt/yql_flatmap_over_join.cpp b/yql/essentials/core/common_opt/yql_flatmap_over_join.cpp index d0f039a8fc7..baa1fe5f05d 100644 --- a/yql/essentials/core/common_opt/yql_flatmap_over_join.cpp +++ b/yql/essentials/core/common_opt/yql_flatmap_over_join.cpp @@ -1177,6 +1177,12 @@ bool IsExtractOrPredicatesOverEquiJoinEnabled(const TTypeAnnotationContext* type return IsOptimizerEnabled<flag>(*types) && !IsOptimizerDisabled<flag>(*types); } +bool IsNormalizeEqualityFilterOverJoinEnabled(const TTypeAnnotationContext* types) { + YQL_ENSURE(types); + static const char flag[] = "NormalizeEqualityFilterOverJoin"; + return IsOptimizerEnabled<flag>(*types) && !IsOptimizerDisabled<flag>(*types); +} + struct TExtraInputPredicates { TExprNode::TPtr Row; TExprNodeList Preds; @@ -1343,6 +1349,222 @@ TJoinEqRebuildResult RebuildJoinTreeForEquality(TVector<TMap<ui32, TEqColumn>>& } +TExprBase NormalizeEqualityFilterOverJoin(const TCoFlatMapBase& node, const TJoinLabels& labels, + const THashMap<TString, TString>& backRenameMap, const TParentsMap& parentsMap, TExprContext& ctx) +{ + TCoEquiJoin equiJoin = node.Input().Cast<TCoEquiJoin>(); + const TExprNode::TPtr joinTree = equiJoin.Arg(equiJoin.ArgCount() - 2).Ptr(); + const TExprNode::TPtr rowPtr = node.Lambda().Args().Arg(0).Ptr(); + const auto& row = *rowPtr; + const auto body = node.Lambda().Body().Cast<TCoConditionalValueBase>(); + + auto predicate = body.Predicate().Ptr(); + + TExprNodeList andComponents; + if (predicate->IsCallable("And")) { + andComponents = predicate->ChildrenList(); + } else { + andComponents.push_back(predicate); + } + + TExprNodeList rest; + TNodeOnNodeOwnedMap remaps; + TVector<TExprNodeList> nodesByInput(equiJoin.ArgCount() - 2); + + for (auto pred : andComponents) { + TExprNodeList sides(2); + // TODO: handle case IsEquality() && !IsMemberEquality() + if (!IsEquality(pred, sides.front(), sides.back()) || HasDependsOn(pred, rowPtr) || !IsStrict(pred)) { + rest.push_back(pred); + continue; + } + + + TVector<ui32> indexes; + TSet<ui32> inputs; + TSet<TStringBuf> usedFields; + + GatherJoinInputs(sides.front(), row, parentsMap, backRenameMap, labels, inputs, usedFields); + if (inputs.size() != 1) { + rest.push_back(pred); + continue; + } + indexes.push_back(*inputs.begin()); + YQL_ENSURE(indexes.back() < nodesByInput.size()); + + inputs.clear(); + GatherJoinInputs(sides.back(), row, parentsMap, backRenameMap, labels, inputs, usedFields); + if (inputs.size() != 1 || *inputs.begin() == indexes.front()) { + rest.push_back(pred); + continue; + } + indexes.push_back(*inputs.begin()); + YQL_ENSURE(indexes.back() < nodesByInput.size()); + + + size_t count = 0; + for (size_t i = 0; i < 2; ++i) { + TExprNode::TPtr side = sides[i]; + if (!side->IsCallable("Member") || &side->Head() != &row) { + ++count; + if (!remaps.contains(side.Get())) { + TExprNode::TPtr output; + TOptimizeExprSettings settings(nullptr); + auto ret = OptimizeExpr(side, output, [&](const TExprNode::TPtr& node, TExprContext& ctx) -> TExprNode::TPtr { + if (node->IsCallable("Member") && &node->Head() == &row) { + const TString originalMemberName{node->Tail().Content()}; + TString memberName = originalMemberName; + if (auto renamed = backRenameMap.FindPtr(memberName)) { + memberName = *renamed; + } + const auto& label = labels.Inputs[indexes[i]]; + TString sourceMemberName = label.MemberName(label.TableName(memberName), label.ColumnName(memberName)); + if (sourceMemberName != originalMemberName) { + return ctx.ChangeChild(*node, TCoMember::idx_Name, ctx.NewAtom(node->Pos(), sourceMemberName)); + } + } + return node; + }, ctx, settings); + YQL_ENSURE(ret != IGraphTransformer::TStatus::Error); + remaps[side.Get()] = output; + } + nodesByInput[indexes[i]].push_back(side); + } + } + + if (!count) { + rest.push_back(pred); + } + } + + if (remaps.empty()) { + return node; + } + + TExprNodeList newJoinArgs; + YQL_ENSURE(labels.Inputs.size() == nodesByInput.size()); + TNodeOnNodeOwnedMap outputRemaps; + TSet<TString> outputMembers; + for (size_t i = 0; i < nodesByInput.size(); ++i) { + TCoEquiJoinInput ejInput = equiJoin.Arg(i).Cast<TCoEquiJoinInput>(); + if (nodesByInput[i].empty()) { + newJoinArgs.push_back(ejInput.Ptr()); + continue; + } + + const TStructExprType* itemType = ejInput.List().Ref().GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>(); + const TStructExprType* castedType = itemType; + if (!IsRequiredSide(joinTree, labels, i).first) { + TVector<const TItemExprType *> items; + for (auto item : itemType->GetItems()) { + auto type = item->GetItemType(); + if (type->IsOptionalOrNull()) { + items.push_back(item); + } else { + items.push_back(ctx.MakeType<TItemExprType>(item->GetName(), ctx.MakeType<TOptionalExprType>(type))); + } + } + castedType = ctx.MakeType<TStructExprType>(items); + } + TString prefix; + bool isMultiTableInput = false; + if (labels.Inputs[i].Tables.size() > 1) { + prefix = TStringBuilder() << labels.Inputs[i].Tables[0] << "."; + isMultiTableInput = true; + } + prefix += YqlJoinKeyColumnName; + TVector<TString> remappedNames = GenNoClashColumns(*itemType, prefix, nodesByInput[i].size()); + newJoinArgs.push_back(ctx.Builder(equiJoin.Arg(i).Pos()) + .List() + .Callable(0, "OrderedMap") + .Add(0, ejInput.List().Ptr()) + .Lambda(1) + .Param("row") + .Callable("FlattenMembers") + .List(0) + .Atom(0, "") + .Arg(1, "row") + .Seal() + .List(1) + .Atom(0, "") + .Callable(1, "AsStruct") + .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& { + for (size_t j = 0; j < nodesByInput[i].size(); ++j) { + auto calcLambda = ctx.ChangeChild(node.Lambda().Ref(), TCoLambda::idx_Body, TExprNode::TPtr(remaps[nodesByInput[i][j].Get()])); + parent + .List(j) + .Atom(0, remappedNames[j]) + .Apply(1, calcLambda) + .With(0) + .Callable("SafeCast") + .Arg(0, "row") + .Add(1, ExpandType(equiJoin.Arg(i).Pos(), *castedType, ctx)) + .Seal() + .Done() + .Seal() + .Seal(); + TString outputMember = isMultiTableInput ? remappedNames[j] : labels.Inputs[i].FullName(remappedNames[j]); + outputRemaps[nodesByInput[i][j].Get()] = ctx.Builder(nodesByInput[i][j]->Pos()) + .Callable("Member") + .Add(0, rowPtr) + .Atom(1, outputMember) + .Seal() + .Build(); + YQL_ENSURE(outputMembers.insert(outputMember).second); + } + return parent; + }) + .Seal() + .Seal() + .Seal() + .Seal() + .Seal() + .Add(1, ejInput.Scope().Ptr()) + .Seal() + .Build()); + } + newJoinArgs.push_back(equiJoin.Arg(equiJoin.ArgCount() - 2).Ptr()); + newJoinArgs.push_back(equiJoin.Arg(equiJoin.ArgCount() - 1).Ptr()); + + auto newJoin = ctx.ChangeChildren(equiJoin.Ref(), std::move(newJoinArgs)); + TExprNode::TPtr newPredicate; + auto status = RemapExpr(body.Predicate().Ptr(), newPredicate, outputRemaps, ctx, TOptimizeExprSettings{nullptr}); + YQL_ENSURE(status != IGraphTransformer::TStatus::Error); + + TExprNode::TPtr newPredicateLambda = ctx.ChangeChild(node.Lambda().Ref(), TCoLambda::idx_Body, std::move(newPredicate)); + TExprNode::TPtr valueLambda = ctx.ChangeChild(node.Lambda().Ref(), TCoLambda::idx_Body, body.Value().Ptr()); + + return TExprBase(ctx.Builder(node.Pos()) + .Callable(node.CallableName()) + .Add(0, newJoin) + .Lambda(1) + .Param("row") + .Callable(body.CallableName()) + .Apply(0, newPredicateLambda) + .With(0, "row") + .Seal() + .Apply(1, valueLambda) + .With(0) + .Callable("RemoveMembers") + .Arg(0, "row") + .List(1) + .Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& { + size_t i = 0; + for (auto name : outputMembers) { + parent.Atom(i++, name); + } + return parent; + }) + .Seal() + .Seal() + .Done() + .Seal() + .Seal() + .Seal() + .Seal() + .Build()); +} + TExprBase HandleEqualityFilterOverJoin(const TCoFlatMapBase& node, const TJoinLabels& labels, const THashMap<TString, TString>& backRenameMap, TExprContext& ctx) { @@ -1671,6 +1893,14 @@ TExprBase FlatMapOverEquiJoin( } } + if (IsNormalizeEqualityFilterOverJoinEnabled(types)) { + auto newNode = NormalizeEqualityFilterOverJoin(node, labels, backRenameMap, parentsMap, ctx); + if (newNode.Raw() != node.Raw()) { + YQL_CLOG(DEBUG, Core) << "NormalizeEqualityFilterOverJoin"; + return newNode; + } + } + if (IsEqualityFilterOverJoinEnabled(types)) { auto newNode = HandleEqualityFilterOverJoin(node, labels, backRenameMap, ctx); if (newNode.Raw() != node.Raw()) { diff --git a/yql/essentials/core/yql_join.cpp b/yql/essentials/core/yql_join.cpp index d08f1b9515f..085d3316c5b 100644 --- a/yql/essentials/core/yql_join.cpp +++ b/yql/essentials/core/yql_join.cpp @@ -2280,4 +2280,13 @@ TExprNode::TPtr DropAnyOverJoinInputs(TExprNode::TPtr joinTree, const TJoinLabel return joinTree; } +bool IsNoPullColumn(TStringBuf columnName) { + if (columnName.Contains('.')) { + TStringBuf table, column; + SplitTableName(columnName, table, column); + columnName = column; + } + return columnName.StartsWith(YqlCanaryColumnName) || columnName.StartsWith(YqlJoinKeyColumnName); +} + } // namespace NYql diff --git a/yql/essentials/core/yql_join.h b/yql/essentials/core/yql_join.h index 347ad9b9039..43db2beaec5 100644 --- a/yql/essentials/core/yql_join.h +++ b/yql/essentials/core/yql_join.h @@ -197,4 +197,8 @@ void GetPruneKeysColumnsForJoinLeaves(const NNodes::TCoEquiJoinTuple& joinTree, TExprNode::TPtr DropAnyOverJoinInputs(TExprNode::TPtr joinTree, const TJoinLabels& labels, const THashMap<TStringBuf, THashSet<TStringBuf>>& keyColumnsByLabel, TExprContext& ctx); +static constexpr TStringBuf YqlCanaryColumnName = "_yql_canary_"; +static constexpr TStringBuf YqlJoinKeyColumnName = "_yql_join_key"; + +bool IsNoPullColumn(TStringBuf columnName); } diff --git a/yql/essentials/core/yql_opt_utils.cpp b/yql/essentials/core/yql_opt_utils.cpp index 8038d15af80..91f41b80bf4 100644 --- a/yql/essentials/core/yql_opt_utils.cpp +++ b/yql/essentials/core/yql_opt_utils.cpp @@ -1,6 +1,7 @@ #include "yql_opt_utils.h" #include "yql_expr_optimize.h" #include "yql_expr_type_annotation.h" +#include "yql_join.h" #include "yql_type_annotation.h" #include "yql_type_helpers.h" @@ -2393,7 +2394,12 @@ template TPartOfConstraintBase::TSetType GetPathsToKeys<true>(const TExprNode& b template TPartOfConstraintBase::TSetType GetPathsToKeys<false>(const TExprNode& body, const TExprNode& arg); TVector<TString> GenNoClashColumns(const TStructExprType& source, TStringBuf prefix, size_t count) { - YQL_ENSURE(prefix.StartsWith("_yql")); + if (!prefix.StartsWith("_yql")) { + YQL_ENSURE(prefix.Contains('.')); + TStringBuf table, column; + SplitTableName(prefix, table, column); + YQL_ENSURE(column.StartsWith("_yql")); + } TSet<size_t> existing; for (auto& item : source.GetItems()) { TStringBuf column = item->GetName(); diff --git a/yql/essentials/core/yql_opt_utils.h b/yql/essentials/core/yql_opt_utils.h index 1444a45170f..ed7e07bde12 100644 --- a/yql/essentials/core/yql_opt_utils.h +++ b/yql/essentials/core/yql_opt_utils.h @@ -181,7 +181,7 @@ template<bool Ordered = false> TPartOfConstraintBase::TSetType GetPathsToKeys(const TExprNode& body, const TExprNode& arg); // generates column names with pattern "prefixN" that do not clash with source columns -// prefix should start with "_yql" +// prefix should start with "_yql" or "<alias>._yql" TVector<TString> GenNoClashColumns(const TStructExprType& source, TStringBuf prefix, size_t count); bool CheckSupportedTypes( diff --git a/yql/essentials/tests/s-expressions/minirun/part2/canondata/result.json b/yql/essentials/tests/s-expressions/minirun/part2/canondata/result.json index 32fe3083d83..30cf9201c51 100644 --- a/yql/essentials/tests/s-expressions/minirun/part2/canondata/result.json +++ b/yql/essentials/tests/s-expressions/minirun/part2/canondata/result.json @@ -391,6 +391,20 @@ "uri": "https://{canondata_backend}/1775319/1b77f6a2ced99cb43171907f99fa05c5976782c8/resource.tar.gz#test.test_Optimizers-MoveExtractMembersOverSort1-default.txt-Results_/results.txt" } ], + "test.test[Optimizers-NormalizeEqualityFilterOverJoinOptional-default.txt-Debug]": [ + { + "checksum": "ca31dd8effb927d18a48d099fdb1e18b", + "size": 1940, + "uri": "https://{canondata_backend}/1881367/22c06fc51a47708d8e7ffe921b07d495560d2fa1/resource.tar.gz#test.test_Optimizers-NormalizeEqualityFilterOverJoinOptional-default.txt-Debug_/opt.yql" + } + ], + "test.test[Optimizers-NormalizeEqualityFilterOverJoinOptional-default.txt-Results]": [ + { + "checksum": "df92f713466329443519bc6326e864a7", + "size": 2163, + "uri": "https://{canondata_backend}/1881367/22c06fc51a47708d8e7ffe921b07d495560d2fa1/resource.tar.gz#test.test_Optimizers-NormalizeEqualityFilterOverJoinOptional-default.txt-Results_/results.txt" + } + ], "test.test[Optimizers-PushDownToFlow-default.txt-Debug]": [ { "checksum": "ed25dda0682388131b4840cbc338e1d3", diff --git a/yql/essentials/tests/s-expressions/minirun/part3/canondata/result.json b/yql/essentials/tests/s-expressions/minirun/part3/canondata/result.json index f941367312d..751f6b466a0 100644 --- a/yql/essentials/tests/s-expressions/minirun/part3/canondata/result.json +++ b/yql/essentials/tests/s-expressions/minirun/part3/canondata/result.json @@ -447,6 +447,20 @@ "uri": "https://{canondata_backend}/1031349/de1275760530cb37ff3ca2d7aa0c26a1e7476b02/resource.tar.gz#test.test_Optimizers-LengthAfterList-default.txt-Results_/results.txt" } ], + "test.test[Optimizers-NormalizeEqualityFilterOverJoin-default.txt-Debug]": [ + { + "checksum": "a06e166f1839a293a236237622423b8c", + "size": 1612, + "uri": "https://{canondata_backend}/1903885/fdf3d9a48d438f90af2da132ed3e94bcfa6909fe/resource.tar.gz#test.test_Optimizers-NormalizeEqualityFilterOverJoin-default.txt-Debug_/opt.yql" + } + ], + "test.test[Optimizers-NormalizeEqualityFilterOverJoin-default.txt-Results]": [ + { + "checksum": "75a63145bd73f1ceeb0abee9bd7e25db", + "size": 2200, + "uri": "https://{canondata_backend}/1903885/fdf3d9a48d438f90af2da132ed3e94bcfa6909fe/resource.tar.gz#test.test_Optimizers-NormalizeEqualityFilterOverJoin-default.txt-Results_/results.txt" + } + ], "test.test[Result-Pickle-default.txt-Debug]": [ { "checksum": "ff0f81f90265c0afaa74303e7198d058", diff --git a/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoin.yqls b/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoin.yqls new file mode 100644 index 00000000000..4b98074ab5a --- /dev/null +++ b/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoin.yqls @@ -0,0 +1,37 @@ +( +(let config (DataSource 'config)) +(let world (Configure! world config '"OptimizerFlags" '"NormalizeEqualityFilterOverJoin")) +(let world (Configure! world config '"OptimizerFlags" '"EqualityFilterOverJoin")) + +(let list1 (AsList + (AsStruct '('a.key1 (Just (Int32 '1))) '('c.value1 (Just (String 'A)))) + (AsStruct '('a.key1 (Just (Int32 '7))) '('c.value1 (Just (String 'B)))) + (AsStruct '('a.key1 (Just (Int32 '4))) '('c.value1 (Just (String 'C)))) + (AsStruct '('a.key1 (Just (Int32 '4))) '('c.value1 (Just (String 'D)))) +)) + +(let list2 (AsList + (AsStruct '('key2 (Int32 '9)) '('value2 (String 'Z))) + (AsStruct '('key2 (Int32 '4)) '('value2 (String 'Y))) + (AsStruct '('key2 (Int32 '3)) '('value2 (String 'X))) + (AsStruct '('key2 (Int32 '4)) '('value2 (String 'W))) + (AsStruct '('key2 (Int32 '8)) '('value2 (String 'V))) +)) + +(let joinCross (EquiJoin '(list1 '('a 'c)) '(list2 'b) '('Cross 'c 'b '() '() '()) +'( +'('rename '"a.key1" '"akey1") +'('rename '"c.value1" '"cvalue1") + +) +)) + +(let filtered (Filter joinCross +(lambda '(row) (Coalesce (== (+ (Member row 'akey1) (Int32 '2)) (+ (Member row 'b.key2) (Int32 '0))) (Bool 'false)) +))) + +(let res_sink (DataSink 'result)) +(let world (Write! world res_sink (Key) filtered '('('type)))) +(let world (Commit! world res_sink)) +(return world) +) diff --git a/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoinOptional.yqls b/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoinOptional.yqls new file mode 100644 index 00000000000..84a77e7652f --- /dev/null +++ b/yql/essentials/tests/s-expressions/suites/Optimizers/NormalizeEqualityFilterOverJoinOptional.yqls @@ -0,0 +1,54 @@ +( +(let config (DataSource 'config)) +(let world (Configure! world config '"OptimizerFlags" '"NormalizeEqualityFilterOverJoin")) +(let world (Configure! world config '"OptimizerFlags" '"EqualityFilterOverJoin")) + +(let list1 (AsList + (AsStruct '('key1 (Int32 '9)) '('value1 (String 'Z))) + (AsStruct '('key1 (Int32 '4)) '('value1 (String 'Y))) + (AsStruct '('key1 (Int32 '3)) '('value1 (String 'X))) + (AsStruct '('key1 (Int32 '4)) '('value1 (String 'W))) + (AsStruct '('key1 (Int32 '8)) '('value1 (String 'V))) +)) + +(let list2 (AsList + (AsStruct '('key2 (Int32 '9)) '('value2 (String 'Z))) + (AsStruct '('key2 (Int32 '4)) '('value2 (String 'Y))) + (AsStruct '('key2 (Int32 '3)) '('value2 (String 'X))) + (AsStruct '('key2 (Int32 '4)) '('value2 (String 'W))) + (AsStruct '('key2 (Int32 '8)) '('value2 (String 'V))) +)) + +(let list3 (AsList + (AsStruct '('key3 (Int32 '9)) '('value3 (String 'Z))) + (AsStruct '('key3 (Int32 '4)) '('value3 (String 'Y))) + (AsStruct '('key3 (Int32 '3)) '('value3 (String 'X))) + (AsStruct '('key3 (Int32 '4)) '('value3 (String 'W))) + (AsStruct '('key3 (Int32 '8)) '('value3 (String 'V))) +)) + +(let joinLeftCross (EquiJoin '(list1 'a) '(list2 'b) '(list3 'c) + +'('Cross '('Left 'a 'b '('a 'key1) '('b 'key2) '()) 'c '() '() '()) + +'( +'('rename '"a.key1" '"akey1") +'('rename '"b.key2" '"bkey2") + +))) + + +(let filtered (Filter joinLeftCross +(lambda '(row) (Coalesce (== (+ (Member row 'bkey2) (Int32 '1)) (+ (Member row 'c.key3) (Int32 '0))) (Bool 'false)) +))) + +(let filtered (Map filtered +(lambda '(row) (AsStruct '('ckey3 (Member row 'c.key3)) '('akey1 (Member row 'akey1)) '('bkey2 (Member row 'bkey2)) '('avalue1 (Member row 'a.value1))) +))) + + +(let res_sink (DataSink 'result)) +(let world (Write! world res_sink (Key) filtered '('('type)))) +(let world (Commit! world res_sink)) +(return world) +) |
