Skip to content

Commit

Permalink
Drop unused fields from GraceJoinCore inputs. (ydb-platform#7134)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Tony-Romanov
Tony-Romanov authored Jul 26, 2024
1 parent 6f2e7f0 commit 28816eb
Showing 1 changed file with 90 additions and 0 deletions.
90 changes: 90 additions & 0 deletions ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5073,6 +5073,19 @@ TExprNode::TListType&& DropUnused(TExprNode::TListType&& list, const std::vector
return std::move(list);
}

std::set<ui32> FillAllIndexes(const TTypeAnnotationNode& rowType) {
std::set<ui32> set;
for (ui32 i = 0U; i < rowType.Cast<TMultiExprType>()->GetSize(); ++i)
set.emplace(i);
return set;
}

template<bool EvenOnly>
void RemoveUsedIndexes(const TExprNode& list, std::set<ui32>& unused) {
for (auto i = 0U; i < list.ChildrenSize() >> (EvenOnly ? 1U : 0U); ++i)
unused.erase(FromString<ui32>(list.Child(EvenOnly ? i << 1U : i)->Content()));
}

TExprNode::TPtr DropUnusedArgs(const TExprNode& lambda, const std::vector<ui32>& unused, TExprContext& ctx, ui32 skip = 0U) {
const auto& copy = ctx.DeepCopyLambda(lambda);
return ctx.ChangeChild(*copy, 0U, ctx.NewArguments(copy->Head().Pos(), DropUnused(copy->Head().ChildrenList(), unused, skip)));
Expand All @@ -5095,6 +5108,28 @@ void DropUnusedRenames(TExprNode::TPtr& renames, const std::vector<ui32>& unused
renames = ctx.ChangeChildren(*renames, std::move(children));
}

template<bool EvenOnly>
void UpdateInputIndexes(TExprNode::TPtr& indexes, const std::vector<ui32>& unused, TExprContext& ctx) {
TExprNode::TListType children;
children.reserve(indexes->ChildrenSize());
for (auto i = 0U; i < indexes->ChildrenSize() >> (EvenOnly ? 1U : 0U); ++i) {
const auto idx = i << (EvenOnly ? 1U : 0U);
const auto outIndex = indexes->Child(idx);
const auto oldValue = FromString<ui32>(outIndex->Content());
const auto newValue = oldValue - ui32(std::distance(unused.cbegin(), std::lower_bound(unused.cbegin(), unused.cend(), oldValue)));
if (oldValue == newValue) {
children.emplace_back(indexes->ChildPtr(idx));
} else {
children.emplace_back(ctx.NewAtom(indexes->Child(idx)->Pos(), newValue));
}
if constexpr (EvenOnly) {
children.emplace_back(indexes->ChildPtr(1U + idx));
}
}

indexes = ctx.ChangeChildren(*indexes, std::move(children));
}

TExprNode::TPtr DropUnusedStateFromUpdate(const TExprNode& lambda, const std::vector<ui32>& unused, TExprContext& ctx) {
const auto& copy = ctx.DeepCopyLambda(lambda, DropUnused(GetLambdaBody(lambda), unused));
return ctx.ChangeChild(*copy, 0U, ctx.NewArguments(copy->Head().Pos(), DropUnused(copy->Head().ChildrenList(), unused, lambda.Head().ChildrenSize() - lambda.ChildrenSize() + 1U)));
Expand Down Expand Up @@ -6698,6 +6733,59 @@ TExprNode::TPtr OptimizeMapJoinCore(const TExprNode::TPtr& node, TExprContext& c
return node;
}

TExprNode::TPtr OptimizeGraceJoinCore(const TExprNode::TPtr& node, TExprContext& ctx) {
const TCoGraceJoinCore grace(node);
auto leftUnused = FillAllIndexes(GetSeqItemType(*grace.LeftInput().Ref().GetTypeAnn()));
auto rightUnused = FillAllIndexes(GetSeqItemType(*grace.RightInput().Ref().GetTypeAnn()));
RemoveUsedIndexes<false>(grace.LeftKeysColumns().Ref(), leftUnused);
RemoveUsedIndexes<false>(grace.RightKeysColumns().Ref(), rightUnused);
RemoveUsedIndexes<true>(grace.LeftRenames().Ref(), leftUnused);
RemoveUsedIndexes<true>(grace.RightRenames().Ref(), rightUnused);
if (!(leftUnused.empty() && rightUnused.empty())) {
YQL_CLOG(DEBUG, CorePeepHole) << node->Content() << " with " << leftUnused.size() << " left and " << rightUnused.size() << " right unused columns.";

auto children = node->ChildrenList();
if (!leftUnused.empty()) {
const std::vector<ui32> unused(leftUnused.cbegin(), leftUnused.cend());
children[TCoGraceJoinCore::idx_LeftInput] = MakeWideMapForDropUnused(std::move(children[TCoGraceJoinCore::idx_LeftInput]), unused, ctx);
UpdateInputIndexes<false>(children[TCoGraceJoinCore::idx_LeftKeysColumns], unused, ctx);
UpdateInputIndexes<true>(children[TCoGraceJoinCore::idx_LeftRenames], unused, ctx);
}
if (!rightUnused.empty()) {
const std::vector<ui32> unused(rightUnused.cbegin(), rightUnused.cend());
children[TCoGraceJoinCore::idx_RightInput] = MakeWideMapForDropUnused(std::move(children[TCoGraceJoinCore::idx_RightInput]), unused, ctx);
UpdateInputIndexes<false>(children[TCoGraceJoinCore::idx_RightKeysColumns], unused, ctx);
UpdateInputIndexes<true>(children[TCoGraceJoinCore::idx_RightRenames], unused, ctx);
}
return ctx.ChangeChildren(*node, std::move(children));
}

return node;
}

TExprNode::TPtr OptimizeGraceSelfJoinCore(const TExprNode::TPtr& node, TExprContext& ctx) {
const TCoGraceSelfJoinCore grace(node);
auto unused = FillAllIndexes(GetSeqItemType(*grace.Input().Ref().GetTypeAnn()));
RemoveUsedIndexes<false>(grace.LeftKeysColumns().Ref(), unused);
RemoveUsedIndexes<false>(grace.RightKeysColumns().Ref(), unused);
RemoveUsedIndexes<true>(grace.LeftRenames().Ref(), unused);
RemoveUsedIndexes<true>(grace.RightRenames().Ref(), unused);
if (!unused.empty()) {
YQL_CLOG(DEBUG, CorePeepHole) << node->Content() << " with " << unused.size() << " unused columns.";

auto children = node->ChildrenList();
const std::vector<ui32> vector(unused.cbegin(), unused.cend());
children[TCoGraceSelfJoinCore::idx_Input] = MakeWideMapForDropUnused(std::move(children[TCoGraceSelfJoinCore::idx_Input]), vector, ctx);
UpdateInputIndexes<false>(children[TCoGraceSelfJoinCore::idx_LeftKeysColumns], vector, ctx);
UpdateInputIndexes<false>(children[TCoGraceSelfJoinCore::idx_RightKeysColumns], vector, ctx);
UpdateInputIndexes<true>(children[TCoGraceSelfJoinCore::idx_LeftRenames], vector, ctx);
UpdateInputIndexes<true>(children[TCoGraceSelfJoinCore::idx_RightRenames], vector, ctx);
return ctx.ChangeChildren(*node, std::move(children));
}

return node;
}

TExprNode::TPtr OptimizeCommonJoinCore(const TExprNode::TPtr& node, TExprContext& ctx) {
if (const auto& input = node->Head(); input.IsCallable("NarrowMap") && input.Tail().Tail().IsCallable("AsStruct")) {
YQL_CLOG(DEBUG, CorePeepHole) << "Swap " << node->Content() << " with " << input.Content();
Expand Down Expand Up @@ -8338,6 +8426,8 @@ struct TPeepHoleRules {
{"WideCondense1", &OptimizeWideCondense1},
{"WideChopper", &OptimizeWideChopper},
{"MapJoinCore", &OptimizeMapJoinCore},
{"GraceJoinCore", &OptimizeGraceJoinCore},
{"GraceSelfJoinCore", &OptimizeGraceSelfJoinCore},
{"CommonJoinCore", &OptimizeCommonJoinCore},
{"BuildTablePath", &DoBuildTablePath},
{"Exists", &OptimizeExists},
Expand Down

0 comments on commit 28816eb

Please sign in to comment.