Fix failed unit tests after enabling the EnableColumnStatistics feature flag (#11409)

Co-authored-by: pilik <pudge1000-7@ydb.tech>

Author: andrewstalin

ydb/core/grpc_services/rpc_read_rows.cpp

@@ -82,8 +82,8 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
 
     static constexpr TDuration DEFAULT_TIMEOUT = TDuration::Seconds(60);
 public:
-    explicit TReadRowsRPC(std::unique_ptr<IRequestNoOpCtx> request)
-        : Request(std::move(request))
+    explicit TReadRowsRPC(IRequestNoOpCtx* request)
+        : Request(request)
         , PipeCache(MakePipePerNodeCacheID(true))
         , Span(TWilsonGrpc::RequestActor, Request->GetWilsonTraceId(), "ReadRowsRpc")
     {}
@@ -730,8 +730,13 @@ class TReadRowsRPC : public TActorBootstrapped<TReadRowsRPC> {
     NWilson::TSpan Span;
 };
 
-void DoReadRowsRequest(std::unique_ptr<IRequestNoOpCtx> p, const IFacilityProvider& f) {
-    f.RegisterActor(new TReadRowsRPC(std::move(p)));
+ void DoReadRowsRequest(std::unique_ptr<IRequestNoOpCtx> p, const IFacilityProvider& f) {
+    f.RegisterActor(new TReadRowsRPC(p.release()));
+}
+
+template<>
+IActor* TEvReadRowsRequest::CreateRpcActor(NKikimr::NGRpcService::IRequestNoOpCtx* msg) {
+    return new TReadRowsRPC(msg);
 }
 
 } // namespace NKikimr::NGRpcService

ydb/core/kqp/opt/kqp_column_statistics_requester.cpp

@@ -93,8 +93,9 @@ IGraphTransformer::TStatus TKqpColumnStatisticsRequester::DoTransform(TExprNode:
                 continue;
             }
 
-            if (!columns.contains(column)) {
+            if (!columnsMeta.contains(column)) {
                 YQL_CLOG(DEBUG, ProviderKikimr) << "Table: " + table + " doesn't contain " + column + " to request for column statistics";
+                continue;
             }
 
             NKikimr::NStat::TRequest req;

ydb/core/statistics/database/database.cpp

@@ -218,17 +218,15 @@ class TLoadStatisticsQuery : public NKikimr::TQueryBase {
     const TPathId PathId;
     const ui64 StatType;
     const ui32 ColumnTag;
-    const ui64 Cookie;
 
     std::optional<TString> Data;
 
 public:
-    TLoadStatisticsQuery(const TString& database, const TPathId& pathId, ui64 statType, ui32 columnTag, ui64 cookie)
+    TLoadStatisticsQuery(const TString& database, const TPathId& pathId, ui64 statType, ui32 columnTag)
         : NKikimr::TQueryBase(NKikimrServices::STATISTICS, {}, database, true)
         , PathId(pathId)
         , StatType(statType)
         , ColumnTag(columnTag)
-        , Cookie(cookie)
     {}
 
     void OnRunQuery() override {
@@ -287,7 +285,6 @@ class TLoadStatisticsQuery : public NKikimr::TQueryBase {
         response->Status = status;
         response->Issues = std::move(issues);
         response->Success = (status == Ydb::StatusIds::SUCCESS);
-        response->Cookie = Cookie;
         if (response->Success) {
             response->Data = Data;
         }
@@ -302,21 +299,19 @@ class TLoadStatisticsRetryingQuery : public TActorBootstrapped<TLoadStatisticsRe
     const TPathId PathId;
     const ui64 StatType;
     const ui32 ColumnTag;
-    const ui64 Cookie;
 
 public:
     using TLoadRetryingQuery = TQueryRetryActor<
         TLoadStatisticsQuery, TEvStatistics::TEvLoadStatisticsQueryResponse,
-        const TString&, const TPathId&, ui64, ui32, ui64>;
+        const TString&, const TPathId&, ui64, ui32>;
 
     TLoadStatisticsRetryingQuery(const NActors::TActorId& replyActorId, const TString& database,
-        const TPathId& pathId, ui64 statType, ui32 columnTag, ui64 cookie)
+        const TPathId& pathId, ui64 statType, ui32 columnTag)
         : ReplyActorId(replyActorId)
         , Database(database)
         , PathId(pathId)
         , StatType(statType)
         , ColumnTag(columnTag)
-        , Cookie(cookie)
     {}
 
     void Bootstrap() {
@@ -326,7 +321,7 @@ class TLoadStatisticsRetryingQuery : public TActorBootstrapped<TLoadStatisticsRe
                 TLoadRetryingQuery::Retryable, TDuration::MilliSeconds(10),
                 TDuration::MilliSeconds(200), TDuration::Seconds(1),
                 std::numeric_limits<size_t>::max(), TDuration::Seconds(1)),
-            Database, PathId, StatType, ColumnTag, Cookie
+            Database, PathId, StatType, ColumnTag
         ));
         Become(&TLoadStatisticsRetryingQuery::StateFunc);
     }
@@ -342,9 +337,9 @@ class TLoadStatisticsRetryingQuery : public TActorBootstrapped<TLoadStatisticsRe
 };
 
 NActors::IActor* CreateLoadStatisticsQuery(const NActors::TActorId& replyActorId,
-    const TString& database, const TPathId& pathId, ui64 statType, ui32 columnTag, ui64 cookie)
+    const TString& database, const TPathId& pathId, ui64 statType, ui32 columnTag)
 {
-    return new TLoadStatisticsRetryingQuery(replyActorId, database, pathId, statType, columnTag, cookie);
+    return new TLoadStatisticsRetryingQuery(replyActorId, database, pathId, statType, columnTag);
 }
 
 

ydb/core/statistics/database/database.h

@@ -10,10 +10,10 @@ NActors::IActor* CreateStatisticsTableCreator(std::unique_ptr<NActors::IEventBas
 NActors::IActor* CreateSaveStatisticsQuery(const NActors::TActorId& replyActorId, const TString& database,
     const TPathId& pathId, ui64 statType, std::vector<ui32>&& columnTags, std::vector<TString>&& data);
 
-NActors::IActor* CreateLoadStatisticsQuery(const NActors::TActorId& replyActorId,  const TString& database,
-    const TPathId& pathId, ui64 statType, ui32 columnTag, ui64 cookie);
+NActors::IActor* CreateLoadStatisticsQuery(const NActors::TActorId& replyActorId, const TString& database,
+    const TPathId& pathId, ui64 statType, ui32 columnTag);
 
-NActors::IActor* CreateDeleteStatisticsQuery(const NActors::TActorId& replyActorId,  const TString& database,
+NActors::IActor* CreateDeleteStatisticsQuery(const NActors::TActorId& replyActorId, const TString& database,
     const TPathId& pathId);
 
 };

ydb/core/statistics/database/ut/ut_database.cpp

@@ -33,14 +33,14 @@ Y_UNIT_TEST_SUITE(StatisticsSaveLoad) {
         auto saveResponse = runtime.GrabEdgeEventRethrow<TEvStatistics::TEvSaveStatisticsQueryResponse>(sender);
         UNIT_ASSERT(saveResponse->Get()->Success);
 
-        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 1, 1),
+        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 1),
             0, 0, TMailboxType::Simple, 0, sender);
         auto loadResponseA = runtime.GrabEdgeEventRethrow<TEvStatistics::TEvLoadStatisticsQueryResponse>(sender);
         UNIT_ASSERT(loadResponseA->Get()->Success);
         UNIT_ASSERT(loadResponseA->Get()->Data);
         UNIT_ASSERT_VALUES_EQUAL(*loadResponseA->Get()->Data, "dataA");
 
-        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 2, 1),
+        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 2),
             0, 0, TMailboxType::Simple, 0, sender);
         auto loadResponseB = runtime.GrabEdgeEventRethrow<TEvStatistics::TEvLoadStatisticsQueryResponse>(sender);
         UNIT_ASSERT(loadResponseB->Get()->Success);
@@ -76,7 +76,7 @@ Y_UNIT_TEST_SUITE(StatisticsSaveLoad) {
         auto deleteResponse = runtime.GrabEdgeEvent<TEvStatistics::TEvDeleteStatisticsQueryResponse>(sender);
         UNIT_ASSERT(deleteResponse->Get()->Success);
 
-        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 1, 1),
+        runtime.Register(CreateLoadStatisticsQuery(sender, "/Root/Database", pathId, statType, 1),
             0, 0, TMailboxType::Simple, 0, sender);
         auto loadResponseA = runtime.GrabEdgeEvent<TEvStatistics::TEvLoadStatisticsQueryResponse>(sender);
         UNIT_ASSERT(!loadResponseA->Get()->Success);

ydb/core/statistics/database/ya.make

@@ -2,7 +2,7 @@ LIBRARY()
 
 SRCS(
     database.h
-    database.cpp    
+    database.cpp
 )
 
 PEERDIR(

ydb/core/statistics/events.h

@@ -203,7 +203,6 @@ struct TEvStatistics {
         Ydb::StatusIds::StatusCode Status;
         NYql::TIssues Issues;
         bool Success = true;
-        ui64 Cookie = 0;
         std::optional<TString> Data;
     };
 

ydb/core/statistics/service/http_request.cpp

@@ -1,5 +1,6 @@
 #include "http_request.h"
 
+#include <ydb/core/statistics/service/service.h>
 #include <ydb/core/base/path.h>
 #include <ydb/core/io_formats/cell_maker/cell_maker.h>
 #include <ydb/core/statistics/database/database.h>
@@ -118,15 +119,17 @@ void THttpRequest::Handle(TEvStatistics::TEvAnalyzeStatusResponse::TPtr& ev) {
     }
 }
 
-void THttpRequest::Handle(TEvStatistics::TEvLoadStatisticsQueryResponse::TPtr& ev) {
+void THttpRequest::Handle(TEvStatistics::TEvGetStatisticsResult::TPtr& ev) {
     const auto msg = ev->Get();
-    if (!msg->Success || !msg->Data) {
-        const auto status = std::to_string(static_cast<ui32>(msg->Status));
-        HttpReply("Error occurred while loading statistics. Status: " + status);
+
+    if (!msg->Success
+            || msg->StatResponses.empty() || !msg->StatResponses[0].Success
+            || msg->StatResponses[0].CountMinSketch.CountMin == nullptr) {
+        HttpReply("Error occurred while loading statistics.");
         return;
     }
 
-    const auto typeId = static_cast<NScheme::TTypeId>(msg->Cookie);
+    const auto typeId = static_cast<NScheme::TTypeId>(ev->Cookie);
     const NScheme::TTypeInfo typeInfo(typeId);
     const TStringBuf value(Params[EParamType::CELL_VALUE]);
     TMemoryPool pool(64);
@@ -138,7 +141,7 @@ void THttpRequest::Handle(TEvStatistics::TEvLoadStatisticsQueryResponse::TPtr& e
         return;
     }
 
-    auto countMinSketch = std::unique_ptr<TCountMinSketch>(TCountMinSketch::FromString(msg->Data->Data(), msg->Data->Size()));
+    const auto countMinSketch = msg->StatResponses[0].CountMinSketch.CountMin.get();
     const auto probe = countMinSketch->Probe(cell.Data(), cell.Size());
     HttpReply(Params[EParamType::PATH] + "[" + Params[EParamType::COLUMN_NAME] + "]=" + std::to_string(probe));
 }
@@ -217,10 +220,16 @@ void THttpRequest::DoCountMinSketchProbe(const TNavigate::TEntry& entry) {
 
     for (const auto& [_, tableInfo]: entry.Columns) {
         if (tableInfo.Name == columnName) {
-            const auto columnTag = tableInfo.Id;
+            auto request = std::make_unique<TEvStatistics::TEvGetStatistics>();
+            request->StatType = EStatType::COUNT_MIN_SKETCH;
+            TRequest req;
+            req.PathId = entry.TableId.PathId;
+            req.ColumnTag = tableInfo.Id;
+            request->StatRequests.emplace_back(std::move(req));
+
             const auto typeId = tableInfo.PType.GetTypeId();
-            const auto& pathId = entry.TableId.PathId;
-            Register(CreateLoadStatisticsQuery(SelfId(), Params[EParamType::DATABASE], pathId, EStatType::COUNT_MIN_SKETCH, columnTag, typeId));
+            const auto statService = MakeStatServiceID(SelfId().NodeId());
+            Send(statService, request.release(), 0, typeId);
             return;
         }
     }

ydb/core/statistics/service/http_request.h

@@ -46,7 +46,7 @@ class THttpRequest : public NActors::TActorBootstrapped<THttpRequest> {
             hFunc(TEvTxProxySchemeCache::TEvNavigateKeySetResult, Handle);
             hFunc(TEvStatistics::TEvAnalyzeStatusResponse, Handle);
             hFunc(TEvPipeCache::TEvDeliveryProblem, Handle);
-            hFunc(TEvStatistics::TEvLoadStatisticsQueryResponse, Handle);
+            hFunc(TEvStatistics::TEvGetStatisticsResult, Handle);
             IgnoreFunc(TEvStatistics::TEvAnalyzeResponse);
             default:
                 LOG_CRIT_S(TlsActivationContext->AsActorContext(), NKikimrServices::STATISTICS,
@@ -57,7 +57,7 @@ class THttpRequest : public NActors::TActorBootstrapped<THttpRequest> {
     void Handle(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev);
     void Handle(TEvStatistics::TEvAnalyzeStatusResponse::TPtr& ev);
     void Handle(TEvPipeCache::TEvDeliveryProblem::TPtr& ev);
-    void Handle(TEvStatistics::TEvLoadStatisticsQueryResponse::TPtr& ev);
+    void Handle(TEvStatistics::TEvGetStatisticsResult::TPtr& ev);
 
     void DoRequest(const TNavigate::TEntry& entry);
     void DoAnalyze(const TNavigate::TEntry& entry);