Skip to content

Commit 9d3feac

Browse files
Tony-RomanovTony-Romanov
authored
Fix Unique/Distinct constraints from nested filtering FlatMap. (#584)
* Add KDevelop * Fix Unique/Distinct constraints from nested filtering FlatMap. * A little fix.
1 parent a6e0cde commit 9d3feac

File tree

2 files changed

+43
-16
lines changed

2 files changed

+43
-16
lines changed

‎ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp

+26
Original file line numberDiff line numberDiff line change
@@ -687,6 +687,32 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
687687
CheckConstraint<TDistinctConstraintNode>(exprRoot, "OrderedMap", "Distinct((v))");
688688
}
689689

690+
Y_UNIT_TEST(NestedFlatMapByOptional) {
691+
const auto s = R"((
692+
(let mr_sink (DataSink 'yt (quote plato)))
693+
(let list (AsList
694+
(AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'v)))
695+
(AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'w)))
696+
(AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'u)))
697+
))
698+
(let list (AssumeUnique list '('key 'subkey) '('value)))
699+
(let list (AssumeDistinct list '('key) '('subkey 'value)))
700+
(let sorted (Sort list '((Bool 'False) (Bool 'True)) (lambda '(item) '((Member item 'value) (Member item 'subkey)))))
701+
(let mapped (OrderedFlatMap sorted (lambda '(item) (FlatMap (StrictCast (Member item 'key) (OptionalType (DataType 'Uint8)))
702+
(lambda '(key) (Just (AsStruct '('k key) '('s (Member item 'subkey)) '('v (Member item 'value)))))
703+
))))
704+
(let world (Write! world mr_sink (Key '('table (String 'Output))) mapped '('('mode 'renew))))
705+
(let world (Commit! world mr_sink))
706+
(return world)
707+
))";
708+
709+
TExprContext exprCtx;
710+
const auto exprRoot = ParseAndAnnotate(s, exprCtx);
711+
CheckConstraint<TSortedConstraintNode>(exprRoot, "OrderedFlatMap", "Sorted(v[desc];s[asc])");
712+
CheckConstraint<TUniqueConstraintNode>(exprRoot, "OrderedFlatMap", "Unique((k,s)(v))");
713+
CheckConstraint<TDistinctConstraintNode>(exprRoot, "OrderedFlatMap", "Distinct((s,v))");
714+
}
715+
690716
Y_UNIT_TEST(FlattenMembers) {
691717
const auto s = R"((
692718
(let mr_sink (DataSink 'yt (quote plato)))

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

+17-16
Original file line numberDiff line numberDiff line change
@@ -964,7 +964,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
964964
}
965965

966966
template<class TConstraint, bool OrderedMap, class TInput>
967-
static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, TExprContext& ctx) {
967+
static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, bool isSingleItem, TExprContext& ctx) {
968968
constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
969969
if (const auto lambda = handler.GetConstraint<TConstraint>()) {
970970
const auto original = input.template GetConstraint<typename TConstraint::TMainConstraint>();
@@ -986,15 +986,15 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
986986
if (!mapping.empty()) {
987987
output.AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
988988
}
989-
} else if constexpr (isOrderConstraint) {
989+
} else if (isOrderConstraint || isSingleItem) {
990990
if (const auto filtered = lambda->RemoveOriginal(ctx, original))
991991
output.AddConstraint(filtered);
992992
}
993993
}
994994
}
995995

996996
template<class TConstraint, bool OrderedMap, bool WideOutput>
997-
static void GetFromMapLambda(const TExprNode::TPtr& input, TExprContext& ctx) {
997+
static void GetFromMapLambda(const TExprNode::TPtr& input, bool isSingleItem, TExprContext& ctx) {
998998
constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
999999
if (const auto lambda = GetConstraintFromLambda<TConstraint, WideOutput>(input->Tail(), ctx)) {
10001000
const auto original = GetDetailed(input->Head().GetConstraint<typename TConstraint::TMainConstraint>(), *input->Head().GetTypeAnn(), ctx);
@@ -1016,7 +1016,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
10161016
if (!mapping.empty()) {
10171017
input->AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
10181018
}
1019-
} else if constexpr (isOrderConstraint) {
1019+
} else if (isOrderConstraint || isSingleItem) {
10201020
if (const auto filtered = lambda->RemoveOriginal(ctx, original))
10211021
input->AddConstraint(filtered);
10221022
}
@@ -1051,10 +1051,11 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
10511051
}
10521052
}
10531053

1054-
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered, WideOutput>(input, ctx);
1055-
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered, WideOutput>(input, ctx);
1056-
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered, WideOutput>(input, ctx);
1057-
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered, WideOutput>(input, ctx);
1054+
const bool singleItem = ETypeAnnotationKind::Optional == input->GetTypeAnn()->GetKind();
1055+
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
1056+
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
1057+
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
1058+
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
10581059

10591060
const auto lambdaVarIndex = GetConstraintFromLambda<TVarIndexConstraintNode, WideOutput>(input->Tail(), ctx);
10601061
const auto lambdaMulti = GetConstraintFromLambda<TMultiConstraintNode, WideOutput>(input->Tail(), ctx);
@@ -1101,10 +1102,10 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
11011102
}
11021103
}
11031104
}
1104-
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
1105-
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
1106-
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
1107-
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
1105+
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
1106+
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
1107+
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
1108+
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
11081109

11091110
if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>()) {
11101111
remappedItems.pop_back();
@@ -1127,10 +1128,10 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
11271128
}
11281129
}
11291130
}
1130-
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
1131-
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
1132-
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
1133-
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
1131+
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
1132+
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
1133+
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
1134+
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
11341135

11351136
if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>()) {
11361137
remappedItems.pop_back();

0 commit comments

Comments
 (0)