This is an automated email from the ASF dual-hosted git repository.

taiyang-li pushed a commit to branch fake_add_bolt_backend
in repository https://gitbox.apache.org/repos/asf/gluten.git

commit 91d67dc30d65697ff0456b50a8043e06aea091d7
Author: liyang.127 <[email protected]>
AuthorDate: Mon Jun 29 21:37:32 2026 +0800

    [GLUTEN][CORE] Extract shared Substrait extension registry
    
    Move the duplicated Substrait extension-function registry logic from the 
Velox
    collector into cpp/core so the first-tier Bolt/Velox shared patch can land 
as a
    single common implementation.
    
    Also keep the collector's VeloxSubstraitSignature include after the 
extraction
    so TypePtr remains visible in the header and the patch commit records the
    follow-up compile fix together with the main refactor.
---
 cpp/core/CMakeLists.txt                            |  1 +
 .../substrait/SubstraitExtensionRegistry.cc}       | 43 ++++++--------
 .../substrait/SubstraitExtensionRegistry.h}        | 24 +++-----
 cpp/velox/substrait/SubstraitExtensionCollector.cc | 39 ++----------
 cpp/velox/substrait/SubstraitExtensionCollector.h  | 69 ++--------------------
 5 files changed, 37 insertions(+), 139 deletions(-)

diff --git a/cpp/core/CMakeLists.txt b/cpp/core/CMakeLists.txt
index 9530d603f5..d976e92bd0 100644
--- a/cpp/core/CMakeLists.txt
+++ b/cpp/core/CMakeLists.txt
@@ -121,6 +121,7 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS
     compute/Runtime.cc
     compute/ProtobufUtils.cc
     compute/ResultIterator.cc
+    substrait/SubstraitExtensionRegistry.cc
     config/GlutenConfig.cc
     jni/JniWrapper.cc
     memory/AllocationListener.cc
diff --git a/cpp/velox/substrait/SubstraitExtensionCollector.cc 
b/cpp/core/substrait/SubstraitExtensionRegistry.cc
similarity index 67%
copy from cpp/velox/substrait/SubstraitExtensionCollector.cc
copy to cpp/core/substrait/SubstraitExtensionRegistry.cc
index 472ef04e3f..1dd00933cf 100644
--- a/cpp/velox/substrait/SubstraitExtensionCollector.cc
+++ b/cpp/core/substrait/SubstraitExtensionRegistry.cc
@@ -15,21 +15,12 @@
  * limitations under the License.
  */
 
-#include "SubstraitExtensionCollector.h"
+#include "SubstraitExtensionRegistry.h"
 
 namespace gluten {
 
-int SubstraitExtensionCollector::getReferenceNumber(
-    const std::string& functionName,
-    const std::vector<TypePtr>& arguments) {
-  const auto& substraitFunctionSignature = 
VeloxSubstraitSignature::toSubstraitSignature(functionName, arguments);
-  // TODO: Currently we treat all velox registry based function signatures as
-  // custom substrait extension, so no uri link and leave it as empty.
-  return getReferenceNumber({"", substraitFunctionSignature});
-}
-
 template <typename T>
-bool SubstraitExtensionCollector::BiDirectionHashMap<T>::putIfAbsent(const 
int& key, const T& value) {
+bool SubstraitExtensionRegistry::BiDirectionHashMap<T>::putIfAbsent(const int& 
key, const T& value) {
   if (forwardMap_.find(key) == forwardMap_.end() && reverseMap_.find(value) == 
reverseMap_.end()) {
     forwardMap_[key] = value;
     reverseMap_[value] = key;
@@ -38,7 +29,21 @@ bool 
SubstraitExtensionCollector::BiDirectionHashMap<T>::putIfAbsent(const int&
   return false;
 }
 
-void SubstraitExtensionCollector::addExtensionsToPlan(::substrait::Plan* plan) 
const {
+SubstraitExtensionRegistry::SubstraitExtensionRegistry() {
+  extensionFunctions_ = 
std::make_shared<BiDirectionHashMap<ExtensionFunctionId>>();
+}
+
+int SubstraitExtensionRegistry::getReferenceNumber(const ExtensionFunctionId& 
extensionFunctionId) {
+  const auto& extensionFunctionAnchorIt = 
extensionFunctions_->reverseMap().find(extensionFunctionId);
+  if (extensionFunctionAnchorIt != extensionFunctions_->reverseMap().end()) {
+    return extensionFunctionAnchorIt->second;
+  }
+  ++functionReferenceNumber_;
+  extensionFunctions_->putIfAbsent(functionReferenceNumber_, 
extensionFunctionId);
+  return functionReferenceNumber_;
+}
+
+void SubstraitExtensionRegistry::addExtensionsToPlan(::substrait::Plan* plan) 
const {
   using SimpleExtensionURI = ::substrait::extensions::SimpleExtensionURI;
   // Currently we don't introduce any substrait extension YAML files, so always
   // only have one URI.
@@ -53,18 +58,4 @@ void 
SubstraitExtensionCollector::addExtensionsToPlan(::substrait::Plan* plan) c
   }
 }
 
-SubstraitExtensionCollector::SubstraitExtensionCollector() {
-  extensionFunctions_ = 
std::make_shared<BiDirectionHashMap<ExtensionFunctionId>>();
-}
-
-int SubstraitExtensionCollector::getReferenceNumber(const ExtensionFunctionId& 
extensionFunctionId) {
-  const auto& extensionFunctionAnchorIt = 
extensionFunctions_->reverseMap().find(extensionFunctionId);
-  if (extensionFunctionAnchorIt != extensionFunctions_->reverseMap().end()) {
-    return extensionFunctionAnchorIt->second;
-  }
-  ++functionReferenceNumber;
-  extensionFunctions_->putIfAbsent(functionReferenceNumber, 
extensionFunctionId);
-  return functionReferenceNumber;
-}
-
 } // namespace gluten
diff --git a/cpp/velox/substrait/SubstraitExtensionCollector.h 
b/cpp/core/substrait/SubstraitExtensionRegistry.h
similarity index 82%
copy from cpp/velox/substrait/SubstraitExtensionCollector.h
copy to cpp/core/substrait/SubstraitExtensionRegistry.h
index 32cb25c92a..3e023ee930 100644
--- a/cpp/velox/substrait/SubstraitExtensionCollector.h
+++ b/cpp/core/substrait/SubstraitExtensionRegistry.h
@@ -17,13 +17,12 @@
 
 #pragma once
 
-#include <optional>
-#include "VeloxSubstraitSignature.h"
+#include <memory>
+#include <string>
+#include <unordered_map>
+
 #include "substrait/algebra.pb.h"
 #include "substrait/plan.pb.h"
-#include "velox/core/Expressions.h"
-#include "velox/core/PlanNode.h"
-#include "velox/type/Type.h"
 
 namespace gluten {
 
@@ -45,13 +44,11 @@ struct ExtensionFunctionId {
 };
 
 /// Assigns unique IDs to function signatures using ExtensionFunctionId.
-class SubstraitExtensionCollector {
+class SubstraitExtensionRegistry {
  public:
-  SubstraitExtensionCollector();
+  SubstraitExtensionRegistry();
 
-  /// Given a scalar function name and argument types, return the functionId
-  /// using ExtensionFunctionId.
-  int getReferenceNumber(const std::string& functionName, const 
std::vector<TypePtr>& arguments);
+  int getReferenceNumber(const ExtensionFunctionId& extensionFunctionId);
 
   /// Add extension functions to Substrait plan.
   void addExtensionsToPlan(::substrait::Plan* plan) const;
@@ -81,15 +78,10 @@ class SubstraitExtensionCollector {
     std::unordered_map<T, int> reverseMap_;
   };
 
-  /// Assigns unique IDs to function signatures using ExtensionFunctionId.
-  int getReferenceNumber(const ExtensionFunctionId& extensionFunctionId);
-
-  int functionReferenceNumber = -1;
+  int functionReferenceNumber_ = -1;
   std::shared_ptr<BiDirectionHashMap<ExtensionFunctionId>> extensionFunctions_;
 };
 
-using SubstraitExtensionCollectorPtr = 
std::shared_ptr<SubstraitExtensionCollector>;
-
 } // namespace gluten
 
 namespace std {
diff --git a/cpp/velox/substrait/SubstraitExtensionCollector.cc 
b/cpp/velox/substrait/SubstraitExtensionCollector.cc
index 472ef04e3f..1a4a971a91 100644
--- a/cpp/velox/substrait/SubstraitExtensionCollector.cc
+++ b/cpp/velox/substrait/SubstraitExtensionCollector.cc
@@ -17,6 +17,8 @@
 
 #include "SubstraitExtensionCollector.h"
 
+#include "VeloxSubstraitSignature.h"
+
 namespace gluten {
 
 int SubstraitExtensionCollector::getReferenceNumber(
@@ -25,46 +27,15 @@ int SubstraitExtensionCollector::getReferenceNumber(
   const auto& substraitFunctionSignature = 
VeloxSubstraitSignature::toSubstraitSignature(functionName, arguments);
   // TODO: Currently we treat all velox registry based function signatures as
   // custom substrait extension, so no uri link and leave it as empty.
-  return getReferenceNumber({"", substraitFunctionSignature});
-}
-
-template <typename T>
-bool SubstraitExtensionCollector::BiDirectionHashMap<T>::putIfAbsent(const 
int& key, const T& value) {
-  if (forwardMap_.find(key) == forwardMap_.end() && reverseMap_.find(value) == 
reverseMap_.end()) {
-    forwardMap_[key] = value;
-    reverseMap_[value] = key;
-    return true;
-  }
-  return false;
+  return extensionRegistry_->getReferenceNumber({"", 
substraitFunctionSignature});
 }
 
 void SubstraitExtensionCollector::addExtensionsToPlan(::substrait::Plan* plan) 
const {
-  using SimpleExtensionURI = ::substrait::extensions::SimpleExtensionURI;
-  // Currently we don't introduce any substrait extension YAML files, so always
-  // only have one URI.
-  SimpleExtensionURI* extensionUri = plan->add_extension_uris();
-  extensionUri->set_extension_uri_anchor(1);
-
-  for (const auto& [referenceNum, functionId] : 
extensionFunctions_->forwardMap()) {
-    auto extensionFunction = 
plan->add_extensions()->mutable_extension_function();
-    
extensionFunction->set_extension_uri_reference(extensionUri->extension_uri_anchor());
-    extensionFunction->set_function_anchor(referenceNum);
-    extensionFunction->set_name(functionId.signature);
-  }
+  extensionRegistry_->addExtensionsToPlan(plan);
 }
 
 SubstraitExtensionCollector::SubstraitExtensionCollector() {
-  extensionFunctions_ = 
std::make_shared<BiDirectionHashMap<ExtensionFunctionId>>();
-}
-
-int SubstraitExtensionCollector::getReferenceNumber(const ExtensionFunctionId& 
extensionFunctionId) {
-  const auto& extensionFunctionAnchorIt = 
extensionFunctions_->reverseMap().find(extensionFunctionId);
-  if (extensionFunctionAnchorIt != extensionFunctions_->reverseMap().end()) {
-    return extensionFunctionAnchorIt->second;
-  }
-  ++functionReferenceNumber;
-  extensionFunctions_->putIfAbsent(functionReferenceNumber, 
extensionFunctionId);
-  return functionReferenceNumber;
+  extensionRegistry_ = std::make_shared<SubstraitExtensionRegistry>();
 }
 
 } // namespace gluten
diff --git a/cpp/velox/substrait/SubstraitExtensionCollector.h 
b/cpp/velox/substrait/SubstraitExtensionCollector.h
index 32cb25c92a..4faa9e6f67 100644
--- a/cpp/velox/substrait/SubstraitExtensionCollector.h
+++ b/cpp/velox/substrait/SubstraitExtensionCollector.h
@@ -17,33 +17,18 @@
 
 #pragma once
 
-#include <optional>
+#include <memory>
+#include <string>
+#include <vector>
+
 #include "VeloxSubstraitSignature.h"
-#include "substrait/algebra.pb.h"
-#include "substrait/plan.pb.h"
+#include "substrait/SubstraitExtensionRegistry.h"
 #include "velox/core/Expressions.h"
 #include "velox/core/PlanNode.h"
 #include "velox/type/Type.h"
 
 namespace gluten {
 
-struct ExtensionFunctionId {
-  /// Substrait extension YAML file uri.
-  std::string uri;
-
-  /// Substrait signature used in the function extension declaration is a
-  /// combination of the name of the function along with a list of input
-  /// argument types.The format is as follows : <function
-  /// name>:<short_arg_type0>_<short_arg_type1>_..._<short_arg_typeN> for more
-  /// detail information about the argument type please refer to link
-  /// https://substrait.io/extensions/#function-signature-compound-names.
-  std::string signature;
-
-  bool operator==(const ExtensionFunctionId& other) const {
-    return (uri == other.uri && signature == other.signature);
-  }
-};
-
 /// Assigns unique IDs to function signatures using ExtensionFunctionId.
 class SubstraitExtensionCollector {
  public:
@@ -57,51 +42,9 @@ class SubstraitExtensionCollector {
   void addExtensionsToPlan(::substrait::Plan* plan) const;
 
  private:
-  /// A bi-direction hash map to keep the relation between reference number and
-  /// either function or type signature.
-  template <class T>
-  class BiDirectionHashMap {
-   public:
-    /// Add (key, value) pair if doesn't exist already, i.e. forwardMap doesn't
-    /// contain the key and reverseMap doesn't contain the value.
-    ///
-    /// @return True if the values were added successfully. False, otherwise.
-    bool putIfAbsent(const int& key, const T& value);
-
-    const std::unordered_map<int, ExtensionFunctionId> forwardMap() const {
-      return forwardMap_;
-    }
-
-    const std::unordered_map<T, int>& reverseMap() const {
-      return reverseMap_;
-    }
-
-   private:
-    std::unordered_map<int, T> forwardMap_;
-    std::unordered_map<T, int> reverseMap_;
-  };
-
-  /// Assigns unique IDs to function signatures using ExtensionFunctionId.
-  int getReferenceNumber(const ExtensionFunctionId& extensionFunctionId);
-
-  int functionReferenceNumber = -1;
-  std::shared_ptr<BiDirectionHashMap<ExtensionFunctionId>> extensionFunctions_;
+  std::shared_ptr<SubstraitExtensionRegistry> extensionRegistry_;
 };
 
 using SubstraitExtensionCollectorPtr = 
std::shared_ptr<SubstraitExtensionCollector>;
 
 } // namespace gluten
-
-namespace std {
-
-/// Hash function of gluten::ExtensionFunctionId.
-template <>
-struct hash<gluten::ExtensionFunctionId> {
-  size_t operator()(const gluten::ExtensionFunctionId& k) const {
-    size_t val = hash<std::string>()(k.uri);
-    val = val * 31 + hash<std::string>()(k.signature);
-    return val;
-  }
-};
-
-}; // namespace std


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to