[OLAP Projection pushdown] Do not push projection if kqp expects a full json document from CS (#28026)

Author: denis0x0D

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

@@ -807,17 +807,20 @@ bool IsSuitableToCollectProjection(TExprBase node) {
 // Collects all operations for projections and returns a vector of pair - [columName, olap operation].
 TVector<std::pair<TString, TExprNode::TPtr>> CollectOlapOperationsForProjections(const TExprNode::TPtr& node, const TExprNode& arg,
                                                                                  TNodeOnNodeOwnedMap& replaces,
-                                                                                 const THashSet<TString>& predicateMembers, TExprContext& ctx) {
+                                                                                 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;
+    THashSet<TString> notSuitableToPushMembers;
     ui32 nextMemberId = 0;
 
     TVector<std::pair<TString, TExprNode::TPtr>> olapOperationsForProjections;
     // Expressions for projections are placed in `AsStruct` callable.
     if (auto asStruct = FindNode(node, asStructPred)) {
         // Process each child for `AsStruct` callable.
         for (auto child : TExprBase(asStruct).Cast<TCoAsStruct>()) {
+            bool memberCollected = false;
             if (IsSuitableToCollectProjection(child.Item(1))) {
                 // Search for the `TCoMember` in expression, we need expression with only one `TCoMember`.
                 if (auto originalMembers = FindNodes(child.Item(1).Ptr(), memberPred); originalMembers.size() == 1) {
@@ -834,24 +837,39 @@ TVector<std::pair<TString, TExprNode::TPtr>> CollectOlapOperationsForProjections
                                 projectionMembers.insert(originalMemberName);
                             }
 
+                            // clang-format off
                             auto newMember = Build<TCoMember>(ctx, node->Pos())
                                 .Struct(originalMember.Struct())
                                 .Name<TCoAtom>()
                                     .Value(originalMemberName)
                                     .Build()
                             .Done();
+                            // clang-format on
 
                             auto olapOperation = olapOperations.front();
                             // Replace full expression with only member.
                             replaces[child.Item(1).Raw()] = newMember.Ptr();
                             olapOperationsForProjections.emplace_back(TString(newMember.Name()), olapOperation.Ptr());
+                            memberCollected = true;
 
                             YQL_CLOG(TRACE, ProviderKqp)
                                 << "[OLAP PROJECTION] Operation in olap dialect: " << KqpExprToPrettyString(olapOperation, ctx);
                         }
                     }
                 }
             }
+            if (!memberCollected) {
+                auto members = FindNodes(child.Item(1).Ptr(), memberPred);
+                for (const auto& member : members) {
+                    notSuitableToPushMembers.insert(TString(TExprBase(member).Cast<TCoMember>().Name()));
+                }
+            }
+        }
+    }
+
+    for (const auto& [colName, expr] : olapOperationsForProjections) {
+        if (notSuitableToPushMembers.count(colName)) {
+            return {};
         }
     }
 

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

@@ -1861,6 +1861,29 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
             TString output = FormatResultSetYson(result.GetResultSet(0));
             CompareYson(output, results[i]);
         }
+
+        std::vector<TString> notPushedQueries = {
+            R"(
+                PRAGMA Kikimr.OptEnableOlapPushdownProjections = "true";
+
+                SELECT jsonDoc, JSON_VALUE(jsonDoc, "$.\"a.b.c\"")
+                FROM `/Root/foo`
+                where b == 1;
+            )"
+        };
+
+        for (ui32 i = 0; i < notPushedQueries.size(); ++i) {
+            const auto query = notPushedQueries[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. Query: " << query);
+            UNIT_ASSERT_C(ast.find("KqpOlapProjection") == std::string::npos, TStringBuilder() << "Projection pushed down. Query: " << query);
+        }
     }
 
     // Unit tests for datetime pushdowns in query service