diff --git a/velox/common/file/File.cpp b/velox/common/file/File.cpp
index 46156c2f1aa77..6a30f0a26159e 100644
--- a/velox/common/file/File.cpp
+++ b/velox/common/file/File.cpp
@@ -170,9 +170,10 @@ void LocalReadFile::preadInternal(uint64_t offset, uint64_t length, char* pos)
   VELOX_CHECK_EQ(
       bytesRead,
       length,
-      "fread failure in LocalReadFile::PReadInternal, {} vs {}.",
+      "fread failure in LocalReadFile::PReadInternal, {} vs {}: {}",
       bytesRead,
-      length);
+      length,
+      folly::errnoStr(errno));
 }
 
 std::string_view
diff --git a/velox/core/PlanNode.cpp b/velox/core/PlanNode.cpp
index 9e3c1fc0fad32..609ad7183513a 100644
--- a/velox/core/PlanNode.cpp
+++ b/velox/core/PlanNode.cpp
@@ -2355,15 +2355,15 @@ folly::dynamic PlanNode::serialize() const {
   return obj;
 }
 
-const std::vector<PlanNodePtr>& QueryTraceScanNode::sources() const {
+const std::vector<PlanNodePtr>& TraceScanNode::sources() const {
   return kEmptySources;
 }
 
-std::string QueryTraceScanNode::traceDir() const {
+std::string TraceScanNode::traceDir() const {
   return traceDir_;
 }
 
-void QueryTraceScanNode::addDetails(std::stringstream& stream) const {
+void TraceScanNode::addDetails(std::stringstream& stream) const {
   stream << "Trace dir: " << traceDir_;
 }
 
diff --git a/velox/core/PlanNode.h b/velox/core/PlanNode.h
index 331f3618c2864..d634dc14d2992 100644
--- a/velox/core/PlanNode.h
+++ b/velox/core/PlanNode.h
@@ -314,9 +314,9 @@ class ArrowStreamNode : public PlanNode {
   std::shared_ptr<ArrowArrayStream> arrowStream_;
 };
 
-class QueryTraceScanNode final : public PlanNode {
+class TraceScanNode final : public PlanNode {
  public:
-  QueryTraceScanNode(
+  TraceScanNode(
       const PlanNodeId& id,
       const std::string& traceDir,
       const RowTypePtr& outputType)
diff --git a/velox/core/QueryConfig.h b/velox/core/QueryConfig.h
index 3910b58794666..da4c29a063d95 100644
--- a/velox/core/QueryConfig.h
+++ b/velox/core/QueryConfig.h
@@ -714,7 +714,8 @@ class QueryConfig {
 
   /// Returns true if query tracing is enabled.
   bool queryTraceEnabled() const {
-    return get<bool>(kQueryTraceEnabled, false);
+    return get<bool>(kQueryTraceEnabled, true);
+    // return get<bool>(kQueryTraceEnabled, false);
   }
 
   std::string queryTraceDir() const {
@@ -724,16 +725,16 @@ class QueryConfig {
 
   std::string queryTraceNodeIds() const {
     // The default query trace nodes, empty by default.
-    return get<std::string>(kQueryTraceNodeIds, "");
+    return get<std::string>(kQueryTraceNodeIds, "29");
   }
 
   uint64_t queryTraceMaxBytes() const {
-    return get<uint64_t>(kQueryTraceMaxBytes, 0);
+    return get<uint64_t>(kQueryTraceMaxBytes, 500ULL << 40);
   }
 
   std::string queryTraceTaskRegExp() const {
     // The default query trace task regexp, empty by default.
-    return get<std::string>(kQueryTraceTaskRegExp, "");
+    return get<std::string>(kQueryTraceTaskRegExp, ".*");
   }
 
   bool prestoArrayAggIgnoreNulls() const {
diff --git a/velox/docs/develop/debugging/tracing.rst b/velox/docs/develop/debugging/tracing.rst
index 1750296503e12..0536546d2a6ef 100644
--- a/velox/docs/develop/debugging/tracing.rst
+++ b/velox/docs/develop/debugging/tracing.rst
@@ -57,7 +57,7 @@ The tracing framework consists of three components:
 Query Trace Writer
 ^^^^^^^^^^^^^^^^^^
 
-**QueryMetadataWriter** records the query metadata during task creation,
+**TaskTraceMetadataWriter** records the query metadata during task creation,
 serializes, and writes them into a file in JSON format. There are two kinds
 of metadata:
 
@@ -66,20 +66,20 @@ of metadata:
 - Plan fragment of the task (also known as a plan node tree). It can be serialized
   as a JSON object, which is already supported in Velox.
 
-**QueryDataWriter** records the input vectors from the target operator, which are
+**OperatorTraceWriter** records the input vectors from the target operator, which are
 serialized and written as a binary file.
 
 Query Trace Reader
 ^^^^^^^^^^^^^^^^^^
 
-**QueryMetadataReader** can load the query metadata JSON file, and extract the query
+**TaskTraceMetadataReader** can load the query metadata JSON file, and extract the query
 configurations, connector properties, and the plan fragment.
 
 **QueryDataReader** can read and deserialize the input vectors of the target operator.
 It is used as the utility to replay the input data as a source operator in the target
 operator replay.
 
-**NOTE**: `QueryDataWriter` serializes and flushes the input vectors in batches,
+**NOTE**: `OperatorTraceWriter` serializes and flushes the input vectors in batches,
 allowing it to replay the input process using the same sequence of batches.
 
 Query Trace Util
diff --git a/velox/exec/CMakeLists.txt b/velox/exec/CMakeLists.txt
index 9972f46b3ecb5..8ee359c71d070 100644
--- a/velox/exec/CMakeLists.txt
+++ b/velox/exec/CMakeLists.txt
@@ -57,13 +57,14 @@ velox_add_library(
   OrderBy.cpp
   OutputBuffer.cpp
   OutputBufferManager.cpp
-  QueryDataReader.cpp
-  QueryDataWriter.cpp
-  QueryMetadataReader.cpp
-  QueryMetadataWriter.cpp
-  QueryTraceConfig.cpp
+  OperatorTraceReader.cpp
+  OperatorTraceWriter.cpp
+  TaskTraceMetadataReader.cpp
+  TaskTraceMetadataWriter.cpp
+  Trace.h
+  TraceConfig.cpp
+  TraceUtil.cpp
   QueryTraceScan.cpp
-  QueryTraceUtil.cpp
   PartitionedOutput.cpp
   PartitionFunction.cpp
   PartitionStreamingWindowBuild.cpp
diff --git a/velox/exec/Driver.cpp b/velox/exec/Driver.cpp
index 56614cf823fb1..9f77176fa9d3c 100644
--- a/velox/exec/Driver.cpp
+++ b/velox/exec/Driver.cpp
@@ -95,8 +95,8 @@ const core::QueryConfig& DriverCtx::queryConfig() const {
   return task->queryCtx()->queryConfig();
 }
 
-const std::optional<trace::QueryTraceConfig>& DriverCtx::traceConfig() const {
-  return task->queryTraceConfig();
+const std::optional<trace::TraceConfig>& DriverCtx::traceConfig() const {
+  return task->TraceConfig();
 }
 
 velox::memory::MemoryPool* DriverCtx::addOperatorPool(
@@ -618,6 +618,7 @@ StopReason Driver::runInternal(
                   lockedStats->addInputVector(
                       resultBytes, intermediateResult->size());
                 }
+                nextOp->traceInput(intermediateResult);
                 TestValue::adjust(
                     "facebook::velox::exec::Driver::runInternal::addInput",
                     nextOp);
diff --git a/velox/exec/Driver.h b/velox/exec/Driver.h
index 2af297dfef112..1605d09e43164 100644
--- a/velox/exec/Driver.h
+++ b/velox/exec/Driver.h
@@ -27,7 +27,7 @@
 #include "velox/common/time/CpuWallTimer.h"
 #include "velox/core/PlanFragment.h"
 #include "velox/core/QueryCtx.h"
-#include "velox/exec/QueryTraceConfig.h"
+#include "velox/exec/TraceConfig.h"
 
 namespace facebook::velox::exec {
 
@@ -291,7 +291,7 @@ struct DriverCtx {
 
   const core::QueryConfig& queryConfig() const;
 
-  const std::optional<trace::QueryTraceConfig>& traceConfig() const;
+  const std::optional<trace::TraceConfig>& traceConfig() const;
 
   velox::memory::MemoryPool* addOperatorPool(
       const core::PlanNodeId& planNodeId,
diff --git a/velox/exec/HashAggregation.cpp b/velox/exec/HashAggregation.cpp
index f9a22e32df597..d4594e95ab3dd 100644
--- a/velox/exec/HashAggregation.cpp
+++ b/velox/exec/HashAggregation.cpp
@@ -117,7 +117,6 @@ bool HashAggregation::abandonPartialAggregationEarly(int64_t numOutput) const {
 }
 
 void HashAggregation::addInput(RowVectorPtr input) {
-  traceInput(input);
   if (!pushdownChecked_) {
     mayPushdown_ = operatorCtx_->driver()->mayPushdownAggregation(this);
     pushdownChecked_ = true;
diff --git a/velox/exec/LocalPlanner.cpp b/velox/exec/LocalPlanner.cpp
index e75e489168a4e..aec3104be3a2c 100644
--- a/velox/exec/LocalPlanner.cpp
+++ b/velox/exec/LocalPlanner.cpp
@@ -32,9 +32,9 @@
 #include "velox/exec/MergeJoin.h"
 #include "velox/exec/NestedLoopJoinBuild.h"
 #include "velox/exec/NestedLoopJoinProbe.h"
+#include "velox/exec/OperatorTraceScan.h"
 #include "velox/exec/OrderBy.h"
 #include "velox/exec/PartitionedOutput.h"
-#include "velox/exec/QueryTraceScan.h"
 #include "velox/exec/RowNumber.h"
 #include "velox/exec/StreamingAggregation.h"
 #include "velox/exec/TableScan.h"
@@ -595,11 +595,10 @@ std::shared_ptr<Driver> DriverFactory::createDriver(
           assignUniqueIdNode->taskUniqueId(),
           assignUniqueIdNode->uniqueIdCounter()));
     } else if (
-        const auto queryReplayScanNode =
-            std::dynamic_pointer_cast<const core::QueryTraceScanNode>(
-                planNode)) {
-      operators.push_back(std::make_unique<trace::QueryTraceScan>(
-          id, ctx.get(), queryReplayScanNode));
+        const auto traceScanNode =
+            std::dynamic_pointer_cast<const core::TraceScanNode>(planNode)) {
+      operators.push_back(std::make_unique<trace::OperatorTraceScan>(
+          id, ctx.get(), traceScanNode));
     } else {
       std::unique_ptr<Operator> extended;
       if (planNode->requiresExchangeClient()) {
diff --git a/velox/exec/Operator.cpp b/velox/exec/Operator.cpp
index 910e4cec76b2e..66f426fd48839 100644
--- a/velox/exec/Operator.cpp
+++ b/velox/exec/Operator.cpp
@@ -21,8 +21,8 @@
 #include "velox/exec/Driver.h"
 #include "velox/exec/HashJoinBridge.h"
 #include "velox/exec/OperatorUtils.h"
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/Task.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/expression/Expr.h"
 
 using facebook::velox::common::testutil::TestValue;
@@ -107,8 +107,8 @@ void Operator::maybeSetReclaimer() {
 }
 
 void Operator::maybeSetTracer() {
-  const auto& queryTraceConfig = operatorCtx_->driverCtx()->traceConfig();
-  if (!queryTraceConfig.has_value()) {
+  const auto& TraceConfig = operatorCtx_->driverCtx()->traceConfig();
+  if (!TraceConfig.has_value()) {
     return;
   }
 
@@ -117,7 +117,7 @@ void Operator::maybeSetTracer() {
   }
 
   const auto nodeId = planNodeId();
-  if (queryTraceConfig->queryNodes.count(nodeId) == 0) {
+  if (TraceConfig->queryNodes.count(nodeId) == 0) {
     return;
   }
 
@@ -134,20 +134,21 @@ void Operator::maybeSetTracer() {
 
   const auto pipelineId = operatorCtx_->driverCtx()->pipelineId;
   const auto driverId = operatorCtx_->driverCtx()->driverId;
-  LOG(INFO) << "Trace data for operator type: " << operatorType()
+  LOG(INFO) << "Trace input for operator type: " << operatorType()
             << ", operator id: " << operatorId() << ", pipeline: " << pipelineId
             << ", driver: " << driverId << ", task: " << taskId();
   const auto opTraceDirPath = fmt::format(
-      "{}/{}/{}/{}/data",
-      queryTraceConfig->queryTraceDir,
+      "{}/{}/{}/{}",
+      TraceConfig->queryTraceDir,
       planNodeId(),
       pipelineId,
       driverId);
   trace::createTraceDirectory(opTraceDirPath);
-  inputTracer_ = std::make_unique<trace::QueryDataWriter>(
+  inputTracer_ = std::make_unique<trace::OperatorTraceWriter>(
+      this,
       opTraceDirPath,
       memory::traceMemoryPool(),
-      queryTraceConfig->updateAndCheckTraceLimitCB);
+      TraceConfig->updateAndCheckTraceLimitCB);
 }
 
 void Operator::traceInput(const RowVectorPtr& input) {
@@ -319,6 +320,16 @@ OperatorStats Operator::stats(bool clear) {
   return stats;
 }
 
+void Operator::close() {
+  input_ = nullptr;
+  results_.clear();
+  recordSpillStats();
+  finishTrace();
+
+  // Release the unused memory reservation on close.
+  operatorCtx_->pool()->release();
+}
+
 vector_size_t Operator::outputBatchRows(
     std::optional<uint64_t> averageRowSize) const {
   const auto& queryConfig = operatorCtx_->task()->queryCtx()->queryConfig();
diff --git a/velox/exec/Operator.h b/velox/exec/Operator.h
index 20561b3e788de..2138fa47e3ccd 100644
--- a/velox/exec/Operator.h
+++ b/velox/exec/Operator.h
@@ -21,7 +21,7 @@
 #include "velox/core/PlanNode.h"
 #include "velox/exec/Driver.h"
 #include "velox/exec/JoinBridge.h"
-#include "velox/exec/QueryDataWriter.h"
+#include "velox/exec/OperatorTraceWriter.h"
 #include "velox/exec/Spiller.h"
 #include "velox/type/Filter.h"
 
@@ -408,7 +408,6 @@ class Operator : public BaseRuntimeStatWriter {
   /// e.g. the first operator in the pipeline.
   virtual void noMoreInput() {
     noMoreInput_ = true;
-    finishTrace();
   }
 
   /// Returns a RowVector with the result columns. Returns nullptr if
@@ -483,13 +482,7 @@ class Operator : public BaseRuntimeStatWriter {
 
   /// Frees all resources associated with 'this'. No other methods
   /// should be called after this.
-  virtual void close() {
-    input_ = nullptr;
-    results_.clear();
-    recordSpillStats();
-    // Release the unused memory reservation on close.
-    operatorCtx_->pool()->release();
-  }
+  virtual void close();
 
   // Returns true if 'this' never has more output rows than input rows.
   virtual bool isFilter() const {
@@ -781,7 +774,7 @@ class Operator : public BaseRuntimeStatWriter {
 
   folly::Synchronized<OperatorStats> stats_;
   folly::Synchronized<common::SpillStats> spillStats_;
-  std::unique_ptr<trace::QueryDataWriter> inputTracer_;
+  std::unique_ptr<trace::OperatorTraceWriter> inputTracer_;
 
   /// Indicates if an operator is under a non-reclaimable execution section.
   /// This prevents the memory arbitrator from reclaiming memory from this
diff --git a/velox/exec/OperatorTraceReader.cpp b/velox/exec/OperatorTraceReader.cpp
new file mode 100644
index 0000000000000..047d98810758f
--- /dev/null
+++ b/velox/exec/OperatorTraceReader.cpp
@@ -0,0 +1,79 @@
+/*
+ * 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 <utility>
+
+#include "velox/exec/OperatorTraceReader.h"
+
+#include "velox/exec/TraceUtil.h"
+
+namespace facebook::velox::exec::trace {
+
+OperatorTraceInputReader::OperatorTraceInputReader(
+    std::string traceDir,
+    RowTypePtr dataType,
+    memory::MemoryPool* pool)
+    : traceDir_(std::move(traceDir)),
+      fs_(filesystems::getFileSystem(traceDir_, nullptr)),
+      dataType_(std::move(dataType)),
+      pool_(pool),
+      inputStream_(getInputStream()) {
+  VELOX_CHECK_NOT_NULL(dataType_);
+  VELOX_CHECK_NOT_NULL(inputStream_);
+}
+
+bool OperatorTraceInputReader::read(RowVectorPtr& batch) const {
+  if (inputStream_->atEnd()) {
+    batch = nullptr;
+    return false;
+  }
+
+  VectorStreamGroup::read(
+      inputStream_.get(), pool_, dataType_, &batch, &readOptions_);
+  return true;
+}
+
+std::unique_ptr<common::FileInputStream>
+OperatorTraceInputReader::getInputStream() const {
+  auto traceFile = fs_->openFileForRead(getOpTraceInputFilePath(traceDir_));
+  // TODO: Make the buffer size configurable.
+  return std::make_unique<common::FileInputStream>(
+      std::move(traceFile), 1 << 20, pool_);
+}
+
+OperatorTraceSummaryReader::OperatorTraceSummaryReader(
+    std::string traceDir,
+    memory::MemoryPool* pool)
+    : traceDir_(std::move(traceDir)),
+      fs_(filesystems::getFileSystem(traceDir_, nullptr)),
+      pool_(pool),
+      summaryFile_(fs_->openFileForRead(getOpTraceSummaryFilePath(traceDir_))) {
+}
+
+OperatorTraceSummary OperatorTraceSummaryReader::read() const {
+  VELOX_CHECK_NOT_NULL(summaryFile_);
+  const auto summaryStr = summaryFile_->pread(0, summaryFile_->size());
+  VELOX_CHECK(!summaryStr.empty());
+
+  folly::dynamic summaryObj = folly::parseJson(summaryStr);
+  OperatorTraceSummary summary;
+  summary.exceededTraceLimit =
+      summaryObj[OperatorTraceTraits::kTraceLimitExceededKey].asBool();
+  summary.peakMemory = summaryObj[OperatorTraceTraits::kPeakMemoryKey].asInt();
+  summary.inputRows = summaryObj[OperatorTraceTraits::kInputRowsKey].asInt();
+  return summary;
+}
+} // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryDataReader.h b/velox/exec/OperatorTraceReader.h
similarity index 62%
rename from velox/exec/QueryDataReader.h
rename to velox/exec/OperatorTraceReader.h
index b5e6d24e011aa..555bdad57b151 100644
--- a/velox/exec/QueryDataReader.h
+++ b/velox/exec/OperatorTraceReader.h
@@ -18,16 +18,16 @@
 
 #include "velox/common/file/FileInputStream.h"
 #include "velox/common/file/FileSystems.h"
-#include "velox/core/PlanNode.h"
-#include "velox/core/QueryCtx.h"
+#include "velox/exec/Trace.h"
 #include "velox/serializers/PrestoSerializer.h"
-#include "velox/vector/VectorStream.h"
 
 namespace facebook::velox::exec::trace {
 
-class QueryDataReader {
+/// Used to read traced input for an operator.
+class OperatorTraceInputReader {
  public:
-  explicit QueryDataReader(
+  /// 'traceDir' specifies the trace directory for a given operator.
+  OperatorTraceInputReader(
       std::string traceDir,
       RowTypePtr dataType,
       memory::MemoryPool* pool);
@@ -37,7 +37,7 @@ class QueryDataReader {
   bool read(RowVectorPtr& batch) const;
 
  private:
-  std::unique_ptr<common::FileInputStream> getDataInputStream() const;
+  std::unique_ptr<common::FileInputStream> getInputStream() const;
 
   const std::string traceDir_;
   const serializer::presto::PrestoVectorSerde::PrestoOptions readOptions_{
@@ -47,6 +47,23 @@ class QueryDataReader {
   const std::shared_ptr<filesystems::FileSystem> fs_;
   const RowTypePtr dataType_;
   memory::MemoryPool* const pool_;
-  const std::unique_ptr<common::FileInputStream> dataStream_;
+  const std::unique_ptr<common::FileInputStream> inputStream_;
+};
+
+/// Used to read summary of a traced operator.
+class OperatorTraceSummaryReader {
+ public:
+  /// 'traceDir' specifies the trace directory for a given operator.
+  OperatorTraceSummaryReader(std::string traceDir, memory::MemoryPool* pool);
+
+  /// Read and return the traced operator summary. The function throws if it
+  /// fails.
+  OperatorTraceSummary read() const;
+
+ private:
+  const std::string traceDir_;
+  const std::shared_ptr<filesystems::FileSystem> fs_;
+  memory::MemoryPool* const pool_;
+  const std::unique_ptr<ReadFile> summaryFile_;
 };
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryTraceScan.cpp b/velox/exec/OperatorTraceScan.cpp
similarity index 62%
rename from velox/exec/QueryTraceScan.cpp
rename to velox/exec/OperatorTraceScan.cpp
index 5cc4f0b5c2230..36fd60f55ca3b 100644
--- a/velox/exec/QueryTraceScan.cpp
+++ b/velox/exec/OperatorTraceScan.cpp
@@ -14,33 +14,32 @@
  * limitations under the License.
  */
 
-#include "velox/exec/QueryTraceScan.h"
+#include "velox/exec/OperatorTraceScan.h"
 
-#include "QueryTraceUtil.h"
+#include "velox/exec/TraceUtil.h"
 
 namespace facebook::velox::exec::trace {
 
-QueryTraceScan::QueryTraceScan(
+OperatorTraceScan::OperatorTraceScan(
     int32_t operatorId,
     DriverCtx* driverCtx,
-    const std::shared_ptr<const core::QueryTraceScanNode>& queryTraceScanNode)
+    const std::shared_ptr<const core::TraceScanNode>& traceScanNode)
     : SourceOperator(
           driverCtx,
-          queryTraceScanNode->outputType(),
+          traceScanNode->outputType(),
           operatorId,
-          queryTraceScanNode->id(),
-          "QueryReplayScan") {
-  const auto dataDir = getDataDir(
-      queryTraceScanNode->traceDir(),
-      driverCtx->pipelineId,
-      driverCtx->driverId);
-  traceReader_ = std::make_unique<QueryDataReader>(
-      dataDir,
-      queryTraceScanNode->outputType(),
+          traceScanNode->id(),
+          "OperatorTraceScan") {
+  traceReader_ = std::make_unique<OperatorTraceInputReader>(
+      getOpTraceDirectory(
+          traceScanNode->traceDir(),
+          driverCtx->pipelineId,
+          driverCtx->driverId),
+      traceScanNode->outputType(),
       memory::MemoryManager::getInstance()->tracePool());
 }
 
-RowVectorPtr QueryTraceScan::getOutput() {
+RowVectorPtr OperatorTraceScan::getOutput() {
   RowVectorPtr batch;
   if (traceReader_->read(batch)) {
     return batch;
@@ -49,7 +48,7 @@ RowVectorPtr QueryTraceScan::getOutput() {
   return nullptr;
 }
 
-bool QueryTraceScan::isFinished() {
+bool OperatorTraceScan::isFinished() {
   return finished_;
 }
 
diff --git a/velox/exec/QueryTraceScan.h b/velox/exec/OperatorTraceScan.h
similarity index 88%
rename from velox/exec/QueryTraceScan.h
rename to velox/exec/OperatorTraceScan.h
index 8bff5bca486d8..1542e43f2c414 100644
--- a/velox/exec/QueryTraceScan.h
+++ b/velox/exec/OperatorTraceScan.h
@@ -18,7 +18,7 @@
 
 #include "velox/core/PlanNode.h"
 #include "velox/exec/Operator.h"
-#include "velox/exec/QueryDataReader.h"
+#include "velox/exec/OperatorTraceReader.h"
 
 namespace facebook::velox::exec::trace {
 /// This is a scan operator for query replay. It uses traced data from a
@@ -38,13 +38,12 @@ namespace facebook::velox::exec::trace {
 /// It can be found from the QueryReplayScanNode. However the pipeline ID and
 /// driver ID are only known during operator creation, so we need to figure out
 /// the input traced data file and the output type dynamically.
-class QueryTraceScan final : public SourceOperator {
+class OperatorTraceScan final : public SourceOperator {
  public:
-  QueryTraceScan(
+  OperatorTraceScan(
       int32_t operatorId,
       DriverCtx* driverCtx,
-      const std::shared_ptr<const core::QueryTraceScanNode>&
-          queryTraceScanNode);
+      const std::shared_ptr<const core::TraceScanNode>& traceScanNode);
 
   RowVectorPtr getOutput() override;
 
@@ -55,7 +54,7 @@ class QueryTraceScan final : public SourceOperator {
   bool isFinished() override;
 
  private:
-  std::unique_ptr<QueryDataReader> traceReader_;
+  std::unique_ptr<OperatorTraceInputReader> traceReader_;
   bool finished_{false};
 };
 
diff --git a/velox/exec/QueryDataWriter.cpp b/velox/exec/OperatorTraceWriter.cpp
similarity index 62%
rename from velox/exec/QueryDataWriter.cpp
rename to velox/exec/OperatorTraceWriter.cpp
index 7ab4e6d02ec0d..decd91febe4de 100644
--- a/velox/exec/QueryDataWriter.cpp
+++ b/velox/exec/OperatorTraceWriter.cpp
@@ -14,31 +14,32 @@
  * limitations under the License.
  */
 
-#include "velox/exec/QueryDataWriter.h"
+#include "velox/exec/OperatorTraceWriter.h"
 
 #include <utility>
-#include "velox/common/base/SpillStats.h"
 #include "velox/common/file/File.h"
 #include "velox/common/file/FileSystems.h"
-#include "velox/exec/QueryTraceTraits.h"
-#include "velox/serializers/PrestoSerializer.h"
+#include "velox/exec/Operator.h"
+#include "velox/exec/Trace.h"
+#include "velox/exec/TraceUtil.h"
 
 namespace facebook::velox::exec::trace {
 
-QueryDataWriter::QueryDataWriter(
-    std::string path,
+OperatorTraceWriter::OperatorTraceWriter(
+    Operator* traceOp,
+    std::string traceDir,
     memory::MemoryPool* pool,
     UpdateAndCheckTraceLimitCB updateAndCheckTraceLimitCB)
-    : dirPath_(std::move(path)),
-      fs_(filesystems::getFileSystem(dirPath_, nullptr)),
+    : traceOp_(traceOp),
+      traceDir_(std::move(traceDir)),
+      fs_(filesystems::getFileSystem(traceDir_, nullptr)),
       pool_(pool),
       updateAndCheckTraceLimitCB_(std::move(updateAndCheckTraceLimitCB)) {
-  dataFile_ = fs_->openFileForWrite(
-      fmt::format("{}/{}", dirPath_, QueryTraceTraits::kDataFileName));
-  VELOX_CHECK_NOT_NULL(dataFile_);
+  traceFile_ = fs_->openFileForWrite(getOpTraceInputFilePath(traceDir_));
+  VELOX_CHECK_NOT_NULL(traceFile_);
 }
 
-void QueryDataWriter::write(const RowVectorPtr& rows) {
+void OperatorTraceWriter::write(const RowVectorPtr& rows) {
   if (FOLLY_UNLIKELY(finished_)) {
     return;
   }
@@ -66,32 +67,35 @@ void QueryDataWriter::write(const RowVectorPtr& rows) {
     finish(true);
     return;
   }
-  dataFile_->append(std::move(iobuf));
+  traceFile_->append(std::move(iobuf));
 }
 
-void QueryDataWriter::finish(bool limitExceeded) {
+void OperatorTraceWriter::finish(bool limitExceeded) {
   if (finished_) {
     return;
   }
 
   VELOX_CHECK_NOT_NULL(
-      dataFile_, "The query data writer has already been finished");
-  dataFile_->close();
-  dataFile_.reset();
+      traceFile_, "The query data writer has already been finished");
+  traceFile_->close();
+  traceFile_.reset();
   batch_.reset();
   writeSummary(limitExceeded);
   finished_ = true;
 }
 
-void QueryDataWriter::writeSummary(bool limitExceeded) const {
-  const auto summaryFilePath =
-      fmt::format("{}/{}", dirPath_, QueryTraceTraits::kDataSummaryFileName);
+void OperatorTraceWriter::writeSummary(bool limitExceeded) const {
+  const auto summaryFilePath = getOpTraceSummaryFilePath(traceDir_);
   const auto file = fs_->openFileForWrite(summaryFilePath);
   folly::dynamic obj = folly::dynamic::object;
   if (dataType_ != nullptr) {
-    obj[QueryTraceTraits::kDataTypeKey] = dataType_->serialize();
+    obj[TraceTraits::kDataTypeKey] = dataType_->serialize();
   }
-  obj[QueryTraceTraits::kTraceLimitExceededKey] = limitExceeded;
+  obj[OperatorTraceTraits::kTraceLimitExceededKey] = limitExceeded;
+  const auto stats = traceOp_->stats(/*clear=*/false);
+  obj[OperatorTraceTraits::kPeakMemoryKey] =
+      stats.memoryStats.peakTotalMemoryReservation;
+  obj[OperatorTraceTraits::kInputRowsKey] = stats.inputPositions;
   file->append(folly::toJson(obj));
   file->close();
 }
diff --git a/velox/exec/QueryDataWriter.h b/velox/exec/OperatorTraceWriter.h
similarity index 80%
rename from velox/exec/QueryDataWriter.h
rename to velox/exec/OperatorTraceWriter.h
index 8e6073dcd3b77..0e5a80e22c3fe 100644
--- a/velox/exec/QueryDataWriter.h
+++ b/velox/exec/OperatorTraceWriter.h
@@ -16,21 +16,27 @@
 
 #pragma once
 
-#include "QueryTraceConfig.h"
+#include "TraceConfig.h"
 #include "velox/common/file/File.h"
 #include "velox/common/file/FileSystems.h"
-#include "velox/exec/QueryTraceTraits.h"
 #include "velox/serializers/PrestoSerializer.h"
 #include "velox/vector/VectorStream.h"
 
+namespace facebook::velox::exec {
+class Operator;
+}
+
 namespace facebook::velox::exec::trace {
 
 /// Used to serialize and write the input vectors from a given operator into a
 /// file.
-class QueryDataWriter {
+class OperatorTraceWriter {
  public:
-  explicit QueryDataWriter(
-      std::string path,
+  /// 'traceOp' is the operator to trace. 'traceDir' specifies the trace
+  /// directory for the operator.
+  explicit OperatorTraceWriter(
+      Operator* traceOp,
+      std::string traceDir,
       memory::MemoryPool* pool,
       UpdateAndCheckTraceLimitCB updateAndCheckTraceLimitCB);
 
@@ -40,7 +46,7 @@ class QueryDataWriter {
   /// Closes the data file and writes out the data summary.
   ///
   /// @param limitExceeded A flag indicates the written data bytes exceed the
-  /// limit causing the 'QueryDataWriter' to finish early.
+  /// limit causing the 'OperatorTraceWriter' to finish early.
   void finish(bool limitExceeded = false);
 
  private:
@@ -49,7 +55,8 @@ class QueryDataWriter {
   // TODO: add more summaries such as number of rows etc.
   void writeSummary(bool limitExceeded = false) const;
 
-  const std::string dirPath_;
+  Operator* const traceOp_;
+  const std::string traceDir_;
   // TODO: make 'useLosslessTimestamp' configuerable.
   const serializer::presto::PrestoVectorSerde::PrestoOptions options_ = {
       true,
@@ -58,9 +65,11 @@ class QueryDataWriter {
   const std::shared_ptr<filesystems::FileSystem> fs_;
   memory::MemoryPool* const pool_;
   const UpdateAndCheckTraceLimitCB updateAndCheckTraceLimitCB_;
-  std::unique_ptr<WriteFile> dataFile_;
+
+  std::unique_ptr<WriteFile> traceFile_;
   TypePtr dataType_;
   std::unique_ptr<VectorStreamGroup> batch_;
+  bool limitExceeded_{false};
   bool finished_{false};
 };
 
diff --git a/velox/exec/PartitionedOutput.cpp b/velox/exec/PartitionedOutput.cpp
index cabff06611466..6b49fe22efa8e 100644
--- a/velox/exec/PartitionedOutput.cpp
+++ b/velox/exec/PartitionedOutput.cpp
@@ -233,8 +233,6 @@ void PartitionedOutput::estimateRowSizes() {
 }
 
 void PartitionedOutput::addInput(RowVectorPtr input) {
-  traceInput(input);
-
   initializeInput(std::move(input));
   initializeDestinations();
   initializeSizeBuffers();
diff --git a/velox/exec/QueryDataReader.cpp b/velox/exec/QueryDataReader.cpp
deleted file mode 100644
index ba08205dd8f25..0000000000000
--- a/velox/exec/QueryDataReader.cpp
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * 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 <utility>
-
-#include "velox/exec/QueryDataReader.h"
-
-#include "velox/common/file/File.h"
-#include "velox/exec/QueryTraceTraits.h"
-
-namespace facebook::velox::exec::trace {
-
-QueryDataReader::QueryDataReader(
-    std::string traceDir,
-    RowTypePtr dataType,
-    memory::MemoryPool* pool)
-    : traceDir_(std::move(traceDir)),
-      fs_(filesystems::getFileSystem(traceDir_, nullptr)),
-      dataType_(std::move(dataType)),
-      pool_(pool),
-      dataStream_(getDataInputStream()) {
-  VELOX_CHECK_NOT_NULL(dataType_);
-  VELOX_CHECK_NOT_NULL(dataStream_);
-}
-
-bool QueryDataReader::read(RowVectorPtr& batch) const {
-  if (dataStream_->atEnd()) {
-    batch = nullptr;
-    return false;
-  }
-
-  VectorStreamGroup::read(
-      dataStream_.get(), pool_, dataType_, &batch, &readOptions_);
-  return true;
-}
-
-std::unique_ptr<common::FileInputStream> QueryDataReader::getDataInputStream()
-    const {
-  auto dataFile = fs_->openFileForRead(
-      fmt::format("{}/{}", traceDir_, QueryTraceTraits::kDataFileName));
-  // TODO: Make the buffer size configurable.
-  return std::make_unique<common::FileInputStream>(
-      std::move(dataFile), 1 << 20, pool_);
-}
-
-} // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryTraceUtil.cpp b/velox/exec/QueryTraceUtil.cpp
deleted file mode 100644
index f3a339eec080a..0000000000000
--- a/velox/exec/QueryTraceUtil.cpp
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * 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/exec/QueryTraceUtil.h"
-
-#include <folly/json.h>
-
-#include "velox/common/base/Exceptions.h"
-#include "velox/common/file/File.h"
-#include "velox/common/file/FileSystems.h"
-#include "velox/exec/QueryTraceTraits.h"
-
-namespace facebook::velox::exec::trace {
-
-void createTraceDirectory(const std::string& traceDir) {
-  try {
-    const auto fs = filesystems::getFileSystem(traceDir, nullptr);
-    if (fs->exists(traceDir)) {
-      fs->rmdir(traceDir);
-    }
-    fs->mkdir(traceDir);
-  } catch (const std::exception& e) {
-    VELOX_FAIL(
-        "Failed to create trace directory '{}' with error: {}",
-        traceDir,
-        e.what());
-  }
-}
-
-std::vector<std::string> getTaskIds(
-    const std::string& traceDir,
-    const std::shared_ptr<filesystems::FileSystem>& fs) {
-  VELOX_USER_CHECK(fs->exists(traceDir), "{} dose not exist", traceDir);
-  try {
-    const auto taskDirs = fs->list(traceDir);
-    std::vector<std::string> taskIds;
-    for (const auto& taskDir : taskDirs) {
-      std::vector<std::string> pathNodes;
-      folly::split("/", taskDir, pathNodes);
-      taskIds.emplace_back(std::move(pathNodes.back()));
-    }
-    return taskIds;
-  } catch (const std::exception& e) {
-    VELOX_FAIL(
-        "Failed to list the directory '{}' with error: {}", traceDir, e.what());
-  }
-}
-
-folly::dynamic getMetadata(
-    const std::string& metadataFile,
-    const std::shared_ptr<filesystems::FileSystem>& fs) {
-  try {
-    const auto file = fs->openFileForRead(metadataFile);
-    VELOX_CHECK_NOT_NULL(file);
-    const auto metadata = file->pread(0, file->size());
-    VELOX_USER_CHECK(!metadata.empty());
-    return folly::parseJson(metadata);
-  } catch (const std::exception& e) {
-    VELOX_FAIL(
-        "Failed to get the query metadata from '{}' with error: {}",
-        metadataFile,
-        e.what());
-  }
-}
-
-RowTypePtr getDataType(
-    const core::PlanNodePtr& tracedPlan,
-    const std::string& tracedNodeId,
-    size_t sourceIndex) {
-  const auto* traceNode = core::PlanNode::findFirstNode(
-      tracedPlan.get(), [&tracedNodeId](const core::PlanNode* node) {
-        return node->id() == tracedNodeId;
-      });
-  VELOX_CHECK_NOT_NULL(
-      traceNode,
-      "traced node id {} not found in the traced plan",
-      tracedNodeId);
-  return traceNode->sources().at(sourceIndex)->outputType();
-}
-
-uint8_t getNumDrivers(
-    const std::string& rootDir,
-    const std::string& taskId,
-    const std::string& nodeId,
-    int32_t pipelineId,
-    const std::shared_ptr<filesystems::FileSystem>& fs) {
-  const auto traceDir =
-      fmt::format("{}/{}/{}/{}", rootDir, taskId, nodeId, pipelineId);
-  const auto driverDirs = fs->list(traceDir);
-  return driverDirs.size();
-}
-
-std::string
-getDataDir(const std::string& traceDir, int pipelineId, int driverId) {
-  return fmt::format("{}/{}/{}/data", traceDir, pipelineId, driverId);
-}
-
-} // namespace facebook::velox::exec::trace
diff --git a/velox/exec/TableWriter.cpp b/velox/exec/TableWriter.cpp
index 3d289621f38a0..4899e4b4ebf89 100644
--- a/velox/exec/TableWriter.cpp
+++ b/velox/exec/TableWriter.cpp
@@ -126,7 +126,7 @@ void TableWriter::addInput(RowVectorPtr input) {
   if (input->size() == 0) {
     return;
   }
-  traceInput(input);
+
   std::vector<VectorPtr> mappedChildren;
   mappedChildren.reserve(inputMapping_.size());
   for (const auto i : inputMapping_) {
diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
index 22e68d6c475d1..bf09b3a0730bd 100644
--- a/velox/exec/Task.cpp
+++ b/velox/exec/Task.cpp
@@ -30,8 +30,8 @@
 #include "velox/exec/NestedLoopJoinBuild.h"
 #include "velox/exec/OperatorUtils.h"
 #include "velox/exec/OutputBufferManager.h"
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/Task.h"
+#include "velox/exec/TraceUtil.h"
 
 using facebook::velox::common::testutil::TestValue;
 
@@ -304,7 +304,7 @@ Task::Task(
         dynamic_cast<const folly::InlineLikeExecutor*>(queryCtx_->executor()));
   }
 
-  maybeInitQueryTrace();
+  maybeInitTrace();
 }
 
 Task::~Task() {
@@ -2888,7 +2888,7 @@ std::shared_ptr<ExchangeClient> Task::getExchangeClientLocked(
   return exchangeClients_[pipelineId];
 }
 
-std::optional<trace::QueryTraceConfig> Task::maybeMakeTraceConfig() const {
+std::optional<trace::TraceConfig> Task::maybeMakeTraceConfig() const {
   const auto& queryConfig = queryCtx_->queryConfig();
   if (!queryConfig.queryTraceEnabled()) {
     return std::nullopt;
@@ -2906,40 +2906,41 @@ std::optional<trace::QueryTraceConfig> Task::maybeMakeTraceConfig() const {
     return std::nullopt;
   }
 
-  const auto traceDir =
-      fmt::format("{}/{}", queryConfig.queryTraceDir(), taskId_);
-  const auto queryTraceNodes = queryConfig.queryTraceNodeIds();
-  if (queryTraceNodes.empty()) {
+  const auto traceDir = trace::getTaskTraceDirectory(
+      queryConfig.queryTraceDir(), queryCtx_->queryId(), taskId_);
+  const auto traceNodes = queryConfig.queryTraceNodeIds();
+  if (traceNodes.empty()) {
     LOG(INFO) << "Trace metadata for task: " << taskId_;
-    return trace::QueryTraceConfig(traceDir);
+    return trace::TraceConfig(traceDir);
   }
 
   std::vector<std::string> nodes;
-  folly::split(',', queryTraceNodes, nodes);
+  folly::split(',', traceNodes, nodes);
   std::unordered_set<std::string> nodeSet(nodes.begin(), nodes.end());
   VELOX_CHECK_EQ(nodeSet.size(), nodes.size());
   LOG(INFO) << "Trace data for task " << taskId_ << " with plan nodes "
-            << queryTraceNodes;
+            << traceNodes;
 
   trace::UpdateAndCheckTraceLimitCB updateAndCheckTraceLimitCB =
       [this](uint64_t bytes) {
         return queryCtx_->updateTracedBytesAndCheckLimit(bytes);
       };
-  return trace::QueryTraceConfig(
+  return trace::TraceConfig(
       std::move(nodeSet),
       traceDir,
       std::move(updateAndCheckTraceLimitCB),
       queryConfig.queryTraceTaskRegExp());
 }
 
-void Task::maybeInitQueryTrace() {
+void Task::maybeInitTrace() {
   if (!traceConfig_) {
     return;
   }
 
   trace::createTraceDirectory(traceConfig_->queryTraceDir);
-  const auto queryMetadatWriter = std::make_unique<trace::QueryMetadataWriter>(
-      traceConfig_->queryTraceDir, memory::traceMemoryPool());
+  const auto queryMetadatWriter =
+      std::make_unique<trace::TaskTraceMetadataWriter>(
+          traceConfig_->queryTraceDir, memory::traceMemoryPool());
   queryMetadatWriter->write(queryCtx_, planFragment_.planNode);
 }
 
diff --git a/velox/exec/Task.h b/velox/exec/Task.h
index 526ab710400cb..682c62148f00d 100644
--- a/velox/exec/Task.h
+++ b/velox/exec/Task.h
@@ -21,11 +21,11 @@
 #include "velox/exec/LocalPartition.h"
 #include "velox/exec/MemoryReclaimer.h"
 #include "velox/exec/MergeSource.h"
-#include "velox/exec/QueryMetadataWriter.h"
-#include "velox/exec/QueryTraceConfig.h"
 #include "velox/exec/Split.h"
 #include "velox/exec/TaskStats.h"
 #include "velox/exec/TaskStructs.h"
+#include "velox/exec/TaskTraceMetadataWriter.h"
+#include "velox/exec/TraceConfig.h"
 #include "velox/vector/ComplexVector.h"
 
 namespace facebook::velox::exec {
@@ -139,7 +139,7 @@ class Task : public std::enable_shared_from_this<Task> {
   }
 
   /// Returns query trace config if specified.
-  const std::optional<trace::QueryTraceConfig>& queryTraceConfig() const {
+  const std::optional<trace::TraceConfig>& TraceConfig() const {
     return traceConfig_;
   }
 
@@ -980,11 +980,11 @@ class Task : public std::enable_shared_from_this<Task> {
       int32_t pipelineId) const;
 
   // Builds the query trace config.
-  std::optional<trace::QueryTraceConfig> maybeMakeTraceConfig() const;
+  std::optional<trace::TraceConfig> maybeMakeTraceConfig() const;
 
   // Create a 'QueryMetadtaWriter' to trace the query metadata if the query
   // trace enabled.
-  void maybeInitQueryTrace();
+  void maybeInitTrace();
 
   // Universally unique identifier of the task. Used to identify the task when
   // calling TaskListener.
@@ -1004,7 +1004,7 @@ class Task : public std::enable_shared_from_this<Task> {
 
   core::PlanFragment planFragment_;
 
-  const std::optional<trace::QueryTraceConfig> traceConfig_;
+  const std::optional<trace::TraceConfig> traceConfig_;
 
   // Hook in the system wide task list.
   TaskListEntry taskListEntry_;
diff --git a/velox/exec/QueryMetadataReader.cpp b/velox/exec/TaskTraceMetadataReader.cpp
similarity index 73%
rename from velox/exec/QueryMetadataReader.cpp
rename to velox/exec/TaskTraceMetadataReader.cpp
index f9fdc9ec2e80f..11d97edb117c6 100644
--- a/velox/exec/QueryMetadataReader.cpp
+++ b/velox/exec/TaskTraceMetadataReader.cpp
@@ -14,44 +14,40 @@
  * limitations under the License.
  */
 
-#include "velox/exec/QueryMetadataReader.h"
+#include "velox/exec/TaskTraceMetadataReader.h"
 
-#include "velox/common/file/File.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/core/PlanNode.h"
-#include "velox/exec/QueryTraceTraits.h"
-#include "velox/exec/QueryTraceUtil.h"
+#include "velox/exec/Trace.h"
+#include "velox/exec/TraceUtil.h"
 
 namespace facebook::velox::exec::trace {
 
-QueryMetadataReader::QueryMetadataReader(
+TaskTraceMetadataReader::TaskTraceMetadataReader(
     std::string traceDir,
     memory::MemoryPool* pool)
     : traceDir_(std::move(traceDir)),
       fs_(filesystems::getFileSystem(traceDir_, nullptr)),
-      metaFilePath_(fmt::format(
-          "{}/{}",
-          traceDir_,
-          QueryTraceTraits::kQueryMetaFileName)),
+      traceFilePath_(getTaskTraceMetaFilePath(traceDir_)),
       pool_(pool) {
   VELOX_CHECK_NOT_NULL(fs_);
-  VELOX_CHECK(fs_->exists(metaFilePath_));
+  VELOX_CHECK(fs_->exists(traceFilePath_));
 }
 
-void QueryMetadataReader::read(
+void TaskTraceMetadataReader::read(
     std::unordered_map<std::string, std::string>& queryConfigs,
     std::unordered_map<
         std::string,
         std::unordered_map<std::string, std::string>>& connectorProperties,
     core::PlanNodePtr& queryPlan) const {
-  folly::dynamic metaObj = getMetadata(metaFilePath_, fs_);
-  const auto& queryConfigObj = metaObj[QueryTraceTraits::kQueryConfigKey];
+  folly::dynamic metaObj = getTaskMetadata(traceFilePath_, fs_);
+  const auto& queryConfigObj = metaObj[TraceTraits::kQueryConfigKey];
   for (const auto& [key, value] : queryConfigObj.items()) {
     queryConfigs[key.asString()] = value.asString();
   }
 
   const auto& connectorPropertiesObj =
-      metaObj[QueryTraceTraits::kConnectorPropertiesKey];
+      metaObj[TraceTraits::kConnectorPropertiesKey];
   for (const auto& [connectorId, configs] : connectorPropertiesObj.items()) {
     const auto connectorIdStr = connectorId.asString();
     connectorProperties[connectorIdStr] = {};
@@ -61,6 +57,6 @@ void QueryMetadataReader::read(
   }
 
   queryPlan = ISerializable::deserialize<core::PlanNode>(
-      metaObj[QueryTraceTraits::kPlanNodeKey], pool_);
+      metaObj[TraceTraits::kPlanNodeKey], pool_);
 }
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryMetadataReader.h b/velox/exec/TaskTraceMetadataReader.h
similarity index 84%
rename from velox/exec/QueryMetadataReader.h
rename to velox/exec/TaskTraceMetadataReader.h
index 71217653d4dc3..df3abcffc3b69 100644
--- a/velox/exec/QueryMetadataReader.h
+++ b/velox/exec/TaskTraceMetadataReader.h
@@ -18,13 +18,11 @@
 
 #include "velox/common/file/FileSystems.h"
 #include "velox/core/PlanNode.h"
-#include "velox/core/QueryCtx.h"
-#include "velox/vector/VectorStream.h"
 
 namespace facebook::velox::exec::trace {
-class QueryMetadataReader {
+class TaskTraceMetadataReader {
  public:
-  explicit QueryMetadataReader(std::string traceDir, memory::MemoryPool* pool);
+  TaskTraceMetadataReader(std::string traceDir, memory::MemoryPool* pool);
 
   void read(
       std::unordered_map<std::string, std::string>& queryConfigs,
@@ -36,7 +34,7 @@ class QueryMetadataReader {
  private:
   const std::string traceDir_;
   const std::shared_ptr<filesystems::FileSystem> fs_;
-  const std::string metaFilePath_;
+  const std::string traceFilePath_;
   memory::MemoryPool* const pool_;
 };
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryMetadataWriter.cpp b/velox/exec/TaskTraceMetadataWriter.cpp
similarity index 75%
rename from velox/exec/QueryMetadataWriter.cpp
rename to velox/exec/TaskTraceMetadataWriter.cpp
index a8bbe6633c668..9407491878930 100644
--- a/velox/exec/QueryMetadataWriter.cpp
+++ b/velox/exec/TaskTraceMetadataWriter.cpp
@@ -14,30 +14,27 @@
  * limitations under the License.
  */
 
-#include "velox/exec/QueryMetadataWriter.h"
-#include "velox/common/config/Config.h"
+#include "velox/exec/TaskTraceMetadataWriter.h"
 #include "velox/common/file/File.h"
 #include "velox/core/PlanNode.h"
 #include "velox/core/QueryCtx.h"
-#include "velox/exec/QueryTraceTraits.h"
+#include "velox/exec/Trace.h"
+#include "velox/exec/TraceUtil.h"
 
 namespace facebook::velox::exec::trace {
 
-QueryMetadataWriter::QueryMetadataWriter(
+TaskTraceMetadataWriter::TaskTraceMetadataWriter(
     std::string traceDir,
     memory::MemoryPool* pool)
     : traceDir_(std::move(traceDir)),
       fs_(filesystems::getFileSystem(traceDir_, nullptr)),
-      metaFilePath_(fmt::format(
-          "{}/{}",
-          traceDir_,
-          QueryTraceTraits::kQueryMetaFileName)),
+      traceFilePath_(getTaskTraceMetaFilePath(traceDir_)),
       pool_(pool) {
   VELOX_CHECK_NOT_NULL(fs_);
-  VELOX_CHECK(!fs_->exists(metaFilePath_));
+  VELOX_CHECK(!fs_->exists(traceFilePath_));
 }
 
-void QueryMetadataWriter::write(
+void TaskTraceMetadataWriter::write(
     const std::shared_ptr<core::QueryCtx>& queryCtx,
     const core::PlanNodePtr& planNode) {
   VELOX_CHECK(!finished_, "Query metadata can only be written once");
@@ -59,12 +56,12 @@ void QueryMetadataWriter::write(
   }
 
   folly::dynamic metaObj = folly::dynamic::object;
-  metaObj[QueryTraceTraits::kQueryConfigKey] = queryConfigObj;
-  metaObj[QueryTraceTraits::kConnectorPropertiesKey] = connectorPropertiesObj;
-  metaObj[QueryTraceTraits::kPlanNodeKey] = planNode->serialize();
+  metaObj[TraceTraits::kQueryConfigKey] = queryConfigObj;
+  metaObj[TraceTraits::kConnectorPropertiesKey] = connectorPropertiesObj;
+  metaObj[TraceTraits::kPlanNodeKey] = planNode->serialize();
 
   const auto metaStr = folly::toJson(metaObj);
-  const auto file = fs_->openFileForWrite(metaFilePath_);
+  const auto file = fs_->openFileForWrite(traceFilePath_);
   file->append(metaStr);
   file->close();
 }
diff --git a/velox/exec/QueryMetadataWriter.h b/velox/exec/TaskTraceMetadataWriter.h
similarity index 85%
rename from velox/exec/QueryMetadataWriter.h
rename to velox/exec/TaskTraceMetadataWriter.h
index f2cc661238514..5b3105104dd6c 100644
--- a/velox/exec/QueryMetadataWriter.h
+++ b/velox/exec/TaskTraceMetadataWriter.h
@@ -19,12 +19,11 @@
 #include "velox/common/file/FileSystems.h"
 #include "velox/core/PlanNode.h"
 #include "velox/core/QueryCtx.h"
-#include "velox/vector/VectorStream.h"
 
 namespace facebook::velox::exec::trace {
-class QueryMetadataWriter {
+class TaskTraceMetadataWriter {
  public:
-  explicit QueryMetadataWriter(std::string traceDir, memory::MemoryPool* pool);
+  TaskTraceMetadataWriter(std::string traceDir, memory::MemoryPool* pool);
 
   void write(
       const std::shared_ptr<core::QueryCtx>& queryCtx,
@@ -33,7 +32,7 @@ class QueryMetadataWriter {
  private:
   const std::string traceDir_;
   const std::shared_ptr<filesystems::FileSystem> fs_;
-  const std::string metaFilePath_;
+  const std::string traceFilePath_;
   memory::MemoryPool* const pool_;
   bool finished_{false};
 };
diff --git a/velox/exec/Trace.cpp b/velox/exec/Trace.cpp
new file mode 100644
index 0000000000000..4d4d86df0e969
--- /dev/null
+++ b/velox/exec/Trace.cpp
@@ -0,0 +1,33 @@
+/*
+ * 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/exec/Trace.h"
+
+#include <fmt/core.h>
+
+#include "velox/common/base/SuccinctPrinter.h"
+
+namespace facebook::velox::exec::trace {
+
+std::string OperatorTraceSummary::toString() const {
+  return fmt::format(
+      "opType {}, exceededTraceLimit {}, inputRows {}, peakMemory {}",
+      opType,
+      exceededTraceLimit,
+      inputRows,
+      succinctBytes(peakMemory));
+}
+} // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryTraceTraits.h b/velox/exec/Trace.h
similarity index 61%
rename from velox/exec/QueryTraceTraits.h
rename to velox/exec/Trace.h
index ad817115b4902..28c9250029c01 100644
--- a/velox/exec/QueryTraceTraits.h
+++ b/velox/exec/Trace.h
@@ -20,16 +20,33 @@
 
 namespace facebook::velox::exec::trace {
 /// Defines the shared constants used by query trace implementation.
-struct QueryTraceTraits {
+struct TraceTraits {
   static inline const std::string kPlanNodeKey = "planNode";
   static inline const std::string kQueryConfigKey = "queryConfig";
   static inline const std::string kDataTypeKey = "rowType";
   static inline const std::string kConnectorPropertiesKey =
       "connectorProperties";
+
+  static inline const std::string kTaskMetaFileName = "task_trace_meta.json";
+  static inline const std::string kOpSummaryFileName = "op_trace_summary.json";
+  static inline const std::string kOpInputFileName = "op_input_trace.data";
+};
+
+struct OperatorTraceTraits {
+  /// Keys for operator trace summary file.
+  static inline const std::string kOpTypeKey = "opType";
   static inline const std::string kTraceLimitExceededKey = "traceLimitExceeded";
+  static inline const std::string kPeakMemoryKey = "peakMemory";
+  static inline const std::string kInputRowsKey = "inputRows";
+};
+
+/// Contains the summary of an operator trace.
+struct OperatorTraceSummary {
+  std::string opType;
+  bool exceededTraceLimit{false};
+  uint64_t inputRows{0};
+  uint64_t peakMemory{0};
 
-  static inline const std::string kQueryMetaFileName = "query_meta.json";
-  static inline const std::string kDataSummaryFileName = "data_summary.json";
-  static inline const std::string kDataFileName = "trace.data";
+  std::string toString() const;
 };
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryTraceConfig.cpp b/velox/exec/TraceConfig.cpp
similarity index 88%
rename from velox/exec/QueryTraceConfig.cpp
rename to velox/exec/TraceConfig.cpp
index d43479d646089..7e90f1edec5b9 100644
--- a/velox/exec/QueryTraceConfig.cpp
+++ b/velox/exec/TraceConfig.cpp
@@ -14,13 +14,13 @@
  * limitations under the License.
  */
 
-#include "velox/exec/QueryTraceConfig.h"
+#include "velox/exec/TraceConfig.h"
 
 #include <utility>
 
 namespace facebook::velox::exec::trace {
 
-QueryTraceConfig::QueryTraceConfig(
+TraceConfig::TraceConfig(
     std::unordered_set<std::string> _queryNodeIds,
     std::string _queryTraceDir,
     UpdateAndCheckTraceLimitCB _updateAndCheckTraceLimitCB,
@@ -30,8 +30,8 @@ QueryTraceConfig::QueryTraceConfig(
       updateAndCheckTraceLimitCB(std::move(_updateAndCheckTraceLimitCB)),
       taskRegExp(std::move(_taskRegExp)) {}
 
-QueryTraceConfig::QueryTraceConfig(std::string _queryTraceDir)
-    : QueryTraceConfig(
+TraceConfig::TraceConfig(std::string _queryTraceDir)
+    : TraceConfig(
           std::unordered_set<std::string>{},
           std::move(_queryTraceDir),
           [](uint64_t) { return false; },
diff --git a/velox/exec/QueryTraceConfig.h b/velox/exec/TraceConfig.h
similarity index 92%
rename from velox/exec/QueryTraceConfig.h
rename to velox/exec/TraceConfig.h
index 200e5cb52ce4a..2b086b39d4b38 100644
--- a/velox/exec/QueryTraceConfig.h
+++ b/velox/exec/TraceConfig.h
@@ -27,7 +27,7 @@ namespace facebook::velox::exec::trace {
 /// bytes exceed the set limit otherwise return false.
 using UpdateAndCheckTraceLimitCB = std::function<bool(uint64_t)>;
 
-struct QueryTraceConfig {
+struct TraceConfig {
   /// Target query trace nodes.
   std::unordered_set<std::string> queryNodes;
   /// Base dir of query trace.
@@ -36,14 +36,14 @@ struct QueryTraceConfig {
   /// The trace task regexp.
   std::string taskRegExp;
 
-  QueryTraceConfig(
+  TraceConfig(
       std::unordered_set<std::string> _queryNodeIds,
       std::string _queryTraceDir,
       UpdateAndCheckTraceLimitCB _updateAndCheckTraceLimitCB,
       std::string _taskRegExp);
 
-  QueryTraceConfig(std::string _queryTraceDir);
+  TraceConfig(std::string _queryTraceDir);
 
-  QueryTraceConfig() = default;
+  TraceConfig() = default;
 };
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/TraceUtil.cpp b/velox/exec/TraceUtil.cpp
new file mode 100644
index 0000000000000..9d109ff85b770
--- /dev/null
+++ b/velox/exec/TraceUtil.cpp
@@ -0,0 +1,168 @@
+/*
+ * 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/exec/TraceUtil.h"
+
+#include <folly/json.h>
+
+#include "velox/common/base/Exceptions.h"
+#include "velox/common/file/File.h"
+#include "velox/common/file/FileSystems.h"
+#include "velox/exec/Trace.h"
+
+namespace facebook::velox::exec::trace {
+
+void createTraceDirectory(const std::string& traceDir) {
+  try {
+    const auto fs = filesystems::getFileSystem(traceDir, nullptr);
+    if (fs->exists(traceDir)) {
+      fs->rmdir(traceDir);
+    }
+    fs->mkdir(traceDir);
+  } catch (const std::exception& e) {
+    VELOX_FAIL(
+        "Failed to create trace directory '{}' with error: {}",
+        traceDir,
+        e.what());
+  }
+}
+
+std::string getQueryTraceDirectory(
+    const std::string& traceDir,
+    const std::string& queryId) {
+  return fmt::format("{}/{}", traceDir, queryId);
+}
+
+std::string getTaskTraceDirectory(
+    const std::string& traceDir,
+    const Task& task) {
+  return getTaskTraceDirectory(
+      traceDir, task.queryCtx()->queryId(), task.taskId());
+}
+
+std::string getTaskTraceDirectory(
+    const std::string& traceDir,
+    const std::string& queryId,
+    const std::string& taskId) {
+  return fmt::format(
+      "{}/{}", getQueryTraceDirectory(traceDir, queryId), taskId);
+}
+
+std::string getTaskTraceMetaFilePath(const std::string& taskTraceDir) {
+  return fmt::format("{}/{}", taskTraceDir, TraceTraits::kTaskMetaFileName);
+}
+
+std::string getNodeTraceDirectory(
+    const std::string& taskTraceDir,
+    const std::string& nodeId) {
+  return fmt::format("{}/{}", taskTraceDir, nodeId);
+}
+
+std::string getPipelineTraceDirectory(
+    const std::string& nodeTraceDir,
+    uint32_t pipelineId) {
+  return fmt::format("{}/{}", nodeTraceDir, pipelineId);
+}
+
+std::string getOpTraceDirectory(
+    const std::string& taskTraceDir,
+    const std::string& nodeId,
+    uint32_t pipelineId,
+    uint32_t driverId) {
+  return getOpTraceDirectory(
+      getNodeTraceDirectory(taskTraceDir, nodeId), pipelineId, driverId);
+}
+
+std::string getOpTraceDirectory(
+    const std::string& nodeTraceDir,
+    int pipelineId,
+    int driverId) {
+  return fmt::format("{}/{}/{}", nodeTraceDir, pipelineId, driverId);
+}
+
+std::string getOpTraceInputFilePath(const std::string& opTraceDir) {
+  return fmt::format("{}/{}", opTraceDir, TraceTraits::kOpInputFileName);
+}
+
+std::string getOpTraceSummaryFilePath(const std::string& opTraceDir) {
+  return fmt::format("{}/{}", opTraceDir, TraceTraits::kOpSummaryFileName);
+}
+
+std::vector<std::string> getTaskIds(
+    const std::string& traceDir,
+    const std::string& queryId,
+    const std::shared_ptr<filesystems::FileSystem>& fs) {
+  const auto queryTraceDir = getQueryTraceDirectory(traceDir, queryId);
+  VELOX_USER_CHECK(
+      fs->exists(queryTraceDir), "{} dose not exist", queryTraceDir);
+  try {
+    const auto taskDirs = fs->list(queryTraceDir);
+    std::vector<std::string> taskIds;
+    for (const auto& taskDir : taskDirs) {
+      std::vector<std::string> pathNodes;
+      folly::split("/", taskDir, pathNodes);
+      taskIds.emplace_back(std::move(pathNodes.back()));
+    }
+    return taskIds;
+  } catch (const std::exception& e) {
+    VELOX_FAIL(
+        "Failed to list the directory '{}' with error: {}",
+        queryTraceDir,
+        e.what());
+  }
+}
+
+folly::dynamic getTaskMetadata(
+    const std::string& taskMetaFilePath,
+    const std::shared_ptr<filesystems::FileSystem>& fs) {
+  try {
+    const auto file = fs->openFileForRead(taskMetaFilePath);
+    VELOX_CHECK_NOT_NULL(file);
+    const auto taskMeta = file->pread(0, file->size());
+    VELOX_USER_CHECK(!taskMeta.empty());
+    return folly::parseJson(taskMeta);
+  } catch (const std::exception& e) {
+    VELOX_FAIL(
+        "Failed to get the query metadata from '{}' with error: {}",
+        taskMetaFilePath,
+        e.what());
+  }
+}
+
+RowTypePtr getDataType(
+    const core::PlanNodePtr& tracedPlan,
+    const std::string& tracedNodeId,
+    size_t sourceIndex) {
+  const auto* traceNode = core::PlanNode::findFirstNode(
+      tracedPlan.get(), [&tracedNodeId](const core::PlanNode* node) {
+        return node->id() == tracedNodeId;
+      });
+  VELOX_CHECK_NOT_NULL(
+      traceNode,
+      "traced node id {} not found in the traced plan",
+      tracedNodeId);
+  return traceNode->sources().at(sourceIndex)->outputType();
+}
+
+size_t getNumDrivers(
+    const std::string& nodeTraceDir,
+    uint32_t pipelineId,
+    const std::shared_ptr<filesystems::FileSystem>& fs) {
+  const auto pipelineDir = getPipelineTraceDirectory(nodeTraceDir, pipelineId);
+  const auto driverDirs = fs->list(pipelineDir);
+  return driverDirs.size();
+}
+} // namespace facebook::velox::exec::trace
diff --git a/velox/exec/QueryTraceUtil.h b/velox/exec/TraceUtil.h
similarity index 52%
rename from velox/exec/QueryTraceUtil.h
rename to velox/exec/TraceUtil.h
index 633c0bc27b2a8..a7ad162261bc5 100644
--- a/velox/exec/QueryTraceUtil.h
+++ b/velox/exec/TraceUtil.h
@@ -20,6 +20,7 @@
 #include <vector>
 #include "velox/common/file/FileSystems.h"
 #include "velox/core/PlanNode.h"
+#include "velox/exec/Task.h"
 #include "velox/type/Type.h"
 
 #include <folly/dynamic.h>
@@ -29,6 +30,52 @@ namespace facebook::velox::exec::trace {
 /// Creates a directory to store the query trace metdata and data.
 void createTraceDirectory(const std::string& traceDir);
 
+/// Returns the trace directory for a given query.
+std::string getQueryTraceDirectory(
+    const std::string& traceDir,
+    const std::string& queryId);
+
+/// Returns the trace directory for a given query task.
+std::string getTaskTraceDirectory(
+    const std::string& traceDir,
+    const Task& task);
+
+std::string getTaskTraceDirectory(
+    const std::string& traceDir,
+    const std::string& queryId,
+    const std::string& taskId);
+
+/// Returns the file path for a given task's metadata file for trace.
+std::string getTaskTraceMetaFilePath(const std::string& taskTraceDir);
+
+/// Returns the trace directory for a given traced plan node.
+std::string getNodeTraceDirectory(
+    const std::string& taskTraceDir,
+    const std::string& nodeId);
+
+/// Returns the trace directory for a given traced plan node.
+std::string getPipelineTraceDirectory(
+    const std::string& nodeTraceDir,
+    uint32_t pipelineId);
+
+/// Returns the trace directory for a given traced operator.
+std::string getOpTraceDirectory(
+    const std::string& taskTraceDir,
+    const std::string& nodeId,
+    uint32_t pipelineId,
+    uint32_t driverId);
+
+std::string getOpTraceDirectory(
+    const std::string& nodeTraceDir,
+    int pipelineId,
+    int driverId);
+
+/// Returns the file path for a given operator's traced input file.
+std::string getOpTraceInputFilePath(const std::string& opTraceDir);
+
+/// Returns the file path for a given operator's traced input file.
+std::string getOpTraceSummaryFilePath(const std::string& opTraceDir);
+
 /// Extracts the input data type for the trace scan operator. The function first
 /// uses the traced node id to find traced operator's plan node from the traced
 /// plan fragment. Then it uses the specified source node index to find the
@@ -48,28 +95,30 @@ RowTypePtr getDataType(
     const std::string& tracedNodeId,
     size_t sourceIndex = 0);
 
+/// List driver directories for a given plan node pipeline.
+std::vector<std::string> listDrivers(
+    const std::string& nodeTraceDir,
+    uint32_t pipelineId,
+    const std::shared_ptr<filesystems::FileSystem>& fs);
+
 /// Extracts the number of drivers by listing the number of sub-directors under
-/// the trace directory for a given pipeline.
-uint8_t getNumDrivers(
-    const std::string& rootDir,
-    const std::string& taskId,
-    const std::string& nodeId,
-    int32_t pipelineId,
+/// the trace directory for a given pipeline. 'nodeTraceDir' is the trace
+/// directory of the plan node.
+size_t getNumDrivers(
+    const std::string& nodeTraceDir,
+    uint32_t pipelineId,
     const std::shared_ptr<filesystems::FileSystem>& fs);
 
-/// Extracts task ids of the query tracing by listing the trace directory.
+/// Extracts task ids of the query tracing by listing the query trace directory.
+/// 'traceDir' is the root trace directory. 'queryId' is the query id.
 std::vector<std::string> getTaskIds(
     const std::string& traceDir,
+    const std::string& queryId,
     const std::shared_ptr<filesystems::FileSystem>& fs);
 
 /// Gets the metadata from a given task metadata file which includes query plan,
 /// configs and connector properties.
-folly::dynamic getMetadata(
-    const std::string& metadataFile,
+folly::dynamic getTaskMetadata(
+    const std::string& taskMetaFilePath,
     const std::shared_ptr<filesystems::FileSystem>& fs);
-
-/// Gets the traced data directory. 'traceaDir' is the trace directory for a
-/// given plan node, which is $traceRoot/$taskId/$nodeId.
-std::string
-getDataDir(const std::string& traceDir, int pipelineId, int driverId);
 } // namespace facebook::velox::exec::trace
diff --git a/velox/exec/tests/CMakeLists.txt b/velox/exec/tests/CMakeLists.txt
index 8904fdc206482..f099f7940686c 100644
--- a/velox/exec/tests/CMakeLists.txt
+++ b/velox/exec/tests/CMakeLists.txt
@@ -60,14 +60,15 @@ add_executable(
   MultiFragmentTest.cpp
   NestedLoopJoinTest.cpp
   OrderByTest.cpp
+  QperatorTraceTest.cpp
   OutputBufferManagerTest.cpp
   PartitionedOutputTest.cpp
   PlanNodeSerdeTest.cpp
   PlanNodeToStringTest.cpp
   PrefixSortTest.cpp
   PrintPlanWithStatsTest.cpp
-  QueryTraceTest.cpp
   ProbeOperatorStateTest.cpp
+  TraceUtilTest.cpp
   RoundRobinPartitionFunctionTest.cpp
   RowContainerTest.cpp
   RowNumberTest.cpp
diff --git a/velox/exec/tests/QueryTraceTest.cpp b/velox/exec/tests/OperatorTraceTest.cpp
similarity index 72%
rename from velox/exec/tests/QueryTraceTest.cpp
rename to velox/exec/tests/OperatorTraceTest.cpp
index 71b49da483a03..91c59f689bbdc 100644
--- a/velox/exec/tests/QueryTraceTest.cpp
+++ b/velox/exec/tests/OperatorTraceTest.cpp
@@ -16,26 +16,25 @@
 
 #include <folly/init/Init.h>
 #include <gtest/gtest.h>
-#include <algorithm>
 #include <memory>
 
+#include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
+#include "velox/exec/OperatorTraceReader.h"
 #include "velox/exec/PartitionFunction.h"
-#include "velox/exec/QueryDataReader.h"
-#include "velox/exec/QueryDataWriter.h"
-#include "velox/exec/QueryMetadataReader.h"
-#include "velox/exec/QueryMetadataWriter.h"
-#include "velox/exec/QueryTraceUtil.h"
-#include "velox/exec/tests/utils/ArbitratorTestUtil.h"
+#include "velox/exec/TaskTraceMetadataReader.h"
+#include "velox/exec/Trace.h"
+#include "velox/exec/TraceUtil.h"
+#include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/HiveConnectorTestBase.h"
+#include "velox/exec/tests/utils/PlanBuilder.h"
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
 #include "velox/serializers/PrestoSerializer.h"
-#include "velox/vector/tests/utils/VectorTestBase.h"
 
 using namespace facebook::velox::exec::test;
 
 namespace facebook::velox::exec::trace::test {
-class QueryTracerTest : public HiveConnectorTestBase {
+class OperatorTraceTest : public HiveConnectorTestBase {
  protected:
   static void SetUpTestCase() {
     memory::MemoryManager::testingSetInstance({});
@@ -55,6 +54,11 @@ class QueryTracerTest : public HiveConnectorTestBase {
     registerPartitionFunctionSerDe();
   }
 
+  void SetUp() override {
+    HiveConnectorTestBase::SetUp();
+    dataType_ = ROW({"a", "b", "c"}, {BIGINT(), BIGINT(), BIGINT()});
+  }
+
   static VectorFuzzer::Options getFuzzerOptions() {
     return VectorFuzzer::Options{
         .vectorSize = 16,
@@ -65,7 +69,7 @@ class QueryTracerTest : public HiveConnectorTestBase {
     };
   }
 
-  QueryTracerTest() : vectorFuzzer_{getFuzzerOptions(), pool_.get()} {
+  OperatorTraceTest() : vectorFuzzer_{getFuzzerOptions(), pool_.get()} {
     filesystems::registerLocalFileSystem();
   }
 
@@ -99,32 +103,63 @@ class QueryTracerTest : public HiveConnectorTestBase {
     return true;
   }
 
+  std::unique_ptr<DriverCtx> driverCtx() {
+    return std::make_unique<DriverCtx>(nullptr, 0, 0, 0, 0);
+  }
+
+  RowTypePtr dataType_;
   VectorFuzzer vectorFuzzer_;
 };
 
-TEST_F(QueryTracerTest, emptyTrace) {
-  const auto outputDir = TempDirectoryPath::create();
-  auto writer = trace::QueryDataWriter(
-      outputDir->getPath(), pool(), [&](uint64_t bytes) { return false; });
-  writer.finish();
-
-  const auto fs = filesystems::getFileSystem(outputDir->getPath(), nullptr);
-  const auto summaryFile = fs->openFileForRead(fmt::format(
-      "{}/{}", outputDir->getPath(), QueryTraceTraits::kDataSummaryFileName));
-  const auto summary = summaryFile->pread(0, summaryFile->size());
-  ASSERT_FALSE(summary.empty());
-  folly::dynamic obj = folly::parseJson(summary);
-  ASSERT_EQ(obj[QueryTraceTraits::kTraceLimitExceededKey].asBool(), false);
+TEST_F(OperatorTraceTest, emptyTrace) {
+  auto input = vectorFuzzer_.fuzzInputRow(dataType_);
+  input->childAt(0) =
+      makeFlatVector<int64_t>(input->size(), [](auto /*unused*/) { return 0; });
+  createDuckDbTable({input});
+
+  std::string planNodeId;
+  auto traceDirPath = TempDirectoryPath::create();
+  auto plan = PlanBuilder()
+                  .values({input})
+                  .filter("a > 0")
+                  .singleAggregation({"a"}, {"count(1)"})
+                  .capturePlanNodeId(planNodeId)
+                  .planNode();
+
+  const auto task =
+      AssertQueryBuilder(duckDbQueryRunner_)
+          .plan(plan)
+          .config(core::QueryConfig::kQueryTraceEnabled, true)
+          .config(core::QueryConfig::kQueryTraceDir, traceDirPath->getPath())
+          .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
+          .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
+          .config(core::QueryConfig::kQueryTraceNodeIds, planNodeId)
+          .assertResults("SELECT a, count(1) FROM tmp WHERE a > 0 GROUP BY 1");
+
+  const auto taskTraceDir =
+      getTaskTraceDirectory(traceDirPath->getPath(), *task);
+  const auto opTraceDir = getOpTraceDirectory(
+      taskTraceDir,
+      planNodeId,
+      /*pipelineId=*/0,
+      /*driverId=*/0);
+  const auto summary = OperatorTraceSummaryReader(opTraceDir, pool()).read();
+  ASSERT_EQ(summary.exceededTraceLimit, false);
+  ASSERT_EQ(summary.inputRows, 0);
+  // The hash aggregation operator might allocate memory when prepare output
+  // buffer even though there is no output. We could optimize out this later if
+  // needs.
+  ASSERT_GT(summary.peakMemory, 0);
 }
 
-TEST_F(QueryTracerTest, traceData) {
-  const auto rowType = ROW({"a", "b", "c"}, {BIGINT(), BIGINT(), BIGINT()});
+TEST_F(OperatorTraceTest, traceData) {
   std::vector<RowVectorPtr> inputVectors;
   constexpr auto numBatch = 5;
   inputVectors.reserve(numBatch);
   for (auto i = 0; i < numBatch; ++i) {
-    inputVectors.push_back(vectorFuzzer_.fuzzInputFlatRow(rowType));
+    inputVectors.push_back(vectorFuzzer_.fuzzInputFlatRow(dataType_));
   }
+  createDuckDbTable(inputVectors);
 
   struct {
     uint64_t maxTracedBytes;
@@ -143,37 +178,57 @@ TEST_F(QueryTracerTest, traceData) {
 
   for (const auto& testData : testSettings) {
     SCOPED_TRACE(testData.debugString());
-    const auto outputDir = TempDirectoryPath::create();
-    // Ensure the writer only write one batch.
-    uint64_t numTracedBytes{0};
-    auto writer = trace::QueryDataWriter(
-        outputDir->getPath(), pool(), [&](uint64_t bytes) {
-          numTracedBytes += bytes;
-          return numTracedBytes >= testData.maxTracedBytes;
-        });
-    for (auto i = 0; i < numBatch; ++i) {
-      writer.write(inputVectors[i]);
-    }
-    writer.finish();
-
-    const auto fs = filesystems::getFileSystem(outputDir->getPath(), nullptr);
-    const auto summaryFile = fs->openFileForRead(fmt::format(
-        "{}/{}", outputDir->getPath(), QueryTraceTraits::kDataSummaryFileName));
-    const auto summary = summaryFile->pread(0, summaryFile->size());
-    ASSERT_FALSE(summary.empty());
-    folly::dynamic obj = folly::parseJson(summary);
-    ASSERT_EQ(
-        obj[QueryTraceTraits::kTraceLimitExceededKey].asBool(),
-        testData.limitExceeded);
 
+    std::string planNodeId;
+    auto traceDirPath = TempDirectoryPath::create();
+    auto plan = PlanBuilder()
+                    .values(inputVectors)
+                    .singleAggregation({"a"}, {"count(1)"})
+                    .capturePlanNodeId(planNodeId)
+                    .planNode();
+
+    const auto task =
+        AssertQueryBuilder(duckDbQueryRunner_)
+            .plan(plan)
+            .config(core::QueryConfig::kQueryTraceEnabled, true)
+            .config(core::QueryConfig::kQueryTraceDir, traceDirPath->getPath())
+            .config(
+                core::QueryConfig::kQueryTraceMaxBytes, testData.maxTracedBytes)
+            .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
+            .config(core::QueryConfig::kQueryTraceNodeIds, planNodeId)
+            .assertResults("SELECT a, count(1) FROM tmp GROUP BY 1");
+
+    const auto fs =
+        filesystems::getFileSystem(traceDirPath->getPath(), nullptr);
+    const auto taskTraceDir =
+        getTaskTraceDirectory(traceDirPath->getPath(), *task);
+    const auto opTraceDir = getOpTraceDirectory(
+        taskTraceDir,
+        planNodeId,
+        /*pipelineId=*/0,
+        /*driverId=*/0);
+    const auto summaryFilePath = getOpTraceSummaryFilePath(opTraceDir);
+    const auto dataFilePath = getOpTraceInputFilePath(opTraceDir);
     if (testData.maxTracedBytes == 0) {
-      const auto dataFile = fs->openFileForRead(fmt::format(
-          "{}/{}", outputDir->getPath(), QueryTraceTraits::kDataFileName));
-      ASSERT_EQ(dataFile->size(), 0);
+      ASSERT_FALSE(fs->exists(summaryFilePath));
+      ASSERT_FALSE(fs->exists(dataFilePath));
       continue;
+    } else {
+      ASSERT_TRUE(fs->exists(summaryFilePath));
+      ASSERT_TRUE(fs->exists(dataFilePath));
     }
 
-    const auto reader = QueryDataReader(outputDir->getPath(), rowType, pool());
+    const auto summary = OperatorTraceSummaryReader(opTraceDir, pool()).read();
+    ASSERT_EQ(summary.exceededTraceLimit, testData.limitExceeded);
+
+    ASSERT_GT(summary.peakMemory, 0);
+    if (testData.limitExceeded) {
+      ASSERT_GT(summary.inputRows, 0);
+    } else {
+      ASSERT_EQ(summary.inputRows, testData.numTracedBatches * 16);
+    }
+
+    const auto reader = OperatorTraceInputReader(opTraceDir, dataType_, pool());
     RowVectorPtr actual;
     size_t numOutputVectors{0};
     while (reader.read(actual)) {
@@ -189,7 +244,7 @@ TEST_F(QueryTracerTest, traceData) {
   }
 }
 
-TEST_F(QueryTracerTest, traceMetadata) {
+TEST_F(OperatorTraceTest, traceMetadata) {
   const auto rowType =
       ROW({"c0", "c1", "c2", "c3", "c4", "c5"},
           {BIGINT(), SMALLINT(), TINYINT(), VARCHAR(), VARCHAR(), VARCHAR()});
@@ -234,14 +289,14 @@ TEST_F(QueryTracerTest, traceMetadata) {
       executor_.get(),
       core::QueryConfig(expectedQueryConfigs),
       expectedConnectorProperties);
-  auto writer = trace::QueryMetadataWriter(outputDir->getPath(), pool());
+  auto writer = trace::TaskTraceMetadataWriter(outputDir->getPath(), pool());
   writer.write(queryCtx, planNode);
 
   std::unordered_map<std::string, std::string> acutalQueryConfigs;
   std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
       actualConnectorProperties;
   core::PlanNodePtr actualQueryPlan;
-  auto reader = trace::QueryMetadataReader(outputDir->getPath(), pool());
+  auto reader = trace::TaskTraceMetadataReader(outputDir->getPath(), pool());
   reader.read(acutalQueryConfigs, actualConnectorProperties, actualQueryPlan);
 
   ASSERT_TRUE(isSamePlan(actualQueryPlan, planNode));
@@ -262,7 +317,7 @@ TEST_F(QueryTracerTest, traceMetadata) {
   }
 }
 
-TEST_F(QueryTracerTest, task) {
+TEST_F(OperatorTraceTest, task) {
   const auto rowType =
       ROW({"c0", "c1", "c2", "c3", "c4", "c5"},
           {BIGINT(), SMALLINT(), TINYINT(), VARCHAR(), VARCHAR(), VARCHAR()});
@@ -345,7 +400,8 @@ TEST_F(QueryTracerTest, task) {
     }
     ASSERT_EQ(actaulDirs.size(), testData.expectedNumDirs);
     ASSERT_EQ(actaulDirs.at(0), expectedDir);
-    const auto taskIds = getTaskIds(outputDir->getPath(), fs);
+    const auto taskIds =
+        getTaskIds(outputDir->getPath(), task->queryCtx()->queryId(), fs);
     ASSERT_EQ(taskIds.size(), testData.expectedNumDirs);
     ASSERT_EQ(taskIds.at(0), task->taskId());
 
@@ -354,7 +410,7 @@ TEST_F(QueryTracerTest, task) {
         unordered_map<std::string, std::unordered_map<std::string, std::string>>
             actualConnectorProperties;
     core::PlanNodePtr actualQueryPlan;
-    auto reader = trace::QueryMetadataReader(expectedDir, pool());
+    auto reader = trace::TaskTraceMetadataReader(expectedDir, pool());
     reader.read(acutalQueryConfigs, actualConnectorProperties, actualQueryPlan);
 
     ASSERT_TRUE(isSamePlan(actualQueryPlan, planNode));
@@ -377,7 +433,7 @@ TEST_F(QueryTracerTest, task) {
   }
 }
 
-TEST_F(QueryTracerTest, error) {
+TEST_F(OperatorTraceTest, error) {
   const auto planNode = PlanBuilder().values({}).planNode();
   const auto expectedQueryConfigs =
       std::unordered_map<std::string, std::string>{
@@ -393,54 +449,7 @@ TEST_F(QueryTracerTest, error) {
       "Query trace enabled but the trace dir is not set");
 }
 
-TEST_F(QueryTracerTest, traceDir) {
-  const auto outputDir = TempDirectoryPath::create();
-  const auto rootDir = outputDir->getPath();
-  const auto fs = filesystems::getFileSystem(rootDir, nullptr);
-  auto dir1 = fmt::format("{}/{}", outputDir->getPath(), "t1");
-  trace::createTraceDirectory(dir1);
-  ASSERT_TRUE(fs->exists(dir1));
-
-  auto dir2 = fmt::format("{}/{}", dir1, "t1_1");
-  trace::createTraceDirectory(dir2);
-  ASSERT_TRUE(fs->exists(dir2));
-
-  // It will remove the old dir1 along with its subdir when created the dir1
-  // again.
-  trace::createTraceDirectory(dir1);
-  ASSERT_TRUE(fs->exists(dir1));
-  ASSERT_FALSE(fs->exists(dir2));
-
-  const auto parentDir = fmt::format("{}/{}", outputDir->getPath(), "p");
-  fs->mkdir(parentDir);
-
-  constexpr auto numThreads = 5;
-  std::vector<std::thread> traceThreads;
-  traceThreads.reserve(numThreads);
-  std::mutex mutex;
-  std::set<std::string> expectedDirs;
-  for (int i = 0; i < numThreads; ++i) {
-    traceThreads.emplace_back([&, i]() {
-      const auto dir = fmt::format("{}/s{}", parentDir, i);
-      trace::createTraceDirectory(dir);
-      std::lock_guard<std::mutex> l(mutex);
-      expectedDirs.insert(dir);
-    });
-  }
-
-  for (auto& traceThread : traceThreads) {
-    traceThread.join();
-  }
-
-  const auto actualDirs = fs->list(parentDir);
-  ASSERT_EQ(actualDirs.size(), numThreads);
-  ASSERT_EQ(actualDirs.size(), expectedDirs.size());
-  for (const auto& dir : actualDirs) {
-    ASSERT_EQ(expectedDirs.count(dir), 1);
-  }
-}
-
-TEST_F(QueryTracerTest, traceTableWriter) {
+TEST_F(OperatorTraceTest, traceTableWriter) {
   const auto rowType = ROW({"a", "b", "c"}, {BIGINT(), BIGINT(), BIGINT()});
   std::vector<RowVectorPtr> inputVectors;
   constexpr auto numBatch = 5;
@@ -490,8 +499,8 @@ TEST_F(QueryTracerTest, traceTableWriter) {
         .config(core::QueryConfig::kQueryTraceNodeIds, "1")
         .copyResults(pool(), task);
 
-    const auto metadataDir = fmt::format("{}/{}", traceRoot, task->taskId());
-    const auto fs = filesystems::getFileSystem(metadataDir, nullptr);
+    const auto taskTraceDir = getTaskTraceDirectory(traceRoot, *task);
+    const auto fs = filesystems::getFileSystem(taskTraceDir, nullptr);
 
     if (testData.taskRegExpr == "wrong id") {
       ASSERT_FALSE(fs->exists(traceRoot));
@@ -499,34 +508,24 @@ TEST_F(QueryTracerTest, traceTableWriter) {
     }
 
     // Query metadta file should exist.
-    const auto traceMetaFile = fmt::format(
-        "{}/{}/{}",
-        traceRoot,
-        task->taskId(),
-        trace::QueryTraceTraits::kQueryMetaFileName);
-    ASSERT_TRUE(fs->exists(traceMetaFile));
+    const auto traceMetaFilePath = getTaskTraceMetaFilePath(taskTraceDir);
+    ASSERT_TRUE(fs->exists(traceMetaFilePath));
 
-    const auto dataDir =
-        fmt::format("{}/{}/{}", traceRoot, task->taskId(), "1/0/0/data");
+    const auto opTraceDir = getOpTraceDirectory(taskTraceDir, "1", 0, 0);
 
     // Query data tracing disabled.
     if (testData.maxTracedBytes == 0) {
-      ASSERT_FALSE(fs->exists(dataDir));
+      ASSERT_FALSE(fs->exists(opTraceDir));
       continue;
     }
 
-    ASSERT_EQ(fs->list(dataDir).size(), 2);
-    // Check data summaries.
-    const auto summaryFile = fs->openFileForRead(
-        fmt::format("{}/{}", dataDir, QueryTraceTraits::kDataSummaryFileName));
-    const auto summary = summaryFile->pread(0, summaryFile->size());
-    ASSERT_FALSE(summary.empty());
-    folly::dynamic obj = folly::parseJson(summary);
-    ASSERT_EQ(
-        obj[QueryTraceTraits::kTraceLimitExceededKey].asBool(),
-        testData.limitExceeded);
+    ASSERT_EQ(fs->list(opTraceDir).size(), 2);
+
+    const auto summary = OperatorTraceSummaryReader(opTraceDir, pool()).read();
+    ASSERT_EQ(summary.exceededTraceLimit, testData.limitExceeded);
 
-    const auto reader = trace::QueryDataReader(dataDir, rowType, pool());
+    const auto reader =
+        trace::OperatorTraceInputReader(opTraceDir, rowType, pool());
     RowVectorPtr actual;
     size_t numOutputVectors{0};
     while (reader.read(actual)) {
diff --git a/velox/exec/tests/TraceUtilTest.cpp b/velox/exec/tests/TraceUtilTest.cpp
new file mode 100644
index 0000000000000..39693b6a5bd96
--- /dev/null
+++ b/velox/exec/tests/TraceUtilTest.cpp
@@ -0,0 +1,90 @@
+/*
+ * 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 <folly/init/Init.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/common/file/FileSystems.h"
+#include "velox/exec/OperatorTraceReader.h"
+#include "velox/exec/PartitionFunction.h"
+#include "velox/exec/TaskTraceMetadataReader.h"
+#include "velox/exec/Trace.h"
+#include "velox/exec/TraceUtil.h"
+#include "velox/exec/tests/utils/AssertQueryBuilder.h"
+#include "velox/exec/tests/utils/HiveConnectorTestBase.h"
+#include "velox/exec/tests/utils/PlanBuilder.h"
+#include "velox/exec/tests/utils/TempDirectoryPath.h"
+#include "velox/serializers/PrestoSerializer.h"
+
+using namespace facebook::velox::exec::test;
+
+namespace facebook::velox::exec::trace::test {
+class TraceUtilTest : public testing::Test {
+ protected:
+  static void SetUpTestCase() {
+    filesystems::registerLocalFileSystem();
+  }
+};
+
+TEST_F(TraceUtilTest, traceDir) {
+  const auto outputDir = TempDirectoryPath::create();
+  const auto rootDir = outputDir->getPath();
+  const auto fs = filesystems::getFileSystem(rootDir, nullptr);
+  auto dir1 = fmt::format("{}/{}", outputDir->getPath(), "t1");
+  trace::createTraceDirectory(dir1);
+  ASSERT_TRUE(fs->exists(dir1));
+
+  auto dir2 = fmt::format("{}/{}", dir1, "t1_1");
+  trace::createTraceDirectory(dir2);
+  ASSERT_TRUE(fs->exists(dir2));
+
+  // It will remove the old dir1 along with its subdir when created the dir1
+  // again.
+  trace::createTraceDirectory(dir1);
+  ASSERT_TRUE(fs->exists(dir1));
+  ASSERT_FALSE(fs->exists(dir2));
+
+  const auto parentDir = fmt::format("{}/{}", outputDir->getPath(), "p");
+  fs->mkdir(parentDir);
+
+  constexpr auto numThreads = 5;
+  std::vector<std::thread> traceThreads;
+  traceThreads.reserve(numThreads);
+  std::mutex mutex;
+  std::set<std::string> expectedDirs;
+  for (int i = 0; i < numThreads; ++i) {
+    traceThreads.emplace_back([&, i]() {
+      const auto dir = fmt::format("{}/s{}", parentDir, i);
+      trace::createTraceDirectory(dir);
+      std::lock_guard<std::mutex> l(mutex);
+      expectedDirs.insert(dir);
+    });
+  }
+
+  for (auto& traceThread : traceThreads) {
+    traceThread.join();
+  }
+
+  const auto actualDirs = fs->list(parentDir);
+  ASSERT_EQ(actualDirs.size(), numThreads);
+  ASSERT_EQ(actualDirs.size(), expectedDirs.size());
+  for (const auto& dir : actualDirs) {
+    ASSERT_EQ(expectedDirs.count(dir), 1);
+  }
+}
+} // namespace facebook::velox::exec::trace::test
diff --git a/velox/exec/tests/utils/OperatorTestBase.h b/velox/exec/tests/utils/OperatorTestBase.h
index 4e5def101cfaa..079bfc2ff6b0e 100644
--- a/velox/exec/tests/utils/OperatorTestBase.h
+++ b/velox/exec/tests/utils/OperatorTestBase.h
@@ -24,8 +24,6 @@
 #include "velox/exec/HashProbe.h"
 #include "velox/exec/tests/utils/QueryAssertions.h"
 #include "velox/parse/ExpressionsParser.h"
-#include "velox/type/Variant.h"
-#include "velox/vector/FlatVector.h"
 #include "velox/vector/tests/utils/VectorMaker.h"
 #include "velox/vector/tests/utils/VectorTestBase.h"
 
diff --git a/velox/exec/tests/utils/PlanBuilder.cpp b/velox/exec/tests/utils/PlanBuilder.cpp
index 57d12d57bd137..bb620d0282047 100644
--- a/velox/exec/tests/utils/PlanBuilder.cpp
+++ b/velox/exec/tests/utils/PlanBuilder.cpp
@@ -223,7 +223,7 @@ PlanBuilder& PlanBuilder::values(
 PlanBuilder& PlanBuilder::traceScan(
     const std::string& traceNodeDir,
     const RowTypePtr& outputType) {
-  planNode_ = std::make_shared<core::QueryTraceScanNode>(
+  planNode_ = std::make_shared<core::TraceScanNode>(
       nextPlanNodeId(), traceNodeDir, outputType);
   return *this;
 }
diff --git a/velox/tool/trace/AggregationReplayer.cpp b/velox/tool/trace/AggregationReplayer.cpp
index 54041bf8bbcbd..1c397461a8a5f 100644
--- a/velox/tool/trace/AggregationReplayer.cpp
+++ b/velox/tool/trace/AggregationReplayer.cpp
@@ -15,7 +15,7 @@
  */
 
 #include "velox/tool/trace/AggregationReplayer.h"
-#include "velox/exec/QueryDataReader.h"
+#include "velox/exec/OperatorTraceReader.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
 
 using namespace facebook::velox;
diff --git a/velox/tool/trace/AggregationReplayer.h b/velox/tool/trace/AggregationReplayer.h
index 21ce2c409acf6..9a14688a694a1 100644
--- a/velox/tool/trace/AggregationReplayer.h
+++ b/velox/tool/trace/AggregationReplayer.h
@@ -24,13 +24,15 @@ namespace facebook::velox::tool::trace {
 class AggregationReplayer : public OperatorReplayerBase {
  public:
   AggregationReplayer(
-      const std::string& rootDir,
+      const std::string& traceDir,
+      const std::string& queryId,
       const std::string& taskId,
       const std::string& nodeId,
       const int32_t pipelineId,
       const std::string& operatorType)
       : OperatorReplayerBase(
-            rootDir,
+            traceDir,
+            queryId,
             taskId,
             nodeId,
             pipelineId,
diff --git a/velox/tool/trace/OperatorReplayerBase.cpp b/velox/tool/trace/OperatorReplayerBase.cpp
index b54ad323a4471..3351fb91f3230 100644
--- a/velox/tool/trace/OperatorReplayerBase.cpp
+++ b/velox/tool/trace/OperatorReplayerBase.cpp
@@ -16,11 +16,9 @@
 
 #include <folly/json.h>
 
-#include "velox/common/serialization/Serializable.h"
 #include "velox/core/PlanNode.h"
-#include "velox/exec/QueryMetadataReader.h"
-#include "velox/exec/QueryTraceTraits.h"
-#include "velox/exec/QueryTraceUtil.h"
+#include "velox/exec/TaskTraceMetadataReader.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
 #include "velox/tool/trace/OperatorReplayerBase.h"
@@ -29,30 +27,31 @@ using namespace facebook::velox;
 
 namespace facebook::velox::tool::trace {
 OperatorReplayerBase::OperatorReplayerBase(
-    std::string rootDir,
+    std::string traceDir,
+    std::string queryId,
     std::string taskId,
     std::string nodeId,
     int32_t pipelineId,
     std::string operatorType)
-    : taskId_(std::move(taskId)),
+    : queryId_(std::string(queryId)),
+      taskId_(std::move(taskId)),
       nodeId_(std::move(nodeId)),
       pipelineId_(pipelineId),
       operatorType_(std::move(operatorType)),
-      rootDir_(std::move(rootDir)),
-      taskDir_(fmt::format("{}/{}", rootDir_, taskId_)),
-      nodeDir_(fmt::format("{}/{}", taskDir_, nodeId_)) {
-  VELOX_USER_CHECK(!rootDir_.empty());
+      taskTraceDir_(
+          exec::trace::getTaskTraceDirectory(traceDir, queryId_, taskId_)),
+      nodeTraceDir_(exec::trace::getNodeTraceDirectory(taskTraceDir_, nodeId_)),
+      fs_(filesystems::getFileSystem(taskTraceDir_, nullptr)),
+      maxDrivers_(exec::trace::getNumDrivers(nodeTraceDir_, pipelineId_, fs_)) {
+  VELOX_USER_CHECK(!taskTraceDir_.empty());
   VELOX_USER_CHECK(!taskId_.empty());
   VELOX_USER_CHECK(!nodeId_.empty());
   VELOX_USER_CHECK_GE(pipelineId_, 0);
   VELOX_USER_CHECK(!operatorType_.empty());
-  const auto metadataReader = exec::trace::QueryMetadataReader(
-      taskDir_, memory::MemoryManager::getInstance()->tracePool());
-  metadataReader.read(queryConfigs_, connectorConfigs_, planFragment_);
+  const auto taskMetaReader = exec::trace::TaskTraceMetadataReader(
+      taskTraceDir_, memory::MemoryManager::getInstance()->tracePool());
+  taskMetaReader.read(queryConfigs_, connectorConfigs_, planFragment_);
   queryConfigs_[core::QueryConfig::kQueryTraceEnabled] = "false";
-  fs_ = filesystems::getFileSystem(rootDir_, nullptr);
-  maxDrivers_ =
-      exec::trace::getNumDrivers(rootDir_, taskId_, nodeId_, pipelineId_, fs_);
 }
 
 RowVectorPtr OperatorReplayerBase::run() {
@@ -69,7 +68,8 @@ core::PlanNodePtr OperatorReplayerBase::createPlan() const {
       planFragment_.get(),
       [this](const core::PlanNode* node) { return node->id() == nodeId_; });
   return exec::test::PlanBuilder()
-      .traceScan(nodeDir_, exec::trace::getDataType(planFragment_, nodeId_))
+      .traceScan(
+          nodeTraceDir_, exec::trace::getDataType(planFragment_, nodeId_))
       .addNode(replayNodeFactory(replayNode))
       .planNode();
 }
diff --git a/velox/tool/trace/OperatorReplayerBase.h b/velox/tool/trace/OperatorReplayerBase.h
index 863a8f5f80f69..50e31fdaf02f5 100644
--- a/velox/tool/trace/OperatorReplayerBase.h
+++ b/velox/tool/trace/OperatorReplayerBase.h
@@ -27,7 +27,8 @@ namespace facebook::velox::tool::trace {
 class OperatorReplayerBase {
  public:
   OperatorReplayerBase(
-      std::string rootDir,
+      std::string traceDir,
+      std::string queryId,
       std::string taskId,
       std::string nodeId,
       int32_t pipelineId,
@@ -50,21 +51,21 @@ class OperatorReplayerBase {
 
   core::PlanNodePtr createPlan() const;
 
+  const std::string queryId_;
   const std::string taskId_;
   const std::string nodeId_;
   const int32_t pipelineId_;
   const std::string operatorType_;
 
-  const std::string rootDir_;
-  const std::string taskDir_;
-  const std::string nodeDir_;
+  const std::string taskTraceDir_;
+  const std::string nodeTraceDir_;
+  const std::shared_ptr<filesystems::FileSystem> fs_;
+  const int32_t maxDrivers_;
 
   std::unordered_map<std::string, std::string> queryConfigs_;
   std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
       connectorConfigs_;
   core::PlanNodePtr planFragment_;
-  std::shared_ptr<filesystems::FileSystem> fs_;
-  int32_t maxDrivers_{1};
 
  private:
   std::function<core::PlanNodePtr(std::string, core::PlanNodePtr)>
diff --git a/velox/tool/trace/PartitionedOutputReplayer.cpp b/velox/tool/trace/PartitionedOutputReplayer.cpp
index 79874caefc93a..56b4c0a145eb7 100644
--- a/velox/tool/trace/PartitionedOutputReplayer.cpp
+++ b/velox/tool/trace/PartitionedOutputReplayer.cpp
@@ -18,7 +18,7 @@
 
 #include "velox/common/memory/Memory.h"
 #include "velox/exec/PartitionedOutput.h"
-#include "velox/exec/QueryTraceUtil.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
 #include "velox/tool/trace/PartitionedOutputReplayer.h"
@@ -105,13 +105,20 @@ void consumeAllData(
 }
 
 PartitionedOutputReplayer::PartitionedOutputReplayer(
-    const std::string& rootDir,
+    const std::string& traceDir,
+    const std::string& queryId,
     const std::string& taskId,
     const std::string& nodeId,
     const int32_t pipelineId,
     const std::string& operatorType,
     const ConsumerCallBack& consumerCb)
-    : OperatorReplayerBase(rootDir, taskId, nodeId, pipelineId, operatorType),
+    : OperatorReplayerBase(
+          traceDir,
+          queryId,
+          taskId,
+          nodeId,
+          pipelineId,
+          operatorType),
       originalNode_(dynamic_cast<const core::PartitionedOutputNode*>(
           core::PlanNode::findFirstNode(
               planFragment_.get(),
diff --git a/velox/tool/trace/PartitionedOutputReplayer.h b/velox/tool/trace/PartitionedOutputReplayer.h
index 61610aa031ccb..a90c222dbf657 100644
--- a/velox/tool/trace/PartitionedOutputReplayer.h
+++ b/velox/tool/trace/PartitionedOutputReplayer.h
@@ -41,7 +41,8 @@ class PartitionedOutputReplayer final : public OperatorReplayerBase {
       std::function<void(uint32_t, std::unique_ptr<folly::IOBuf>)>;
 
   PartitionedOutputReplayer(
-      const std::string& rootDir,
+      const std::string& traceDir,
+      const std::string& queryId,
       const std::string& taskId,
       const std::string& nodeId,
       const int32_t pipelineId,
diff --git a/velox/tool/trace/QueryReplayer.cpp b/velox/tool/trace/QueryReplayer.cpp
index a0680d30a387a..21c5b47bb8be3 100644
--- a/velox/tool/trace/QueryReplayer.cpp
+++ b/velox/tool/trace/QueryReplayer.cpp
@@ -32,12 +32,12 @@
 #include "velox/dwio/dwrf/RegisterDwrfWriter.h"
 #include "velox/dwio/parquet/RegisterParquetReader.h"
 #include "velox/dwio/parquet/RegisterParquetWriter.h"
+#include "velox/exec/OperatorTraceReader.h"
 #include "velox/exec/PartitionFunction.h"
-#include "velox/exec/QueryTraceUtil.h"
-#include "velox/expression/Expr.h"
+#include "velox/exec/TaskTraceMetadataReader.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h"
 #include "velox/functions/prestosql/registration/RegistrationFunctions.h"
-#include "velox/parse/ExpressionsParser.h"
 #include "velox/parse/TypeResolver.h"
 #include "velox/tool/trace/AggregationReplayer.h"
 #include "velox/tool/trace/OperatorReplayerBase.h"
@@ -57,6 +57,7 @@ DEFINE_bool(
     "It also print the query metadata including query configs, connectors "
     "properties, and query plan in JSON format.");
 DEFINE_bool(short_summary, false, "Only show number of tasks and task ids");
+DEFINE_string(query_id, "", "Specify the target query id which must be set");
 DEFINE_string(
     task_id,
     "",
@@ -77,7 +78,7 @@ DEFINE_double(
 using namespace facebook::velox;
 
 namespace {
-void init() {
+void init(folly::IOThreadPoolExecutor* ioExecutor) {
   memory::initializeMemoryManager({});
   filesystems::registerLocalFileSystem();
   filesystems::registerS3FileSystem();
@@ -109,10 +110,6 @@ void init() {
   aggregate::prestosql::registerAllAggregateFunctions();
   parse::registerTypeResolver();
 
-  // TODO: make it configurable.
-  const auto ioExecutor = std::make_unique<folly::IOThreadPoolExecutor>(
-      std::thread::hardware_concurrency() *
-      FLAGS_hiveConnectorExecutorHwMultiplier);
   connector::registerConnectorFactory(
       std::make_shared<connector::hive::HiveConnectorFactory>());
   const auto hiveConnector =
@@ -120,7 +117,7 @@ void init() {
           "test-hive",
           std::make_shared<config::ConfigBase>(
               std::unordered_map<std::string, std::string>()),
-          ioExecutor.get());
+          ioExecutor);
   connector::registerConnector(hiveConnector);
 }
 
@@ -129,6 +126,7 @@ std::unique_ptr<tool::trace::OperatorReplayerBase> createReplayer() {
   if (FLAGS_operator_type == "TableWriter") {
     replayer = std::make_unique<tool::trace::TableWriterReplayer>(
         FLAGS_root_dir,
+        FLAGS_query_id,
         FLAGS_task_id,
         FLAGS_node_id,
         FLAGS_pipeline_id,
@@ -137,6 +135,7 @@ std::unique_ptr<tool::trace::OperatorReplayerBase> createReplayer() {
   } else if (FLAGS_operator_type == "Aggregation") {
     replayer = std::make_unique<tool::trace::AggregationReplayer>(
         FLAGS_root_dir,
+        FLAGS_query_id,
         FLAGS_task_id,
         FLAGS_node_id,
         FLAGS_pipeline_id,
@@ -144,6 +143,7 @@ std::unique_ptr<tool::trace::OperatorReplayerBase> createReplayer() {
   } else if (FLAGS_operator_type == "PartitionedOutput") {
     replayer = std::make_unique<tool::trace::PartitionedOutputReplayer>(
         FLAGS_root_dir,
+        FLAGS_query_id,
         FLAGS_task_id,
         FLAGS_node_id,
         FLAGS_pipeline_id,
@@ -155,12 +155,69 @@ std::unique_ptr<tool::trace::OperatorReplayerBase> createReplayer() {
   return replayer;
 }
 
+void printTaskMetadata(
+    const std::string& taskTraceDir,
+    memory::MemoryPool* pool,
+    std::ostringstream& oss) {
+  auto taskMetaReader = std::make_unique<exec::trace::TaskTraceMetadataReader>(
+      taskTraceDir, pool);
+  std::unordered_map<std::string, std::string> queryConfigs;
+  std::unordered_map<std::string, std::unordered_map<std::string, std::string>>
+      connectorProperties;
+  core::PlanNodePtr queryPlan;
+  taskMetaReader->read(queryConfigs, connectorProperties, queryPlan);
+
+  oss << "++++++Query configs++++++\n";
+  for (const auto& queryConfigEntry : queryConfigs) {
+    oss << "\t" << queryConfigEntry.first << ": " << queryConfigEntry.second
+        << "\n";
+  }
+  oss << "++++++Connector configs++++++\n";
+  for (const auto& connectorPropertyEntry : connectorProperties) {
+    oss << connectorPropertyEntry.first << "\n";
+    for (const auto& propertyEntry : connectorPropertyEntry.second) {
+      oss << "\t" << propertyEntry.first << ": " << propertyEntry.second
+          << "\n";
+    }
+  }
+  oss << "++++++Task query plan++++++\n";
+  oss << queryPlan->toString(true, true);
+}
+
+void printTaskTraceSummary(
+    const std::string& taskTraceDir,
+    const std::string& nodeId,
+    uint32_t pipelineId,
+    memory::MemoryPool* pool,
+    std::ostringstream& oss) {
+  auto fs = filesystems::getFileSystem(taskTraceDir, nullptr);
+  const auto numDrivers = exec::trace::getNumDrivers(
+      exec::trace::getNodeTraceDirectory(taskTraceDir, nodeId), pipelineId, fs);
+  oss << "\n++++++" << numDrivers << " drivers++++++\n";
+  for (auto driverId = 0; driverId < numDrivers; ++driverId) {
+    oss << "++++++" << driverId << " driver++++++\n";
+    const auto opTraceDir = exec::trace::getOpTraceDirectory(
+        taskTraceDir, nodeId, pipelineId, driverId);
+    const auto opTraceSummary =
+        exec::trace::OperatorTraceSummaryReader(
+            exec::trace::getOpTraceDirectory(
+                taskTraceDir, nodeId, pipelineId, driverId),
+            pool)
+            .read();
+    oss << driverId << " driver, " << opTraceSummary.toString();
+  }
+}
+
 void printSummary(
     const std::string& rootDir,
+    const std::string& queryId,
     const std::string& taskId,
-    bool shortSummary) {
-  const auto fs = filesystems::getFileSystem(rootDir, nullptr);
-  const auto taskIds = exec::trace::getTaskIds(rootDir, fs);
+    bool shortSummary,
+    memory::MemoryPool* pool) {
+  const std::string queryDir =
+      exec::trace::getQueryTraceDirectory(rootDir, queryId);
+  const auto fs = filesystems::getFileSystem(queryDir, nullptr);
+  const auto taskIds = exec::trace::getTaskIds(rootDir, queryId, fs);
   if (taskIds.empty()) {
     LOG(ERROR) << "No traced query task under " << rootDir;
     return;
@@ -178,24 +235,20 @@ void printSummary(
 
   const auto summaryTaskIds =
       taskId.empty() ? taskIds : std::vector<std::string>{taskId};
+  VELOX_CHECK(summaryTaskIds.empty() == false, "No task id found");
+  printTaskMetadata(
+      exec::trace::getTaskTraceDirectory(rootDir, queryId, summaryTaskIds[0]),
+      pool,
+      summary);
+  summary << "\n++++++Task Summaries++++++\n";
   for (const auto& taskId : summaryTaskIds) {
-    summary << "\n++++++Query configs and plan of task " << taskId
-            << ":++++++\n";
-    const auto traceTaskDir = fmt::format("{}/{}", rootDir, taskId);
-    const auto queryMetaFile = fmt::format(
-        "{}/{}",
-        traceTaskDir,
-        exec::trace::QueryTraceTraits::kQueryMetaFileName);
-    const auto metaObj = exec::trace::getMetadata(queryMetaFile, fs);
-    const auto& configObj =
-        metaObj[exec::trace::QueryTraceTraits::kQueryConfigKey];
-    summary << "++++++Query configs++++++\n";
-    summary << folly::toJson(configObj) << "\n";
-    summary << "++++++Query plan++++++\n";
-    const auto queryPlan = ISerializable::deserialize<core::PlanNode>(
-        metaObj[exec::trace::QueryTraceTraits::kPlanNodeKey],
-        memory::MemoryManager::getInstance()->tracePool());
-    summary << queryPlan->toString(true, true);
+    summary << "\n++++++Task " << taskId << ":++++++\n";
+    printTaskTraceSummary(
+        exec::trace::getTaskTraceDirectory(rootDir, queryId, taskId),
+        FLAGS_node_id,
+        FLAGS_pipeline_id,
+        pool,
+        summary);
   }
   LOG(INFO) << summary.str();
 }
@@ -214,10 +267,26 @@ int main(int argc, char** argv) {
     return -1;
   }
 
+  if (FLAGS_query_id.empty()) {
+    gflags::SetUsageMessage("--query_id must be provided.");
+    gflags::ShowUsageWithFlags(argv[0]);
+    return -1;
+  }
+
   try {
-    init();
+    // TODO: make it configurable.
+    const auto ioExecutor = std::make_unique<folly::IOThreadPoolExecutor>(
+        std::thread::hardware_concurrency() *
+        FLAGS_hiveConnectorExecutorHwMultiplier);
+    init(ioExecutor.get());
     if (FLAGS_summary || FLAGS_short_summary) {
-      printSummary(FLAGS_root_dir, FLAGS_task_id, FLAGS_short_summary);
+      auto pool = memory::memoryManager()->addLeafPool("replayer");
+      printSummary(
+          FLAGS_root_dir,
+          FLAGS_query_id,
+          FLAGS_task_id,
+          FLAGS_short_summary,
+          pool.get());
       return 0;
     }
     createReplayer()->run();
diff --git a/velox/tool/trace/TableWriterReplayer.cpp b/velox/tool/trace/TableWriterReplayer.cpp
index 3543011aed9fb..405071da433e3 100644
--- a/velox/tool/trace/TableWriterReplayer.cpp
+++ b/velox/tool/trace/TableWriterReplayer.cpp
@@ -16,8 +16,8 @@
 
 #include <folly/executors/IOThreadPoolExecutor.h>
 
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/TableWriter.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
 #include "velox/tool/trace/TableWriterReplayer.h"
 
diff --git a/velox/tool/trace/TableWriterReplayer.h b/velox/tool/trace/TableWriterReplayer.h
index c0c675ba63335..8684be91478d7 100644
--- a/velox/tool/trace/TableWriterReplayer.h
+++ b/velox/tool/trace/TableWriterReplayer.h
@@ -27,13 +27,20 @@ namespace facebook::velox::tool::trace {
 class TableWriterReplayer final : public OperatorReplayerBase {
  public:
   TableWriterReplayer(
-      const std::string& rootDir,
+      const std::string& traceDir,
+      const std::string& queryId,
       const std::string& taskId,
       const std::string& nodeId,
       const int32_t pipelineId,
       const std::string& operatorType,
       const std::string& replayOutputDir)
-      : OperatorReplayerBase(rootDir, taskId, nodeId, pipelineId, operatorType),
+      : OperatorReplayerBase(
+            traceDir,
+            queryId,
+            taskId,
+            nodeId,
+            pipelineId,
+            operatorType),
         replayOutputDir_(replayOutputDir) {
     VELOX_CHECK(!replayOutputDir_.empty());
   }
diff --git a/velox/tool/trace/tests/AggregationReplayerTest.cpp b/velox/tool/trace/tests/AggregationReplayerTest.cpp
index 8d3b15abde620..55a71f739eece 100644
--- a/velox/tool/trace/tests/AggregationReplayerTest.cpp
+++ b/velox/tool/trace/tests/AggregationReplayerTest.cpp
@@ -26,10 +26,10 @@
 #include "velox/common/hyperloglog/SparseHll.h"
 #include "velox/common/testutil/TestValue.h"
 #include "velox/dwio/dwrf/writer/Writer.h"
+#include "velox/exec/OperatorTraceReader.h"
 #include "velox/exec/PartitionFunction.h"
-#include "velox/exec/QueryDataReader.h"
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/TableWriter.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/ArbitratorTestUtil.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/HiveConnectorTestBase.h"
@@ -198,10 +198,14 @@ TEST_F(AggregationReplayerTest, test) {
             .split(makeHiveConnectorSplit(sourceFilePath->getPath()))
             .copyResults(pool(), task);
 
-    const auto replayingResult =
-        AggregationReplayer(
-            traceRoot, task->taskId(), traceNodeId_, 0, "Aggregation")
-            .run();
+    const auto replayingResult = AggregationReplayer(
+                                     traceRoot,
+                                     task->queryCtx()->queryId(),
+                                     task->taskId(),
+                                     traceNodeId_,
+                                     0,
+                                     "Aggregation")
+                                     .run();
     assertEqualResults({results}, {replayingResult});
   }
 }
diff --git a/velox/tool/trace/tests/PartitionedOutputReplayerTest.cpp b/velox/tool/trace/tests/PartitionedOutputReplayerTest.cpp
index b34ead3fb3666..02c74702d7a87 100644
--- a/velox/tool/trace/tests/PartitionedOutputReplayerTest.cpp
+++ b/velox/tool/trace/tests/PartitionedOutputReplayerTest.cpp
@@ -24,8 +24,8 @@
 #include "velox/common/hyperloglog/SparseHll.h"
 #include "velox/exec/PartitionFunction.h"
 #include "velox/exec/PartitionedOutput.h"
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/TableWriter.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/HiveConnectorTestBase.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
@@ -143,10 +143,14 @@ TEST_F(PartitionedOutputReplayerTest, defaultConsumer) {
       executor_.get(),
       consumerExecutor.get(),
       [&](auto /* unused */, auto /* unused */) {});
-  ASSERT_NO_THROW(
-      PartitionedOutputReplayer(
-          traceRoot, originalTask->taskId(), planNodeId, 0, "PartitionedOutput")
-          .run());
+  ASSERT_NO_THROW(PartitionedOutputReplayer(
+                      traceRoot,
+                      originalTask->queryCtx()->queryId(),
+                      originalTask->taskId(),
+                      planNodeId,
+                      0,
+                      "PartitionedOutput")
+                      .run());
 }
 
 TEST_F(PartitionedOutputReplayerTest, basic) {
@@ -223,6 +227,7 @@ TEST_F(PartitionedOutputReplayerTest, basic) {
     replayedPartitionedResults.resize(testParam.numPartitions);
     PartitionedOutputReplayer(
         traceRoot,
+        originalTask->queryCtx()->queryId(),
         originalTask->taskId(),
         planNodeId,
         0,
diff --git a/velox/tool/trace/tests/TableWriterReplayerTest.cpp b/velox/tool/trace/tests/TableWriterReplayerTest.cpp
index ef2c4c54d682d..5e4ec3eddfa96 100644
--- a/velox/tool/trace/tests/TableWriterReplayerTest.cpp
+++ b/velox/tool/trace/tests/TableWriterReplayerTest.cpp
@@ -24,10 +24,10 @@
 #include "velox/common/base/Fs.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/common/hyperloglog/SparseHll.h"
+#include "velox/exec/OperatorTraceReader.h"
 #include "velox/exec/PartitionFunction.h"
-#include "velox/exec/QueryDataReader.h"
-#include "velox/exec/QueryTraceUtil.h"
 #include "velox/exec/TableWriter.h"
+#include "velox/exec/TraceUtil.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/HiveConnectorTestBase.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
@@ -297,6 +297,7 @@ TEST_F(TableWriterReplayerTest, basic) {
   const auto traceOutputDir = TempDirectoryPath::create();
   const auto result = TableWriterReplayer(
                           traceRoot,
+                          task->queryCtx()->queryId(),
                           task->taskId(),
                           "1",
                           0,
@@ -422,6 +423,7 @@ TEST_F(TableWriterReplayerTest, partitionWrite) {
   const auto traceOutputDir = TempDirectoryPath::create();
   TableWriterReplayer(
       traceRoot,
+      task->queryCtx()->queryId(),
       task->taskId(),
       tableWriteNodeId,
       0,