Add TableWriterBuilder::insertHandle() and handling.

Author: mhaseeb123

velox/exec/tests/utils/PlanBuilder.cpp

@@ -234,46 +234,50 @@ core::PlanNodePtr PlanBuilder::TableWriterBuilder::build(core::PlanNodeId id) {
   // upstream operator.
   auto outputType = outputType_ ? outputType_ : upstreamNode->outputType();
 
-  // Create column handles.
-  std::vector<std::shared_ptr<const connector::hive::HiveColumnHandle>>
-      columnHandles;
-  for (auto i = 0; i < outputType->size(); ++i) {
-    const auto column = outputType->nameOf(i);
-    const bool isPartitionKey =
-        std::find(partitionBy_.begin(), partitionBy_.end(), column) !=
-        partitionBy_.end();
-    columnHandles.push_back(std::make_shared<connector::hive::HiveColumnHandle>(
-        column,
-        isPartitionKey
-            ? connector::hive::HiveColumnHandle::ColumnType::kPartitionKey
-            : connector::hive::HiveColumnHandle::ColumnType::kRegular,
-        outputType->childAt(i),
-        outputType->childAt(i)));
-  }
+  // If insertHandle_ is not specified, build a HiveInsertTableHandle along with
+  // columnHandles, bucketProperty and locationHandle.
+  if (!insertHandle_) {
+    // Create column handles.
+    std::vector<std::shared_ptr<const connector::hive::HiveColumnHandle>>
+        columnHandles;
+    for (auto i = 0; i < outputType->size(); ++i) {
+      const auto column = outputType->nameOf(i);
+      const bool isPartitionKey =
+          std::find(partitionBy_.begin(), partitionBy_.end(), column) !=
+          partitionBy_.end();
+      columnHandles.push_back(std::make_shared<connector::hive::HiveColumnHandle>(
+          column,
+          isPartitionKey
+              ? connector::hive::HiveColumnHandle::ColumnType::kPartitionKey
+              : connector::hive::HiveColumnHandle::ColumnType::kRegular,
+          outputType->childAt(i),
+          outputType->childAt(i)));
+    }
 
-  auto locationHandle = std::make_shared<connector::hive::LocationHandle>(
-      outputDirectoryPath_,
-      outputDirectoryPath_,
-      connector::hive::LocationHandle::TableType::kNew,
-      outputFileName_);
+    auto locationHandle = std::make_shared<connector::hive::LocationHandle>(
+        outputDirectoryPath_,
+        outputDirectoryPath_,
+        connector::hive::LocationHandle::TableType::kNew,
+        outputFileName_);
 
-  std::shared_ptr<HiveBucketProperty> bucketProperty;
-  if (bucketCount_ != 0) {
-    bucketProperty =
-        buildHiveBucketProperty(outputType, bucketCount_, bucketedBy_, sortBy_);
-  }
+    std::shared_ptr<HiveBucketProperty> bucketProperty;
+    if (bucketCount_ != 0) {
+      bucketProperty =
+          buildHiveBucketProperty(outputType, bucketCount_, bucketedBy_, sortBy_);
+    }
 
-  auto hiveHandle = std::make_shared<connector::hive::HiveInsertTableHandle>(
-      columnHandles,
-      locationHandle,
-      fileFormat_,
-      bucketProperty,
-      compressionKind_,
-      serdeParameters_,
-      options_);
+    auto hiveHandle = std::make_shared<connector::hive::HiveInsertTableHandle>(
+        columnHandles,
+        locationHandle,
+        fileFormat_,
+        bucketProperty,
+        compressionKind_,
+        serdeParameters_,
+        options_);
 
-  auto insertHandle =
+    insertHandle_ =
       std::make_shared<core::InsertTableHandle>(connectorId_, hiveHandle);
+  }
 
   std::shared_ptr<core::AggregationNode> aggregationNode;
   if (!aggregates_.empty()) {

velox/exec/tests/utils/PlanBuilder.h

@@ -347,6 +347,14 @@ class PlanBuilder {
       return *this;
     }
 
+    /// @param insertHandle TableInsertHandle (optional). Other builder arguments such as
+    /// the `connectorId`, `outputDirectoryPath`, `fileFormat` and so on will be ignored.
+    TableWriterBuilder& insertHandle(
+        std::shared_ptr<core::InsertTableHandle> insertHandle) {
+      insertHandle_ = std::move(insertHandle);
+      return *this;
+    }
+
     /// @param partitionBy Specifies the partition key columns.
     TableWriterBuilder& partitionBy(std::vector<std::string> partitionBy) {
       partitionBy_ = std::move(partitionBy);
@@ -423,6 +431,7 @@ class PlanBuilder {
     std::string outputDirectoryPath_;
     std::string outputFileName_;
     std::string connectorId_{kHiveDefaultConnectorId};
+    shared_ptr<core::InsertTableHandle> insertHandle_;
 
     std::vector<std::string> partitionBy_;
     int32_t bucketCount_{0};