Make dwio::common::WriterOptions polymorphic (facebookincubator#10380)

Summary:
Pull Request resolved: facebookincubator#10380

Making writer options a virtual class that could be specialized for
different file formats to enable Velox users to provided them as part of a
query plan.

Today, there are a few different ways to pass writer parameters (serdeParams and
hive config). They end up being serialized to strings, so can't be used with
non-serializable configs. Moreover, file format specific configurations and code
end-up in generic parts of the code, like HiveConnector.

This will allow us to better organize file format specific options; for now only
creating the framework to contain the changes.

Reviewed By: kunalkataria, kparichay

Differential Revision: D59302873

velox/connectors/hive/HiveDataSink.cpp

@@ -677,57 +677,107 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   ioStats_.emplace_back(std::make_shared<io::IoStatistics>());
   setMemoryReclaimers(writerInfo_.back().get(), ioStats_.back().get());
 
-  dwio::common::WriterOptions options;
+  // Take the one provided by the user as a starting point, or allocate a new
+  // one.
+  const auto& writerOptions = insertTableHandle_->writerOptions();
+  std::shared_ptr<dwio::common::WriterOptions> options = writerOptions
+      ? writerOptions
+      : std::make_shared<dwio::common::WriterOptions>();
+
   const auto* connectorSessionProperties =
       connectorQueryCtx_->sessionProperties();
-  options.schema = getNonPartitionTypes(dataChannels_, inputType_);
-
-  options.memoryPool = writerInfo_.back()->writerPool.get();
-  options.compressionKind = insertTableHandle_->compressionKind();
-  if (canReclaim()) {
-    options.spillConfig = spillConfig_;
-  }
-  options.nonReclaimableSection =
-      writerInfo_.back()->nonReclaimableSectionHolder.get();
-  options.maxStripeSize = std::optional(
-      hiveConfig_->orcWriterMaxStripeSize(connectorSessionProperties));
-  options.maxDictionaryMemory = std::optional(
-      hiveConfig_->orcWriterMaxDictionaryMemory(connectorSessionProperties));
-  options.orcWriterIntegerDictionaryEncodingEnabled =
-      hiveConfig_->isOrcWriterIntegerDictionaryEncodingEnabled(
-          connectorSessionProperties);
-  options.orcWriterStringDictionaryEncodingEnabled =
-      hiveConfig_->isOrcWriterStringDictionaryEncodingEnabled(
-          connectorSessionProperties);
-  options.parquetWriteTimestampUnit =
-      hiveConfig_->parquetWriteTimestampUnit(connectorSessionProperties);
-  options.orcMinCompressionSize = std::optional(
-      hiveConfig_->orcWriterMinCompressionSize(connectorSessionProperties));
-  options.orcLinearStripeSizeHeuristics =
-      std::optional(hiveConfig_->orcWriterLinearStripeSizeHeuristics(
-          connectorSessionProperties));
-  options.serdeParameters = std::map<std::string, std::string>(
-      insertTableHandle_->serdeParameters().begin(),
-      insertTableHandle_->serdeParameters().end());
+
+  // Only overwrite options in case they were not already provided.
+  if (!options->schema) {
+    options->schema = getNonPartitionTypes(dataChannels_, inputType_);
+  }
+
+  if (!options->memoryPool) {
+    options->memoryPool = writerInfo_.back()->writerPool.get();
+  }
+
+  if (!options->compressionKind) {
+    options->compressionKind = insertTableHandle_->compressionKind();
+  }
+
+  if (!options->spillConfig && canReclaim()) {
+    options->spillConfig = spillConfig_;
+  }
+
+  if (!options->nonReclaimableSection) {
+    options->nonReclaimableSection =
+        writerInfo_.back()->nonReclaimableSectionHolder.get();
+  }
+
+  if (!options->maxStripeSize) {
+    options->maxStripeSize = std::optional(
+        hiveConfig_->orcWriterMaxStripeSize(connectorSessionProperties));
+  }
+
+  if (!options->maxDictionaryMemory) {
+    options->maxDictionaryMemory = std::optional(
+        hiveConfig_->orcWriterMaxDictionaryMemory(connectorSessionProperties));
+  }
+
+  if (!options->orcWriterIntegerDictionaryEncodingEnabled) {
+    options->orcWriterIntegerDictionaryEncodingEnabled =
+        hiveConfig_->isOrcWriterIntegerDictionaryEncodingEnabled(
+            connectorSessionProperties);
+  }
+
+  if (!options->orcWriterStringDictionaryEncodingEnabled) {
+    options->orcWriterStringDictionaryEncodingEnabled =
+        hiveConfig_->isOrcWriterStringDictionaryEncodingEnabled(
+            connectorSessionProperties);
+  }
+
+  if (!options->parquetWriteTimestampUnit) {
+    options->parquetWriteTimestampUnit =
+        hiveConfig_->parquetWriteTimestampUnit(connectorSessionProperties);
+  }
+
+  if (!options->orcMinCompressionSize) {
+    options->orcMinCompressionSize = std::optional(
+        hiveConfig_->orcWriterMinCompressionSize(connectorSessionProperties));
+  }
+
+  if (!options->orcLinearStripeSizeHeuristics) {
+    options->orcLinearStripeSizeHeuristics =
+        std::optional(hiveConfig_->orcWriterLinearStripeSizeHeuristics(
+            connectorSessionProperties));
+  }
+
+  if (options->serdeParameters.empty()) {
+    options->serdeParameters = std::map<std::string, std::string>(
+        insertTableHandle_->serdeParameters().begin(),
+        insertTableHandle_->serdeParameters().end());
+  }
 
   auto compressionLevel =
       hiveConfig_->orcWriterCompressionLevel(connectorSessionProperties);
-  options.zlibCompressionLevel =
-      compressionLevel.value_or(kDefaultZlibCompressionLevel);
-  options.zstdCompressionLevel =
-      compressionLevel.value_or(kDefaultZstdCompressionLevel);
+
+  if (!options->zlibCompressionLevel) {
+    options->zlibCompressionLevel =
+        compressionLevel.value_or(kDefaultZlibCompressionLevel);
+  }
+  if (!options->zstdCompressionLevel) {
+    options->zstdCompressionLevel =
+        compressionLevel.value_or(kDefaultZstdCompressionLevel);
+  }
 
   // Prevents the memory allocation during the writer creation.
   WRITER_NON_RECLAIMABLE_SECTION_GUARD(writerInfo_.size() - 1);
   auto writer = writerFactory_->createWriter(
       dwio::common::FileSink::create(
           writePath,
-          {.bufferWrite = false,
-           .connectorProperties = hiveConfig_->config(),
-           .fileCreateConfig = hiveConfig_->writeFileCreateConfig(),
-           .pool = writerInfo_.back()->sinkPool.get(),
-           .metricLogger = dwio::common::MetricsLog::voidLog(),
-           .stats = ioStats_.back().get()}),
+          {
+              .bufferWrite = false,
+              .connectorProperties = hiveConfig_->config(),
+              .fileCreateConfig = hiveConfig_->writeFileCreateConfig(),
+              .pool = writerInfo_.back()->sinkPool.get(),
+              .metricLogger = dwio::common::MetricsLog::voidLog(),
+              .stats = ioStats_.back().get(),
+          }),
       options);
   writer = maybeCreateBucketSortWriter(std::move(writer));
   writers_.emplace_back(std::move(writer));

velox/connectors/hive/HiveDataSink.h

@@ -186,9 +186,7 @@ FOLLY_ALWAYS_INLINE std::ostream& operator<<(
 class HiveInsertTableHandle;
 using HiveInsertTableHandlePtr = std::shared_ptr<HiveInsertTableHandle>;
 
-/**
- * Represents a request for Hive write.
- */
+/// Represents a request for Hive write.
 class HiveInsertTableHandle : public ConnectorInsertTableHandle {
  public:
   HiveInsertTableHandle(
@@ -198,13 +196,16 @@ class HiveInsertTableHandle : public ConnectorInsertTableHandle {
           dwio::common::FileFormat::DWRF,
       std::shared_ptr<HiveBucketProperty> bucketProperty = nullptr,
       std::optional<common::CompressionKind> compressionKind = {},
-      const std::unordered_map<std::string, std::string>& serdeParameters = {})
+      const std::unordered_map<std::string, std::string>& serdeParameters = {},
+      const std::shared_ptr<dwio::common::WriterOptions>& writerOptions =
+          nullptr)
       : inputColumns_(std::move(inputColumns)),
         locationHandle_(std::move(locationHandle)),
         tableStorageFormat_(tableStorageFormat),
         bucketProperty_(std::move(bucketProperty)),
         compressionKind_(compressionKind),
-        serdeParameters_(serdeParameters) {
+        serdeParameters_(serdeParameters),
+        writerOptions_(writerOptions) {
     if (compressionKind.has_value()) {
       VELOX_CHECK(
           compressionKind.value() != common::CompressionKind_MAX,
@@ -235,6 +236,10 @@ class HiveInsertTableHandle : public ConnectorInsertTableHandle {
     return serdeParameters_;
   }
 
+  const std::shared_ptr<dwio::common::WriterOptions>& writerOptions() const {
+    return writerOptions_;
+  }
+
   bool supportsMultiThreading() const override {
     return true;
   }
@@ -262,6 +267,7 @@ class HiveInsertTableHandle : public ConnectorInsertTableHandle {
   const std::shared_ptr<HiveBucketProperty> bucketProperty_;
   const std::optional<common::CompressionKind> compressionKind_;
   const std::unordered_map<std::string, std::string> serdeParameters_;
+  const std::shared_ptr<dwio::common::WriterOptions> writerOptions_;
 };
 
 /// Parameters for Hive writers.

velox/dwio/common/Options.h

@@ -608,6 +608,8 @@ struct WriterOptions {
   std::optional<uint8_t> parquetWriteTimestampUnit;
   std::optional<uint8_t> zlibCompressionLevel;
   std::optional<uint8_t> zstdCompressionLevel;
+
+  virtual ~WriterOptions() = default;
 };
 
 } // namespace facebook::velox::dwio::common

velox/dwio/common/WriterFactory.h

@@ -50,7 +50,19 @@ class WriterFactory {
   /// @return writer object
   virtual std::unique_ptr<Writer> createWriter(
       std::unique_ptr<dwio::common::FileSink> sink,
-      const dwio::common::WriterOptions& options) = 0;
+      std::shared_ptr<dwio::common::WriterOptions> options) = 0;
+
+#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY
+  // TODO: for backward compatibility with old clients. Remove once Prestissimo
+  // code is moved to the new API above.
+  std::unique_ptr<Writer> createWriter(
+      std::unique_ptr<dwio::common::FileSink> sink,
+      const dwio::common::WriterOptions& options) {
+    return createWriter(
+        std::move(sink),
+        std::make_shared<dwio::common::WriterOptions>(options));
+  }
+#endif
 
  private:
   const FileFormat format_;

velox/dwio/dwrf/writer/Writer.cpp

@@ -855,8 +855,8 @@ dwrf::WriterOptions getDwrfOptions(const dwio::common::WriterOptions& options) {
 
 std::unique_ptr<dwio::common::Writer> DwrfWriterFactory::createWriter(
     std::unique_ptr<dwio::common::FileSink> sink,
-    const dwio::common::WriterOptions& options) {
-  auto dwrfOptions = getDwrfOptions(options);
+    std::shared_ptr<dwio::common::WriterOptions> options) {
+  auto dwrfOptions = getDwrfOptions(*options);
   return std::make_unique<Writer>(std::move(sink), dwrfOptions);
 }
 

velox/dwio/dwrf/writer/Writer.h

@@ -219,7 +219,7 @@ class DwrfWriterFactory : public dwio::common::WriterFactory {
 
   std::unique_ptr<dwio::common::Writer> createWriter(
       std::unique_ptr<dwio::common::FileSink> sink,
-      const dwio::common::WriterOptions& options) override;
+      std::shared_ptr<dwio::common::WriterOptions> options) override;
 };
 
 } // namespace facebook::velox::dwrf

velox/dwio/parquet/writer/Writer.cpp

@@ -416,10 +416,10 @@ void Writer::setMemoryReclaimers() {
 
 std::unique_ptr<dwio::common::Writer> ParquetWriterFactory::createWriter(
     std::unique_ptr<dwio::common::FileSink> sink,
-    const dwio::common::WriterOptions& options) {
-  auto parquetOptions = getParquetOptions(options);
+    std::shared_ptr<dwio::common::WriterOptions> options) {
+  auto parquetOptions = getParquetOptions(*options);
   return std::make_unique<Writer>(
-      std::move(sink), parquetOptions, asRowType(options.schema));
+      std::move(sink), parquetOptions, asRowType(options->schema));
 }
 
 } // namespace facebook::velox::parquet

velox/dwio/parquet/writer/Writer.h

@@ -171,7 +171,7 @@ class ParquetWriterFactory : public dwio::common::WriterFactory {
 
   std::unique_ptr<dwio::common::Writer> createWriter(
       std::unique_ptr<dwio::common::FileSink> sink,
-      const dwio::common::WriterOptions& options) override;
+      std::shared_ptr<dwio::common::WriterOptions> options) override;
 };
 
 } // namespace facebook::velox::parquet

velox/exec/fuzzer/PrestoQueryRunner.cpp

@@ -34,7 +34,6 @@
 using namespace facebook::velox;
 
 namespace facebook::velox::exec::test {
-
 namespace {
 
 void writeToFile(
@@ -43,9 +42,9 @@ void writeToFile(
     memory::MemoryPool* pool) {
   VELOX_CHECK_GT(data.size(), 0);
 
-  dwio::common::WriterOptions options;
-  options.schema = data[0]->type();
-  options.memoryPool = pool;
+  auto options = std::make_shared<dwio::common::WriterOptions>();
+  options->schema = data[0]->type();
+  options->memoryPool = pool;
 
   auto writeFile = std::make_unique<LocalWriteFile>(path, true, false);
   auto sink =
@@ -159,6 +158,7 @@ class ServerResponse {
  private:
   folly::dynamic response_;
 };
+
 } // namespace
 
 PrestoQueryRunner::PrestoQueryRunner(

velox/exec/fuzzer/PrestoQueryRunner.h

@@ -23,6 +23,7 @@
 #include "velox/vector/ComplexVector.h"
 
 namespace facebook::velox::exec::test {
+
 template <typename T>
 T extractSingleValue(const std::vector<RowVectorPtr>& data) {
   auto simpleVector = data[0]->childAt(0)->as<SimpleVector<T>>();

velox/exec/tests/utils/PlanBuilder.cpp

@@ -383,7 +383,8 @@ PlanBuilder& PlanBuilder::tableWrite(
     const dwio::common::FileFormat fileFormat,
     const std::vector<std::string>& aggregates,
     const std::string& connectorId,
-    const std::unordered_map<std::string, std::string>& serdeParameters) {
+    const std::unordered_map<std::string, std::string>& serdeParameters,
+    const std::shared_ptr<dwio::common::WriterOptions>& writerOptions) {
   VELOX_CHECK_NOT_NULL(planNode_, "TableWrite cannot be the source node");
   auto rowType = planNode_->outputType();
 
@@ -418,7 +419,8 @@ PlanBuilder& PlanBuilder::tableWrite(
       fileFormat,
       bucketProperty,
       common::CompressionKind_NONE,
-      serdeParameters);
+      serdeParameters,
+      writerOptions);
 
   auto insertHandle =
       std::make_shared<core::InsertTableHandle>(connectorId, hiveHandle);

velox/exec/tests/utils/PlanBuilder.h

@@ -435,6 +435,9 @@ class PlanBuilder {
   /// @param fileFormat File format to use for the written data.
   /// @param aggregates Aggregations for column statistics collection during
   /// write.
+  /// @param connectorId Name used to register the connector.
+  /// @param serdeParameters Additional parameters passed to the writer.
+  /// @param Option objects passed to the writer.
   PlanBuilder& tableWrite(
       const std::string& outputDirectoryPath,
       const std::vector<std::string>& partitionBy,
@@ -446,7 +449,9 @@ class PlanBuilder {
           dwio::common::FileFormat::DWRF,
       const std::vector<std::string>& aggregates = {},
       const std::string& connectorId = "test-hive",
-      const std::unordered_map<std::string, std::string>& serdeParameters = {});
+      const std::unordered_map<std::string, std::string>& serdeParameters = {},
+      const std::shared_ptr<dwio::common::WriterOptions>& writerOptions =
+          nullptr);
 
   /// Add a TableWriteMergeNode.
   PlanBuilder& tableWriteMerge(