Add iceberg partition transforms. Co-authored-by: Chengcheng Jin <[email protected]>

Author: PingLiuPing

velox/connectors/hive/HiveDataSink.cpp

@@ -374,7 +374,18 @@ HiveDataSink::HiveDataSink(
                     *insertTableHandle->bucketProperty(),
                     inputType)
               : nullptr,
-          getNonPartitionChannels(insertTableHandle)) {}
+          getPartitionChannels(insertTableHandle),
+          getNonPartitionChannels(insertTableHandle),
+          !getPartitionChannels(insertTableHandle).empty()
+              ? std::make_unique<PartitionIdGenerator>(
+                    inputType,
+                    getPartitionChannels(insertTableHandle),
+                    hiveConfig->maxPartitionsPerWriters(
+                        connectorQueryCtx->sessionProperties()),
+                    connectorQueryCtx->memoryPool(),
+                    hiveConfig->isPartitionPathAsLowerCase(
+                        connectorQueryCtx->sessionProperties()))
+              : nullptr) {}
 
 HiveDataSink::HiveDataSink(
     RowTypePtr inputType,
@@ -384,7 +395,9 @@ HiveDataSink::HiveDataSink(
     const std::shared_ptr<const HiveConfig>& hiveConfig,
     uint32_t bucketCount,
     std::unique_ptr<core::PartitionFunction> bucketFunction,
-    const std::vector<column_index_t>& dataChannels)
+    const std::vector<column_index_t>& partitionChannels,
+    const std::vector<column_index_t>& dataChannels,
+    std::unique_ptr<PartitionIdGenerator> partitionIdGenerator)
     : inputType_(std::move(inputType)),
       insertTableHandle_(std::move(insertTableHandle)),
       connectorQueryCtx_(connectorQueryCtx),
@@ -393,17 +406,8 @@ HiveDataSink::HiveDataSink(
       updateMode_(getUpdateMode()),
       maxOpenWriters_(hiveConfig_->maxPartitionsPerWriters(
           connectorQueryCtx->sessionProperties())),
-      partitionChannels_(getPartitionChannels(insertTableHandle_)),
-      partitionIdGenerator_(
-          !partitionChannels_.empty()
-              ? std::make_unique<PartitionIdGenerator>(
-                    inputType_,
-                    partitionChannels_,
-                    maxOpenWriters_,
-                    connectorQueryCtx_->memoryPool(),
-                    hiveConfig_->isPartitionPathAsLowerCase(
-                        connectorQueryCtx->sessionProperties()))
-              : nullptr),
+      partitionChannels_(partitionChannels),
+      partitionIdGenerator_(std::move(partitionIdGenerator)),
       dataChannels_(dataChannels),
       bucketCount_(static_cast<int32_t>(bucketCount)),
       bucketFunction_(std::move(bucketFunction)),
@@ -753,6 +757,32 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   ioStats_.emplace_back(std::make_shared<io::IoStatistics>());
   setMemoryReclaimers(writerInfo_.back().get(), ioStats_.back().get());
 
+  auto options = createWriterOptions();
+  // 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(),
+          }),
+      options);
+  writer = maybeCreateBucketSortWriter(std::move(writer));
+  writers_.emplace_back(std::move(writer));
+
+  extendBuffersForPartitionedTables();
+
+  writerIndexMap_.emplace(id, writers_.size() - 1);
+  return writerIndexMap_[id];
+}
+
+std::shared_ptr<dwio::common::WriterOptions> HiveDataSink::createWriterOptions()
+    const {
   // Take the writer options provided by the user as a starting point, or
   // allocate a new one.
   auto options = insertTableHandle_->writerOptions();
@@ -802,28 +832,7 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   options->adjustTimestampToTimezone =
       connectorQueryCtx_->adjustTimestampToTimezone();
   options->processConfigs(*hiveConfig_->config(), *connectorSessionProperties);
-
-  // 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(),
-          }),
-      options);
-  writer = maybeCreateBucketSortWriter(std::move(writer));
-  writers_.emplace_back(std::move(writer));
-
-  extendBuffersForPartitionedTables();
-
-  writerIndexMap_.emplace(id, writers_.size() - 1);
-  return writerIndexMap_[id];
+  return options;
 }
 
 std::optional<std::string> HiveDataSink::getPartitionName(

velox/connectors/hive/HiveDataSink.h

@@ -529,7 +529,9 @@ class HiveDataSink : public DataSink {
       const std::shared_ptr<const HiveConfig>& hiveConfig,
       uint32_t bucketCount,
       std::unique_ptr<core::PartitionFunction> bucketFunction,
-      const std::vector<column_index_t>& dataChannels);
+      const std::vector<column_index_t>& partitionChannels,
+      const std::vector<column_index_t>& dataChannels,
+      std::unique_ptr<PartitionIdGenerator> partitionIdGenerator);
 
   void appendData(RowVectorPtr input) override;
 
@@ -633,6 +635,9 @@ class HiveDataSink : public DataSink {
   // the newly created writer in 'writers_'.
   uint32_t appendWriter(const HiveWriterId& id);
 
+  virtual std::shared_ptr<dwio::common::WriterOptions> createWriterOptions()
+      const;
+
   virtual std::optional<std::string> getPartitionName(
       const HiveWriterId& id) const;
 

velox/connectors/hive/HivePartitionUtil.cpp

@@ -78,11 +78,10 @@ std::pair<std::string, std::string> makePartitionKeyValueString(
     const BaseVector* partitionVector,
     vector_size_t row,
     const std::string& name,
-    bool isDate,
-    const std::string& nullValueName) {
+    bool isDate) {
   using T = typename TypeTraits<Kind>::NativeType;
   if (partitionVector->as<SimpleVector<T>>()->isNullAt(row)) {
-    return std::make_pair(name, nullValueName);
+    return std::make_pair(name, "");
   }
   if (isDate) {
     return std::make_pair(
@@ -100,8 +99,7 @@ std::pair<std::string, std::string> makePartitionKeyValueString(
 
 std::vector<std::pair<std::string, std::string>> extractPartitionKeyValues(
     const RowVectorPtr& partitionsVector,
-    vector_size_t row,
-    const std::string& nullValueName) {
+    vector_size_t row) {
   std::vector<std::pair<std::string, std::string>> partitionKeyValues;
   for (auto i = 0; i < partitionsVector->childrenSize(); i++) {
     partitionKeyValues.push_back(PARTITION_TYPE_DISPATCH(
@@ -110,8 +108,7 @@ std::vector<std::pair<std::string, std::string>> extractPartitionKeyValues(
         partitionsVector->childAt(i)->loadedVector(),
         row,
         asRowType(partitionsVector->type())->nameOf(i),
-        partitionsVector->childAt(i)->type()->isDate(),
-        nullValueName));
+        partitionsVector->childAt(i)->type()->isDate()));
   }
   return partitionKeyValues;
 }

velox/connectors/hive/HivePartitionUtil.h

@@ -21,7 +21,6 @@ namespace facebook::velox::connector::hive {
 
 std::vector<std::pair<std::string, std::string>> extractPartitionKeyValues(
     const RowVectorPtr& partitionsVector,
-    vector_size_t row,
-    const std::string& nullValueName = "");
+    vector_size_t row);
 
 } // namespace facebook::velox::connector::hive

velox/connectors/hive/PartitionIdGenerator.cpp

@@ -29,10 +29,10 @@ PartitionIdGenerator::PartitionIdGenerator(
     uint32_t maxPartitions,
     memory::MemoryPool* pool,
     bool partitionPathAsLowerCase)
-    : pool_(pool),
-      partitionChannels_(std::move(partitionChannels)),
+    : partitionChannels_(std::move(partitionChannels)),
       maxPartitions_(maxPartitions),
-      partitionPathAsLowerCase_(partitionPathAsLowerCase) {
+      partitionPathAsLowerCase_(partitionPathAsLowerCase),
+      pool_(pool) {
   VELOX_USER_CHECK(
       !partitionChannels_.empty(), "There must be at least one partition key.");
   for (auto channel : partitionChannels_) {
@@ -61,6 +61,19 @@ PartitionIdGenerator::PartitionIdGenerator(
   }
 }
 
+PartitionIdGenerator::PartitionIdGenerator(
+    std::vector<column_index_t> partitionChannels,
+    uint32_t maxPartitions,
+    memory::MemoryPool* pool,
+    bool partitionPathAsLowerCase)
+    : partitionChannels_(std::move(partitionChannels)),
+      maxPartitions_(maxPartitions),
+      partitionPathAsLowerCase_(partitionPathAsLowerCase),
+      pool_(pool) {
+  VELOX_USER_CHECK(
+      !partitionChannels_.empty(), "There must be at least one partition key.");
+}
+
 void PartitionIdGenerator::run(
     const RowVectorPtr& input,
     raw_vector<uint64_t>& result) {
@@ -97,11 +110,9 @@ void PartitionIdGenerator::run(
   }
 }
 
-std::string PartitionIdGenerator::partitionName(
-    uint64_t partitionId,
-    const std::string& nullValueName) const {
+std::string PartitionIdGenerator::partitionName(uint64_t partitionId) const {
   return FileUtils::makePartName(
-      extractPartitionKeyValues(partitionValues_, partitionId, nullValueName),
+      extractPartitionKeyValues(partitionValues_, partitionId),
       partitionPathAsLowerCase_);
 }
 
@@ -172,7 +183,7 @@ void PartitionIdGenerator::updateValueToPartitionIdMapping() {
 }
 
 void PartitionIdGenerator::savePartitionValues(
-    uint64_t partitionId,
+    uint32_t partitionId,
     const RowVectorPtr& input,
     vector_size_t row) {
   for (auto i = 0; i < partitionChannels_.size(); ++i) {

velox/connectors/hive/PartitionIdGenerator.h

@@ -38,10 +38,12 @@ class PartitionIdGenerator {
       memory::MemoryPool* pool,
       bool partitionPathAsLowerCase);
 
+  virtual ~PartitionIdGenerator() = default;
+
   /// Generate sequential partition IDs for input vector.
   /// @param input Input RowVector.
   /// @param result Generated integer IDs indexed by input row number.
-  void run(const RowVectorPtr& input, raw_vector<uint64_t>& result);
+  virtual void run(const RowVectorPtr& input, raw_vector<uint64_t>& result);
 
   /// Return the total number of distinct partitions processed so far.
   uint64_t numPartitions() const {
@@ -52,18 +54,38 @@ class PartitionIdGenerator {
   /// style. It is derived from the partitionValues_ at index partitionId.
   /// Partition keys appear in the order of partition columns in the table
   /// schema.
-  std::string partitionName(
-      uint64_t partitionId,
-      const std::string& nullValueName = "") const;
+  virtual std::string partitionName(uint64_t partitionId) const;
 
- private:
-  static constexpr const int32_t kHasherReservePct = 20;
+ protected:
+  PartitionIdGenerator(
+      std::vector<column_index_t> partitionChannels,
+      uint32_t maxPartitions,
+      memory::MemoryPool* pool,
+      bool partitionPathAsLowerCase);
 
   // Computes value IDs using VectorHashers for all rows in 'input'.
   void computeValueIds(
       const RowVectorPtr& input,
       raw_vector<uint64_t>& valueIds);
 
+  const std::vector<column_index_t> partitionChannels_;
+
+  std::vector<std::unique_ptr<exec::VectorHasher>> hashers_;
+
+  // A vector holding unique partition key values. One row per partition. Row
+  // numbers match partition IDs.
+  RowVectorPtr partitionValues_;
+
+  const uint32_t maxPartitions_;
+
+  // A mapping from value ID produced by VectorHashers to a partition ID.
+  std::unordered_map<uint64_t, uint64_t> partitionIds_;
+
+  const bool partitionPathAsLowerCase_;
+
+ private:
+  static constexpr const int32_t kHasherReservePct = 20;
+
   // In case of rehash (when value IDs produced by VectorHashers change), we
   // update value id for pre-existing partitions while keeping partition ids.
   // This method rebuilds 'partitionIds_' by re-calculating the value ids using
@@ -72,29 +94,15 @@ class PartitionIdGenerator {
 
   // Copies partition values of 'row' from 'input' into 'partitionId' row in
   // 'partitionValues_'.
-  void savePartitionValues(
-      uint64_t partitionId,
+  virtual void savePartitionValues(
+      uint32_t partitionId,
       const RowVectorPtr& input,
       vector_size_t row);
 
   memory::MemoryPool* const pool_;
 
-  const std::vector<column_index_t> partitionChannels_;
-
-  const uint32_t maxPartitions_;
-
-  const bool partitionPathAsLowerCase_;
-
-  std::vector<std::unique_ptr<exec::VectorHasher>> hashers_;
   bool hasMultiplierSet_ = false;
 
-  // A mapping from value ID produced by VectorHashers to a partition ID.
-  std::unordered_map<uint64_t, uint64_t> partitionIds_;
-
-  // A vector holding unique partition key values. One row per partition. Row
-  // numbers match partition IDs.
-  RowVectorPtr partitionValues_;
-
   // All rows are set valid to compute partition IDs for all input rows.
   SelectivityVector allRows_;
 };

velox/connectors/hive/SplitReader.cpp

@@ -54,6 +54,18 @@ VectorPtr newConstantFromString(
         pool, size, false, type, std::move(days));
   }
 
+  if constexpr (std::is_same_v<T, int64_t> || std::is_same_v<T, int128_t>) {
+    if (type->isDecimal()) {
+      auto [precision, scale] = getDecimalPrecisionScale(*type);
+      T result;
+      const auto status = DecimalUtil::castFromString<T>(
+          StringView(value.value()), precision, scale, result);
+      VELOX_USER_CHECK(status.ok(), status.message());
+      return std::make_shared<ConstantVector<T>>(
+          pool, size, false, type, std::move(result));
+    }
+  }
+
   if constexpr (std::is_same_v<T, StringView>) {
     return std::make_shared<ConstantVector<StringView>>(
         pool, size, false, type, StringView(value.value()));

velox/connectors/hive/SplitReader.h

@@ -160,18 +160,18 @@ class SplitReader {
       VectorPtr& output,
       const std::vector<BaseVector::CopyRange>& ranges);
 
+  void setPartitionValue(
+      common::ScanSpec* spec,
+      const std::string& partitionKey,
+      const std::optional<std::string>& value) const;
+
  private:
   /// Different table formats may have different meatadata columns.
   /// This function will be used to update the scanSpec for these columns.
-  std::vector<TypePtr> adaptColumns(
+  virtual std::vector<TypePtr> adaptColumns(
       const RowTypePtr& fileType,
       const std::shared_ptr<const velox::RowType>& tableSchema) const;
 
-  void setPartitionValue(
-      common::ScanSpec* spec,
-      const std::string& partitionKey,
-      const std::optional<std::string>& value) const;
-
  protected:
   std::shared_ptr<const HiveConnectorSplit> hiveSplit_;
   const std::shared_ptr<const HiveTableHandle> hiveTableHandle_;

velox/connectors/hive/iceberg/CMakeLists.txt

@@ -14,13 +14,23 @@
 
 velox_add_library(
   velox_hive_iceberg_connector
-  IcebergSplitReader.cpp
+  IcebergDataSink.cpp
+  IcebergPartitionIdGenerator.cpp
   IcebergSplit.cpp
+  IcebergSplitReader.cpp
+  Murmur3.cpp
+  PartitionSpec.cpp
   PositionalDeleteFileReader.cpp
-  IcebergDataSink.cpp)
+  TransformFactory.cpp
+  Transforms.cpp)
 
 velox_link_libraries(velox_hive_iceberg_connector velox_connector Folly::folly)
 
+if(VELOX_ENABLE_PARQUET)
+  velox_link_libraries(velox_hive_iceberg_connector
+                       velox_dwio_arrow_parquet_writer)
+endif()
+
 if(${VELOX_BUILD_TESTING})
   add_subdirectory(tests)
 endif()