ARROW-638: [C++] Complex Number Support via ExtensionTypes

cpp/src/arrow/CMakeLists.txt

@@ -155,6 +155,7 @@ set(ARROW_SRCS
     array/diff.cc
     array/util.cc
     array/validate.cc
+    extensions/complex_type.cc
     builder.cc
     buffer.cc
     chunked_array.cc

cpp/src/arrow/extension_type_test.cc

@@ -27,6 +27,7 @@
 #include "arrow/array/array_nested.h"
 #include "arrow/array/util.h"
 #include "arrow/extension_type.h"
+#include "arrow/extensions/complex_type.h"
 #include "arrow/io/memory.h"
 #include "arrow/ipc/options.h"
 #include "arrow/ipc/reader.h"
@@ -178,7 +179,9 @@ class ExtStructType : public ExtensionType {
 
 class TestExtensionType : public ::testing::Test {
  public:
-  void SetUp() { ASSERT_OK(RegisterExtensionType(std::make_shared<UuidType>())); }
+  void SetUp() {
+    ASSERT_OK(RegisterExtensionType(std::make_shared<UuidType>()));
+  }
 
   void TearDown() {
     if (GetExtensionType("uuid")) {
@@ -187,6 +190,23 @@ class TestExtensionType : public ::testing::Test {
   }
 };
 
+TEST_F(TestExtensionType, ComplexTypeTest) {
+  auto registered_type = GetExtensionType("arrow.complex64");
+  ASSERT_NE(registered_type, nullptr);
+
+  auto type = complex64();
+  ASSERT_EQ(type->id(), Type::EXTENSION);
+
+  const auto& ext_type = static_cast<const ExtensionType&>(*type);
+  std::string serialized = ext_type.Serialize();
+
+  ASSERT_OK_AND_ASSIGN(auto deserialized,
+                       ext_type.Deserialize(fixed_size_list(float32(), 2), serialized));
+
+  ASSERT_TRUE(deserialized->Equals(*type));
+  ASSERT_FALSE(deserialized->Equals(*fixed_size_list(float32(), 2)));
+}
+
 TEST_F(TestExtensionType, ExtensionTypeTest) {
   auto type_not_exist = GetExtensionType("uuid-unknown");
   ASSERT_EQ(type_not_exist, nullptr);

cpp/src/arrow/extensions/complex_type.cc

@@ -0,0 +1,67 @@
+// 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.
+
+// Complex Number Extension Type
+
+#include <mutex>
+#include <thread>
+#include <sstream>
+
+#include "arrow/extensions/complex_type.h"
+
+namespace arrow {
+
+bool ComplexFloatType::ExtensionEquals(const ExtensionType& other) const {
+  const auto& other_ext = static_cast<const ExtensionType&>(other);
+  return other_ext.extension_name() == this->extension_name();
+}
+
+bool ComplexDoubleType::ExtensionEquals(const ExtensionType& other) const {
+  const auto& other_ext = static_cast<const ExtensionType&>(other);
+  return other_ext.extension_name() == this->extension_name();
+}
+
+
+std::shared_ptr<DataType> complex64() {
+  return std::make_shared<ComplexFloatType>();
+}
+
+std::shared_ptr<DataType> complex128() {
+  return std::make_shared<ComplexDoubleType>();
+}
+
+/// NOTE(sjperkins)
+// Suggestions on how to improve this welcome!
+std::once_flag complex_float_registered;
+std::once_flag complex_double_registered;
+
+Status register_complex_types()
+{
+  std::call_once(complex_float_registered,
+                 RegisterExtensionType,
+                 std::make_shared<ComplexFloatType>());
+
+  std::call_once(complex_double_registered,
+                 RegisterExtensionType,
+                 std::make_shared<ComplexDoubleType>());
+
+  return Status::OK();
+}
+
+static Status complex_types_registered = register_complex_types();
+
+};  // namespace arrow

cpp/src/arrow/extensions/complex_type.h

@@ -0,0 +1,107 @@
+// 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.
+
+// Complex Number Extension Type
+#pragma once
+
+#include <complex>
+
+#include "arrow/extension_type.h"
+
+namespace arrow {
+
+
+std::shared_ptr<DataType> complex64();
+std::shared_ptr<DataType> complex128();
+
+
+class ComplexFloatArray : public ExtensionArray {
+ public:
+  using ExtensionArray::ExtensionArray;
+};
+
+class ComplexFloatType : public ExtensionType {
+ public:
+  using c_type = std::complex<float>;
+
+  explicit ComplexFloatType()
+      : ExtensionType(fixed_size_list(float32(), 2)) {}
+
+  std::string name() const override {
+    return "complex64";
+  }
+
+  std::string extension_name() const override {
+    return "arrow.complex64";
+  }
+
+  bool ExtensionEquals(const ExtensionType& other) const override;
+
+  std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override {
+    return std::make_shared<ComplexFloatArray>(data);
+  }
+
+  Result<std::shared_ptr<DataType>> Deserialize(
+      std::shared_ptr<DataType> storage_type,
+      const std::string& serialized) const override {
+    return complex64();
+  };
+
+  std::string Serialize() const override {
+    return "";
+  }
+};
+
+
+class ComplexDoubleArray : public ExtensionArray {
+ public:
+  using ExtensionArray::ExtensionArray;
+};
+
+class ComplexDoubleType : public ExtensionType {
+ public:
+  using c_type = std::complex<double>;
+
+  explicit ComplexDoubleType()
+      : ExtensionType(fixed_size_list(float64(), 2)) {}
+
+  std::string name() const override {
+    return "complex128";
+  }
+
+  std::string extension_name() const override {
+    return "arrow.complex128";
+  }
+
+  bool ExtensionEquals(const ExtensionType& other) const override;
+
+  std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override {
+    return std::make_shared<ComplexFloatArray>(data);
+  }
+
+  Result<std::shared_ptr<DataType>> Deserialize(
+      std::shared_ptr<DataType> storage_type,
+      const std::string& serialized) const override {
+    return complex128();
+  };
+
+  std::string Serialize() const override {
+    return "";
+  }
+};
+
+};  // namespace arrow

cpp/src/arrow/python/CMakeLists.txt

@@ -44,7 +44,8 @@ set(ARROW_PYTHON_SRCS
     numpy_to_arrow.cc
     python_to_arrow.cc
     pyarrow.cc
-    serialize.cc)
+    serialize.cc
+    type_traits.cc)
 
 set_source_files_properties(init.cc PROPERTIES SKIP_PRECOMPILE_HEADERS ON
                                                SKIP_UNITY_BUILD_INCLUSION ON)