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,106 @@ 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());
+
+ 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,7 @@ 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;
 
  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(