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 #pragma once
19 
20 #include <memory>
21 #include <string>
22 
23 #include "arrow/extension_type.h"
24 #include "arrow/testing/visibility.h"
25 #include "arrow/util/macros.h"
26 
27 namespace arrow {
28 
29 class ARROW_TESTING_EXPORT UuidArray : public ExtensionArray {
30  public:
31   using ExtensionArray::ExtensionArray;
32 };
33 
34 class ARROW_TESTING_EXPORT UuidType : public ExtensionType {
35  public:
UuidType()36   UuidType() : ExtensionType(fixed_size_binary(16)) {}
37 
extension_name()38   std::string extension_name() const override { return "uuid"; }
39 
40   bool ExtensionEquals(const ExtensionType& other) const override;
41 
42   std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
43 
44   Result<std::shared_ptr<DataType>> Deserialize(
45       std::shared_ptr<DataType> storage_type,
46       const std::string& serialized) const override;
47 
Serialize()48   std::string Serialize() const override { return "uuid-serialized"; }
49 };
50 
51 class ARROW_TESTING_EXPORT SmallintArray : public ExtensionArray {
52  public:
53   using ExtensionArray::ExtensionArray;
54 };
55 
56 class ARROW_TESTING_EXPORT SmallintType : public ExtensionType {
57  public:
SmallintType()58   SmallintType() : ExtensionType(int16()) {}
59 
extension_name()60   std::string extension_name() const override { return "smallint"; }
61 
62   bool ExtensionEquals(const ExtensionType& other) const override;
63 
64   std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
65 
66   Result<std::shared_ptr<DataType>> Deserialize(
67       std::shared_ptr<DataType> storage_type,
68       const std::string& serialized) const override;
69 
Serialize()70   std::string Serialize() const override { return "smallint"; }
71 };
72 
73 class ARROW_TESTING_EXPORT DictExtensionType : public ExtensionType {
74  public:
DictExtensionType()75   DictExtensionType() : ExtensionType(dictionary(int8(), utf8())) {}
76 
extension_name()77   std::string extension_name() const override { return "dict-extension"; }
78 
79   bool ExtensionEquals(const ExtensionType& other) const override;
80 
81   std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
82 
83   Result<std::shared_ptr<DataType>> Deserialize(
84       std::shared_ptr<DataType> storage_type,
85       const std::string& serialized) const override;
86 
Serialize()87   std::string Serialize() const override { return "dict-extension-serialized"; }
88 };
89 
90 ARROW_TESTING_EXPORT
91 std::shared_ptr<DataType> uuid();
92 
93 ARROW_TESTING_EXPORT
94 std::shared_ptr<DataType> smallint();
95 
96 ARROW_TESTING_EXPORT
97 std::shared_ptr<DataType> dict_extension_type();
98 
99 ARROW_TESTING_EXPORT
100 std::shared_ptr<Array> ExampleUuid();
101 
102 ARROW_TESTING_EXPORT
103 std::shared_ptr<Array> ExampleSmallint();
104 
105 // A RAII class that registers an extension type on construction
106 // and unregisters it on destruction.
107 class ARROW_TESTING_EXPORT ExtensionTypeGuard {
108  public:
109   explicit ExtensionTypeGuard(const std::shared_ptr<DataType>& type);
110   ~ExtensionTypeGuard();
111   ARROW_DEFAULT_MOVE_AND_ASSIGN(ExtensionTypeGuard);
112 
113  protected:
114   ARROW_DISALLOW_COPY_AND_ASSIGN(ExtensionTypeGuard);
115 
116   std::string extension_name_;
117 };
118 
119 }  // namespace arrow
120