1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17
18 #include <memory>
19 #include <string>
20 #include <vector>
21
22 #include <gtest/gtest.h>
23
24 #include "arrow/compute/function.h"
25 #include "arrow/compute/registry.h"
26 #include "arrow/result.h"
27 #include "arrow/status.h"
28 #include "arrow/testing/gtest_util.h"
29 #include "arrow/util/macros.h"
30
31 namespace arrow {
32 namespace compute {
33
34 class TestRegistry : public ::testing::Test {
35 public:
SetUp()36 void SetUp() { registry_ = FunctionRegistry::Make(); }
37
38 protected:
39 std::unique_ptr<FunctionRegistry> registry_;
40 };
41
TEST_F(TestRegistry,CreateBuiltInRegistry)42 TEST_F(TestRegistry, CreateBuiltInRegistry) {
43 // This does DCHECK_OK internally for now so this will fail in debug builds
44 // if there is a problem initializing the global function registry
45 FunctionRegistry* registry = GetFunctionRegistry();
46 ARROW_UNUSED(registry);
47 }
48
TEST_F(TestRegistry,Basics)49 TEST_F(TestRegistry, Basics) {
50 ASSERT_EQ(0, registry_->num_functions());
51
52 std::shared_ptr<Function> func =
53 std::make_shared<ScalarFunction>("f1", Arity::Unary(), /*doc=*/nullptr);
54 ASSERT_OK(registry_->AddFunction(func));
55 ASSERT_EQ(1, registry_->num_functions());
56
57 func = std::make_shared<VectorFunction>("f0", Arity::Binary(), /*doc=*/nullptr);
58 ASSERT_OK(registry_->AddFunction(func));
59 ASSERT_EQ(2, registry_->num_functions());
60
61 ASSERT_OK_AND_ASSIGN(std::shared_ptr<const Function> f1, registry_->GetFunction("f1"));
62 ASSERT_EQ("f1", f1->name());
63
64 // Non-existent function
65 ASSERT_RAISES(KeyError, registry_->GetFunction("f2"));
66
67 // Try adding a function with name collision
68 func = std::make_shared<ScalarAggregateFunction>("f1", Arity::Unary(), /*doc=*/nullptr);
69 ASSERT_RAISES(KeyError, registry_->AddFunction(func));
70
71 // Allow overwriting by flag
72 ASSERT_OK(registry_->AddFunction(func, /*allow_overwrite=*/true));
73 ASSERT_OK_AND_ASSIGN(f1, registry_->GetFunction("f1"));
74 ASSERT_EQ(Function::SCALAR_AGGREGATE, f1->kind());
75
76 std::vector<std::string> expected_names = {"f0", "f1"};
77 ASSERT_EQ(expected_names, registry_->GetFunctionNames());
78
79 // Aliases
80 ASSERT_RAISES(KeyError, registry_->AddAlias("f33", "f3"));
81 ASSERT_OK(registry_->AddAlias("f11", "f1"));
82 ASSERT_OK_AND_ASSIGN(std::shared_ptr<const Function> f2, registry_->GetFunction("f11"));
83 ASSERT_EQ(func, f2);
84 }
85
86 } // namespace compute
87 } // namespace arrow
88