YQ-3697 Add partition count to dqrun (#9837)

Author: kardymonds

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp

@@ -53,8 +53,8 @@ NThreading::TFuture<IPqGateway::TListStreams> TDummyPqGateway::ListStreams(const
 TDummyPqGateway& TDummyPqGateway::AddDummyTopic(const TDummyTopic& topic) {
     with_lock (Mutex) {
         Y_ENSURE(topic.Cluster);
-        Y_ENSURE(topic.Path);
-        const auto key = std::make_pair(topic.Cluster, topic.Path);
+        Y_ENSURE(topic.TopicName);
+        const auto key = std::make_pair(topic.Cluster, topic.TopicName);
         Y_ENSURE(Topics.emplace(key, topic).second, "Already inserted dummy topic {" << topic.Cluster << ", " << topic.Path << "}");
         return *this;
     }

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h

@@ -9,10 +9,11 @@
 namespace NYql {
 
 struct TDummyTopic {
-    TDummyTopic(const TString& cluster, const TString& path, const TMaybe<TString>& filePath = {})
+    TDummyTopic(const TString& cluster, const TString& topicName, const TMaybe<TString>& path = {}, size_t partitionCount = 1)
         : Cluster(cluster)
+        , TopicName(topicName)
         , Path(path)
-        , FilePath(filePath)
+        , PartitionsCount(partitionCount)
     {
     }
 
@@ -22,9 +23,9 @@ struct TDummyTopic {
     }
 
     TString Cluster;
-    TString Path;
-    TMaybe<TString> FilePath;
-    size_t PartitionsCount = 1;
+    TString TopicName;
+    TMaybe<TString> Path;
+    size_t PartitionsCount;
 };
 
 // Dummy Pq gateway for tests.

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp

@@ -5,6 +5,7 @@
 
 #include <library/cpp/threading/future/async.h>
 
+#include <util/folder/path.h>
 #include <util/system/file.h>
 #include "yql_pq_blocking_queue.h"
 
@@ -337,17 +338,24 @@ struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession {
 
 std::shared_ptr<NYdb::NTopic::IReadSession> TFileTopicClient::CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) {
     Y_ENSURE(!settings.Topics_.empty());
-    auto topicPath = settings.Topics_.front().Path_;
-
+    const auto& topic = settings.Topics_.front();
+    auto topicPath = topic.Path_;
+    Y_ENSURE(topic.PartitionIds_.size() >= 1);
+    ui64 partitionId = topic.PartitionIds_.front();
     auto topicsIt = Topics_.find(make_pair("pq", topicPath));
     Y_ENSURE(topicsIt != Topics_.end());
-    auto filePath = topicsIt->second.FilePath;
+    auto filePath = topicsIt->second.Path;
     Y_ENSURE(filePath);
 
+    TFsPath fsPath(*filePath);
+    if (fsPath.IsDirectory()) {
+        filePath = TStringBuilder() << *filePath << "/" << ToString(partitionId);
+    } else if (!fsPath.Exists()) {
+        filePath = TStringBuilder() << *filePath << "_" << partitionId;
+    }
+
     // TODO
     ui64 sessionId = 0;
-    ui64 partitionId = 0;
-
     return std::make_shared<TFileTopicReadSession>(
         TFile(*filePath, EOpenMode::TEnum::RdOnly),
         MakeIntrusive<TDummyPartitionSession>(sessionId, TString{topicPath}, partitionId)
@@ -411,7 +419,7 @@ std::shared_ptr<NYdb::NTopic::IWriteSession> TFileTopicClient::CreateWriteSessio
     auto topicPath = TString{settings.Path_};
     auto topicsIt = Topics_.find(make_pair("pq", topicPath));
     Y_ENSURE(topicsIt != Topics_.end());
-    auto filePath = topicsIt->second.FilePath;
+    auto filePath = topicsIt->second.Path;
     Y_ENSURE(filePath);
 
     return std::make_shared<TFileTopicWriteSession>(TFile(*filePath, EOpenMode::TEnum::RdWr));

ydb/tests/tools/fqrun/fqrun.cpp

@@ -163,13 +163,15 @@ class TMain : public TMainBase {
         options.AddLongOption("emulate-pq", "Emulate YDS with local file, accepts list of tables to emulate with following format: topic@file (can be used in query from cluster `pq`)")
             .RequiredArgument("topic@file")
             .Handler1([this](const NLastGetopt::TOptsParser* option) {
-                TStringBuf topicName;
-                TStringBuf filePath;
-                TStringBuf(option->CurVal()).Split('@', topicName, filePath);
-                if (topicName.empty() || filePath.empty()) {
-                    ythrow yexception() << "Incorrect PQ file mapping, expected form topic@file";
+                TStringBuf topicName, others;
+                TStringBuf(option->CurVal()).Split('@', topicName, others);
+                TStringBuf path, partitionCountStr;
+                TStringBuf(others).Split(':', path, partitionCountStr);
+                size_t partitionCount = !partitionCountStr.empty() ? FromString<size_t>(partitionCountStr) : 1;
+                if (topicName.empty() || path.empty()) {
+                    ythrow yexception() << "Incorrect PQ file mapping, expected form topic@path[:partitions_count]" << Endl;
                 }
-                if (!PqFilesMapping.emplace(topicName, filePath).second) {
+                if (!PqFilesMapping.emplace(topicName, NYql::TDummyTopic("pq", TString(topicName), TString(path), partitionCount)).second) {
                     ythrow yexception() << "Got duplicated topic name: " << topicName;
                 }
             });
@@ -222,8 +224,8 @@ class TMain : public TMainBase {
 
         if (!PqFilesMapping.empty()) {
             auto fileGateway = MakeIntrusive<NYql::TDummyPqGateway>();
-            for (const auto& [topic, file] : PqFilesMapping) {
-                fileGateway->AddDummyTopic(NYql::TDummyTopic("pq", TString(topic), TString(file)));
+            for (const auto& [_, topic] : PqFilesMapping) {
+                fileGateway->AddDummyTopic(topic);
             }
             RunnerOptions.FqSettings.PqGateway = std::move(fileGateway);
         }
@@ -246,7 +248,7 @@ class TMain : public TMainBase {
 private:
     TExecutionOptions ExecutionOptions;
     TRunnerOptions RunnerOptions;
-    std::unordered_map<TString, TString> PqFilesMapping;
+    std::unordered_map<TString, NYql::TDummyTopic> PqFilesMapping;
 };
 
 }  // anonymous namespace