Skip to content

Commit 6709bae

Browse files
Tony-RomanovTony-Romanov
authored
Fix constrains after expand PartitionsByKeys. (#685)
* Fix constrains after expand PartitionsByKeys. * A little fix and cosmetic.
1 parent 8fea8a7 commit 6709bae

File tree

4 files changed

+127
-100
lines changed

4 files changed

+127
-100
lines changed

‎ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp‎

Lines changed: 39 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2445,23 +2445,17 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
24452445
const bool isStream = node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Flow ||
24462446
node->Head().GetTypeAnn()->GetKind() == ETypeAnnotationKind::Stream;
24472447
TExprNode::TPtr sort;
2448-
auto keyExtractor = node->ChildPtr(1);
2448+
const auto keyExtractor = node->Child(TCoPartitionsByKeys::idx_KeySelectorLambda);
24492449
const bool isConstKey = !IsDepended(keyExtractor->Tail(), keyExtractor->Head().Head());
2450-
const bool haveSort = !node->Child(2)->IsCallable("Void");
2451-
auto idLambda = ctx.Builder(node->Pos())
2452-
.Lambda()
2453-
.Param("x")
2454-
.Arg("x")
2455-
.Seal()
2456-
.Build();
2450+
const bool haveSort = node->Child(TCoPartitionsByKeys::idx_SortKeySelectorLambda)->IsLambda();
24572451

24582452
auto sortLambda = ctx.Builder(node->Pos())
24592453
.Lambda()
24602454
.Param("x")
24612455
.Callable("Sort")
24622456
.Arg(0, "x")
2463-
.Add(1, node->ChildPtr(2))
2464-
.Add(2, node->ChildPtr(3))
2457+
.Add(1, node->ChildPtr(TCoPartitionsByKeys::idx_SortDirections))
2458+
.Add(2, node->ChildPtr(TCoPartitionsByKeys::idx_SortKeySelectorLambda))
24652459
.Seal()
24662460
.Seal()
24672461
.Build();
@@ -2480,7 +2474,7 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
24802474
.Callable(0, "DictPayloads")
24812475
.Arg(0, "dict")
24822476
.Seal()
2483-
.Add(1, haveSort ? sortLambda : idLambda)
2477+
.Add(1, haveSort ? sortLambda : MakeIdentityLambda(node->Pos(), ctx))
24842478
.Seal()
24852479
.Seal()
24862480
.Build();
@@ -2490,8 +2484,8 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
24902484
sort = ctx.Builder(node->Pos())
24912485
.Callable("Sort")
24922486
.Add(0, node->HeadPtr())
2493-
.Add(1, node->ChildPtr(2))
2494-
.Add(2, node->ChildPtr(3))
2487+
.Add(1, node->ChildPtr(TCoPartitionsByKeys::idx_SortDirections))
2488+
.Add(2, node->ChildPtr(TCoPartitionsByKeys::idx_SortKeySelectorLambda))
24952489
.Seal()
24962490
.Build();
24972491
} else {
@@ -2503,8 +2497,8 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
25032497
.Callable("OrderedFlatMap")
25042498
.Callable(0, "SqueezeToDict")
25052499
.Add(0, node->HeadPtr())
2506-
.Add(1, std::move(keyExtractor))
2507-
.Add(2, std::move(idLambda))
2500+
.Add(1, node->ChildPtr(TCoPartitionsByKeys::idx_KeySelectorLambda))
2501+
.Add(2, MakeIdentityLambda(node->Pos(), ctx))
25082502
.Add(3, std::move(settings))
25092503
.Seal()
25102504
.Add(1, std::move(flatten))
@@ -2516,17 +2510,44 @@ TExprNode::TPtr ExpandPartitionsByKeys(const TExprNode::TPtr& node, TExprContext
25162510
.With(0)
25172511
.Callable("ToDict")
25182512
.Add(0, node->HeadPtr())
2519-
.Add(1, std::move(keyExtractor))
2520-
.Add(2, std::move(idLambda))
2513+
.Add(1, node->ChildPtr(TCoPartitionsByKeys::idx_KeySelectorLambda))
2514+
.Add(2, MakeIdentityLambda(node->Pos(), ctx))
25212515
.Add(3, std::move(settings))
25222516
.Seal()
25232517
.Done()
25242518
.Seal()
25252519
.Build();
25262520
}
2521+
2522+
if (auto keys = GetPathsToKeys(keyExtractor->Tail(), keyExtractor->Head().Head()); !keys.empty()) {
2523+
if (const auto sortKeySelector = node->Child(TCoPartitionsByKeys::idx_SortKeySelectorLambda); sortKeySelector->IsLambda()) {
2524+
auto sortKeys = GetPathsToKeys(sortKeySelector->Tail(), sortKeySelector->Head().Head());
2525+
std::move(sortKeys.begin(), sortKeys.end(), std::back_inserter(keys));
2526+
std::sort(keys.begin(), keys.end());
2527+
}
2528+
2529+
TExprNode::TListType columns;
2530+
columns.reserve(keys.size());
2531+
for (const auto& path : keys) {
2532+
if (1U == path.size())
2533+
columns.emplace_back(ctx.NewAtom(node->Pos(), path.front()));
2534+
else {
2535+
TExprNode::TListType atoms(path.size());
2536+
std::transform(path.cbegin(), path.cend(), atoms.begin(), [&](const std::string_view& name) { return ctx.NewAtom(node->Pos(), name); });
2537+
columns.emplace_back(ctx.NewList(node->Pos(), std::move(atoms)));
2538+
}
2539+
}
2540+
2541+
sort = ctx.Builder(node->Pos())
2542+
.Callable("AssumeChopped")
2543+
.Add(0, std::move(sort))
2544+
.List(1).Add(std::move(columns)).Seal()
2545+
.Seal()
2546+
.Build();
2547+
}
25272548
}
25282549

2529-
return KeepConstraints(ctx.ReplaceNode(node->Tail().TailPtr(), node->Tail().Head().Head(), std::move(sort)), *node, ctx);
2550+
return ctx.ReplaceNode(node->Tail().TailPtr(), node->Tail().Head().Head(), std::move(sort));
25302551
}
25312552

25322553
TExprNode::TPtr ExpandIsKeySwitch(const TExprNode::TPtr& node, TExprContext& ctx) {

‎ydb/library/yql/core/yql_expr_constraint.cpp‎

Lines changed: 4 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -753,38 +753,6 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
753753
return std::nullopt;
754754
}
755755

756-
static std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& body, const TExprNode& arg) {
757-
if (&body == &arg)
758-
return TPartOfConstraintBase::TPathType();
759-
760-
if (body.IsCallable({"Member","Nth"})) {
761-
if (auto path = GetPathToKey(body.Head(), arg)) {
762-
path->emplace_back(body.Tail().Content());
763-
return path;
764-
}
765-
}
766-
767-
if (body.IsCallable({"CastStruct","FilterMembers","Just","Unwrap"}))
768-
return GetPathToKey(body.Head(), arg);
769-
if (body.IsCallable("Member") && body.Head().IsCallable("AsStruct"))
770-
return GetPathToKey(GetLiteralStructMember(body.Head(), body.Tail()), arg);
771-
if (body.IsCallable("Nth") && body.Head().IsList())
772-
return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), arg);
773-
if (body.IsList() && 1U == body.ChildrenSize() && body.Head().IsCallable("Nth") && body.Head().Tail().IsAtom("0") &&
774-
1U == RemoveOptionality(*body.Head().Head().GetTypeAnn()).Cast<TTupleExprType>()->GetSize())
775-
// Especialy for "Extract single item tuple from Condense1" optimizer.
776-
return GetPathToKey(body.Head().Head(), arg);
777-
if (body.IsCallable("AsStruct") && 1U == body.ChildrenSize() && body.Head().Tail().IsCallable("Member") &&
778-
body.Head().Head().Content() == body.Head().Tail().Tail().Content() &&
779-
1U == RemoveOptionality(*body.Head().Tail().Head().GetTypeAnn()).Cast<TStructExprType>()->GetSize())
780-
// Especialy for "Extract single item struct from Condense1" optimizer.
781-
return GetPathToKey(body.Head().Tail().Head(), arg);
782-
if (IsTransparentIfPresent(body) && &body.Head() == &arg)
783-
return GetPathToKey(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
784-
785-
return std::nullopt;
786-
}
787-
788756
static std::vector<std::pair<TPartOfConstraintBase::TPathType, bool>>
789757
ExtractSimpleSortTraits(const TExprNode& sortDirections, const TExprNode& keySelectorLambda) {
790758
const auto& keySelectorBody = keySelectorLambda.Tail();
@@ -2684,34 +2652,6 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
26842652
.Combine(UpdateLambdaConstraints(input->ChildRef(TCoIsKeySwitch::idx_StateKeyExtractor), ctx, stateConstraints));
26852653
}
26862654

2687-
static const TExprNode& GetLiteralStructMember(const TExprNode& literal, const TExprNode& member) {
2688-
for (const auto& child : literal.Children())
2689-
if (&child->Head() == &member || child->Head().Content() == member.Content())
2690-
return child->Tail();
2691-
ythrow yexception() << "Member '" << member.Content() << "' not found in literal struct.";
2692-
}
2693-
2694-
static std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(const TExprNode& body, const TExprNode::TChildrenType& args) {
2695-
if (body.IsArgument()) {
2696-
for (auto i = 0U; i < args.size(); ++i)
2697-
if (&body == args[i].Get())
2698-
return std::make_pair(TPartOfConstraintBase::TPathType(), i);
2699-
} else if (body.IsCallable({"Member","Nth"})) {
2700-
if (auto path = GetPathToKey(body.Head(), args)) {
2701-
path->first.emplace_back(body.Tail().Content());
2702-
return path;
2703-
} else if (const auto& head = SkipCallables(body.Head(), {"CastStruct","FilterMembers"}); head.IsCallable("AsStruct") && body.IsCallable("Member")) {
2704-
return GetPathToKey(GetLiteralStructMember(head, body.Tail()), args);
2705-
} else if (body.IsCallable("Nth") && body.Head().IsList()) {
2706-
return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), args);
2707-
} else if (body.IsCallable({"CastStruct","FilterMembers"})) {
2708-
return GetPathToKey(body.Head(), args);
2709-
}
2710-
}
2711-
2712-
return std::nullopt;
2713-
}
2714-
27152655
template<bool Wide>
27162656
static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& body, const TExprNode::TChildrenType& args, TExprContext& ctx) {
27172657
TPartOfConstraintBase::TSetType keys;
@@ -2749,33 +2689,15 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
27492689
return keys;
27502690
}
27512691

2752-
static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& body, const TExprNode& arg) {
2753-
TPartOfConstraintBase::TSetType keys;
2754-
if (body.IsList()) {
2755-
if (const auto size = body.ChildrenSize()) {
2756-
keys.reserve(size);
2757-
for (auto i = 0U; i < size; ++i)
2758-
if (auto path = GetPathToKey(*body.Child(i), arg))
2759-
keys.insert_unique(std::move(*path));
2760-
}
2761-
} else if (body.IsCallable("StablePickle")) {
2762-
return GetSimpleKeys(body.Head(), arg);
2763-
} else if (auto path = GetPathToKey(body, arg)) {
2764-
keys.insert_unique(std::move(*path));
2765-
}
2766-
2767-
return keys;
2768-
}
2769-
27702692
template<bool Wide>
27712693
static TPartOfConstraintBase::TSetType GetSimpleKeys(const TExprNode& selector, TExprContext& ctx) {
27722694
YQL_ENSURE(selector.IsLambda() && 2U == selector.ChildrenSize());
27732695
const auto& body = selector.Tail();
27742696
if constexpr (!Wide) {
27752697
if (TCoIsKeySwitch::Match(&body)) {
27762698
const TCoIsKeySwitch keySwitch(&body);
2777-
const auto& i = GetSimpleKeys(*ctx.ReplaceNode(keySwitch.ItemKeyExtractor().Body().Ptr(), keySwitch.ItemKeyExtractor().Args().Arg(0).Ref(), keySwitch.Item().Ptr()), keySwitch.Item().Ref());
2778-
const auto& s = GetSimpleKeys(*ctx.ReplaceNode(keySwitch.StateKeyExtractor().Body().Ptr(), keySwitch.StateKeyExtractor().Args().Arg(0).Ref(), keySwitch.State().Ptr()), keySwitch.Item().Ref());
2699+
const auto& i = GetPathsToKeys(*ctx.ReplaceNode(keySwitch.ItemKeyExtractor().Body().Ptr(), keySwitch.ItemKeyExtractor().Args().Arg(0).Ref(), keySwitch.Item().Ptr()), keySwitch.Item().Ref());
2700+
const auto& s = GetPathsToKeys(*ctx.ReplaceNode(keySwitch.StateKeyExtractor().Body().Ptr(), keySwitch.StateKeyExtractor().Args().Arg(0).Ref(), keySwitch.State().Ptr()), keySwitch.Item().Ref());
27792701
return i == s ? i : TPartOfConstraintBase::TSetType();
27802702
}
27812703
}
@@ -3005,13 +2927,13 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
30052927

30062928
TPartOfConstraintBase::TSetType keys;
30072929
if constexpr (Partitions) {
3008-
keys = GetSimpleKeys(input->Child(TCoBase::idx_KeySelectorLambda)->Tail(), input->Child(TCoBase::idx_KeySelectorLambda)->Head().Head());
2930+
keys = GetPathsToKeys(input->Child(TCoBase::idx_KeySelectorLambda)->Tail(), input->Child(TCoBase::idx_KeySelectorLambda)->Head().Head());
30092931
if (const auto sortKeySelector = input->Child(TCoBase::idx_SortKeySelectorLambda); sortKeySelector->IsLambda()) {
30102932
if (const auto status = UpdateLambdaConstraints(*sortKeySelector); status != TStatus::Ok) {
30112933
return status;
30122934
}
30132935

3014-
auto sortKeys = GetSimpleKeys(sortKeySelector->Tail(), sortKeySelector->Head().Head());
2936+
auto sortKeys = GetPathsToKeys(sortKeySelector->Tail(), sortKeySelector->Head().Head());
30152937
std::move(sortKeys.begin(), sortKeys.end(), std::back_inserter(keys));
30162938
std::sort(keys.begin(), keys.end());
30172939
}

‎ydb/library/yql/core/yql_opt_utils.cpp‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,13 @@ TExprNodeBuilder GetterBuilder(TExprNodeBuilder parent, ui32 index, const TTypeA
8282
return GetterBuilder(parent.Callable(index, ETypeAnnotationKind::Struct == parentType->GetKind() ? "Member" : "Nth"), 0U, type, path).Atom(1, name).Seal();
8383
}
8484

85+
const TExprNode& GetLiteralStructMember(const TExprNode& literal, const TExprNode& member) {
86+
for (const auto& child : literal.Children())
87+
if (&child->Head() == &member || child->Head().Content() == member.Content())
88+
return child->Tail();
89+
ythrow yexception() << "Member '" << member.Content() << "' not found in literal struct.";
90+
}
91+
8592
}
8693

8794
TExprNode::TPtr MakeBoolNothing(TPositionHandle position, TExprContext& ctx) {
@@ -1994,4 +2001,76 @@ void OptimizeSubsetFieldsForNodeWithMultiUsage(const TExprNode::TPtr& node, cons
19942001
}
19952002
}
19962003

2004+
2005+
std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(const TExprNode& body, const TExprNode::TChildrenType& args) {
2006+
if (body.IsArgument()) {
2007+
for (auto i = 0U; i < args.size(); ++i)
2008+
if (&body == args[i].Get())
2009+
return std::make_pair(TPartOfConstraintBase::TPathType(), i);
2010+
} else if (body.IsCallable({"Member","Nth"})) {
2011+
if (auto path = GetPathToKey(body.Head(), args)) {
2012+
path->first.emplace_back(body.Tail().Content());
2013+
return path;
2014+
} else if (const auto& head = SkipCallables(body.Head(), {"CastStruct","FilterMembers"}); head.IsCallable("AsStruct") && body.IsCallable("Member")) {
2015+
return GetPathToKey(GetLiteralStructMember(head, body.Tail()), args);
2016+
} else if (body.IsCallable("Nth") && body.Head().IsList()) {
2017+
return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), args);
2018+
} else if (body.IsCallable({"CastStruct","FilterMembers"})) {
2019+
return GetPathToKey(body.Head(), args);
2020+
}
2021+
}
2022+
2023+
return std::nullopt;
2024+
}
2025+
2026+
std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& body, const TExprNode& arg) {
2027+
if (&body == &arg)
2028+
return TPartOfConstraintBase::TPathType();
2029+
2030+
if (body.IsCallable({"Member","Nth"})) {
2031+
if (auto path = GetPathToKey(body.Head(), arg)) {
2032+
path->emplace_back(body.Tail().Content());
2033+
return path;
2034+
}
2035+
}
2036+
2037+
if (body.IsCallable({"CastStruct","FilterMembers","Just","Unwrap"}))
2038+
return GetPathToKey(body.Head(), arg);
2039+
if (body.IsCallable("Member") && body.Head().IsCallable("AsStruct"))
2040+
return GetPathToKey(GetLiteralStructMember(body.Head(), body.Tail()), arg);
2041+
if (body.IsCallable("Nth") && body.Head().IsList())
2042+
return GetPathToKey(*body.Head().Child(FromString<ui32>(body.Tail().Content())), arg);
2043+
if (body.IsList() && 1U == body.ChildrenSize() && body.Head().IsCallable("Nth") && body.Head().Tail().IsAtom("0") &&
2044+
1U == RemoveOptionality(*body.Head().Head().GetTypeAnn()).Cast<TTupleExprType>()->GetSize())
2045+
// Especialy for "Extract single item tuple from Condense1" optimizer.
2046+
return GetPathToKey(body.Head().Head(), arg);
2047+
if (body.IsCallable("AsStruct") && 1U == body.ChildrenSize() && body.Head().Tail().IsCallable("Member") &&
2048+
body.Head().Head().Content() == body.Head().Tail().Tail().Content() &&
2049+
1U == RemoveOptionality(*body.Head().Tail().Head().GetTypeAnn()).Cast<TStructExprType>()->GetSize())
2050+
// Especialy for "Extract single item struct from Condense1" optimizer.
2051+
return GetPathToKey(body.Head().Tail().Head(), arg);
2052+
if (IsTransparentIfPresent(body) && &body.Head() == &arg)
2053+
return GetPathToKey(body.Child(1)->Tail().Head(), body.Child(1)->Head().Head());
2054+
2055+
return std::nullopt;
2056+
}
2057+
2058+
TPartOfConstraintBase::TSetType GetPathsToKeys(const TExprNode& body, const TExprNode& arg) {
2059+
TPartOfConstraintBase::TSetType keys;
2060+
if (body.IsList()) {
2061+
if (const auto size = body.ChildrenSize()) {
2062+
keys.reserve(size);
2063+
for (auto i = 0U; i < size; ++i)
2064+
if (auto path = GetPathToKey(*body.Child(i), arg))
2065+
keys.insert_unique(std::move(*path));
2066+
}
2067+
} else if (body.IsCallable("StablePickle")) {
2068+
return GetPathsToKeys(body.Head(), arg);
2069+
} else if (auto path = GetPathToKey(body, arg)) {
2070+
keys.insert_unique(std::move(*path));
2071+
}
2072+
2073+
return keys;
2074+
}
2075+
19972076
}

‎ydb/library/yql/core/yql_opt_utils.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,4 +149,9 @@ void OptimizeSubsetFieldsForNodeWithMultiUsage(const TExprNode::TPtr& node, cons
149149
TNodeOnNodeOwnedMap& toOptimize, TExprContext& ctx,
150150
std::function<TExprNode::TPtr(const TExprNode::TPtr&, const TExprNode::TPtr&, const TParentsMap&, TExprContext&)> handler);
151151

152+
std::optional<TPartOfConstraintBase::TPathType> GetPathToKey(const TExprNode& body, const TExprNode& arg);
153+
std::optional<std::pair<TPartOfConstraintBase::TPathType, ui32>> GetPathToKey(const TExprNode& body, const TExprNode::TChildrenType& args);
154+
155+
TPartOfConstraintBase::TSetType GetPathsToKeys(const TExprNode& body, const TExprNode& arg);
156+
152157
}

0 commit comments

Comments
 (0)