[Olap projection] Do not push projection for the same column (#21702)

2 files changed, 36 insertions, 13 deletions

diff --git a/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp b/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp
index 53494b3f227..30e1d6557c6 100644
--- a/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp
+++ b/ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp
@@ -808,6 +808,7 @@ TVector<std::pair<TString, TExprNode::TPtr>> CollectOlapOperationsForProjections
                                                                                  const THashSet<TString>& predicateMembers, TExprContext& ctx) {
     auto asStructPred = [](const TExprNode::TPtr& node) -> bool { return !!TMaybeNode<TCoAsStruct>(node); };
     auto memberPred = [](const TExprNode::TPtr& node) { return !!TMaybeNode<TCoMember>(node); };
+    THashSet<TString> projectionMembers;
 
     TVector<std::pair<TString, TExprNode::TPtr>> olapOperationsForProjections;
     // Expressions for projections are placed in `AsStruct` callable.
@@ -821,14 +822,19 @@ TVector<std::pair<TString, TExprNode::TPtr>> CollectOlapOperationsForProjections
                     if (auto olapOperations = ConvertComparisonNode(TExprBase(child.Item(1)), arg, ctx, node->Pos(), false);
                         olapOperations.size() == 1) {
                         auto originalMember = TExprBase(originalMembers.front()).Cast<TCoMember>();
-
                         auto originalMemberName = TString(originalMember.Name());
+                        // Do not push projection for the same column, we are not ready for this.
+                        if (projectionMembers.contains(originalMemberName)) {
+                            return {};
+                        }
+                        projectionMembers.insert(originalMemberName);
+
                         // We cannot push projection if some predicate for the same column still not pushed.
                         if (!predicateMembers.contains(originalMemberName)) {
                             auto newMember = Build<TCoMember>(ctx, node->Pos())
                                 .Struct(originalMember.Struct())
                                 .Name(originalMember.Name())
-                                .Done();
+                            .Done();
 
                             auto olapOperation = olapOperations.front();
                             // Replace full expression with only member.
diff --git a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
index bb0895f1bb3..85ccd42c258 100644
--- a/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
+++ b/ydb/core/kqp/ut/olap/kqp_olap_ut.cpp
@@ -1737,8 +1737,6 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
 
         for (ui32 i = 0; i < queries.size(); ++i) {
             const auto query = queries[i];
-            Cerr << "QUERY: " << Endl;
-            Cerr << query << Endl;
 
             auto result =
                 session2
@@ -1747,11 +1745,6 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
             UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
 
             auto ast = *result.GetStats()->GetAst();
-            // TString plan = *result.GetStats()->GetPlan();
-            // Cerr << "PLAN " << plan << Endl;
-            // TODO: Add pushed projection to explain.
-            // NYdb::NConsoleClient::TQueryPlanPrinter queryPlanPrinter(NYdb::NConsoleClient::EDataFormat::PrettyTable, true, Cout, 0);
-            // queryPlanPrinter.Print(plan);
 
             UNIT_ASSERT_C(ast.find("KqpOlapProjections") != std::string::npos, TStringBuilder() << "Projections not pushed down. Query: " << query);
             UNIT_ASSERT_C(ast.find("KqpOlapProjection") != std::string::npos, TStringBuilder() << "Projection not pushed down. Query: " << query);
@@ -1760,12 +1753,36 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
             UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
 
             TString output = FormatResultSetYson(result.GetResultSet(0));
-            Cerr << "QUERY: " << Endl;
-            Cerr << query << Endl;
-
-            Cerr << "RESULT: " << output << Endl;
             CompareYson(output, results[i]);
         }
+
+        TVector<TString> notForPushdown = {
+        R"(
+            PRAGMA Kikimr.OptEnableOlapPushdownProjections = "true";
+
+            SELECT JSON_VALUE(jsonDoc, "$.a"), JSON_VALUE(jsonDoc, "$.b") FROM `/Root/foo`
+            where JSON_VALUE(jsonDoc, "$.c") = "b";
+        )"};
+
+       for (ui32 i = 0; i < notForPushdown.size(); ++i) {
+            const auto query = notForPushdown[i];
+            auto result =
+                session2
+                    .ExecuteQuery(query, NYdb::NQuery::TTxControl::NoTx(), NYdb::NQuery::TExecuteQuerySettings().ExecMode(NQuery::EExecMode::Explain))
+                    .ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+
+            auto ast = *result.GetStats()->GetAst();
+            UNIT_ASSERT_C(ast.find("KqpOlapProjections") == std::string::npos,
+                          TStringBuilder() << "Projections pushed down, but should not. Query: " << query);
+            UNIT_ASSERT_C(ast.find("KqpOlapFilter") != std::string::npos, TStringBuilder() << "Filter not pushed down. Query: " << query);
+
+            result =
+                session2.ExecuteQuery(query, NYdb::NQuery::TTxControl::NoTx(), NYdb::NQuery::TExecuteQuerySettings()).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::SUCCESS);
+
+            TString output = FormatResultSetYson(result.GetResultSet(0));
+        }
     }
 
     // Unit tests for datetime pushdowns in query service