Add file-create-config to hive config and propagate to file sink

facebookincubator · Dec 13, 2023 · f485f11 · f485f11
1 parent e57c47f
commit f485f11
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git a/velox/connectors/hive/HiveConfig.cpp b/velox/connectors/hive/HiveConfig.cpp
@@ -214,4 +214,11 @@ uint64_t HiveConfig::sortWriterMaxOutputBytes(const Config* session) const {
   return 10UL << 20;
 }
 
+std::string HiveConfig::fileCreateConfig(const Config* session) const {
+  if (session->isValueExists(kFileCreateConfig)) {
+    return session->get<std::string>(kFileCreateConfig).value();
+  }
+  return config_->get<std::string>(kFileCreateConfig, "");
+}
+
 } // namespace facebook::velox::connector::hive
diff --git a/velox/connectors/hive/HiveConfig.h b/velox/connectors/hive/HiveConfig.h
@@ -148,6 +148,11 @@ class HiveConfig {
   static constexpr const char* kSortWriterMaxOutputBytesSession =
       "sort_writer_max_output_bytes";
 
+  /// Config used to create sink files. This config is provided to underlying
+  /// file system and the config is free form. The form should be defined by
+  /// the underlying file system.
+  static constexpr const char* kFileCreateConfig = "file-create-config";
+
   InsertExistingPartitionsBehavior insertExistingPartitionsBehavior(
       const Config* session) const;
 
@@ -201,6 +206,8 @@ class HiveConfig {
 
   uint64_t sortWriterMaxOutputBytes(const Config* session) const;
 
+  std::string fileCreateConfig(const Config* session) const;
+
   HiveConfig(std::shared_ptr<const Config> config) {
     VELOX_CHECK_NOT_NULL(
         config, "Config is null for HiveConfig initialization");

diff --git a/velox/connectors/hive/HiveDataSink.cpp b/velox/connectors/hive/HiveDataSink.cpp
@@ -17,6 +17,7 @@
 #include "velox/connectors/hive/HiveDataSink.h"
 
 #include "velox/common/base/Counters.h"
+#include "velox/common/file/FileSystems.h
 #include "velox/common/base/Fs.h"
 #include "velox/common/base/StatsReporter.h"
 #include "velox/common/testutil/TestValue.h"
@@ -595,6 +596,8 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   setMemoryReclaimers(writerInfo_.back().get());
 
   dwio::common::WriterOptions options;
+  const auto connectorSessionProperties =
+      connectorQueryCtx_->sessionProperties();
   options.schema = inputType_;
   options.memoryPool = writerInfo_.back()->writerPool.get();
   options.compressionKind = insertTableHandle_->compressionKind();
@@ -603,11 +606,10 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   }
   options.nonReclaimableSection =
       writerInfo_.back()->nonReclaimableSectionHolder.get();
-  options.maxStripeSize = std::optional(hiveConfig_->getOrcWriterMaxStripeSize(
-      connectorQueryCtx_->sessionProperties()));
-  options.maxDictionaryMemory =
-      std::optional(hiveConfig_->getOrcWriterMaxDictionaryMemory(
-          connectorQueryCtx_->sessionProperties()));
+  options.maxStripeSize = std::optional(
+      hiveConfig_->getOrcWriterMaxStripeSize(connectorSessionProperties));
+  options.maxDictionaryMemory = std::optional(
+      hiveConfig_->getOrcWriterMaxDictionaryMemory(connectorSessionProperties));
   ioStats_.emplace_back(std::make_shared<io::IoStatistics>());
 
   // Prevents the memory allocation during the writer creation.
@@ -617,6 +619,8 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
           writePath,
           {.bufferWrite = false,
            .connectorProperties = hiveConfig_->config(),
+           .fileCreateConfig =
+               hiveConfig_->fileCreateConfig(connectorSessionProperties),
            .pool = writerInfo_.back()->sinkPool.get(),
            .metricLogger = dwio::common::MetricsLog::voidLog(),
            .stats = ioStats_.back().get()}),

diff --git a/velox/dwio/common/FileSink.h b/velox/dwio/common/FileSink.h
@@ -41,6 +41,10 @@ class FileSink : public Closeable {
     /// Connector properties are required to create a FileSink on FileSystems
     /// such as S3.
     const std::shared_ptr<const Config>& connectorProperties{nullptr};
+    /// Config used to create sink files. This config is provided to underlying
+    /// file system and the config is free form. The form should be defined by
+    /// the underlying file system.
+    const std::string fileCreateConfig;
     memory::MemoryPool* pool{nullptr};
     MetricsLogPtr metricLogger{MetricsLog::voidLog()};
     IoStatistics* stats{nullptr};