support config store decimal as integer for write parquet format

Author: lifulong-zhihu

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

@@ -682,6 +682,58 @@ TEST_F(ParquetWriterTest, updateWriterOptionsFromHiveConfig) {
       TimestampPrecision::kMilliseconds);
 }
 
+TEST_F(ParquetWriterTest, storeDecimalAsIntegerFromConfig) {
+  const config::ConfigBase emptySession({});
+
+  {
+    std::unordered_map<std::string, std::string> connectorMap = {
+        {parquet::WriterOptions::kParquetHiveConnectorStoreDecimalAsInteger,
+         "false"}};
+    const config::ConfigBase connectorConfig(std::move(connectorMap));
+    parquet::WriterOptions options;
+    options.processConfigs(connectorConfig, emptySession);
+    ASSERT_TRUE(options.storeDecimalAsInteger.has_value());
+    ASSERT_FALSE(options.storeDecimalAsInteger.value());
+  }
+
+  {
+    std::unordered_map<std::string, std::string> connectorMap = {
+        {parquet::WriterOptions::kParquetHiveConnectorStoreDecimalAsInteger,
+         "true"}};
+    const config::ConfigBase connectorConfig(std::move(connectorMap));
+    parquet::WriterOptions options;
+    options.processConfigs(connectorConfig, emptySession);
+    ASSERT_TRUE(options.storeDecimalAsInteger.has_value());
+    ASSERT_TRUE(options.storeDecimalAsInteger.value());
+  }
+
+  {
+    std::unordered_map<std::string, std::string> connectorMap = {
+        {parquet::WriterOptions::kParquetHiveConnectorStoreDecimalAsInteger,
+         "true"}};
+    const config::ConfigBase connectorConfig(std::move(connectorMap));
+    std::unordered_map<std::string, std::string> sessionMap = {
+        {parquet::WriterOptions::kParquetSessionStoreDecimalAsInteger,
+         "false"}};
+    const config::ConfigBase session(std::move(sessionMap));
+    parquet::WriterOptions options;
+    options.processConfigs(connectorConfig, session);
+    ASSERT_TRUE(options.storeDecimalAsInteger.has_value());
+    ASSERT_FALSE(options.storeDecimalAsInteger.value());
+  }
+
+  {
+    parquet::WriterOptions options;
+    options.storeDecimalAsInteger = true;
+    std::unordered_map<std::string, std::string> connectorMap = {
+        {parquet::WriterOptions::kParquetHiveConnectorStoreDecimalAsInteger,
+         "false"}};
+    const config::ConfigBase connectorConfig(std::move(connectorMap));
+    options.processConfigs(connectorConfig, emptySession);
+    ASSERT_TRUE(options.storeDecimalAsInteger.value());
+  }
+}
+
 #ifdef VELOX_ENABLE_PARQUET
 DEBUG_ONLY_TEST_F(ParquetWriterTest, timestampUnitAndTimeZone) {
   SCOPED_TESTVALUE_SET(

velox/dwio/parquet/writer/Writer.cpp

@@ -161,7 +161,11 @@ std::shared_ptr<WriterProperties> getArrowParquetWriterOptions(
   properties = properties->maxRowGroupLength(
       static_cast<int64_t>(flushPolicy->rowsInRowGroup()));
   properties = properties->codecOptions(options.codecOptions);
-  properties = properties->enableStoreDecimalAsInteger();
+  if (options.storeDecimalAsInteger.value_or(true)) {
+    properties = properties->enableStoreDecimalAsInteger();
+  } else {
+    properties = properties->disableStoreDecimalAsInteger();
+  }
   if (options.useParquetDataPageV2.value_or(false)) {
     properties = properties->dataPageVersion(arrow::ParquetDataPageVersion::V2);
   } else {
@@ -349,6 +353,20 @@ std::optional<int64_t> toParquetBatchSize(
   }
 }
 
+std::optional<bool> isParquetStoreDecimalAsInteger(
+    const config::ConfigBase& config,
+    const char* configKey) {
+  try {
+    if (const auto storeAsInt = config.get<bool>(configKey)) {
+      return storeAsInt.value();
+    }
+  } catch (const folly::ConversionError& e) {
+    VELOX_USER_FAIL(
+        "Invalid parquet writer store decimal as integer option: {}", e.what());
+  }
+  return std::nullopt;
+}
+
 } // namespace
 
 Writer::Writer(
@@ -639,6 +657,16 @@ void WriterOptions::processConfigs(
             kParquetEnableDictionary, connectorConfig));
   }
 
+  if (!storeDecimalAsInteger) {
+    storeDecimalAsInteger = isParquetStoreDecimalAsInteger(
+                                session, kParquetSessionStoreDecimalAsInteger)
+                                .has_value()
+        ? isParquetStoreDecimalAsInteger(
+              session, kParquetSessionStoreDecimalAsInteger)
+        : isParquetStoreDecimalAsInteger(
+              connectorConfig, kParquetHiveConnectorStoreDecimalAsInteger);
+  }
+
   if (!dictionaryPageSizeLimit) {
     dictionaryPageSizeLimit =
         toParquetPageSize(session.getWithFallback<std::string>(

velox/dwio/parquet/writer/Writer.h

@@ -134,6 +134,9 @@ struct WriterOptions : public dwio::common::WriterOptions, public WriterConfig {
   std::optional<int64_t> dataPageSize;
   std::optional<int64_t> dictionaryPageSizeLimit;
   std::optional<bool> enableDictionary;
+  /// If unset, Writer uses true (INT32/INT64 for short DECIMAL); false forces
+  /// FIXED_LEN_BYTE_ARRAY for those precisions.
+  std::optional<bool> storeDecimalAsInteger;
   std::optional<bool> useParquetDataPageV2;
   std::optional<std::string> createdBy;
 
@@ -163,6 +166,10 @@ struct WriterOptions : public dwio::common::WriterOptions, public WriterConfig {
       "hive.parquet.writer.batch_size";
   static constexpr const char* kParquetCreatedBy =
       "hive.parquet.writer.created_by";
+  static constexpr const char* kParquetSessionStoreDecimalAsInteger =
+      "hive.parquet.writer.store_decimal_as_integer";
+  static constexpr const char* kParquetHiveConnectorStoreDecimalAsInteger =
+      "hive.parquet.writer.store-decimal-as-integer";
   static constexpr const char* kParquetMaxTargetFileSize =
       "max_target_file_size";
   // Serde parameter keys for timestamp settings. These can be set via