Skip to content

Commit

Permalink
Fix Unique/Distinct constraints from nested filtering FlatMap. (#584)
Browse files Browse the repository at this point in the history 
* Add KDevelop

* Fix Unique/Distinct constraints from nested filtering FlatMap.

* A little fix.
  • Loading branch information
@Tony-Romanov
Tony-Romanov authored Dec 22, 2023
1 parent a6e0cde commit 9d3feac
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 16 deletions.
26 changes: 26 additions & 0 deletions ydb/library/yql/core/ut/yql_expr_constraint_ut.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -687,6 +687,32 @@ Y_UNIT_TEST_SUITE(TYqlExprConstraints) {
CheckConstraint<TDistinctConstraintNode>(exprRoot, "OrderedMap", "Distinct((v))");
}

Y_UNIT_TEST(NestedFlatMapByOptional) {
const auto s = R"((
(let mr_sink (DataSink 'yt (quote plato)))
(let list (AsList
(AsStruct '('key (String '4)) '('subkey (String 'c)) '('value (String 'v)))
(AsStruct '('key (String '1)) '('subkey (String 'd)) '('value (String 'w)))
(AsStruct '('key (String '3)) '('subkey (String 'b)) '('value (String 'u)))
))
(let list (AssumeUnique list '('key 'subkey) '('value)))
(let list (AssumeDistinct list '('key) '('subkey 'value)))
(let sorted (Sort list '((Bool 'False) (Bool 'True)) (lambda '(item) '((Member item 'value) (Member item 'subkey)))))
(let mapped (OrderedFlatMap sorted (lambda '(item) (FlatMap (StrictCast (Member item 'key) (OptionalType (DataType 'Uint8)))
(lambda '(key) (Just (AsStruct '('k key) '('s (Member item 'subkey)) '('v (Member item 'value)))))
))))
(let world (Write! world mr_sink (Key '('table (String 'Output))) mapped '('('mode 'renew))))
(let world (Commit! world mr_sink))
(return world)
))";

TExprContext exprCtx;
const auto exprRoot = ParseAndAnnotate(s, exprCtx);
CheckConstraint<TSortedConstraintNode>(exprRoot, "OrderedFlatMap", "Sorted(v[desc];s[asc])");
CheckConstraint<TUniqueConstraintNode>(exprRoot, "OrderedFlatMap", "Unique((k,s)(v))");
CheckConstraint<TDistinctConstraintNode>(exprRoot, "OrderedFlatMap", "Distinct((s,v))");
}

Y_UNIT_TEST(FlattenMembers) {
const auto s = R"((
(let mr_sink (DataSink 'yt (quote plato)))
Expand Down
33 changes: 17 additions & 16 deletions ydb/library/yql/core/yql_expr_constraint.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -964,7 +964,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
}

template<class TConstraint, bool OrderedMap, class TInput>
static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, TExprContext& ctx) {
static void GetFromMapLambda(const TInput& input, const TConstraintSet& handler, TConstraintSet& output, bool isSingleItem, TExprContext& ctx) {
constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
if (const auto lambda = handler.GetConstraint<TConstraint>()) {
const auto original = input.template GetConstraint<typename TConstraint::TMainConstraint>();
Expand All @@ -986,15 +986,15 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
if (!mapping.empty()) {
output.AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
}
} else if constexpr (isOrderConstraint) {
} else if (isOrderConstraint || isSingleItem) {
if (const auto filtered = lambda->RemoveOriginal(ctx, original))
output.AddConstraint(filtered);
}
}
}

template<class TConstraint, bool OrderedMap, bool WideOutput>
static void GetFromMapLambda(const TExprNode::TPtr& input, TExprContext& ctx) {
static void GetFromMapLambda(const TExprNode::TPtr& input, bool isSingleItem, TExprContext& ctx) {
constexpr bool isOrderConstraint = std::is_same<typename TConstraint::TMainConstraint, TSortedConstraintNode>() || std::is_same<typename TConstraint::TMainConstraint, TChoppedConstraintNode>();
if (const auto lambda = GetConstraintFromLambda<TConstraint, WideOutput>(input->Tail(), ctx)) {
const auto original = GetDetailed(input->Head().GetConstraint<typename TConstraint::TMainConstraint>(), *input->Head().GetTypeAnn(), ctx);
Expand All @@ -1016,7 +1016,7 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
if (!mapping.empty()) {
input->AddConstraint(ctx.MakeConstraint<TConstraint>(std::move(mapping)));
}
} else if constexpr (isOrderConstraint) {
} else if (isOrderConstraint || isSingleItem) {
if (const auto filtered = lambda->RemoveOriginal(ctx, original))
input->AddConstraint(filtered);
}
Expand Down Expand Up @@ -1051,10 +1051,11 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
}
}

GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered, WideOutput>(input, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered, WideOutput>(input, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered, WideOutput>(input, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered, WideOutput>(input, ctx);
const bool singleItem = ETypeAnnotationKind::Optional == input->GetTypeAnn()->GetKind();
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered, WideOutput>(input, singleItem, ctx);

const auto lambdaVarIndex = GetConstraintFromLambda<TVarIndexConstraintNode, WideOutput>(input->Tail(), ctx);
const auto lambdaMulti = GetConstraintFromLambda<TMultiConstraintNode, WideOutput>(input->Tail(), ctx);
Expand Down Expand Up @@ -1101,10 +1102,10 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
}
}
}
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(input->Head(), item.second, remappedItems.back().second, singleItem, ctx);

if (const auto empty = item.second.template GetConstraint<TEmptyConstraintNode>()) {
remappedItems.pop_back();
Expand All @@ -1127,10 +1128,10 @@ class TCallableConstraintTransformer : public TCallableTransformerBase<TCallable
}
}
}
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, ctx);
GetFromMapLambda<TPartOfUniqueConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfDistinctConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfSortedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);
GetFromMapLambda<TPartOfChoppedConstraintNode, Ordered>(*origConstr, item.second, remappedItems.back().second, singleItem, ctx);

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

0 comments on commit 9d3feac

Please sign in to comment.