ARROW-16657: [C++] Support nesting of extension-id-registries

cpp/src/arrow/engine/CMakeLists.txt

@@ -66,6 +66,7 @@ endif()
 
 add_arrow_test(substrait_test
                SOURCES
+               substrait/ext_test.cc
                substrait/serde_test.cc
                EXTRA_LINK_LIBS
                ${ARROW_SUBSTRAIT_TEST_LINK_LIBS}

cpp/src/arrow/engine/substrait/ext_test.cc

@@ -0,0 +1,271 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/engine/substrait/extension_set.h"
+#include "arrow/engine/substrait/util.h"
+
+#include <google/protobuf/descriptor.h>
+#include <google/protobuf/util/json_util.h>
+#include <google/protobuf/util/type_resolver_util.h>
+#include <gtest/gtest.h>
+
+#include "arrow/testing/gtest_util.h"
+#include "arrow/testing/matchers.h"
+
+using testing::ElementsAre;
+using testing::Eq;
+using testing::HasSubstr;
+using testing::UnorderedElementsAre;
+
+namespace arrow {
+
+using internal::checked_cast;
+
+namespace engine {
+
+// an extension-id-registry provider to be used as a test parameter
+//
+// we cannot pass a pointer to a nested registry as a test parameter because the
+// shared_ptr in which it is made would not be held and get destructed too early,
+// nor can we pass a shared_ptr to the default nested registry as a test parameter
+// because it is global and must never be cleaned up, so we pass a shared_ptr to a
+// provider that either owns or does not own the registry it provides, depending
+// on the case.
+struct ExtensionIdRegistryProvider {
+  virtual ExtensionIdRegistry* get() const = 0;
+};
+
+struct DefaultExtensionIdRegistryProvider : public ExtensionIdRegistryProvider {
+  virtual ~DefaultExtensionIdRegistryProvider() {}
+  ExtensionIdRegistry* get() const override { return default_extension_id_registry(); }
+};
+
+struct NestedExtensionIdRegistryProvider : public ExtensionIdRegistryProvider {
+  virtual ~NestedExtensionIdRegistryProvider() {}
+  std::shared_ptr<ExtensionIdRegistry> registry_ = substrait::MakeExtensionIdRegistry();
+  ExtensionIdRegistry* get() const override { return &*registry_; }
+};
+
+using Id = ExtensionIdRegistry::Id;
+
+bool operator==(const Id& id1, const Id& id2) {
+  return id1.uri == id2.uri && id1.name == id2.name;
+}
+
+bool operator!=(const Id& id1, const Id& id2) { return !(id1 == id2); }
+
+struct TypeName {
+  std::shared_ptr<DataType> type;
+  util::string_view name;
+};
+
+static const std::vector<TypeName> kTypeNames = {
+    TypeName{uint8(), "u8"},
+    TypeName{uint16(), "u16"},
+    TypeName{uint32(), "u32"},
+    TypeName{uint64(), "u64"},
+    TypeName{float16(), "fp16"},
+    TypeName{null(), "null"},
+    TypeName{month_interval(), "interval_month"},
+    TypeName{day_time_interval(), "interval_day_milli"},
+    TypeName{month_day_nano_interval(), "interval_month_day_nano"},
+};
+
+static const std::vector<util::string_view> kFunctionNames = {
+    "add",
+};
+
+static const std::vector<util::string_view> kTempFunctionNames = {
+    "temp_func_1",
+    "temp_func_2",
+};
+
+static const std::vector<TypeName> kTempTypeNames = {
+    TypeName{timestamp(TimeUnit::SECOND, "temp_tz_1"), "temp_type_1"},
+    TypeName{timestamp(TimeUnit::SECOND, "temp_tz_2"), "temp_type_2"},
+};
+
+static Id kNonExistentId{kArrowExtTypesUri, "non_existent"};
+static TypeName kNonExistentTypeName{timestamp(TimeUnit::SECOND, "non_existent_tz_1"),
+                                     "non_existent_type_1"};
+
+using ExtensionIdRegistryParams =
+    std::tuple<std::shared_ptr<ExtensionIdRegistryProvider>, std::string>;
+
+struct ExtensionIdRegistryTest
+    : public testing::TestWithParam<ExtensionIdRegistryParams> {};
+
+TEST_P(ExtensionIdRegistryTest, GetTypes) {
+  auto provider = std::get<0>(GetParam());
+  auto registry = provider->get();
+
+  for (TypeName e : kTypeNames) {
+    auto id = Id{kArrowExtTypesUri, e.name};
+    for (auto typerec_opt : {registry->GetType(id), registry->GetType(*e.type)}) {
+      ASSERT_TRUE(typerec_opt);
+      auto typerec = typerec_opt.value();
+      ASSERT_EQ(id, typerec.id);
+      ASSERT_EQ(*e.type, *typerec.type);
+    }
+  }
+  ASSERT_FALSE(registry->GetType(kNonExistentId));
+  ASSERT_FALSE(registry->GetType(*kNonExistentTypeName.type));
+}
+
+TEST_P(ExtensionIdRegistryTest, ReregisterTypes) {
+  auto provider = std::get<0>(GetParam());
+  auto registry = provider->get();
+
+  for (TypeName e : kTypeNames) {
+    auto id = Id{kArrowExtTypesUri, e.name};
+    ASSERT_RAISES(Invalid, registry->CanRegisterType(id, e.type));
+    ASSERT_RAISES(Invalid, registry->RegisterType(id, e.type));
+  }
+}
+
+TEST_P(ExtensionIdRegistryTest, GetFunctions) {
+  auto provider = std::get<0>(GetParam());
+  auto registry = provider->get();
+
+  for (util::string_view name : kFunctionNames) {
+    auto id = Id{kArrowExtTypesUri, name};
+    for (auto funcrec_opt : {registry->GetFunction(id), registry->GetFunction(name)}) {
+      ASSERT_TRUE(funcrec_opt);
+      auto funcrec = funcrec_opt.value();
+      ASSERT_EQ(id, funcrec.id);
+      ASSERT_EQ(name, funcrec.function_name);
+    }
+  }
+  ASSERT_FALSE(registry->GetType(kNonExistentId));
+  ASSERT_FALSE(registry->GetType(*kNonExistentTypeName.type));
+}
+
+TEST_P(ExtensionIdRegistryTest, ReregisterFunctions) {
+  auto provider = std::get<0>(GetParam());
+  auto registry = provider->get();
+
+  for (util::string_view name : kFunctionNames) {
+    auto id = Id{kArrowExtTypesUri, name};
+    ASSERT_RAISES(Invalid, registry->CanRegisterFunction(id, name.to_string()));
+    ASSERT_RAISES(Invalid, registry->RegisterFunction(id, name.to_string()));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Substrait, ExtensionIdRegistryTest,
+    testing::Values(
+        std::make_tuple(std::make_shared<DefaultExtensionIdRegistryProvider>(),
+                        "default"),
+        std::make_tuple(std::make_shared<NestedExtensionIdRegistryProvider>(),
+                        "nested")));
+
+TEST(ExtensionIdRegistryTest, RegisterTempTypes) {
+  auto default_registry = default_extension_id_registry();
+  constexpr int rounds = 3;
+  for (int i = 0; i < rounds; i++) {
+    auto registry = substrait::MakeExtensionIdRegistry();
+
+    for (TypeName e : kTempTypeNames) {
+      auto id = Id{kArrowExtTypesUri, e.name};
+      ASSERT_OK(registry->CanRegisterType(id, e.type));
+      ASSERT_OK(registry->RegisterType(id, e.type));
+      ASSERT_RAISES(Invalid, registry->CanRegisterType(id, e.type));
+      ASSERT_RAISES(Invalid, registry->RegisterType(id, e.type));
+      ASSERT_OK(default_registry->CanRegisterType(id, e.type));
+    }
+  }
+}
+
+TEST(ExtensionIdRegistryTest, RegisterTempFunctions) {
+  auto default_registry = default_extension_id_registry();
+  constexpr int rounds = 3;
+  for (int i = 0; i < rounds; i++) {
+    auto registry = substrait::MakeExtensionIdRegistry();
+
+    for (util::string_view name : kTempFunctionNames) {
+      auto id = Id{kArrowExtTypesUri, name};
+      ASSERT_OK(registry->CanRegisterFunction(id, name.to_string()));
+      ASSERT_OK(registry->RegisterFunction(id, name.to_string()));
+      ASSERT_RAISES(Invalid, registry->CanRegisterFunction(id, name.to_string()));
+      ASSERT_RAISES(Invalid, registry->RegisterFunction(id, name.to_string()));
+      ASSERT_OK(default_registry->CanRegisterFunction(id, name.to_string()));
+    }
+  }
+}
+
+TEST(ExtensionIdRegistryTest, RegisterNestedTypes) {
+  std::shared_ptr<DataType> type1 = kTempTypeNames[0].type;
+  std::shared_ptr<DataType> type2 = kTempTypeNames[1].type;
+  auto id1 = Id{kArrowExtTypesUri, kTempTypeNames[0].name};
+  auto id2 = Id{kArrowExtTypesUri, kTempTypeNames[1].name};
+
+  auto default_registry = default_extension_id_registry();
+  constexpr int rounds = 3;
+  for (int i = 0; i < rounds; i++) {
+    auto registry1 = nested_extension_id_registry(default_registry);
+
+    ASSERT_OK(registry1->CanRegisterType(id1, type1));
+    ASSERT_OK(registry1->RegisterType(id1, type1));
+
+    for (int j = 0; j < rounds; j++) {
+      auto registry2 = nested_extension_id_registry(&*registry1);
+
+      ASSERT_OK(registry2->CanRegisterType(id2, type2));
+      ASSERT_OK(registry2->RegisterType(id2, type2));
+      ASSERT_RAISES(Invalid, registry2->CanRegisterType(id2, type2));
+      ASSERT_RAISES(Invalid, registry2->RegisterType(id2, type2));
+      ASSERT_OK(default_registry->CanRegisterType(id2, type2));
+    }
+
+    ASSERT_RAISES(Invalid, registry1->CanRegisterType(id1, type1));
+    ASSERT_RAISES(Invalid, registry1->RegisterType(id1, type1));
+    ASSERT_OK(default_registry->CanRegisterType(id1, type1));
+  }
+}
+
+TEST(ExtensionIdRegistryTest, RegisterNestedFunctions) {
+  util::string_view name1 = kTempFunctionNames[0];
+  util::string_view name2 = kTempFunctionNames[1];
+  auto id1 = Id{kArrowExtTypesUri, name1};
+  auto id2 = Id{kArrowExtTypesUri, name2};
+
+  auto default_registry = default_extension_id_registry();
+  constexpr int rounds = 3;
+  for (int i = 0; i < rounds; i++) {
+    auto registry1 = substrait::MakeExtensionIdRegistry();
+
+    ASSERT_OK(registry1->CanRegisterFunction(id1, name1.to_string()));
+    ASSERT_OK(registry1->RegisterFunction(id1, name1.to_string()));
+
+    for (int j = 0; j < rounds; j++) {
+      auto registry2 = substrait::MakeExtensionIdRegistry();
+
+      ASSERT_OK(registry2->CanRegisterFunction(id2, name2.to_string()));
+      ASSERT_OK(registry2->RegisterFunction(id2, name2.to_string()));
+      ASSERT_RAISES(Invalid, registry2->CanRegisterFunction(id2, name2.to_string()));
+      ASSERT_RAISES(Invalid, registry2->RegisterFunction(id2, name2.to_string()));
+      ASSERT_OK(default_registry->CanRegisterFunction(id2, name2.to_string()));
+    }
+
+    ASSERT_RAISES(Invalid, registry1->CanRegisterFunction(id1, name1.to_string()));
+    ASSERT_RAISES(Invalid, registry1->RegisterFunction(id1, name1.to_string()));
+    ASSERT_OK(default_registry->CanRegisterFunction(id1, name1.to_string()));
+  }
+}
+
+}  // namespace engine
+}  // namespace arrow