Skip to content

Commit 52b880c

Browse files
author
vokayndzop
committed
MR: support projection of aggregations in MEASURES
commit_hash:33457c2eb4b38c16a6b4c7dc3c2a90c588275f2f
1 parent 504c507 commit 52b880c

File tree

19 files changed

+439
-371
lines changed

19 files changed

+439
-371
lines changed

‎yql/essentials/core/type_ann/type_ann_core.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13018,7 +13018,8 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1301813018
Functions["NextValue"] = &NextValueWrapper;
1301913019

1302013020
Functions["MatchRecognize"] = &MatchRecognizeWrapper;
13021-
Functions["MatchRecognizeMeasuresAggregates"] = &MatchRecognizeMeasuresAggregatesWrapper;
13021+
Functions["MatchRecognizeMeasuresCallables"] = &MatchRecognizeMeasuresCallablesWrapper;
13022+
Functions["MatchRecognizeMeasuresCallable"] = &MatchRecognizeMeasuresCallableWrapper;
1302213023
Functions["MatchRecognizeParams"] = &MatchRecognizeParamsWrapper;
1302313024
Functions["MatchRecognizeMeasures"] = &MatchRecognizeMeasuresWrapper;
1302413025
Functions["MatchRecognizePattern"] = &MatchRecognizePatternWrapper;

‎yql/essentials/core/type_ann/type_ann_match_recognize.cpp

Lines changed: 89 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
#include "type_ann_match_recognize.h"
22

3+
#include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
34
#include <yql/essentials/core/sql_types/match_recognize.h>
45
#include <yql/essentials/core/yql_match_recognize.h>
56

67
namespace NYql::NTypeAnnImpl {
78

9+
using namespace NNodes;
10+
811
namespace {
912

1013
const TStructExprType* GetMatchedRowsRangesType(const TExprNode::TPtr& patternVars, TContext &ctx) {
@@ -72,15 +75,14 @@ IGraphTransformer::TStatus MatchRecognizeWrapper(const TExprNode::TPtr& input, T
7275
return IGraphTransformer::TStatus::Ok;
7376
}
7477

75-
IGraphTransformer::TStatus MatchRecognizeMeasuresAggregatesWrapper(const TExprNode::TPtr& input, TExprNode::TPtr&, TContext& ctx) {
76-
if (!EnsureMinArgsCount(*input, 4, ctx.Expr)) {
78+
IGraphTransformer::TStatus MatchRecognizeMeasuresCallablesWrapper(const TExprNode::TPtr& input, TExprNode::TPtr&, TContext& ctx) {
79+
if (!EnsureArgsCount(*input, 4, ctx.Expr)) {
7780
return IGraphTransformer::TStatus::Error;
7881
}
7982
const auto inputRowType = input->Child(0);
8083
const auto patternVars = input->Child(1);
8184
const auto names = input->Child(2);
82-
const auto vars = input->Child(3);
83-
constexpr size_t FirstLambdaIndex = 4;
85+
const auto callablesItems = input->Child(3);
8486

8587
if (!EnsureType(*inputRowType, ctx.Expr)) {
8688
return IGraphTransformer::TStatus::Error;
@@ -91,42 +93,85 @@ IGraphTransformer::TStatus MatchRecognizeMeasuresAggregatesWrapper(const TExprNo
9193
if (!EnsureTupleOfAtoms(*names, ctx.Expr)) {
9294
return IGraphTransformer::TStatus::Error;
9395
}
94-
if (!EnsureTupleOfAtoms(*vars, ctx.Expr)) {
96+
if (!EnsureTupleSize(*callablesItems, names->ChildrenSize(), ctx.Expr)) {
9597
return IGraphTransformer::TStatus::Error;
9698
}
97-
if (!EnsureArgsCount(*vars, names->ChildrenSize(), ctx.Expr)) {
99+
100+
TVector<const TItemExprType*> items;
101+
for (size_t i = 0; i < callablesItems->ChildrenSize(); ++i) {
102+
const auto name = names->Child(i)->Content();
103+
const auto type = callablesItems->Child(i)->GetTypeAnn();
104+
items.push_back(ctx.Expr.MakeType<TItemExprType>(name, type));
105+
}
106+
input->SetTypeAnn(ctx.Expr.MakeType<TStructExprType>(items));
107+
return IGraphTransformer::TStatus::Ok;
108+
}
109+
110+
IGraphTransformer::TStatus MatchRecognizeMeasuresCallableWrapper(const TExprNode::TPtr& input, TExprNode::TPtr&, TContext& ctx) {
111+
if (!EnsureArgsCount(*input, 3, ctx.Expr)) {
98112
return IGraphTransformer::TStatus::Error;
99113
}
100-
if (!EnsureArgsCount(*input, FirstLambdaIndex + names->ChildrenSize(), ctx.Expr)) {
114+
auto& lambda = input->ChildRef(0);
115+
const auto vars = input->Child(1);
116+
const auto aggregates = input->Child(2);
117+
118+
if (!EnsureTuple(*aggregates, ctx.Expr)) {
119+
return IGraphTransformer::TStatus::Error;
120+
}
121+
if (!EnsureTupleOfAtoms(*vars, ctx.Expr)) {
122+
return IGraphTransformer::TStatus::Error;
123+
}
124+
if (!EnsureTupleSize(*aggregates, vars->ChildrenSize(), ctx.Expr)) {
101125
return IGraphTransformer::TStatus::Error;
102126
}
103127

104128
TVector<const TItemExprType*> items;
105-
for (size_t i = 0; i < names->ChildrenSize(); ++i) {
106-
auto lambda = input->Child(FirstLambdaIndex + i);
107-
if (const auto varName = vars->Child(i)->Content()) {
108-
const auto traits = input->Child(FirstLambdaIndex + i);
109-
if (!EnsureTupleMinSize(*traits, 2, ctx.Expr)) {
110-
return IGraphTransformer::TStatus::Error;
111-
}
112-
if (!EnsureTupleMaxSize(*traits, 3, ctx.Expr)) {
113-
return IGraphTransformer::TStatus::Error;
114-
}
115-
if (!EnsureAtom(traits->Head(), ctx.Expr)) {
116-
return IGraphTransformer::TStatus::Error;
117-
}
118-
lambda = traits->Child(1)->Child(NNodes::TCoAggregationTraits::idx_FinishHandler);
129+
for (const auto& aggregate : aggregates->Children()) {
130+
if (!EnsureTupleMinSize(*aggregate, 2, ctx.Expr)) {
131+
return IGraphTransformer::TStatus::Error;
119132
}
120-
if (auto type = lambda->GetTypeAnn()) {
121-
if (type->GetKind() != ETypeAnnotationKind::Optional) {
122-
type = ctx.Expr.MakeType<TOptionalExprType>(type);
123-
}
124-
items.push_back(ctx.Expr.MakeType<TItemExprType>(names->Child(i)->Content(), type));
125-
} else {
126-
return IGraphTransformer::TStatus::Repeat;
133+
if (!EnsureTupleMaxSize(*aggregate, 3, ctx.Expr)) {
134+
return IGraphTransformer::TStatus::Error;
135+
}
136+
137+
const auto key = aggregate->Child(0);
138+
if (!EnsureAtom(*key, ctx.Expr)) {
139+
return IGraphTransformer::TStatus::Error;
127140
}
141+
142+
const auto traits = aggregate->Child(1);
143+
if (!traits->IsCallable(TCoAggregationTraits::CallableName())) {
144+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(aggregate->Pos()), TStringBuilder()
145+
<< "Expected AggregationTraits, but got: " << aggregate->Content()));
146+
return IGraphTransformer::TStatus::Error;
147+
}
148+
149+
if (aggregate->ChildrenSize() == 3 && !EnsureAtom(*aggregate->Child(2), ctx.Expr)) {
150+
return IGraphTransformer::TStatus::Error;
151+
}
152+
153+
auto finishType = traits->Child(TCoAggregationTraits::idx_DefVal)->IsCallable("Null")
154+
? traits->Child(TCoAggregationTraits::idx_FinishHandler)->GetTypeAnn()
155+
: traits->Child(TCoAggregationTraits::idx_DefVal)->GetTypeAnn();
156+
if (!finishType->IsOptionalOrNull()) {
157+
finishType = ctx.Expr.MakeType<TOptionalExprType>(finishType);
158+
}
159+
items.push_back(ctx.Expr.MakeType<TItemExprType>(key->Content(), finishType));
128160
}
129-
input->SetTypeAnn(ctx.Expr.MakeType<TStructExprType>(items));
161+
const auto status = ConvertToLambda(lambda, ctx.Expr, 1, 1);
162+
if (status != IGraphTransformer::TStatus::Ok) {
163+
return status;
164+
}
165+
if (!UpdateLambdaAllArgumentsTypes(lambda, {ctx.Expr.MakeType<TStructExprType>(items)}, ctx.Expr)) {
166+
return IGraphTransformer::TStatus::Error;
167+
}
168+
if (!lambda->GetTypeAnn()) {
169+
return IGraphTransformer::TStatus::Repeat;
170+
}
171+
if (!EnsureComputableType(lambda->Pos(), *lambda->GetTypeAnn(), ctx.Expr)) {
172+
return IGraphTransformer::TStatus::Error;
173+
}
174+
input->SetTypeAnn(lambda->GetTypeAnn());
130175
return IGraphTransformer::TStatus::Ok;
131176
}
132177

@@ -140,13 +185,13 @@ IGraphTransformer::TStatus MatchRecognizeParamsWrapper(const TExprNode::TPtr& in
140185
}
141186

142187
IGraphTransformer::TStatus MatchRecognizeMeasuresWrapper(const TExprNode::TPtr& input, TExprNode::TPtr&, TContext& ctx) {
143-
if (!EnsureMinArgsCount(*input, 3, ctx.Expr)) {
188+
constexpr size_t FirstLambdaIndex = 3;
189+
if (!EnsureMinArgsCount(*input, FirstLambdaIndex, ctx.Expr)) {
144190
return IGraphTransformer::TStatus::Error;
145191
}
146192
const auto inputRowType = input->Child(0);
147193
const auto patternVars = input->Child(1);
148194
const auto names = input->Child(2);
149-
constexpr size_t FirstLambdaIndex = 3;
150195

151196
if (!EnsureType(*inputRowType, ctx.Expr)) {
152197
return IGraphTransformer::TStatus::Error;
@@ -190,11 +235,13 @@ IGraphTransformer::TStatus MatchRecognizeMeasuresWrapper(const TExprNode::TPtr&
190235
ctx.Expr)) {
191236
return IGraphTransformer::TStatus::Error;
192237
}
193-
if (auto type = lambda->GetTypeAnn()) {
194-
items.push_back(ctx.Expr.MakeType<TItemExprType>(names->Child(i)->Content(), type));
195-
} else {
238+
if (!lambda->GetTypeAnn()) {
196239
return IGraphTransformer::TStatus::Repeat;
197240
}
241+
if (!EnsureComputableType(lambda->Pos(), *lambda->GetTypeAnn(), ctx.Expr)) {
242+
return IGraphTransformer::TStatus::Error;
243+
}
244+
items.push_back(ctx.Expr.MakeType<TItemExprType>(names->Child(i)->Content(), lambda->GetTypeAnn()));
198245
}
199246
input->SetTypeAnn(ctx.Expr.MakeType<TStructExprType>(items));
200247
return IGraphTransformer::TStatus::Ok;
@@ -206,13 +253,13 @@ IGraphTransformer::TStatus MatchRecognizePatternWrapper(const TExprNode::TPtr& i
206253
}
207254

208255
IGraphTransformer::TStatus MatchRecognizeDefinesWrapper(const TExprNode::TPtr& input, TExprNode::TPtr&, TContext &ctx) {
209-
if (!EnsureMinArgsCount(*input, 3, ctx.Expr)) {
256+
constexpr size_t FirstLambdaIndex = 3;
257+
if (!EnsureMinArgsCount(*input, FirstLambdaIndex, ctx.Expr)) {
210258
return IGraphTransformer::TStatus::Error;
211259
}
212260
const auto inputRowType = input->Child(0);
213261
const auto patternVars = input->Child(1);
214262
const auto names = input->Child(2);
215-
constexpr size_t FirstLambdaIndex = 3;
216263

217264
if (!EnsureType(*inputRowType, ctx.Expr)) {
218265
return IGraphTransformer::TStatus::Error;
@@ -246,16 +293,14 @@ IGraphTransformer::TStatus MatchRecognizeDefinesWrapper(const TExprNode::TPtr& i
246293
ctx.Expr)) {
247294
return IGraphTransformer::TStatus::Error;
248295
}
249-
if (auto type = lambda->GetTypeAnn()) {
250-
if (IsBoolLike(*type)) {
251-
items.push_back(ctx.Expr.MakeType<TItemExprType>(names->Child(i)->Content(), type));
252-
} else {
253-
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambda->Pos()), "DEFINE expression must be a predicate"));
254-
return IGraphTransformer::TStatus::Error;
255-
}
256-
} else {
296+
if (!lambda->GetTypeAnn()) {
257297
return IGraphTransformer::TStatus::Repeat;
258298
}
299+
if (!IsBoolLike(*lambda->GetTypeAnn())) {
300+
ctx.Expr.AddError(TIssue(ctx.Expr.GetPosition(lambda->Pos()), "DEFINE expression must be a predicate"));
301+
return IGraphTransformer::TStatus::Error;
302+
}
303+
items.push_back(ctx.Expr.MakeType<TItemExprType>(names->Child(i)->Content(), lambda->GetTypeAnn()));
259304
}
260305
input->SetTypeAnn(ctx.Expr.MakeType<TStructExprType>(items));
261306
return IGraphTransformer::TStatus::Ok;

‎yql/essentials/core/type_ann/type_ann_match_recognize.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
namespace NYql::NTypeAnnImpl {
99

1010
IGraphTransformer::TStatus MatchRecognizeWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
11-
IGraphTransformer::TStatus MatchRecognizeMeasuresAggregatesWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
11+
IGraphTransformer::TStatus MatchRecognizeMeasuresCallablesWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
12+
IGraphTransformer::TStatus MatchRecognizeMeasuresCallableWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
1213
IGraphTransformer::TStatus MatchRecognizeParamsWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
1314
IGraphTransformer::TStatus MatchRecognizeMeasuresWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
1415
IGraphTransformer::TStatus MatchRecognizePatternWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

0 commit comments

Comments
 (0)