3
3
#include < ydb/core/formats/arrow/ssa_runtime_version.h>
4
4
#include < yql/essentials/core/yql_opt_utils.h>
5
5
#include < yql/essentials/core/yql_expr_optimize.h>
6
+ #include < yql/essentials/utils/log/log.h>
6
7
7
8
namespace NKikimr ::NKqp::NOpt {
8
9
@@ -28,24 +29,28 @@ bool IsSupportedDataType(const TCoDataCtor& node, bool allowOlapApply) {
28
29
node.Maybe <TCoUint32>() ||
29
30
node.Maybe <TCoUint64>() ||
30
31
node.Maybe <TCoUtf8>() ||
31
- node.Maybe <TCoString>()) {
32
- return true ;
33
- }
34
-
35
- if (node.Maybe <TCoTimestamp>()) {
32
+ node.Maybe <TCoString>()
33
+ ) {
36
34
return true ;
37
35
}
38
36
39
37
if (allowOlapApply) {
40
- if (node.Maybe <TCoDate32>() || node.Maybe <TCoDatetime64>() || node.Maybe <TCoTimestamp64>() || node.Maybe <TCoInterval64>()) {
38
+ if (node.Maybe <TCoDate>() ||
39
+ node.Maybe <TCoDate32>() ||
40
+ node.Maybe <TCoDatetime>() ||
41
+ node.Maybe <TCoDatetime64>() ||
42
+ node.Maybe <TCoTimestamp>() ||
43
+ node.Maybe <TCoTimestamp64>() ||
44
+ node.Maybe <TCoInterval>() ||
45
+ node.Maybe <TCoInterval64>()) {
41
46
return true ;
42
47
}
43
48
}
44
49
45
50
return false ;
46
51
}
47
52
48
- bool IsSupportedCast (const TCoSafeCast& cast) {
53
+ bool IsSupportedCast (const TCoSafeCast& cast, bool allowOlapApply= false ) {
49
54
auto maybeDataType = cast.Type ().Maybe <TCoDataType>();
50
55
if (!maybeDataType) {
51
56
if (const auto maybeOptionalType = cast.Type ().Maybe <TCoOptionalType>()) {
@@ -54,6 +59,10 @@ bool IsSupportedCast(const TCoSafeCast& cast) {
54
59
}
55
60
YQL_ENSURE (maybeDataType.IsValid ());
56
61
62
+ if (allowOlapApply) {
63
+ return true ;
64
+ }
65
+
57
66
const auto dataType = maybeDataType.Cast ();
58
67
if (dataType.Type ().Value () == " Int32" ) { // TODO: Support any numeric casts.
59
68
return cast.Value ().Maybe <TCoString>() || cast.Value ().Maybe <TCoUtf8>();
@@ -87,21 +96,25 @@ bool IsMemberColumn(const TExprBase& node, const TExprNode* lambdaArg) {
87
96
}
88
97
89
98
bool IsGoodTypeForArithmeticPushdown (const TTypeAnnotationNode& type, bool allowOlapApply) {
90
- const auto fatures = NUdf::GetDataTypeInfo (RemoveOptionality (type).Cast <TDataExprType>()->GetSlot ()).Features ;
91
- return NUdf::EDataTypeFeatures::NumericType & fatures
92
- || (allowOlapApply && (NUdf::EDataTypeFeatures::BigDateType & fatures) && !(NUdf::EDataTypeFeatures::TzDateType & fatures));
99
+ const auto features = NUdf::GetDataTypeInfo (RemoveOptionality (type).Cast <TDataExprType>()->GetSlot ()).Features ;
100
+ return NUdf::EDataTypeFeatures::NumericType & features
101
+ || (allowOlapApply && ((NUdf::EDataTypeFeatures::BigDateType |
102
+ NUdf::EDataTypeFeatures::DateType |
103
+ NUdf::EDataTypeFeatures::TimeIntervalType) & features) && !(NUdf::EDataTypeFeatures::TzDateType & features));
93
104
}
94
105
95
106
bool IsGoodTypeForComparsionPushdown (const TTypeAnnotationNode& type, bool allowOlapApply) {
96
- const auto fatures = NUdf::GetDataTypeInfo (RemoveOptionality (type).Cast <TDataExprType>()->GetSlot ()).Features ;
97
- return (NUdf::EDataTypeFeatures::CanCompare & fatures)
98
- && (((NUdf::EDataTypeFeatures::NumericType | NUdf::EDataTypeFeatures::StringType) & fatures) ||
99
- (allowOlapApply && (NUdf::EDataTypeFeatures::BigDateType & fatures) && !(NUdf::EDataTypeFeatures::TzDateType & fatures)));
107
+ const auto features = NUdf::GetDataTypeInfo (RemoveOptionality (type).Cast <TDataExprType>()->GetSlot ()).Features ;
108
+ return (NUdf::EDataTypeFeatures::CanCompare & features)
109
+ && (((NUdf::EDataTypeFeatures::NumericType | NUdf::EDataTypeFeatures::StringType) & features) ||
110
+ (allowOlapApply && ((NUdf::EDataTypeFeatures::BigDateType |
111
+ NUdf::EDataTypeFeatures::DateType |
112
+ NUdf::EDataTypeFeatures::TimeIntervalType) & features) && !(NUdf::EDataTypeFeatures::TzDateType & features)));
100
113
}
101
114
102
115
[[maybe_unused]]
103
116
bool AbstractTreeCanBePushed (const TExprBase& expr, const TExprNode* ) {
104
- if (!expr.Ref ().IsCallable ({" Apply" , " NamedApply" , " IfPresent" , " Visit" })) {
117
+ if (!expr.Ref ().IsCallable ({" Apply" , " Coalesce " , " NamedApply" , " IfPresent" , " Visit" })) {
105
118
return false ;
106
119
}
107
120
@@ -139,6 +152,14 @@ bool AbstractTreeCanBePushed(const TExprBase& expr, const TExprNode* ) {
139
152
return true ;
140
153
}
141
154
155
+ bool IfPresentCanBePushed (const TCoIfPresent& ifPresent, const TExprNode* lambdaArg, bool allowOlapApply) {
156
+
157
+ Y_UNUSED (ifPresent);
158
+ Y_UNUSED (lambdaArg);
159
+
160
+ return allowOlapApply;
161
+ }
162
+
142
163
bool CheckExpressionNodeForPushdown (const TExprBase& node, const TExprNode* lambdaArg, bool allowOlapApply) {
143
164
if (allowOlapApply) {
144
165
if (node.Maybe <TCoJust>() || node.Maybe <TCoCoalesce>()) {
@@ -154,7 +175,7 @@ bool CheckExpressionNodeForPushdown(const TExprBase& node, const TExprNode* lamb
154
175
}
155
176
156
177
if (const auto maybeSafeCast = node.Maybe <TCoSafeCast>()) {
157
- return IsSupportedCast (maybeSafeCast.Cast ());
178
+ return IsSupportedCast (maybeSafeCast.Cast (), allowOlapApply );
158
179
} else if (const auto maybeData = node.Maybe <TCoDataCtor>()) {
159
180
return IsSupportedDataType (maybeData.Cast (), allowOlapApply);
160
181
} else if (const auto maybeMember = node.Maybe <TCoMember>()) {
@@ -174,6 +195,9 @@ bool CheckExpressionNodeForPushdown(const TExprBase& node, const TExprNode* lamb
174
195
}
175
196
176
197
if (allowOlapApply) {
198
+ if (const auto maybeIfPresent = node.Maybe <TCoIfPresent>()) {
199
+ return IfPresentCanBePushed (maybeIfPresent.Cast (), lambdaArg, allowOlapApply);
200
+ }
177
201
return AbstractTreeCanBePushed (node, lambdaArg);
178
202
}
179
203
@@ -208,8 +232,9 @@ bool IsGoodTypesForPushdownCompare(const TTypeAnnotationNode& typeOne, const TTy
208
232
}
209
233
return true ;
210
234
}
211
- case ETypeAnnotationKind::Data:
235
+ case ETypeAnnotationKind::Data: {
212
236
return IsGoodTypeForComparsionPushdown (typeOne, allowOlapApply) && IsGoodTypeForComparsionPushdown (typeTwo, allowOlapApply);
237
+ }
213
238
default :
214
239
break ;
215
240
}
@@ -327,6 +352,8 @@ bool CoalesceCanBePushed(const TCoCoalesce& coalesce, const TExprNode* lambdaArg
327
352
return SafeCastCanBePushed (maybeFlatmap.Cast (), lambdaArg, allowOlapApply);
328
353
} else if (const auto maybeJsonExists = predicate.Maybe <TCoJsonExists>()) {
329
354
return JsonExistsCanBePushed (maybeJsonExists.Cast (), lambdaArg);
355
+ } else if (const auto maybeIfPresent = predicate.Maybe <TCoIfPresent>()) {
356
+ return IfPresentCanBePushed (maybeIfPresent.Cast (), lambdaArg, allowOlapApply);
330
357
}
331
358
332
359
return false ;
@@ -339,39 +366,49 @@ bool ExistsCanBePushed(const TCoExists& exists, const TExprNode* lambdaArg) {
339
366
void CollectChildrenPredicates (const TExprNode& opNode, TOLAPPredicateNode& predicateTree, const TExprNode* lambdaArg, const TExprBase& lambdaBody, bool allowOlapApply) {
340
367
predicateTree.Children .reserve (opNode.ChildrenSize ());
341
368
predicateTree.CanBePushed = true ;
369
+ predicateTree.CanBePushedApply = true ;
370
+
342
371
for (const auto & childNodePtr: opNode.Children ()) {
343
372
TOLAPPredicateNode child;
344
373
child.ExprNode = childNodePtr;
345
- if (const auto maybeCtor = TMaybeNode<TCoDataCtor>(child.ExprNode ))
346
- child.CanBePushed = IsSupportedDataType (maybeCtor.Cast (), allowOlapApply);
374
+ if (const auto maybeCtor = TMaybeNode<TCoDataCtor>(child.ExprNode )) {
375
+ child.CanBePushed = IsSupportedDataType (maybeCtor.Cast (), false );
376
+ child.CanBePushedApply = IsSupportedDataType (maybeCtor.Cast (), true );
377
+ }
347
378
else
348
379
CollectPredicates (TExprBase (child.ExprNode ), child, lambdaArg, lambdaBody, allowOlapApply);
349
380
predicateTree.Children .emplace_back (child);
350
381
predicateTree.CanBePushed &= child.CanBePushed ;
382
+ predicateTree.CanBePushedApply &= child.CanBePushedApply ;
351
383
}
352
384
}
353
385
354
386
} // namespace
355
387
356
388
void CollectPredicates (const TExprBase& predicate, TOLAPPredicateNode& predicateTree, const TExprNode* lambdaArg, const TExprBase& lambdaBody, bool allowOlapApply) {
357
389
if (predicate.Maybe <TCoNot>() || predicate.Maybe <TCoAnd>() || predicate.Maybe <TCoOr>() || predicate.Maybe <TCoXor>()) {
358
- return CollectChildrenPredicates (predicate.Ref (), predicateTree, lambdaArg, lambdaBody, allowOlapApply);
390
+ CollectChildrenPredicates (predicate.Ref (), predicateTree, lambdaArg, lambdaBody, allowOlapApply);
359
391
} else if (const auto maybeCoalesce = predicate.Maybe <TCoCoalesce>()) {
360
- predicateTree.CanBePushed = CoalesceCanBePushed (maybeCoalesce.Cast (), lambdaArg, lambdaBody, allowOlapApply);
392
+ predicateTree.CanBePushed = CoalesceCanBePushed (maybeCoalesce.Cast (), lambdaArg, lambdaBody, false );
393
+ predicateTree.CanBePushedApply = CoalesceCanBePushed (maybeCoalesce.Cast (), lambdaArg, lambdaBody, true );
361
394
} else if (const auto maybeCompare = predicate.Maybe <TCoCompare>()) {
362
- predicateTree.CanBePushed = CompareCanBePushed (maybeCompare.Cast (), lambdaArg, lambdaBody, allowOlapApply);
395
+ predicateTree.CanBePushed = CompareCanBePushed (maybeCompare.Cast (), lambdaArg, lambdaBody, false );
396
+ predicateTree.CanBePushedApply = CompareCanBePushed (maybeCompare.Cast (), lambdaArg, lambdaBody, true );
363
397
} else if (const auto maybeExists = predicate.Maybe <TCoExists>()) {
364
398
predicateTree.CanBePushed = ExistsCanBePushed (maybeExists.Cast (), lambdaArg);
399
+ predicateTree.CanBePushedApply = predicateTree.CanBePushed ;
365
400
} else if (const auto maybeJsonExists = predicate.Maybe <TCoJsonExists>()) {
366
401
predicateTree.CanBePushed = JsonExistsCanBePushed (maybeJsonExists.Cast (), lambdaArg);
402
+ predicateTree.CanBePushedApply = predicateTree.CanBePushed ;
367
403
}
368
- if (predicateTree.CanBePushed ) {
369
- return ;
370
- } else if (allowOlapApply){
404
+
405
+ if (allowOlapApply && !predicateTree.CanBePushedApply ){
371
406
if (predicate.Maybe <TCoIf>() || predicate.Maybe <TCoJust>() || predicate.Maybe <TCoCoalesce>()) {
372
- return CollectChildrenPredicates (predicate.Ref (), predicateTree, lambdaArg, lambdaBody, allowOlapApply);
407
+ CollectChildrenPredicates (predicate.Ref (), predicateTree, lambdaArg, lambdaBody, true );
408
+ }
409
+ if (!predicateTree.CanBePushedApply ) {
410
+ predicateTree.CanBePushedApply = AbstractTreeCanBePushed (predicate, lambdaArg);
373
411
}
374
- predicateTree.CanBePushed = AbstractTreeCanBePushed (predicate, lambdaArg);
375
412
}
376
413
}
377
414
} // namespace NKikimr::NKqp::NOpt
0 commit comments