Moving VariantExtensionType to parquet::arrow namespace

apache · Feb 25, 2025 · 8c5ec3e · 8c5ec3e
1 parent 5096d74
commit 8c5ec3e
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 11 deletions.
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
@@ -377,7 +377,6 @@ set(ARROW_SRCS
     extension/bool8.cc
     extension/json.cc
     extension/uuid.cc
-    extension/variant.cc
     pretty_print.cc
     record_batch.cc
     result.cc

diff --git a/cpp/src/arrow/extension/CMakeLists.txt b/cpp/src/arrow/extension/CMakeLists.txt
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc uuid_test.cc variant_test.cc)
+set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc uuid_test.cc)
 
 if(ARROW_JSON)
   list(APPEND CANONICAL_EXTENSION_TESTS fixed_shape_tensor_test.cc opaque_test.cc)

diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt
@@ -157,6 +157,7 @@ set(PARQUET_SRCS
     arrow/reader_internal.cc
     arrow/schema.cc
     arrow/schema_internal.cc
+    arrow/variant.cc
     arrow/writer.cc
     bloom_filter.cc
     bloom_filter_reader.cc
@@ -399,7 +400,8 @@ add_parquet_test(arrow-test
                  arrow/arrow_metadata_test.cc
                  arrow/arrow_reader_writer_test.cc
                  arrow/arrow_schema_test.cc
-                 arrow/arrow_statistics_test.cc)
+                 arrow/arrow_statistics_test.cc
+                 arrow/variant_test.cc)
 
 add_parquet_test(arrow-internals-test SOURCES arrow/path_internal_test.cc
                  arrow/reconstruct_internal_test.cc)

diff --git a/cpp/src/arrow/extension/variant.cc → cpp/src/parquet/arrow/variant.cc b/cpp/src/arrow/extension/variant.cc → cpp/src/parquet/arrow/variant.cc
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "arrow/extension/variant.h"
+#include "parquet/arrow/variant.h"
 
 #include <string>
 
@@ -25,18 +25,42 @@
 #include "arrow/type_fwd.h"
 #include "arrow/util/logging.h"
 
-namespace arrow::extension {
+namespace parquet::arrow {
+
+using ::arrow::Status;
+using ::arrow::Type;
 
 bool isBinary(Type::type type) {
   return type == Type::BINARY || type == Type::LARGE_BINARY;
 }
 
+bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const {
+  // TODO(neilechao)
+  return false;
+}
+
+Result<std::shared_ptr<DataType>> VariantExtensionType::Deserialize(
+    std::shared_ptr<DataType> storage_type, const std::string& serialized) const {
+  return VariantExtensionType::Make(std::move(storage_type));
+}
+
+std::string VariantExtensionType::Serialize() const { return ""; }
+
+std::shared_ptr<Array> VariantExtensionType::MakeArray(
+    std::shared_ptr<ArrayData> data) const {
+  DCHECK_EQ(data->type->id(), Type::EXTENSION);
+  DCHECK_EQ("parquet.variant",
+            ::arrow::internal::checked_cast<const ExtensionType&>(*data->type)
+                .extension_name());
+  return std::make_shared<::arrow::ExtensionArray>(data);
+}
+
 bool VariantExtensionType::IsSupportedStorageType(
     std::shared_ptr<DataType> storage_type) {
   if (storage_type->id() == Type::STRUCT) {
     // TODO(neilechao) assertions for binary types, and non-nullable first field for
     // metadata
-    return storage_type->num_fields() == 3;
+    return storage_type->num_fields() >= 2;
   }
 
   return false;
@@ -57,4 +81,4 @@ std::shared_ptr<DataType> variant(std::shared_ptr<DataType> storage_type) {
   return VariantExtensionType::Make(std::move(storage_type)).ValueOrDie();
 }
 
-}  // namespace arrow::extension
+}  // namespace parquet::arrow
diff --git a/cpp/src/arrow/extension/variant.h → cpp/src/parquet/arrow/variant.h b/cpp/src/arrow/extension/variant.h → cpp/src/parquet/arrow/variant.h
@@ -25,14 +25,20 @@
 #include "arrow/type_fwd.h"
 #include "arrow/util/visibility.h"
 
-namespace arrow::extension {
+namespace parquet::arrow {
+
+using ::arrow::Array;
+using ::arrow::ArrayData;
+using ::arrow::DataType;
+using ::arrow::ExtensionType;
+using ::arrow::Result;
 
 class ARROW_EXPORT VariantExtensionType : public ExtensionType {
  public:
   explicit VariantExtensionType(const std::shared_ptr<DataType>& storage_type)
       : ExtensionType(storage_type), storage_type_(storage_type) {}
 
-  std::string extension_name() const override { return "variant.json"; }
+  std::string extension_name() const override { return "parquet.variant"; }
 
   bool ExtensionEquals(const ExtensionType& other) const override;
 
@@ -55,4 +61,4 @@ class ARROW_EXPORT VariantExtensionType : public ExtensionType {
 /// \brief Return a VariantExtensionType instance.
 ARROW_EXPORT std::shared_ptr<DataType> variant(std::shared_ptr<DataType> storage_type);
 
-}  // namespace arrow::extension
+}  // namespace parquet::arrow
diff --git a/cpp/src/arrow/extension/variant_test.cc → cpp/src/parquet/arrow/variant_test.cc b/cpp/src/arrow/extension/variant_test.cc → cpp/src/parquet/arrow/variant_test.cc
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-#include "arrow/extension/variant.h"
+#include "parquet/arrow/variant.h"
 
 #include "arrow/array/validate.h"
 #include "arrow/ipc/test_common.h"