feat: Add a trace file operation tool (#12021)

duanmeng · facebook-github-bot · commit 915f1d8c80ce · 2025-01-16T23:19:54.000-08:00
Summary: Add a trace file operation tool, which provides trace file copy operation. We can specify the target query or task by gflags. Pull Request resolved: #12021 Reviewed By: yuandagits Differential Revision: D68254596 Pulled By: xiaoxmeng fbshipit-source-id: bb1e580898ca17f491ff4ec984efeac24cfbba5a
diff --git a/velox/tool/trace/CMakeLists.txt b/velox/tool/trace/CMakeLists.txt
@@ -52,4 +52,22 @@ target_link_libraries(
   velox_exec_test_lib
   velox_tpch_connector)
 
+add_library(velox_trace_file_tool_base TraceFileToolRunner.cpp)
+
+target_link_libraries(
+  velox_trace_file_tool_base
+  velox_exec
+  velox_type
+  velox_vector
+  velox_exec
+  velox_hive_connector
+  Folly::folly
+  glog::glog
+  gflags::gflags)
+
+add_executable(velox_trace_file_tool TraceFileToolMain.cpp)
+
+target_link_libraries(
+  velox_trace_file_tool velox_trace_file_tool_base)
+
 add_subdirectory(tests)
diff --git a/velox/tool/trace/TraceFileToolMain.cpp b/velox/tool/trace/TraceFileToolMain.cpp
@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <date/date.h>
+#include <gflags/gflags.h>
+#include "velox/tool/trace/TraceFileToolRunner.h"
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  facebook::velox::tool::trace::TraceFileToolRunner runner;
+  runner.init();
+  runner.run();
+  return 0;
+}
diff --git a/velox/tool/trace/TraceFileToolRunner.cpp b/velox/tool/trace/TraceFileToolRunner.cpp
@@ -0,0 +1,120 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/tool/trace/TraceFileToolRunner.h"
+#include "velox/common/file/FileSystems.h"
+
+DEFINE_string(trace_file_op, "copy", "Operation type of this run");
+DEFINE_string(
+    source_root_dir,
+    "",
+    "Source root directory of the tracing data, it must be set");
+DEFINE_string(
+    dest_root_dir,
+    "",
+    "Dest root directory, it must be set if the operation type is copy");
+DEFINE_string(trace_query_id, "", "Specify the trace query id");
+DEFINE_string(trace_task_id, "", "Specify the trace task id");
+
+namespace facebook::velox::tool::trace {
+
+TraceFileToolRunner::TraceFileToolRunner()
+    : sourceRootDir_(FLAGS_source_root_dir), destRootDir_(FLAGS_dest_root_dir) {
+  VELOX_USER_CHECK(!sourceRootDir_.empty());
+}
+
+void TraceFileToolRunner::init() {
+  filesystems::registerLocalFileSystem();
+  sourceFs_ = filesystems::getFileSystem(sourceRootDir_, nullptr);
+  VELOX_CHECK_NOT_NULL(sourceFs_);
+  if (FLAGS_trace_file_op == "copy") {
+    VELOX_USER_CHECK(!destRootDir_.empty());
+    destFs_ = filesystems::getFileSystem(destRootDir_, nullptr);
+    VELOX_CHECK_NOT_NULL(destFs_);
+    std::string copyRootDir;
+    if (FLAGS_trace_query_id.empty()) {
+      VELOX_CHECK(
+          FLAGS_trace_task_id.empty(),
+          "Trace query ID is empty but trace task ID is not empty");
+      copyRootDir = sourceRootDir_;
+    } else if (FLAGS_trace_task_id.empty()) {
+      copyRootDir = fmt::format("{}/{}", sourceRootDir_, FLAGS_trace_query_id);
+    } else {
+      copyRootDir = fmt::format(
+          "{}/{}/{}",
+          sourceRootDir_,
+          FLAGS_trace_query_id,
+          FLAGS_trace_task_id);
+    }
+    listFiles(copyRootDir);
+  } else {
+    VELOX_UNSUPPORTED(
+        "Unsupported trace file operation type {}", FLAGS_trace_file_op);
+  }
+}
+
+void TraceFileToolRunner::run() {
+  if (FLAGS_trace_file_op == "copy") {
+    copyFiles();
+  } else {
+    VELOX_UNSUPPORTED(
+        "Unsupported trace file operation type {}", FLAGS_trace_file_op);
+  }
+}
+
+void TraceFileToolRunner::copyFiles() const {
+  const auto prefixLen = sourceFs_->extractPath(sourceRootDir_).length();
+  for (const auto& source : sourceFiles_) {
+    const auto targetFile =
+        fmt::format("{}/{}", destRootDir_, source.substr(prefixLen));
+    const auto readFile = sourceFs_->openFileForRead(source);
+    const auto writeFile = destFs_->openFileForWrite(
+        targetFile,
+        filesystems::FileOptions{
+            .values = {},
+            .fileSize = std::nullopt,
+            .shouldCreateParentDirectories = true});
+    const auto fileSize = readFile->size();
+    constexpr auto batchSize = 4 << 10;
+    const auto ioBuf = folly::IOBuf::create(batchSize);
+    uint64_t offset = 0;
+    while (offset < fileSize) {
+      const auto curLen = std::min<uint64_t>(fileSize - offset, batchSize);
+      const auto dataView =
+          readFile->pread(offset, curLen, ioBuf->writableData());
+      writeFile->append(dataView);
+      ioBuf->append(curLen);
+      offset += curLen;
+      ioBuf->clear();
+    }
+    writeFile->flush();
+    writeFile->close();
+  }
+}
+
+void TraceFileToolRunner::listFiles(const std::string& path) {
+  VELOX_USER_CHECK(sourceFs_->exists(path), "{} dose not exist", path);
+  if (!sourceFs_->isDirectory(path)) {
+    sourceFiles_.push_back(path);
+    return;
+  }
+
+  for (const auto& p : sourceFs_->list(sourceFs_->extractPath(path))) {
+    listFiles(p);
+  }
+}
+
+} // namespace facebook::velox::tool::trace
diff --git a/velox/tool/trace/TraceFileToolRunner.h b/velox/tool/trace/TraceFileToolRunner.h
@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/io/IOBuf.h>
+#include <gflags/gflags.h>
+#include <string>
+#include "velox/common/file/FileSystems.h"
+#include "velox/common/file/Utils.h"
+
+DECLARE_string(source_root_dir);
+DECLARE_string(dest_root_dir);
+DECLARE_string(trace_file_op);
+DECLARE_string(trace_query_id);
+DECLARE_string(trace_task_id);
+
+namespace facebook::velox::tool::trace {
+
+/// The trace replay runner. It is configured through a set of gflags passed
+/// from replayer tool command line.
+class TraceFileToolRunner {
+ public:
+  TraceFileToolRunner();
+  virtual ~TraceFileToolRunner() = default;
+
+  /// Initializes the trace file tool runner by setting the velox runtime
+  /// environment for the trace file operations. It is invoked before run().
+  virtual void init();
+
+  /// Runs the trace file operations.
+  void run();
+
+ private:
+  // List all the files in the source root dir recursively.
+  void listFiles(const std::string& path);
+
+  void copyFiles() const;
+
+  const std::string sourceRootDir_;
+  const std::string destRootDir_;
+  std::shared_ptr<filesystems::FileSystem> sourceFs_;
+  std::shared_ptr<filesystems::FileSystem> destFs_;
+  std::vector<std::string> sourceFiles_;
+};
+
+} // namespace facebook::velox::tool::trace
diff --git a/velox/tool/trace/tests/CMakeLists.txt b/velox/tool/trace/tests/CMakeLists.txt
@@ -18,6 +18,7 @@ add_executable(
   FilterProjectReplayerTest.cpp
   HashJoinReplayerTest.cpp
   PartitionedOutputReplayerTest.cpp
+  TraceFileToolTest.cpp
   TableScanReplayerTest.cpp
   TableWriterReplayerTest.cpp)
 
@@ -34,6 +35,7 @@ target_link_libraries(
   velox_exec_test_lib
   velox_memory
   velox_query_trace_replayer_base
+  velox_trace_file_tool_base
   velox_vector_fuzzer
   GTest::gtest_main
   GTest::gmock
diff --git a/velox/tool/trace/tests/TraceFileToolTest.cpp b/velox/tool/trace/tests/TraceFileToolTest.cpp
