Skip to content

Commit ebc5e19

Browse files
nepalnepal
committed
Apply MemberNthOverFlatMap only for single use FlatMap
commit_hash:51c45d45c91bce2df4992a8e76ed07bfc4bfbe7b
1 parent 55cdb44 commit ebc5e19

File tree

8 files changed

+506
-69
lines changed

8 files changed

+506
-69
lines changed

‎yql/essentials/core/common_opt/yql_co_extr_members.cpp‎

Lines changed: 83 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ TExprNode::TPtr ApplyExtractMembersToExtend(const TExprNode::TPtr& node, const T
5151
return ctx.NewCallable(node->Pos(), node->Content(), std::move(inputs));
5252
}
5353

54-
TExprNode::TPtr ApplyExtractMembersToSkipNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
54+
namespace {
55+
TExprNode::TPtr ApplyExtractMembersToSkipNullMembersLegacy(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
5556
TCoSkipNullMembers skipNullMembers(node);
5657
const auto& filtered = skipNullMembers.Members();
5758
if (!filtered) {
@@ -85,9 +86,7 @@ TExprNode::TPtr ApplyExtractMembersToSkipNullMembers(const TExprNode::TPtr& node
8586
.Done().Ptr();
8687
}
8788

88-
TExprNode::TPtr ApplyExtractMembersToFilterNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx,
89-
TOptimizeContext& optCtx, TStringBuf logSuffix)
90-
{
89+
TExprNode::TPtr ApplyExtractMembersToFilterNullMembersLegacy(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix) {
9190
TCoFilterNullMembers filterNullMembers(node);
9291
if (!filterNullMembers.Input().Maybe<TCoAssumeAllMembersNullableAtOnce>()) {
9392
return {};
@@ -156,31 +155,100 @@ TExprNode::TPtr ApplyExtractMembersToFilterNullMembers(const TExprNode::TPtr& no
156155

157156
YQL_CLOG(DEBUG, Core) << "Move ExtractMembers over " << node->Content() << logSuffix;
158157

159-
static const char optName[] = "MemberNthOverFlatMap";
160-
YQL_ENSURE(optCtx.Types);
161-
const bool keepAssume = !IsOptimizerDisabled<optName>(*optCtx.Types);
162-
163-
auto innerExtract = Build<TCoExtractMembers>(ctx, filterNullMembers.Pos())
164-
.Input(input)
165-
.Members(extendedMembers ? extendedMembers : members)
166-
.Done().Ptr();
167-
168158
if (extendedMembers) {
169159
return Build<TCoExtractMembers>(ctx, filterNullMembers.Pos())
170160
.Input<TCoFilterNullMembers>()
171-
.Input(ctx.WrapByCallableIf(keepAssume, TCoAssumeAllMembersNullableAtOnce::CallableName(), std::move(innerExtract)))
161+
.Input<TCoExtractMembers>()
162+
.Input(input)
163+
.Members(extendedMembers)
164+
.Build()
172165
.Members(filteredMembers)
173166
.Build()
174167
.Members(members)
175168
.Done().Ptr();
176169
}
177170

178171
return Build<TCoFilterNullMembers>(ctx, filterNullMembers.Pos())
179-
.Input(ctx.WrapByCallableIf(keepAssume, TCoAssumeAllMembersNullableAtOnce::CallableName(), std::move(innerExtract)))
172+
.Input<TCoExtractMembers>()
173+
.Input(input)
174+
.Members(members)
175+
.Build()
180176
.Members(filteredMembers)
181177
.Done().Ptr();
182178
}
183179

180+
} // namespace
181+
182+
TExprNode::TPtr ApplyExtractMembersToFilterSkipNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx,
183+
TOptimizeContext& optCtx, TStringBuf logSuffix)
184+
{
185+
static const char optName[] = "MemberNthOverFlatMap";
186+
YQL_ENSURE(optCtx.Types);
187+
if (IsOptimizerDisabled<optName>(*optCtx.Types)) {
188+
return node->IsCallable("FilterNullMembers") ?
189+
ApplyExtractMembersToFilterNullMembersLegacy(node, members, ctx, logSuffix) :
190+
ApplyExtractMembersToSkipNullMembersLegacy(node, members, ctx, logSuffix);
191+
}
192+
193+
TCoFilterNullMembersBase self(node);
194+
TSet<TStringBuf> filteredMembers = GetFilteredMembers(self);
195+
YQL_ENSURE(!filteredMembers.empty());
196+
TSet<TStringBuf> extractedMembers;
197+
for (const auto& atom : members->ChildrenList()) {
198+
extractedMembers.insert(atom->Content());
199+
}
200+
201+
TSet<TStringBuf> filteredAndExtracted;
202+
std::set_intersection(filteredMembers.begin(), filteredMembers.end(),
203+
extractedMembers.begin(), extractedMembers.end(),
204+
std::inserter(filteredAndExtracted, filteredAndExtracted.end()));
205+
206+
auto filterInput = self.Input();
207+
bool hasAssume = false;
208+
if (auto maybeAssume = filterInput.Maybe<TCoAssumeAllMembersNullableAtOnce>()) {
209+
filterInput = maybeAssume.Cast().Input();
210+
hasAssume = true;
211+
}
212+
213+
if (filteredAndExtracted == filteredMembers || (hasAssume && !filteredAndExtracted.empty())) {
214+
// simple pushdown
215+
YQL_CLOG(DEBUG, Core) << "Move ExtractMembers over " << self.CallableName() << logSuffix;
216+
auto newInput = ctx.NewCallable(self.Pos(), "ExtractMembers", { filterInput.Ptr(), members });
217+
return ctx.NewCallable(self.Pos(), self.CallableName(), {
218+
ctx.WrapByCallableIf(hasAssume, TCoAssumeAllMembersNullableAtOnce::CallableName(), std::move(newInput)),
219+
MakeAtomList(self.Pos(), filteredAndExtracted, ctx)
220+
});
221+
}
222+
223+
TSet<TStringBuf> innerExtracted = extractedMembers;
224+
if (hasAssume) {
225+
YQL_ENSURE(filteredAndExtracted.empty());
226+
// just leave single member
227+
filteredMembers = { *filteredMembers.begin() };
228+
}
229+
innerExtracted.insert(filteredMembers.begin(), filteredMembers.end());
230+
231+
const auto inputType = GetSequenceItemType(filterInput, false);
232+
YQL_ENSURE(inputType);
233+
const size_t inputWidth = inputType->Cast<TStructExprType>()->GetSize();
234+
YQL_ENSURE(inputWidth >= innerExtracted.size());
235+
236+
if (inputWidth == innerExtracted.size()) {
237+
return {};
238+
}
239+
240+
YQL_CLOG(DEBUG, Core) << "Push ExtractMembers over " << self.CallableName() << logSuffix;
241+
auto newInput = ctx.NewCallable(self.Pos(), "ExtractMembers", { filterInput.Ptr(), MakeAtomList(self.Pos(), innerExtracted, ctx) });
242+
auto newFilter = ctx.Builder(self.Pos())
243+
.Callable(self.CallableName())
244+
.Add(0, ctx.WrapByCallableIf(hasAssume, TCoAssumeAllMembersNullableAtOnce::CallableName(), std::move(newInput)))
245+
.Add(1, MakeAtomList(self.Pos(), filteredMembers, ctx))
246+
.Seal()
247+
.Build();
248+
return ctx.NewCallable(self.Pos(), "ExtractMembers", { newFilter, members });
249+
}
250+
251+
184252
TExprNode::TPtr ApplyExtractMembersToSortOrPruneKeys(const TExprNode::TPtr& node, const TExprNode::TPtr& members, const TParentsMap& parentsMap, TExprContext& ctx, TStringBuf logSuffix) {
185253
auto nodeIsPruneKeys = node->IsCallable("PruneKeys") || node->IsCallable("PruneAdjacentKeys");
186254
auto nodeIsSort = !nodeIsPruneKeys;

‎yql/essentials/core/common_opt/yql_co_extr_members.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,7 @@ namespace NYql {
1212
TExprNode::TPtr ApplyExtractMembersToTake(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
1313
TExprNode::TPtr ApplyExtractMembersToSkip(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
1414
TExprNode::TPtr ApplyExtractMembersToExtend(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
15-
TExprNode::TPtr ApplyExtractMembersToSkipNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);
16-
TExprNode::TPtr ApplyExtractMembersToFilterNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx,
15+
TExprNode::TPtr ApplyExtractMembersToFilterSkipNullMembers(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx,
1716
TOptimizeContext& optCtx, TStringBuf logSuffix);
1817
TExprNode::TPtr ApplyExtractMembersToSortOrPruneKeys(const TExprNode::TPtr& node, const TExprNode::TPtr& members, const TParentsMap& parentsMap, TExprContext& ctx, TStringBuf logSuffix);
1918
TExprNode::TPtr ApplyExtractMembersToAssumeUnique(const TExprNode::TPtr& node, const TExprNode::TPtr& members, TExprContext& ctx, TStringBuf logSuffix);

‎yql/essentials/core/common_opt/yql_co_finalizers.cpp‎

Lines changed: 105 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,22 @@ bool AllConsumersAreUnordered(const TExprNode::TPtr& node, const TParentsMap& pa
297297
return true;
298298
}
299299

300+
bool AllConsumersAreMembers(const TExprNode::TPtr& node, const TParentsMap& parents, TNodeSet& memberConsumers) {
301+
auto it = parents.find(node.Get());
302+
if (it == parents.cend() || it->second.empty()) {
303+
return false;
304+
}
305+
306+
for (auto parent : it->second) {
307+
if (!TCoMember::Match(parent)) {
308+
return false;
309+
}
310+
memberConsumers.insert(parent);
311+
}
312+
313+
return true;
314+
}
315+
300316
bool OptimizeForUnorderedConsumers(const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
301317
static const char optName[] = "UnorderedOverSortImproved";
302318
YQL_ENSURE(optCtx.Types);
@@ -332,6 +348,87 @@ bool IsFieldSubsetForOptionalsEnabled(const TOptimizeContext& optCtx) {
332348
return !IsOptimizerDisabled<optName>(*optCtx.Types);
333349
}
334350

351+
void OptimizeForMemberConsumers(const TCoFlatMapBase& self, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
352+
YQL_ENSURE(optCtx.Types);
353+
static const char optName[] = "MemberNthOverFlatMap";
354+
if (IsOptimizerDisabled<optName>(*optCtx.Types)) {
355+
return;
356+
}
357+
358+
auto maybeAsStruct = self.Lambda().Body().Maybe<TCoJust>().Input().Maybe<TCoAsStruct>();
359+
if (!maybeAsStruct) {
360+
return;
361+
}
362+
363+
YQL_ENSURE(optCtx.ParentsMap);
364+
const auto& parentsMap = *optCtx.ParentsMap;
365+
TNodeSet memberConsumers;
366+
if (!AllConsumersAreMembers(self.Ptr(), parentsMap, memberConsumers)) {
367+
return;
368+
}
369+
370+
TMap<TStringBuf, TExprNode::TPtr> structItems;
371+
for (auto item : maybeAsStruct.Cast().Ref().ChildrenList()) {
372+
TCoNameValueTuple tuple(item);
373+
structItems[tuple.Name().Value()] = tuple.Value().Cast().Ptr();
374+
}
375+
376+
TNodeSet restMembers;
377+
size_t separableMembersCount = 0;
378+
for (auto memberNode : memberConsumers) {
379+
TCoMember member(memberNode);
380+
auto name = member.Name().Value();
381+
auto it = structItems.find(name);
382+
YQL_ENSURE(it != structItems.end());
383+
384+
TExprNode::TPtr entry = it->second;
385+
if (IsDepended(*entry, self.Lambda().Args().Arg(0).Ref())) {
386+
restMembers.insert(memberNode);
387+
} else {
388+
++separableMembersCount;
389+
toOptimize[memberNode] = ctx.Builder(memberNode->Pos())
390+
.Callable(self.CallableName())
391+
.Add(0, self.Input().Ptr())
392+
.Lambda(1)
393+
.Param("unused")
394+
.Add(0, ctx.WrapByCallableIf(!entry->GetTypeAnn()->IsOptionalOrNull(), "Just", std::move(entry)))
395+
.Seal()
396+
.Seal()
397+
.Build();
398+
structItems.erase(it);
399+
}
400+
}
401+
402+
if (separableMembersCount && !restMembers.empty()) {
403+
auto restBody = ctx.Builder(self.Lambda().Body().Pos())
404+
.Callable("Just")
405+
.Callable(0, "AsStruct")
406+
.Do([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
407+
ui32 i = 0U;
408+
for (const auto& [name, value] : structItems) {
409+
parent.List(i)
410+
.Atom(0, name)
411+
.Add(1, value)
412+
.Seal();
413+
++i;
414+
}
415+
return parent;
416+
})
417+
.Seal()
418+
.Seal()
419+
.Build();
420+
auto restFlatMap = ctx.ChangeChild(self.Ref(), TCoFlatMapBase::idx_Lambda,
421+
ctx.DeepCopyLambda(*ctx.ChangeChild(self.Lambda().Ref(), TCoLambda::idx_Body, std::move(restBody))));
422+
for (auto restMember : restMembers) {
423+
toOptimize[restMember] = ctx.ChangeChild(*restMember, TCoMember::idx_Struct, TExprNode::TPtr(restFlatMap));
424+
}
425+
}
426+
427+
if (separableMembersCount) {
428+
YQL_CLOG(DEBUG, Core) << separableMembersCount << " separable Members and " << restMembers.size() << " non separable over " << self.CallableName();
429+
}
430+
}
431+
335432
}
336433

337434
void RegisterCoFinalizers(TFinalizingOptimizerMap& map) {
@@ -368,21 +465,12 @@ void RegisterCoFinalizers(TFinalizingOptimizerMap& map) {
368465
return true;
369466
};
370467

371-
map[TCoSkipNullMembers::CallableName()] = [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
372-
OptimizeSubsetFieldsForNodeWithMultiUsage(node, *optCtx.ParentsMap, toOptimize, ctx,
373-
[] (const TExprNode::TPtr& input, const TExprNode::TPtr& members, const TParentsMap&, TExprContext& ctx) {
374-
return ApplyExtractMembersToSkipNullMembers(input, members, ctx, " with multi-usage");
375-
},
376-
IsFieldSubsetForOptionalsEnabled(optCtx)
377-
);
378-
379-
return true;
380-
};
381-
382-
map[TCoFilterNullMembers::CallableName()] = [](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
468+
map[TCoFilterNullMembers::CallableName()] =
469+
map[TCoSkipNullMembers::CallableName()] =
470+
[](const TExprNode::TPtr& node, TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx, TOptimizeContext& optCtx) {
383471
OptimizeSubsetFieldsForNodeWithMultiUsage(node, *optCtx.ParentsMap, toOptimize, ctx,
384472
[&] (const TExprNode::TPtr& input, const TExprNode::TPtr& members, const TParentsMap&, TExprContext& ctx) {
385-
return ApplyExtractMembersToFilterNullMembers(input, members, ctx, optCtx, " with multi-usage");
473+
return ApplyExtractMembersToFilterSkipNullMembers(input, members, ctx, optCtx, " with multi-usage");
386474
},
387475
IsFieldSubsetForOptionalsEnabled(optCtx)
388476
);
@@ -398,6 +486,10 @@ void RegisterCoFinalizers(TFinalizingOptimizerMap& map) {
398486
IsFieldSubsetForOptionalsEnabled(optCtx)
399487
);
400488

489+
if (toOptimize.empty() && TCoFlatMapBase::Match(node.Get())) {
490+
OptimizeForMemberConsumers(TCoFlatMapBase(node), toOptimize, ctx, optCtx);
491+
}
492+
401493
return true;
402494
};
403495

‎yql/essentials/core/common_opt/yql_co_flow1.cpp‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2068,6 +2068,39 @@ void RegisterCoFlowCallables1(TCallableOptimizerMap& map) {
20682068
return node;
20692069
};
20702070

2071+
map[TCoMember::CallableName()] = map[TCoNth::CallableName()] = [](const TExprNode::TPtr& node, TExprContext& ctx, TOptimizeContext& optCtx) {
2072+
YQL_ENSURE(optCtx.Types);
2073+
static const char optName[] = "MemberNthOverFlatMap";
2074+
if (IsOptimizerDisabled<optName>(*optCtx.Types)) {
2075+
return node;
2076+
}
2077+
if (!optCtx.IsSingleUsage(node->Head())) {
2078+
return node;
2079+
}
2080+
if (auto maybeFlatMap = TMaybeNode<TCoFlatMapBase>(node->HeadPtr())) {
2081+
auto flatMap = maybeFlatMap.Cast();
2082+
if (flatMap.Input().Ref().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional &&
2083+
flatMap.Lambda().Ref().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Optional)
2084+
{
2085+
YQL_CLOG(DEBUG, Core) << node->Content() << " over " << node->Head().Content();
2086+
return ctx.Builder(node->Pos())
2087+
.Callable(flatMap.CallableName())
2088+
.Add(0, flatMap.Input().Ptr())
2089+
.Lambda(1)
2090+
.Param("item")
2091+
.Callable(node->Content())
2092+
.Apply(0, flatMap.Lambda().Ptr())
2093+
.With(0, "item")
2094+
.Seal()
2095+
.Add(1, node->Child(1))
2096+
.Seal()
2097+
.Seal()
2098+
.Seal()
2099+
.Build();
2100+
}
2101+
}
2102+
return node;
2103+
};
20712104
}
20722105

20732106
}

0 commit comments

Comments
 (0)