apacheGH-33850: [C++] Allow Substrait's default extension provider to…

… be configured (apache#34042) See apache#33850 * Closes: apache#33850 Authored-by: Yaron Gvili <rtpsw@hotmail.com> Signed-off-by: Weston Pace <weston.pace@gmail.com>
gringasalpastor · Feb 17, 2023 · ed53892 · ed53892
1 parent 4882350
commit ed53892
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/cpp/src/arrow/engine/substrait/options.cc b/cpp/src/arrow/engine/substrait/options.cc
@@ -122,11 +122,23 @@ class DefaultExtensionProvider : public BaseExtensionProvider {
   }
 };
 
-std::shared_ptr<ExtensionProvider> ExtensionProvider::kDefaultExtensionProvider =
+namespace {
+
+std::shared_ptr<ExtensionProvider> g_default_extension_provider =
     std::make_shared<DefaultExtensionProvider>();
 
+std::mutex g_default_extension_provider_mutex;
+
+}  // namespace
+
 std::shared_ptr<ExtensionProvider> default_extension_provider() {
-  return ExtensionProvider::kDefaultExtensionProvider;
+  std::unique_lock<std::mutex> lock(g_default_extension_provider_mutex);
+  return g_default_extension_provider;
+}
+
+void set_default_extension_provider(const std::shared_ptr<ExtensionProvider>& provider) {
+  std::unique_lock<std::mutex> lock(g_default_extension_provider_mutex);
+  g_default_extension_provider = provider;
 }
 
 }  // namespace engine

diff --git a/cpp/src/arrow/engine/substrait/options.h b/cpp/src/arrow/engine/substrait/options.h
@@ -74,14 +74,19 @@ class ARROW_ENGINE_EXPORT ExtensionDetails {
 
 class ARROW_ENGINE_EXPORT ExtensionProvider {
  public:
-  static std::shared_ptr<ExtensionProvider> kDefaultExtensionProvider;
   virtual ~ExtensionProvider() = default;
   virtual Result<RelationInfo> MakeRel(const std::vector<DeclarationInfo>& inputs,
                                        const ExtensionDetails& ext_details,
                                        const ExtensionSet& ext_set) = 0;
 };
 
+/// \brief Get the default extension provider
 ARROW_ENGINE_EXPORT std::shared_ptr<ExtensionProvider> default_extension_provider();
+/// \brief Set the default extension provider
+///
+/// \param[in] provider the new provider to be set as default
+ARROW_ENGINE_EXPORT void set_default_extension_provider(
+    const std::shared_ptr<ExtensionProvider>& provider);
 
 /// Options that control the conversion between Substrait and Acero representations of a
 /// plan.