Skip to content

Commit 8fb8631

Browse files
committed
regularize coalesce pushdown
1 parent f43e9dd commit 8fb8631

File tree

1 file changed

+73
-93
lines changed

1 file changed

+73
-93
lines changed

ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp

Lines changed: 73 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ std::vector<std::pair<TExprBase, TExprBase>> ExtractComparisonParameters(const T
139139
TMaybeNode<TExprBase> ComparisonPushdown(const std::vector<std::pair<TExprBase, TExprBase>>& parameters, const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos);
140140

141141
[[maybe_unused]]
142-
TMaybeNode<TExprBase> YqlCoalescePushdown(const TCoCoalesce& coalesce, const TExprNode& argument, TExprContext& ctx) {
142+
TMaybeNode<TExprBase> CoalescePushdown(const TCoCoalesce& coalesce, const TExprNode& argument, TExprContext& ctx) {
143143
if (const auto params = ExtractBinaryFunctionParameters(coalesce, argument, ctx, coalesce.Pos())) {
144144
return Build<TKqpOlapFilterBinaryOp>(ctx, coalesce.Pos())
145145
.Operator().Value("??", TNodeFlags::Default).Build()
@@ -213,9 +213,55 @@ TMaybeNode<TExprBase> JsonExistsPushdown(const TCoJsonExists& jsonExists, TExprC
213213
.Done();
214214
}
215215

216+
TMaybeNode<TExprBase> SafeCastPredicatePushdown(const TCoFlatMap& inputFlatmap, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
217+
{
218+
/*
219+
* There are three ways of comparison in following format:
220+
*
221+
* FlatMap (LeftArgument, FlatMap(RightArgument(), Just(Predicate))
222+
*
223+
* Examples:
224+
* FlatMap (SafeCast(), FlatMap(Member(), Just(Comparison))
225+
* FlatMap (Member(), FlatMap(SafeCast(), Just(Comparison))
226+
* FlatMap (SafeCast(), FlatMap(SafeCast(), Just(Comparison))
227+
*/
228+
auto left = ConvertComparisonNode(inputFlatmap.Input(), argument, ctx, pos);
229+
if (left.empty()) {
230+
return NullNode;
231+
}
232+
233+
auto flatmap = inputFlatmap.Lambda().Body().Cast<TCoFlatMap>();
234+
auto right = ConvertComparisonNode(flatmap.Input(), argument, ctx, pos);
235+
if (right.empty()) {
236+
return NullNode;
237+
}
238+
239+
auto predicate = flatmap.Lambda().Body().Cast<TCoJust>().Input().Cast<TCoCompare>();
240+
241+
std::vector<std::pair<TExprBase, TExprBase>> parameters;
242+
if (left.size() != right.size()) {
243+
return NullNode;
244+
}
245+
246+
for (ui32 i = 0; i < left.size(); ++i) {
247+
parameters.emplace_back(std::move(std::make_pair(left[i], right[i])));
248+
}
249+
250+
return ComparisonPushdown(parameters, predicate, ctx, pos);
251+
}
252+
253+
TMaybeNode<TExprBase> SimplePredicatePushdown(const TCoCompare& predicate, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
254+
{
255+
const auto parameters = ExtractComparisonParameters(predicate, argument, ctx, pos);
256+
if (parameters.empty()) {
257+
return NullNode;
258+
}
259+
260+
return ComparisonPushdown(parameters, predicate, ctx, pos);
261+
}
262+
216263
std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
217264
{
218-
std::vector<TExprBase> out;
219265
const auto convertNode = [&ctx, &pos, &argument](const TExprBase& node) -> TMaybeNode<TExprBase> {
220266
if (node.Maybe<TCoNull>()) {
221267
return node;
@@ -299,7 +345,7 @@ std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExp
299345
}
300346

301347
if (const auto maybeCoalesce = node.Maybe<TCoCoalesce>()) {
302-
return YqlCoalescePushdown(maybeCoalesce.Cast(), argument, ctx);
348+
return CoalescePushdown(maybeCoalesce.Cast(), argument, ctx);
303349
}
304350

305351
if (const auto maybeCompare = node.Maybe<TCoCompare>()) {
@@ -308,6 +354,11 @@ std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExp
308354
}
309355
}
310356

357+
if (const auto maybeFlatmap = node.Maybe<TCoFlatMap>()) {
358+
return SafeCastPredicatePushdown(maybeFlatmap.Cast(), argument, ctx, pos);
359+
} else if (auto maybePredicate = node.Maybe<TCoCompare>()) {
360+
return SimplePredicatePushdown(maybePredicate.Cast(), argument, ctx, pos);
361+
}
311362

312363
if constexpr (NKikimr::NSsa::RuntimeVersion >= 5U) {
313364
return YqlApplyPushdown(node, argument, ctx);
@@ -316,36 +367,27 @@ std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExp
316367
}
317368
};
318369

319-
// Columns & values may be single element
320-
TMaybeNode<TExprBase> node = convertNode(nodeIn);
321-
322-
if (node.IsValid()) {
323-
out.emplace_back(std::move(node.Cast()));
324-
return out;
325-
}
326-
327-
// Or columns and values can be Tuple
328-
if (!nodeIn.Maybe<TExprList>()) {
329-
// something unusual found, return empty vector
330-
return out;
331-
}
370+
if (const auto& list = nodeIn.Maybe<TExprList>()) {
371+
const auto& tuple = list.Cast();
372+
std::vector<TExprBase> out;
332373

333-
auto tuple = nodeIn.Cast<TExprList>();
374+
out.reserve(tuple.Size());
375+
for (ui32 i = 0; i < tuple.Size(); ++i) {
376+
TMaybeNode<TExprBase> node = convertNode(tuple.Item(i));
334377

335-
out.reserve(tuple.Size());
336-
337-
for (ui32 i = 0; i < tuple.Size(); ++i) {
338-
TMaybeNode<TExprBase> node = convertNode(tuple.Item(i));
378+
if (!node.IsValid()) {
379+
// Return empty vector
380+
return TVector<TExprBase>();
381+
}
339382

340-
if (!node.IsValid()) {
341-
// Return empty vector
342-
return TVector<TExprBase>();
383+
out.emplace_back(node.Cast());
343384
}
344-
345-
out.emplace_back(node.Cast());
385+
return out;
386+
} else if (const auto& node = convertNode(nodeIn); node.IsValid()) {
387+
return {node.Cast()};
388+
} else {
389+
return {};
346390
}
347-
348-
return out;
349391
}
350392

351393
TExprBase BuildOneElementComparison(const std::pair<TExprBase, TExprBase>& parameter, const TCoCompare& predicate,
@@ -480,16 +522,6 @@ TMaybeNode<TExprBase> ComparisonPushdown(const std::vector<std::pair<TExprBase,
480522
.Done();
481523
}
482524

483-
TMaybeNode<TExprBase> SimplePredicatePushdown(const TCoCompare& predicate, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
484-
{
485-
const auto parameters = ExtractComparisonParameters(predicate, argument, ctx, pos);
486-
if (parameters.empty()) {
487-
return NullNode;
488-
}
489-
490-
return ComparisonPushdown(parameters, predicate, ctx, pos);
491-
}
492-
493525
// TODO: Check how to reduce columns if they are not needed. Unfortunately columnshard need columns list
494526
// for every column present in program even if it is not used in result set.
495527
//#define ENABLE_COLUMNS_PRUNING
@@ -528,62 +560,10 @@ TMaybeNode<TExprBase> ExistsPushdown(const TCoExists& exists, TExprContext& ctx,
528560
.Done();
529561
}
530562

531-
TMaybeNode<TExprBase> SafeCastPredicatePushdown(const TCoFlatMap& inputFlatmap, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
532-
{
533-
/*
534-
* There are three ways of comparison in following format:
535-
*
536-
* FlatMap (LeftArgument, FlatMap(RightArgument(), Just(Predicate))
537-
*
538-
* Examples:
539-
* FlatMap (SafeCast(), FlatMap(Member(), Just(Comparison))
540-
* FlatMap (Member(), FlatMap(SafeCast(), Just(Comparison))
541-
* FlatMap (SafeCast(), FlatMap(SafeCast(), Just(Comparison))
542-
*/
543-
auto left = ConvertComparisonNode(inputFlatmap.Input(), argument, ctx, pos);
544-
if (left.empty()) {
545-
return NullNode;
546-
}
547-
548-
auto flatmap = inputFlatmap.Lambda().Body().Cast<TCoFlatMap>();
549-
auto right = ConvertComparisonNode(flatmap.Input(), argument, ctx, pos);
550-
if (right.empty()) {
551-
return NullNode;
552-
}
553-
554-
auto predicate = flatmap.Lambda().Body().Cast<TCoJust>().Input().Cast<TCoCompare>();
555-
556-
std::vector<std::pair<TExprBase, TExprBase>> parameters;
557-
if (left.size() != right.size()) {
558-
return NullNode;
559-
}
560-
561-
for (ui32 i = 0; i < left.size(); ++i) {
562-
parameters.emplace_back(std::move(std::make_pair(left[i], right[i])));
563-
}
564-
565-
return ComparisonPushdown(parameters, predicate, ctx, pos);
566-
}
567-
568-
TMaybeNode<TExprBase> CoalescePushdown(const TCoCoalesce& coalesce, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
569-
{
570-
if (const auto node = YqlCoalescePushdown(coalesce, argument, ctx)) {
571-
return node;
572-
}
573-
574-
auto predicate = coalesce.Predicate();
575-
if (const auto maybeFlatmap = predicate.Maybe<TCoFlatMap>()) {
576-
return SafeCastPredicatePushdown(maybeFlatmap.Cast(), argument, ctx, pos);
577-
} else if (auto maybePredicate = predicate.Maybe<TCoCompare>()) {
578-
return SimplePredicatePushdown(maybePredicate.Cast(), argument, ctx, pos);
579-
}
580-
return NullNode;
581-
}
582-
583563
TFilterOpsLevels PredicatePushdown(const TExprBase& predicate, const TExprNode& argument, TExprContext& ctx, TPositionHandle pos)
584564
{
585565
if (const auto maybeCoalesce = predicate.Maybe<TCoCoalesce>()) {
586-
auto coalescePred = CoalescePushdown(maybeCoalesce.Cast(), argument, ctx, pos);
566+
auto coalescePred = CoalescePushdown(maybeCoalesce.Cast(), argument, ctx);
587567
return TFilterOpsLevels(coalescePred);
588568
}
589569

@@ -753,7 +733,7 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
753733

754734
const auto& lambda = flatmap.Lambda();
755735

756-
YQL_CLOG(TRACE, ProviderKqp) << "Initial OLAP lambda: " << KqpExprToPrettyString(lambda, ctx);
736+
YQL_CLOG(ERROR, ProviderKqp) << "Initial OLAP lambda: " << KqpExprToPrettyString(lambda, ctx);
757737

758738
const auto maybeOptionalIf = lambda.Body().Maybe<TCoOptionalIf>();
759739
if (!maybeOptionalIf.IsValid()) {
@@ -824,7 +804,7 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
824804
.Build()
825805
.Done();
826806

827-
YQL_CLOG(TRACE, ProviderKqp) << "Pushed OLAP lambda: " << KqpExprToPrettyString(newProcessLambda, ctx);
807+
YQL_CLOG(ERROR, ProviderKqp) << "Pushed OLAP lambda: " << KqpExprToPrettyString(newProcessLambda, ctx);
828808

829809
#ifdef ENABLE_COLUMNS_PRUNING
830810
TMaybeNode<TCoAtomList> readColumns = BuildColumnsFromLambda(lambda, ctx, node.Pos());

0 commit comments

Comments
 (0)