ARROW-6187: [C++] Fallback to storage type when writing ExtensionType…

… to Parquet
kszucs · Sep 24, 2019 · fb4b810 · fb4b810
1 parent 35b6d07
commit fb4b810
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc
@@ -20,6 +20,7 @@
 #include <string>
 #include <vector>
 
+#include "arrow/extension_type.h"
 #include "arrow/type.h"
 #include "arrow/util/checked_cast.h"
 
@@ -308,6 +309,12 @@ Status FieldToNode(const std::shared_ptr<Field>& field,
           field->name(), dict_type.value_type(), field->nullable(), field->metadata());
       return FieldToNode(unpacked_field, properties, arrow_properties, out);
     }
+    case ArrowTypeId::EXTENSION: {
+      auto ext_type = std::static_pointer_cast<::arrow::ExtensionType>(field->type());
+      std::shared_ptr<::arrow::Field> storage_field = ::arrow::field(
+          field->name(), ext_type->storage_type(), field->nullable(), field->metadata());
+      return FieldToNode(storage_field, properties, arrow_properties, out);
+    }
     default: {
       // TODO: DENSE_UNION, SPARE_UNION, JSON_SCALAR, DECIMAL_TEXT, VARCHAR
       return Status::NotImplemented(

diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc
@@ -26,6 +26,7 @@
 
 #include "arrow/array.h"
 #include "arrow/buffer_builder.h"
+#include "arrow/extension_type.h"
 #include "arrow/ipc/writer.h"
 #include "arrow/table.h"
 #include "arrow/type.h"
@@ -48,6 +49,7 @@ using arrow::DictionaryArray;
 using arrow::Field;
 using arrow::FixedSizeBinaryArray;
 using Int16BufferBuilder = arrow::TypedBufferBuilder<int16_t>;
+using arrow::ExtensionArray;
 using arrow::ListArray;
 using arrow::MemoryPool;
 using arrow::NumericArray;
@@ -115,6 +117,8 @@ class LevelBuilder {
     return VisitInline(*array.values());
   }
 
+  Status Visit(const ExtensionArray& array) { return VisitInline(*array.storage()); }
+
 #define NOT_IMPLEMENTED_VISIT(ArrowTypePrefix)                             \
   Status Visit(const ::arrow::ArrowTypePrefix##Array& array) {             \
     return Status::NotImplemented("Level generation for " #ArrowTypePrefix \
@@ -126,7 +130,6 @@ class LevelBuilder {
   NOT_IMPLEMENTED_VISIT(FixedSizeList)
   NOT_IMPLEMENTED_VISIT(Struct)
   NOT_IMPLEMENTED_VISIT(Union)
-  NOT_IMPLEMENTED_VISIT(Extension)
 
 #undef NOT_IMPLEMENTED_VISIT