misc: Add compression format suffix to the written Parquet file

velox/common/compression/Compression.cpp

@@ -98,4 +98,24 @@ CompressionKind stringToCompressionKind(const std::string& kind) {
     VELOX_UNSUPPORTED("Not support compression kind {}", kind);
   }
 }
+
+std::string fileNameSuffix(CompressionKind kind) {
+  switch (static_cast<int32_t>(kind)) {
+    case CompressionKind_ZLIB:
+      return ".zlib";
+    case CompressionKind_SNAPPY:
+      return ".snappy";
+    case CompressionKind_LZO:
+      return ".lzo";
+    case CompressionKind_ZSTD:
+      return ".zstd";
+    case CompressionKind_LZ4:
+      return ".lz4";
+    case CompressionKind_GZIP:
+      return ".gz";
+    case CompressionKind_NONE:
+    default:
+      return "";
+  }
+}
 } // namespace facebook::velox::common

velox/common/compression/Compression.h

@@ -42,6 +42,8 @@ std::string compressionKindToString(CompressionKind kind);
 
 CompressionKind stringToCompressionKind(const std::string& kind);
 
+std::string fileNameSuffix(CompressionKind kind);
+
 constexpr uint64_t DEFAULT_COMPRESSION_BLOCK_SIZE = 256 * 1024;
 
 } // namespace facebook::velox::common
@@ -56,3 +58,16 @@ struct fmt::formatter<facebook::velox::common::CompressionKind>
         facebook::velox::common::compressionKindToString(s), ctx);
   }
 };
+
+template <>
+struct fmt::formatter<std::optional<facebook::velox::common::CompressionKind>>
+    : fmt::formatter<std::string> {
+  auto format(
+      const std::optional<facebook::velox::common::CompressionKind>& s,
+      format_context& ctx) {
+    return formatter<std::string>::format(
+        facebook::velox::common::fileNameSuffix(s.value_or(
+            facebook::velox::common::CompressionKind::CompressionKind_NONE)),
+        ctx);
+  }
+};

velox/connectors/hive/HiveDataSink.cpp

@@ -939,8 +939,16 @@ std::pair<std::string, std::string> HiveDataSink::getWriterFileNames(
       insertTableHandle_->storageFormat() ==
           dwio::common::FileFormat::PARQUET) {
     return {
-        fmt::format("{}{}", targetFileName, ".parquet"),
-        fmt::format("{}{}", writeFileName, ".parquet")};
+        fmt::format(
+            "{}{}{}",
+            targetFileName,
+            insertTableHandle_->compressionKind(),
+            ".parquet"),
+        fmt::format(
+            "{}{}{}",
+            writeFileName,
+            insertTableHandle_->compressionKind(),
+            ".parquet")};
   }
   return {targetFileName, writeFileName};
 }
@@ -1185,4 +1193,8 @@ uint64_t HiveDataSink::WriterReclaimer::reclaim(
   }
   return reclaimedBytes;
 }
+
+std::string HiveDataSink::writeFileName() const {
+  return getWriterFileNames(std::nullopt).second;
+}
 } // namespace facebook::velox::connector::hive

velox/connectors/hive/HiveDataSink.h

@@ -467,6 +467,9 @@ class HiveDataSink : public DataSink {
 
   bool canReclaim() const;
 
+  // For testing: check the write file name.
+  std::string writeFileName() const;
+
  private:
   // Validates the state transition from 'oldState' to 'newState'.
   void checkStateTransition(State oldState, State newState);

velox/connectors/hive/tests/HiveDataSinkTest.cpp

@@ -144,7 +144,9 @@ class HiveDataSinkTest : public exec::test::HiveConnectorTestBase {
       const std::shared_ptr<connector::hive::HiveBucketProperty>&
           bucketProperty = nullptr,
       const std::shared_ptr<dwio::common::WriterOptions>& writerOptions =
-          nullptr) {
+          nullptr,
+      const CompressionKind compressionKind =
+          CompressionKind::CompressionKind_ZSTD) {
     return makeHiveInsertTableHandle(
         outputRowType->names(),
         outputRowType->children(),
@@ -155,7 +157,7 @@ class HiveDataSinkTest : public exec::test::HiveConnectorTestBase {
             std::nullopt,
             connector::hive::LocationHandle::TableType::kNew),
         fileFormat,
-        CompressionKind::CompressionKind_ZSTD,
+        compressionKind,
         {},
         writerOptions);
   }
@@ -168,7 +170,9 @@ class HiveDataSinkTest : public exec::test::HiveConnectorTestBase {
       const std::shared_ptr<connector::hive::HiveBucketProperty>&
           bucketProperty = nullptr,
       const std::shared_ptr<dwio::common::WriterOptions>& writerOptions =
-          nullptr) {
+          nullptr,
+      const CompressionKind compressionKind =
+          CompressionKind::CompressionKind_ZSTD) {
     return std::make_shared<HiveDataSink>(
         rowType,
         createHiveInsertTableHandle(
@@ -177,7 +181,8 @@ class HiveDataSinkTest : public exec::test::HiveConnectorTestBase {
             fileFormat,
             partitionedBy,
             bucketProperty,
-            writerOptions),
+            writerOptions,
+            compressionKind),
         connectorQueryCtx_.get(),
         CommitStrategy::kNoCommit,
         connectorConfig_);
@@ -1168,6 +1173,43 @@ TEST_F(HiveDataSinkTest, flushPolicyWithParquet) {
   EXPECT_EQ(fileMeta.numRowGroups(), 10);
   EXPECT_EQ(fileMeta.rowGroup(0).numRows(), 500);
 }
+
+TEST_F(HiveDataSinkTest, parquetCompressionFileName) {
+  auto checkFileSuffix = [this](
+                             const CompressionKind compression,
+                             std::string_view expectedSuffix) {
+    auto endsWith = [](std::string_view str, std::string_view suffix) {
+      if (str.length() < suffix.length()) {
+        return false;
+      }
+      return (
+          0 ==
+          str.compare(str.length() - suffix.length(), suffix.length(), suffix));
+    };
+    auto dataSink = createDataSink(
+        rowType_,
+        "/path/to/test",
+        dwio::common::FileFormat::PARQUET,
+        {},
+        nullptr,
+        nullptr,
+        compression);
+    auto writeFileName = dataSink->writeFileName();
+    VELOX_CHECK(
+        endsWith(writeFileName, expectedSuffix),
+        "{} not end with {}",
+        writeFileName,
+        expectedSuffix);
+  };
+  checkFileSuffix(CompressionKind::CompressionKind_NONE, ".parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_ZLIB, ".zlib.parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_SNAPPY, ".snappy.parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_LZO, ".lzo.parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_ZSTD, ".zstd.parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_LZ4, ".lz4.parquet");
+  checkFileSuffix(CompressionKind::CompressionKind_GZIP, ".gz.parquet");
+}
+
 #endif
 
 TEST_F(HiveDataSinkTest, flushPolicyWithDWRF) {