Skip to content

Commit b38bd0d

Browse files
pavelvelikhovPavel Velikhov
pavelvelikhov
and
Pavel Velikhov
authored
Expanded constant folding with UDFs (#17533)
Co-authored-by: Pavel Velikhov <pavelvelikhov@localhost.localdomain>
1 parent 6f7b911 commit b38bd0d

File tree

11 files changed

+73
-12
lines changed

11 files changed

+73
-12
lines changed

‎ydb/core/kqp/compile_service/kqp_compile_actor.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -649,6 +649,7 @@ void ApplyServiceConfig(TKikimrConfiguration& kqpConfig, const TTableServiceConf
649649
kqpConfig.DefaultCostBasedOptimizationLevel = serviceConfig.GetDefaultCostBasedOptimizationLevel();
650650
kqpConfig.DefaultEnableShuffleElimination = serviceConfig.GetDefaultEnableShuffleElimination();
651651
kqpConfig.EnableConstantFolding = serviceConfig.GetEnableConstantFolding();
652+
kqpConfig.EnableFoldUdfs = serviceConfig.GetEnableFoldUdfs();
652653
kqpConfig.SetDefaultEnabledSpillingNodes(serviceConfig.GetEnableSpillingNodes());
653654
kqpConfig.EnableSpilling = serviceConfig.GetEnableQueryServiceSpilling();
654655
kqpConfig.EnableSnapshotIsolationRW = serviceConfig.GetEnableSnapshotIsolationRW();

‎ydb/core/kqp/compile_service/kqp_compile_service.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
309309

310310
ui64 defaultCostBasedOptimizationLevel = TableServiceConfig.GetDefaultCostBasedOptimizationLevel();
311311
bool enableConstantFolding = TableServiceConfig.GetEnableConstantFolding();
312+
bool enableFoldUdfs = TableServiceConfig.GetEnableFoldUdfs();
312313

313314
bool defaultEnableShuffleElimination = TableServiceConfig.GetDefaultEnableShuffleElimination();
314315

@@ -341,6 +342,7 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
341342
TableServiceConfig.GetEnableSpillingNodes() != enableSpillingNodes ||
342343
TableServiceConfig.GetDefaultCostBasedOptimizationLevel() != defaultCostBasedOptimizationLevel ||
343344
TableServiceConfig.GetEnableConstantFolding() != enableConstantFolding ||
345+
TableServiceConfig.GetEnableFoldUdfs() != enableFoldUdfs ||
344346
TableServiceConfig.GetEnableAstCache() != enableAstCache ||
345347
TableServiceConfig.GetEnableImplicitQueryParameterTypes() != enableImplicitQueryParameterTypes ||
346348
TableServiceConfig.GetEnablePgConstsToParams() != enablePgConstsToParams ||

‎ydb/core/kqp/opt/kqp_constant_folding_transformer.cpp

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ namespace {
1414
* Traverse a lambda and create a mapping from nodes to nodes wrapped in EvaluateExpr callable
1515
* We check for literals specifically, since they shouldn't be evaluated
1616
*/
17-
void ExtractConstantExprs(const TExprNode::TPtr& input, TNodeOnNodeOwnedMap& replaces, TExprContext& ctx) {
17+
void ExtractConstantExprs(const TExprNode::TPtr& input, TNodeOnNodeOwnedMap& replaces, TExprContext& ctx, bool foldUdfs = true) {
1818
if (TCoLambda::Match(input.Get())) {
1919
auto lambda = TExprBase(input).Cast<TCoLambda>();
2020
return ExtractConstantExprs(lambda.Body().Ptr(), replaces, ctx);
@@ -24,7 +24,7 @@ namespace {
2424
return;
2525
}
2626

27-
if (IsConstantExpr(input) && !input->IsCallable("PgConst")) {
27+
if (IsConstantExpr(input, foldUdfs) && !input->IsCallable("PgConst")) {
2828
TNodeOnNodeOwnedMap deepClones;
2929
auto inputClone = ctx.DeepCopy(*input, ctx, deepClones, false, true, true);
3030

@@ -64,6 +64,8 @@ IGraphTransformer::TStatus TKqpConstantFoldingTransformer::DoTransform(TExprNode
6464
return IGraphTransformer::TStatus::Ok;
6565
}
6666

67+
bool foldUdfs = Config->EnableFoldUdfs;
68+
6769
TNodeOnNodeOwnedMap replaces;
6870

6971
VisitExpr(input, [&](const TExprNode::TPtr& node) {
@@ -78,7 +80,7 @@ IGraphTransformer::TStatus TKqpConstantFoldingTransformer::DoTransform(TExprNode
7880
return true;
7981
}
8082

81-
ExtractConstantExprs(flatmap.Lambda().Body().Ptr(), replaces, ctx);
83+
ExtractConstantExprs(flatmap.Lambda().Body().Ptr(), replaces, ctx, foldUdfs);
8284

8385
return replaces.empty();
8486
}

‎ydb/core/kqp/provider/yql_kikimr_settings.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,7 @@ struct TKikimrConfiguration : public TKikimrSettings, public NCommon::TSettingDi
177177
bool EnableSpilling = true;
178178
ui32 DefaultCostBasedOptimizationLevel = 4;
179179
bool EnableConstantFolding = true;
180+
bool EnableFoldUdfs = true;
180181
ui64 DefaultEnableSpillingNodes = 0;
181182
bool EnableAntlr4Parser = false;
182183
bool EnableSnapshotIsolationRW = false;

‎ydb/core/kqp/ut/join/data/queries/udf_constant_fold.sql

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
SELECT *
2+
FROM `/Root/S` as S
3+
WHERE S.payload2 = String::HexDecode("54");

‎ydb/core/kqp/ut/join/kqp_join_order_ut.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,10 @@ Y_UNIT_TEST_SUITE(KqpJoinOrder) {
683683
ExecuteJoinOrderTestGenericQueryWithStats("queries/datetime_constant_fold.sql", "stats/basic.json", false, ColumnStore);
684684
}
685685

686+
Y_UNIT_TEST_TWIN(UdfConstantFold, ColumnStore) {
687+
ExecuteJoinOrderTestGenericQueryWithStats("queries/udf_constant_fold.sql", "stats/basic.json", false, ColumnStore);
688+
}
689+
686690
Y_UNIT_TEST_TWIN(TPCHRandomJoinViewJustWorks, ColumnStore) {
687691
ExecuteJoinOrderTestGenericQueryWithStats("queries/tpch_random_join_view_just_works.sql", "stats/tpch1000s.json", false, ColumnStore);
688692
}

‎ydb/core/protos/table_service_config.proto

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,4 +381,7 @@ message TTableServiceConfig {
381381
}
382382

383383
optional TBatchOperationSettings BatchOperationSettings = 81;
384+
385+
optional bool EnableFoldUdfs = 82 [ default = true ];
386+
384387
};

‎ydb/library/yql/dq/opt/dq_opt_stat.cpp

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,50 @@ namespace {
1515
* We maintain a white list of callables that we consider part of constant expressions
1616
* All other callables will not be evaluated
1717
*/
18-
THashSet<TString> constantFoldingWhiteList = {
18+
THashSet<TString> ConstantFoldingWhiteList = {
1919
"Concat", "Just", "Optional", "SafeCast", "AsList",
2020
"+", "-", "*", "/", "%"};
2121

22-
THashSet<TString> pgConstantFoldingWhiteList = {
22+
THashSet<TString> PgConstantFoldingWhiteList = {
2323
"PgResolvedOp", "PgResolvedCall", "PgCast", "PgConst", "PgArray", "PgType"};
2424

25+
TVector<TString> UdfBlackList = {
26+
"RandomNumber",
27+
"Random",
28+
"RandomUuid",
29+
"Now",
30+
"CurrentUtcDate",
31+
"CurrentUtcDatetime",
32+
"CurrentUtcTimestamp"
33+
};
34+
35+
bool IsConstantUdf(const TExprNode::TPtr& input, bool withParams = false) {
36+
if (!TCoApply::Match(input.Get())) {
37+
return false;
38+
}
39+
40+
if (input->ChildrenSize()!=2) {
41+
return false;
42+
}
43+
if (input->Child(0)->IsCallable("Udf")) {
44+
auto udf = TCoUdf(input->Child(0));
45+
auto udfName = udf.MethodName().StringValue();
46+
47+
for (auto blck : UdfBlackList) {
48+
if (udfName.find(blck) != TString::npos) {
49+
return false;
50+
}
51+
}
52+
53+
if (withParams) {
54+
return IsConstantExprWithParams(input->Child(1));
55+
}
56+
else {
57+
return IsConstantExpr(input->Child(1));
58+
}
59+
}
60+
return false;
61+
}
2562

2663
TString RemoveAliases(TString attributeName) {
2764
if (auto idx = attributeName.find('.'); idx != TString::npos) {
@@ -167,7 +204,7 @@ bool IsConstantExprPg(const TExprNode::TPtr& input) {
167204
return true;
168205
}
169206

170-
if (input->IsCallable(pgConstantFoldingWhiteList) || input->IsList()) {
207+
if (input->IsCallable(PgConstantFoldingWhiteList) || input->IsList()) {
171208
for (size_t i = 0; i < input->ChildrenSize(); i++) {
172209
auto callableInput = input->Child(i);
173210
if (callableInput->IsLambda() && !IsConstantExprPg(callableInput->Child(1))) {
@@ -190,7 +227,7 @@ bool IsConstantExprPg(const TExprNode::TPtr& input) {
190227
* - If its a callable in the while list and all children are constant expressions, then its a constant expression
191228
* - If one of the child is a type expression, it also passes the check
192229
*/
193-
bool IsConstantExpr(const TExprNode::TPtr& input) {
230+
bool IsConstantExpr(const TExprNode::TPtr& input, bool foldUdfs) {
194231
if (input->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Pg) {
195232
return IsConstantExprPg(input);
196233
}
@@ -203,7 +240,7 @@ bool IsConstantExpr(const TExprNode::TPtr& input) {
203240
return true;
204241
}
205242

206-
else if (input->IsCallable(constantFoldingWhiteList)) {
243+
else if (input->IsCallable(ConstantFoldingWhiteList)) {
207244
for (size_t i = 0; i < input->ChildrenSize(); i++) {
208245
auto callableInput = input->Child(i);
209246
if (callableInput->GetTypeAnn()->GetKind() != ETypeAnnotationKind::Type && !IsConstantExpr(callableInput)) {
@@ -213,6 +250,10 @@ bool IsConstantExpr(const TExprNode::TPtr& input) {
213250
return true;
214251
}
215252

253+
else if (foldUdfs && TCoApply::Match(input.Get()) && IsConstantUdf(input)) {
254+
return true;
255+
}
256+
216257
return false;
217258
}
218259

@@ -233,7 +274,7 @@ bool IsConstantExprWithParams(const TExprNode::TPtr& input) {
233274
return true;
234275
}
235276

236-
else if (input->IsCallable(constantFoldingWhiteList)) {
277+
else if (input->IsCallable(ConstantFoldingWhiteList)) {
237278
for (size_t i = 0; i < input->ChildrenSize(); i++) {
238279
auto callableInput = input->Child(i);
239280
if (callableInput->GetTypeAnn()->GetKind() != ETypeAnnotationKind::Type && !IsConstantExprWithParams(callableInput)) {
@@ -243,6 +284,10 @@ bool IsConstantExprWithParams(const TExprNode::TPtr& input) {
243284
return true;
244285
}
245286

287+
else if (TCoApply::Match(input.Get()) && IsConstantUdf(input, true)) {
288+
return true;
289+
}
290+
246291
return false;
247292
}
248293

‎ydb/library/yql/dq/opt/dq_opt_stat.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ class TPredicateSelectivityComputer {
9999
};
100100

101101
bool NeedCalc(NNodes::TExprBase node);
102-
bool IsConstantExpr(const TExprNode::TPtr& input);
102+
bool IsConstantExpr(const TExprNode::TPtr& input, bool foldUdfs = true);
103103
bool IsConstantExprWithParams(const TExprNode::TPtr& input);
104104

105105
} // namespace NYql::NDq {

‎ydb/tests/functional/canonical/canondata/test_sql.TestCanonicalFolder1.test_case_index_topsort_index_with_selector_aliases_2.sql-plan_/index_topsort_index_with_selector_aliases_2.sql.plan

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
"limit": "1000",
1717
"lookup_by": [
1818
"hashOrderDate (9983033094796217818)",
19-
"orderDate (\u00ab2023-03-01\u00bb)"
19+
"orderDate (2023-03-01)"
2020
],
2121
"scan_by": [
2222
"orderId (-\u221e, +\u221e)"

0 commit comments

Comments
 (0)