[native] Prepare for the actual support of distributed procedures

Author: hantangwangd

presto-native-execution/presto_cpp/main/connectors/IcebergPrestoToVeloxConnector.cpp

@@ -269,6 +269,14 @@ IcebergPrestoToVeloxConnector::toVeloxTableHandle(
       typeParser);
 }
 
+std::unique_ptr<velox::connector::ConnectorInsertTableHandle>
+IcebergPrestoToVeloxConnector::toVeloxInsertTableHandle(
+    const protocol::ExecuteProcedureHandle* executeProcedureHandle,
+    const TypeParser& typeParser) const {
+  // TODO: requires data insertion support
+  VELOX_FAIL("Not yet supported, requires data insertion support first");
+}
+
 std::unique_ptr<protocol::ConnectorProtocol>
 IcebergPrestoToVeloxConnector::createConnectorProtocol() const {
   return std::make_unique<protocol::iceberg::IcebergConnectorProtocol>();

presto-native-execution/presto_cpp/main/connectors/IcebergPrestoToVeloxConnector.h

@@ -38,6 +38,10 @@ class IcebergPrestoToVeloxConnector final : public PrestoToVeloxConnector {
       const TypeParser& typeParser,
       velox::connector::ColumnHandleMap& assignments) const final;
 
+  std::unique_ptr<velox::connector::ConnectorInsertTableHandle> toVeloxInsertTableHandle(
+      const protocol::ExecuteProcedureHandle* executeProcedureHandle,
+      const TypeParser& typeParser) const final;
+
   std::unique_ptr<protocol::ConnectorProtocol> createConnectorProtocol()
       const final;
 };

presto-native-execution/presto_cpp/main/connectors/PrestoToVeloxConnector.h

@@ -103,6 +103,14 @@ class PrestoToVeloxConnector {
     return {};
   }
 
+  [[nodiscard]] virtual std::unique_ptr<
+      velox::connector::ConnectorInsertTableHandle>
+  toVeloxInsertTableHandle(
+      const protocol::ExecuteProcedureHandle* executeProcedureHandle,
+      const TypeParser& typeParser) const {
+    return {};
+  }
+
   [[nodiscard]] std::unique_ptr<velox::core::PartitionFunctionSpec>
   createVeloxPartitionFunctionSpec(
       const protocol::ConnectorPartitioningHandle* partitioningHandle,

presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp

@@ -1541,6 +1541,57 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       sourceVeloxPlan);
 }
 
+std::shared_ptr<const core::TableWriteNode>
+VeloxQueryPlanConverterBase::toVeloxQueryPlan(
+    const std::shared_ptr<const protocol::CallDistributedProcedureNode>& node,
+    const std::shared_ptr<protocol::TableWriteInfo>& tableWriteInfo,
+    const protocol::TaskId& taskId) {
+  const auto executeProcedureHandle = std::dynamic_pointer_cast<protocol::ExecuteProcedureHandle>(tableWriteInfo->writerTarget);
+
+  if (!executeProcedureHandle) {
+    VELOX_UNSUPPORTED(
+        "Unsupported execute procedure handle: {}",
+        toJsonString(tableWriteInfo->writerTarget));
+  }
+
+  std::string connectorId = executeProcedureHandle->handle.connectorId;
+  auto& connector =
+      getPrestoToVeloxConnector(executeProcedureHandle->handle.connectorHandle->_type);
+  auto veloxHandle = connector.toVeloxInsertTableHandle(
+      executeProcedureHandle.get(), typeParser_);
+  auto connectorInsertHandle = std::shared_ptr(std::move(veloxHandle));
+
+  if (!connectorInsertHandle) {
+    VELOX_UNSUPPORTED(
+        "Unsupported execute procedure handle: {}",
+        toJsonString(tableWriteInfo->writerTarget));
+  }
+
+  auto insertTableHandle = std::make_shared<core::InsertTableHandle>(
+      connectorId, connectorInsertHandle);
+
+  const auto outputType = toRowType(
+      generateOutputVariables(
+          {node->rowCountVariable,
+           node->fragmentVariable,
+           node->tableCommitContextVariable},
+          nullptr),
+      typeParser_);
+  const auto sourceVeloxPlan =
+      toVeloxQueryPlan(node->source, tableWriteInfo, taskId);
+
+  return std::make_shared<core::TableWriteNode>(
+      node->id,
+      toRowType(node->columns, typeParser_),
+      node->columnNames,
+      std::nullopt,
+      std::move(insertTableHandle),
+      node->partitioningScheme != nullptr,
+      outputType,
+      getCommitStrategy(),
+      sourceVeloxPlan);
+}
+
 std::shared_ptr<const core::TableWriteNode>
 VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     const std::shared_ptr<const protocol::DeleteNode>& node,
@@ -1929,6 +1980,10 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan(
           std::dynamic_pointer_cast<const protocol::TableWriterNode>(node)) {
     return toVeloxQueryPlan(tableWriter, tableWriteInfo, taskId);
   }
+  if (auto callDistributedProcedure =
+          std::dynamic_pointer_cast<const protocol::CallDistributedProcedureNode>(node)) {
+    return toVeloxQueryPlan(callDistributedProcedure, tableWriteInfo, taskId);
+  }
   if (auto deleteNode =
           std::dynamic_pointer_cast<const protocol::DeleteNode>(node)) {
     return toVeloxQueryPlan(deleteNode, tableWriteInfo, taskId);

presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h

@@ -160,6 +160,11 @@ class VeloxQueryPlanConverterBase {
       const std::shared_ptr<protocol::TableWriteInfo>& tableWriteInfo,
       const protocol::TaskId& taskId);
 
+  std::shared_ptr<const velox::core::TableWriteNode> toVeloxQueryPlan(
+    const std::shared_ptr<const protocol::CallDistributedProcedureNode>& node,
+    const std::shared_ptr<protocol::TableWriteInfo>& tableWriteInfo,
+    const protocol::TaskId& taskId);
+
   std::shared_ptr<const velox::core::TableWriteMergeNode> toVeloxQueryPlan(
       const std::shared_ptr<const protocol::TableWriterMergeNode>& node,
       const std::shared_ptr<protocol::TableWriteInfo>& tableWriteInfo,

presto-native-execution/presto_cpp/presto_protocol/connector/hive/HiveConnectorProtocol.h

@@ -26,5 +26,6 @@ using HiveConnectorProtocol = ConnectorProtocolTemplate<
     HivePartitioningHandle,
     HiveTransactionHandle,
     NotImplemented,
+    NotImplemented,
     NotImplemented>;
 } // namespace facebook::presto::protocol::hive

presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/IcebergConnectorProtocol.h

@@ -27,6 +27,7 @@ using IcebergConnectorProtocol = ConnectorProtocolTemplate<
     IcebergSplit,
     NotImplemented,
     hive::HiveTransactionHandle,
+    IcebergDistributedProcedureHandle,
     NotImplemented,
     NotImplemented>;
 

presto-native-execution/presto_cpp/presto_protocol/connector/tpch/TpchConnectorProtocol.h

@@ -29,6 +29,7 @@ using TpchConnectorProtocol = ConnectorProtocolTemplate<
     TpchPartitioningHandle,
     TpchTransactionHandle,
     NotImplemented,
+    NotImplemented,
     NotImplemented>;
 
 } // namespace facebook::presto::protocol::tpch

presto-native-execution/presto_cpp/presto_protocol/core/ConnectorProtocol.h

@@ -75,6 +75,13 @@ class ConnectorProtocol {
     const std::string& thrift,
     std::shared_ptr<ConnectorInsertTableHandle>& proto) const = 0;
 
+  virtual void to_json(
+        json& j,
+        const std::shared_ptr<ConnectorDistributedProcedureHandle>& p) const = 0;
+  virtual void from_json(
+      const json& j,
+      std::shared_ptr<ConnectorDistributedProcedureHandle>& p) const = 0;
+
   virtual void to_json(
       json& j,
       const std::shared_ptr<ConnectorOutputTableHandle>& p) const = 0;
@@ -117,7 +124,7 @@ class ConnectorProtocol {
       std::string& thrift) const = 0;
   virtual void deserialize(
     const std::string& thrift,
-    std::shared_ptr<ConnectorTransactionHandle>& proto) const = 0;    
+    std::shared_ptr<ConnectorTransactionHandle>& proto) const = 0;
 
   virtual void to_json(
       json& j,
@@ -153,6 +160,7 @@ template <
     typename ConnectorSplitType = NotImplemented,
     typename ConnectorPartitioningHandleType = NotImplemented,
     typename ConnectorTransactionHandleType = NotImplemented,
+    typename ConnectorDistributedProcedureHandleType = NotImplemented,
     typename ConnectorDeleteTableHandleType = NotImplemented,
     typename ConnectorIndexHandleType = NotImplemented>
 class ConnectorProtocolTemplate final : public ConnectorProtocol {
@@ -221,6 +229,15 @@ class ConnectorProtocolTemplate final : public ConnectorProtocol {
     deserializeTemplate<ConnectorInsertTableHandleType>(thrift, proto);
   }
 
+  void to_json(json& j, const std::shared_ptr<ConnectorDistributedProcedureHandle>& p)
+    const final {
+    to_json_template<ConnectorDistributedProcedureHandleType>(j, p);
+  }
+  void from_json(const json& j, std::shared_ptr<ConnectorDistributedProcedureHandle>& p)
+      const final {
+    from_json_template<ConnectorDistributedProcedureHandleType>(j, p);
+  }
+
   void to_json(json& j, const std::shared_ptr<ConnectorOutputTableHandle>& p)
       const final {
     to_json_template<ConnectorOutputTableHandleType>(j, p);
@@ -406,6 +423,7 @@ using SystemConnectorProtocol = ConnectorProtocolTemplate<
     SystemPartitioningHandle,
     SystemTransactionHandle,
     NotImplemented,
+    NotImplemented,
     NotImplemented>;
 
 } // namespace facebook::presto::protocol