Do not serialize data type if not initialized

tanjialiang · Oct 12, 2024 · 8c3e197 · 8c3e197
1 parent b00751e
commit 8c3e197
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/velox/exec/QueryDataWriter.cpp b/velox/exec/QueryDataWriter.cpp
@@ -42,6 +42,11 @@ void QueryDataWriter::write(const RowVectorPtr& rows) {
   if (FOLLY_UNLIKELY(finished_)) {
     return;
   }
+  if (dataType_ == nullptr) {
+    dataType_ = rows->type();
+  } else {
+    VELOX_CHECK(dataType_->kindEquals(rows->type()));
+  }
 
   if (batch_ == nullptr) {
     batch_ = std::make_unique<VectorStreamGroup>(pool_);
@@ -51,7 +56,6 @@ void QueryDataWriter::write(const RowVectorPtr& rows) {
         &options_);
   }
   batch_->append(rows);
-  dataType_ = rows->type();
 
   // Serialize and write out each batch.
   IOBufOutputStream out(
@@ -86,7 +90,9 @@ void QueryDataWriter::writeSummary(bool limitExceeded) const {
       fmt::format("{}/{}", dirPath_, QueryTraceTraits::kDataSummaryFileName);
   const auto file = fs_->openFileForWrite(summaryFilePath);
   folly::dynamic obj = folly::dynamic::object;
-  obj[QueryTraceTraits::kDataTypeKey] = dataType_->serialize();
+  if (dataType_ != nullptr) {
+    obj[QueryTraceTraits::kDataTypeKey] = dataType_->serialize();
+  }
   obj[QueryTraceTraits::kTraceLimitExceededKey] = limitExceeded;
   file->append(folly::toJson(obj));
   file->close();

diff --git a/velox/exec/tests/QueryTraceTest.cpp b/velox/exec/tests/QueryTraceTest.cpp
@@ -102,6 +102,21 @@ class QueryTracerTest : public HiveConnectorTestBase {
   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(QueryTracerTest, traceData) {
   const auto rowType = ROW({"a", "b", "c"}, {BIGINT(), BIGINT(), BIGINT()});
   std::vector<RowVectorPtr> inputVectors;