feat(parquet): Add config for datapage version (#11151)

Summary:
Add config and session properties `hive.parquet.writer.datapage-version`, `hive.parquet.writer.datapage_version`, to determine the parquet writer datapage version (V1 or V2).  Defaults to V1.

Pull Request resolved: #11151

Reviewed By: kagamiori

Differential Revision: D69496525

Pulled By: kgpai

fbshipit-source-id: 0a25a7f3383ad779a44e38b0c49825c52080ae92

Author: svm1

velox/docs/configs.rst

@@ -657,6 +657,12 @@ Each query can override the config by setting corresponding query session proper
      - 9
      - Timestamp unit used when writing timestamps into Parquet through Arrow bridge.
        Valid values are 3 (millisecond), 6 (microsecond), and 9 (nanosecond).
+   * - hive.parquet.writer.datapage-version
+     - hive.parquet.writer.datapage_version
+     - string
+     - V1
+     - Data Page version used when writing into Parquet through Arrow bridge.
+       Valid values are "V1" and "V2".
 
 ``Amazon S3 Configuration``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^

velox/dwio/parquet/tests/writer/ParquetWriterTest.cpp

@@ -22,6 +22,7 @@
 #include "velox/connectors/hive/HiveConnector.h" // @manual
 #include "velox/core/QueryCtx.h"
 #include "velox/dwio/parquet/RegisterParquetWriter.h" // @manual
+#include "velox/dwio/parquet/reader/PageReader.h"
 #include "velox/dwio/parquet/tests/ParquetTestBase.h"
 #include "velox/exec/Cursor.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
@@ -146,6 +147,124 @@ TEST_F(ParquetWriterTest, compression) {
   assertReadWithReaderAndExpected(schema, *rowReader, data, *leafPool_);
 };
 
+TEST_F(ParquetWriterTest, toggleDataPageVersion) {
+  auto schema = ROW({"c0"}, {INTEGER()});
+  const int64_t kRows = 1;
+  const auto data = makeRowVector({
+      makeFlatVector<int32_t>(kRows, [](auto row) { return 987; }),
+  });
+
+  // Write Parquet test data, then read and return the DataPage
+  // (thrift::PageType::type) used.
+  const auto testDataPageVersion =
+      [&](std::unordered_map<std::string, std::string> configFromFile,
+          std::unordered_map<std::string, std::string> sessionProperties) {
+        // Create an in-memory writer.
+        auto sink = std::make_unique<MemorySink>(
+            200 * 1024 * 1024,
+            dwio::common::FileSink::Options{.pool = leafPool_.get()});
+        auto sinkPtr = sink.get();
+        parquet::WriterOptions writerOptions;
+        writerOptions.memoryPool = leafPool_.get();
+
+        // Simulate setting of Hive config & connector session properties, then
+        // write test data.
+        auto connectorConfig = config::ConfigBase(std::move(configFromFile));
+        auto connectorSessionProperties =
+            config::ConfigBase(std::move(sessionProperties));
+
+        writerOptions.processConfigs(
+            connectorConfig, connectorSessionProperties);
+        auto writer = std::make_unique<parquet::Writer>(
+            std::move(sink), writerOptions, rootPool_, schema);
+        writer->write(data);
+        writer->close();
+
+        // Read to identify DataPage used.
+        dwio::common::ReaderOptions readerOptions{leafPool_.get()};
+        auto reader = createReaderInMemory(*sinkPtr, readerOptions);
+
+        auto colChunkPtr = reader->fileMetaData().rowGroup(0).columnChunk(0);
+        std::string_view sinkData(sinkPtr->data(), sinkPtr->size());
+
+        auto readFile = std::make_shared<InMemoryReadFile>(sinkData);
+        auto file = std::make_shared<ReadFileInputStream>(std::move(readFile));
+
+        auto inputStream = std::make_unique<SeekableFileInputStream>(
+            std::move(file),
+            colChunkPtr.dataPageOffset(),
+            150,
+            *leafPool_,
+            LogType::TEST);
+        auto pageReader = std::make_unique<PageReader>(
+            std::move(inputStream),
+            *leafPool_,
+            colChunkPtr.compression(),
+            colChunkPtr.totalCompressedSize());
+
+        return pageReader->readPageHeader().type;
+      };
+
+  // Test default behavior - DataPage should be V1.
+  ASSERT_EQ(testDataPageVersion({}, {}), thrift::PageType::type::DATA_PAGE);
+
+  // Simulate setting DataPage version to V2 via Hive config from file.
+  std::unordered_map<std::string, std::string> configFromFile = {
+      {parquet::WriterOptions::kParquetHiveConnectorDataPageVersion, "V2"}};
+
+  ASSERT_EQ(
+      testDataPageVersion(configFromFile, {}),
+      thrift::PageType::type::DATA_PAGE_V2);
+
+  // Simulate setting DataPage version to V1 via Hive config from file.
+  configFromFile = {
+      {parquet::WriterOptions::kParquetHiveConnectorDataPageVersion, "V1"}};
+
+  ASSERT_EQ(
+      testDataPageVersion(configFromFile, {}),
+      thrift::PageType::type::DATA_PAGE);
+
+  // Simulate setting DataPage version to V2 via connector session property.
+  std::unordered_map<std::string, std::string> sessionProperties = {
+      {parquet::WriterOptions::kParquetSessionDataPageVersion, "V2"}};
+
+  ASSERT_EQ(
+      testDataPageVersion({}, sessionProperties),
+      thrift::PageType::type::DATA_PAGE_V2);
+
+  // Simulate setting DataPage version to V1 via connector session property.
+  sessionProperties = {
+      {parquet::WriterOptions::kParquetSessionDataPageVersion, "V1"}};
+
+  ASSERT_EQ(
+      testDataPageVersion({}, sessionProperties),
+      thrift::PageType::type::DATA_PAGE);
+
+  // Simulate setting DataPage version to V1 via connector session property,
+  // and to V2 via Hive config from file. Session property should take
+  // precedence.
+  sessionProperties = {
+      {parquet::WriterOptions::kParquetSessionDataPageVersion, "V1"}};
+  configFromFile = {
+      {parquet::WriterOptions::kParquetHiveConnectorDataPageVersion, "V2"}};
+
+  ASSERT_EQ(
+      testDataPageVersion({}, sessionProperties),
+      thrift::PageType::type::DATA_PAGE);
+
+  // Simulate setting DataPage version to V2 via connector session property,
+  // and to V1 via Hive config from file. Session property should take
+  // precedence.
+  sessionProperties = {
+      {parquet::WriterOptions::kParquetSessionDataPageVersion, "V2"}};
+  configFromFile = {
+      {parquet::WriterOptions::kParquetHiveConnectorDataPageVersion, "V1"}};
+
+  ASSERT_EQ(
+      testDataPageVersion({}, sessionProperties),
+      thrift::PageType::type::DATA_PAGE_V2);
+}
+
 DEBUG_ONLY_TEST_F(ParquetWriterTest, unitFromWriterOptions) {
   SCOPED_TESTVALUE_SET(
       "facebook::velox::parquet::Writer::write",

velox/dwio/parquet/writer/Writer.cpp

@@ -147,6 +147,13 @@ std::shared_ptr<WriterProperties> getArrowParquetWriterOptions(
       static_cast<int64_t>(flushPolicy->rowsInRowGroup()));
   properties = properties->codec_options(options.codecOptions);
   properties = properties->enable_store_decimal_as_integer();
+  if (options.useParquetDataPageV2.value_or(false)) {
+    properties =
+        properties->data_page_version(arrow::ParquetDataPageVersion::V2);
+  } else {
+    properties =
+        properties->data_page_version(arrow::ParquetDataPageVersion::V1);
+  }
   return properties->build();
 }
 
@@ -238,6 +245,21 @@ std::optional<std::string> getTimestampTimeZone(
   return std::nullopt;
 }
 
+std::optional<bool> getParquetDataPageVersion(
+    const config::ConfigBase& config,
+    const char* configKey) {
+  if (const auto version = config.get<std::string>(configKey)) {
+    if (version == "V1") {
+      return false;
+    } else if (version == "V2") {
+      return true;
+    } else {
+      VELOX_FAIL("Unsupported parquet datapage version {}", version.value());
+    }
+  }
+  return std::nullopt;
+}
+
 } // namespace
 
 Writer::Writer(
@@ -467,6 +489,15 @@ void WriterOptions::processConfigs(
   if (!parquetWriteTimestampTimeZone) {
     parquetWriteTimestampTimeZone = parquetWriterOptions->sessionTimezoneName;
   }
+
+  if (!useParquetDataPageV2) {
+    useParquetDataPageV2 =
+        getParquetDataPageVersion(session, kParquetSessionDataPageVersion)
+            .has_value()
+        ? getParquetDataPageVersion(session, kParquetSessionDataPageVersion)
+        : getParquetDataPageVersion(
+              connectorConfig, kParquetHiveConnectorDataPageVersion);
+  }
 }
 
 } // namespace facebook::velox::parquet

velox/dwio/parquet/writer/Writer.h

@@ -108,6 +108,7 @@ struct WriterOptions : public dwio::common::WriterOptions {
   /// Timestamp time zone for Parquet write through Arrow bridge.
   std::optional<std::string> parquetWriteTimestampTimeZone;
   bool writeInt96AsTimestamp = false;
+  std::optional<bool> useParquetDataPageV2;
 
   // Parsing session and hive configs.
 
@@ -117,6 +118,10 @@ struct WriterOptions : public dwio::common::WriterOptions {
       "hive.parquet.writer.timestamp_unit";
   static constexpr const char* kParquetHiveConnectorWriteTimestampUnit =
       "hive.parquet.writer.timestamp-unit";
+  static constexpr const char* kParquetSessionDataPageVersion =
+      "hive.parquet.writer.datapage_version";
+  static constexpr const char* kParquetHiveConnectorDataPageVersion =
+      "hive.parquet.writer.datapage-version";
 
   // Process hive connector and session configs.
   void processConfigs(