Split pushdown for different callable kinds

Author: zverevgeny

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

@@ -685,45 +685,40 @@ TFilterOpsLevels PredicatePushdown(const TExprBase& predicate, const TExprNode&
 }
 
 TOLAPPredicateNode WrapPredicates(const std::vector<TOLAPPredicateNode>& predicates, TExprContext& ctx, TPositionHandle pos) {
+    TOLAPPredicateNode result;
     if (predicates.empty()) {
-        return {};
-    }
-
-    if (const auto predicatesSize = predicates.size(); 1U == predicatesSize) {
-        return predicates.front();
+        result.ExprNode = MakeBool<true>(pos, ctx);;
+    } else if (predicates.size() == 1) {
+        result = predicates.front();
     } else {
         TOLAPPredicateNode result;
         result.Children = predicates;
         result.CanBePushed = true;
 
         TVector<NNodes::TExprBase> exprNodes;
-        exprNodes.reserve(predicatesSize);
+        exprNodes.reserve(predicates.size());
         for (const auto& pred : predicates) {
             exprNodes.emplace_back(pred.ExprNode);
             result.CanBePushed &= pred.CanBePushed;
         }
         result.ExprNode = NNodes::Build<NNodes::TCoAnd>(ctx, pos)
             .Add(exprNodes)
             .Done().Ptr();
-        return result;
     }
+    return result;
 }
 
-void SplitForPartialPushdown(const TOLAPPredicateNode& predicateTree, TOLAPPredicateNode& predicatesToPush, TOLAPPredicateNode& remainingPredicates,
-    TExprContext& ctx, TPositionHandle pos)
+std::pair<std::vector<TOLAPPredicateNode>, std::vector<TOLAPPredicateNode>> SplitForPartialPushdown(const TOLAPPredicateNode& predicateTree)
 {
     if (predicateTree.CanBePushed) {
-        predicatesToPush = predicateTree;
-        remainingPredicates.ExprNode = MakeBool<true>(pos, ctx);
-        return;
+        return {{predicateTree}, {}};
     }
 
     if (!TCoAnd::Match(predicateTree.ExprNode.Get())) {
         // We can partially pushdown predicates from AND operator only.
         // For OR operator we would need to have several read operators which is not acceptable.
         // TODO: Add support for NOT(op1 OR op2), because it expands to (!op1 AND !op2).
-        remainingPredicates = predicateTree;
-        return;
+        return {{}, {predicateTree}};
     }
 
     bool isFoundNotStrictOp = false;
@@ -738,8 +733,7 @@ void SplitForPartialPushdown(const TOLAPPredicateNode& predicateTree, TOLAPPredi
             remaining.emplace_back(predicate);
         }
     }
-    predicatesToPush = WrapPredicates(pushable, ctx, pos);
-    remainingPredicates = WrapPredicates(remaining, ctx, pos);
+    return {pushable, remaining};
 }
 
 } // anonymous namespace end
@@ -763,6 +757,7 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
     }
 
     const auto& lambda = flatmap.Lambda();
+    const auto& lambdaArg = lambda.Args().Arg(0).Ref();
 
     YQL_CLOG(TRACE, ProviderKqp) << "Initial OLAP lambda: " << KqpExprToPrettyString(lambda, ctx);
 
@@ -775,35 +770,53 @@ TExprBase KqpPushOlapFilter(TExprBase node, TExprContext& ctx, const TKqpOptimiz
     auto predicate = optionaIf.Predicate();
     auto value = optionaIf.Value();
 
-    if constexpr (NSsa::RuntimeVersion >= 5U) {
-        TExprNode::TPtr afterPeephole;
-        bool hasNonDeterministicFunctions;
-        if (const auto status = PeepHoleOptimizeNode(optionaIf.Ptr(), afterPeephole, ctx, typesCtx, nullptr, hasNonDeterministicFunctions);
-            status != IGraphTransformer::TStatus::Ok) {
-            YQL_CLOG(ERROR, ProviderKqp) << "Peephole OLAP failed." << Endl << ctx.IssueManager.GetIssues().ToString();
-            return node;
-        }
-
-        const TCoIf simplified(std::move(afterPeephole));
-        predicate = simplified.Predicate();
-        value = simplified.ThenValue().Cast<TCoJust>().Input();
-    }
-
     TOLAPPredicateNode predicateTree;
     predicateTree.ExprNode = predicate.Ptr();
-    const auto& lambdaArg = lambda.Args().Arg(0).Ref();
     CollectPredicates(predicate, predicateTree, &lambdaArg, read.Process().Body());
     YQL_ENSURE(predicateTree.IsValid(), "Collected OLAP predicates are invalid");
 
-    TOLAPPredicateNode predicatesToPush, remainingPredicates;
-    SplitForPartialPushdown(predicateTree, predicatesToPush, remainingPredicates, ctx, node.Pos());
-    if (!predicatesToPush.IsValid()) {
+    auto [pushable, remaining] = SplitForPartialPushdown(predicateTree);
+
+    if constexpr (NSsa::RuntimeVersion >= 5U) {
+        if (!remaining.empty()) {
+            const auto remainingPredicates = WrapPredicates(remaining, ctx, node.Pos());
+            const auto recoveredOptinalIfForNonPushedDownPredicates = Build<TCoOptionalIf>(ctx, node.Pos())
+                .Predicate<TCoLambda>()
+                    .Args(lambda.Args())
+                    .Body(remainingPredicates.ExprNode)
+                .Build()
+            .Build();
+            TExprNode::TPtr afterPeephole;
+            bool hasNonDeterministicFunctions;
+            if (const auto status = PeepHoleOptimizeNode(optionaIf.Ptr(), afterPeephole, ctx, typesCtx, nullptr, hasNonDeterministicFunctions);
+                status != IGraphTransformer::TStatus::Ok) {
+                YQL_CLOG(ERROR, ProviderKqp) << "Peephole OLAP failed." << Endl << ctx.IssueManager.GetIssues().ToString();
+                return node;
+            }
+            const TCoIf simplified(std::move(afterPeephole));
+            predicate = simplified.Predicate();
+            value = simplified.ThenValue().Cast<TCoJust>().Input();
+
+            TOLAPPredicateNode predicateTree;
+            predicateTree.ExprNode = predicate.Ptr();
+            CollectPredicates(predicate, predicateTree, &lambdaArg, read.Process().Body());
+            YQL_ENSURE(predicateTree.IsValid(), "Collected OLAP predicates are invalid");
+        
+            auto [pushable2, remaining2] = SplitForPartialPushdown(predicateTree);
+            for (auto& p: pushable2) {
+                pushable.emplace_back(std::move(p));
+            }
+            remaining = std::move(remaining2);
+        }
+    }
+
+    if (pushable.empty()) {
         return node;
     }
 
-    YQL_ENSURE(predicatesToPush.IsValid(), "Predicates to push is invalid");
-    YQL_ENSURE(remainingPredicates.IsValid(), "Remaining predicates is invalid");
 
+    const auto predicatesToPush = WrapPredicates(pushable, ctx, node.Pos());
+    const auto remainingPredicates = WrapPredicates(remaining, ctx, node.Pos());
     const auto pushedFilters = PredicatePushdown(TExprBase(predicatesToPush.ExprNode), lambdaArg, ctx, node.Pos());
     // Temporary fix for https://st.yandex-team.ru/KIKIMR-22560
     // YQL_ENSURE(pushedFilters.IsValid(), "Pushed predicate should be always valid!");

ydb/core/kqp/query_compiler/kqp_olap_compiler.cpp

@@ -450,6 +450,7 @@ const TTypedColumn CompileJsonExists(const TKqpOlapJsonExists& jsonExistsCallabl
         jsonExistsCallable.Column(),
         jsonExistsCallable.Path(),
         type);
+    jsonExistsFunc->SetKernelName("JsonExists");
     jsonExistsFunc->SetKernelIdx(idx);
 
     if constexpr (NSsa::RuntimeVersion >= 4U) {

ydb/core/kqp/ut/olap/aggregations_ut.cpp

@@ -1286,6 +1286,24 @@ Y_UNIT_TEST_SUITE(KqpOlapAggregations) {
         TestTableWithNulls({testCase});
     }
 
+    Y_UNIT_TEST(Json_ValueAndLike) {
+        TAggregationTestCase testCase;
+        testCase.SetQuery(R"(
+                SELECT id FROM `/Root/tableWithNulls`
+                WHERE 
+                    -- (JSON_VALUE(jsonval, "$.labels.http_status_code") = '500')
+                    -- AND JSON_EXISTS(jsonval, "$.meta.exception.stacktrace")
+                    --AND 
+                        (JSON_VALUE(jsonval, "$.meta.uri") like "%unified%")
+            )")
+            .AddExpectedPlanOptions("KqpOlapFilter");
+            // .SetExpectedReply(R"([[1;["val1"];#]])");
+
+        TestTableWithNulls({testCase});
+    }
+
+
+
     Y_UNIT_TEST(Json_GetValue_Minus) {
         TAggregationTestCase testCase;
         testCase.SetQuery(R"(
@@ -1310,9 +1328,7 @@ Y_UNIT_TEST_SUITE(KqpOlapAggregations) {
                 SELECT id, JSON_VALUE(jsonval, "$.col1" RETURNING String), JSON_VALUE(jsondoc, "$.col1") FROM `/Root/tableWithNulls`
                 WHERE JSON_VALUE(jsonval, "$.col1" RETURNING String) = "val1" AND id = 1;
             )")
-#if SSA_RUNTIME_VERSION >= 5U
-            .AddExpectedPlanOptions("KqpOlapApply")
-#elif SSA_RUNTIME_VERSION >= 3U
+#if SSA_RUNTIME_VERSION >= 3U
             .AddExpectedPlanOptions("KqpOlapJsonValue")
 #else
             .AddExpectedPlanOptions("Udf")