add unshred and shredded reader and roundtrip

Author: HuaHuaY

cpp/src/arrow/extension/parquet_variant.cc

@@ -29,44 +29,6 @@
 
 namespace arrow::extension {
 
-VariantExtensionType::VariantExtensionType(const std::shared_ptr<DataType>& storage_type)
-    : ExtensionType(storage_type) {
-  // IsSupportedStorageType should have been called already, asserting that
-  // metadata is present and at least one of value / typed_value is present.
-  for (const auto& field : storage_type->fields()) {
-    if (field->name() == "metadata") {
-      metadata_ = field;
-    } else if (field->name() == "value") {
-      value_ = field;
-    } else if (field->name() == "typed_value") {
-      typed_value_ = field;
-    }
-  }
-}
-
-bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const {
-  return other.extension_name() == this->extension_name() &&
-         other.storage_type()->Equals(this->storage_type());
-}
-
-Result<std::shared_ptr<DataType>> VariantExtensionType::Deserialize(
-    std::shared_ptr<DataType> storage_type, const std::string& serialized) const {
-  if (!serialized.empty()) {
-    return Status::Invalid("Unexpected serialized metadata: '", serialized, "'");
-  }
-  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(kVariantExtensionName,
-            internal::checked_cast<const ExtensionType&>(*data->type).extension_name());
-  return std::make_shared<VariantArray>(data);
-}
-
 namespace {
 
 bool IsSupportedPrimitiveTypedValue(const std::shared_ptr<DataType>& type) {
@@ -109,8 +71,10 @@ bool IsSupportedPrimitiveTypedValue(const std::shared_ptr<DataType>& type) {
   }
 }
 
+template <bool strict>
 bool IsSupportedTypedValue(const std::shared_ptr<Field>& field);
 
+template <bool strict, bool in_object>
 bool IsVariantFieldGroup(const std::shared_ptr<DataType>& type) {
   if (type->id() != Type::STRUCT) {
     return false;
@@ -126,44 +90,51 @@ bool IsVariantFieldGroup(const std::shared_ptr<DataType>& type) {
       }
       value = field;
     } else if (field->name() == "typed_value") {
-      if (typed_value != nullptr || !IsSupportedTypedValue(field)) {
+      if (typed_value != nullptr || !IsSupportedTypedValue<strict>(field)) {
         return false;
       }
       typed_value = field;
     } else {
       return false;
     }
   }
-  return value != nullptr || typed_value != nullptr;
+
+  if constexpr (strict && in_object) {
+    return value != nullptr;
+  } else {
+    return value != nullptr || typed_value != nullptr;
+  }
 }
 
+template <bool strict>
 bool IsSupportedTypedValue(const std::shared_ptr<Field>& field) {
   if (!field->nullable()) {
     return false;
   }
-  auto is_variant_field_group = [](const auto& field) {
-    return !field->nullable() && IsVariantFieldGroup(field->type());
-  };
 
   switch (field->type()->id()) {
     case Type::STRUCT:
       return field->type()->num_fields() > 0 &&
-             std::ranges::all_of(field->type()->fields(), is_variant_field_group);
+             std::ranges::all_of(field->type()->fields(), [](const auto& field) {
+               return (!strict || !field->nullable()) &&
+                      IsVariantFieldGroup<strict, /*in_object=*/true>(field->type());
+             });
     case Type::LIST:
     case Type::LARGE_LIST:
     case Type::LIST_VIEW:
     case Type::LARGE_LIST_VIEW:
-    case Type::FIXED_SIZE_LIST:
-      return is_variant_field_group(field->type()->field(0));
+    case Type::FIXED_SIZE_LIST: {
+      auto& inner_field = field->type()->field(0);
+      return !inner_field->nullable() &&
+             IsVariantFieldGroup<strict, /*in_object=*/false>(inner_field->type());
+    }
     default:
       return IsSupportedPrimitiveTypedValue(field->type());
   }
 }
 
-}  // namespace
-
-bool VariantExtensionType::IsSupportedStorageType(
-    const std::shared_ptr<DataType>& storage_type) {
+template <bool strict>
+bool IsSupportedStorageTypeImpl(const std::shared_ptr<DataType>& storage_type) {
   if (storage_type->id() != Type::STRUCT) {
     return false;
   }
@@ -185,7 +156,7 @@ bool VariantExtensionType::IsSupportedStorageType(
       }
       value = field;
     } else if (field->name() == "typed_value") {
-      if (typed_value != nullptr || !IsSupportedTypedValue(field)) {
+      if (typed_value != nullptr || !IsSupportedTypedValue<strict>(field)) {
         return false;
       }
       typed_value = field;
@@ -194,16 +165,68 @@ bool VariantExtensionType::IsSupportedStorageType(
     }
   }
 
-  if (metadata == nullptr || (value == nullptr && typed_value == nullptr)) {
+  if (metadata == nullptr) {
     return false;
   }
-  if (value == nullptr) {
-    return true;
+
+  if constexpr (strict) {
+    if (value == nullptr) {
+      return false;
+    }
+    return (typed_value == nullptr) != value->nullable();
+  } else {
+    bool value_nullable = (value == nullptr) || value->nullable();
+    return (typed_value == nullptr) != value_nullable;
   }
-  if (typed_value == nullptr) {
-    return !value->nullable();
+}
+
+}  // namespace
+
+VariantExtensionType::VariantExtensionType(const std::shared_ptr<DataType>& storage_type)
+    : ExtensionType(storage_type) {
+  // IsSupportedStorageType should have been called already, asserting that
+  // metadata is present and at least one of value / typed_value is present.
+  for (const auto& field : storage_type->fields()) {
+    if (field->name() == "metadata") {
+      metadata_ = field;
+    } else if (field->name() == "value") {
+      value_ = field;
+    } else if (field->name() == "typed_value") {
+      typed_value_ = field;
+    }
+  }
+}
+
+bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const {
+  return other.extension_name() == this->extension_name() &&
+         other.storage_type()->Equals(this->storage_type());
+}
+
+Result<std::shared_ptr<DataType>> VariantExtensionType::Deserialize(
+    std::shared_ptr<DataType> storage_type, const std::string& serialized) const {
+  if (!serialized.empty()) {
+    return Status::Invalid("Unexpected serialized metadata: '", serialized, "'");
+  }
+  if (!IsSupportedStorageTypeImpl</*strict=*/false>(storage_type)) {
+    return Status::Invalid("Invalid storage type for VariantExtensionType: ",
+                           storage_type->ToString());
   }
-  return value->nullable();
+  return std::make_shared<VariantExtensionType>(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(kVariantExtensionName,
+            internal::checked_cast<const ExtensionType&>(*data->type).extension_name());
+  return std::make_shared<VariantArray>(data);
+}
+
+bool VariantExtensionType::IsSupportedStorageType(
+    const std::shared_ptr<DataType>& storage_type) {
+  return IsSupportedStorageTypeImpl</*strict=*/true>(storage_type);
 }
 
 Result<std::shared_ptr<DataType>> VariantExtensionType::Make(

cpp/src/parquet/CMakeLists.txt

@@ -192,7 +192,10 @@ set(PARQUET_SRCS
     stream_writer.cc
     types.cc
     variant/builder.cc
-    variant/encoding.cc
+    variant/decoding.cc
+    variant/array_internal.cc
+    variant/slot_internal.cc
+    variant/unshred.cc
     variant/validate.cc
     xxhasher.cc)
 

cpp/src/parquet/arrow/arrow_schema_test.cc

@@ -1092,8 +1092,35 @@ TEST_F(TestConvertParquetSchema, ParquetVariantShreddedWithoutValue) {
       ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false),
                         ::arrow::field("typed_value", ::arrow::int64())});
 
-  ASSERT_NO_FATAL_FAILURE(CheckParquetVariantSchema(
-      "variant_shredded", {metadata, typed_value}, storage_type));
+  {
+    auto arrow_schema = ::arrow::schema({::arrow::field("variant_shredded", storage_type)});
+    ASSERT_OK(ConvertSchema({GroupNode::Make("variant_shredded", Repetition::OPTIONAL,
+                                            {metadata, typed_value},
+                                            LogicalType::Variant())}));
+    ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema, /*check_metadata=*/true));
+  }
+
+  auto registered_storage_type =
+      ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false),
+                        ::arrow::field("value", ::arrow::binary(), false)});
+  auto registered_variant = ::arrow::extension::variant(registered_storage_type);
+  ::arrow::ExtensionTypeGuard guard({registered_variant});
+
+  ArrowReaderProperties props;
+  props.set_arrow_extensions_enabled(true);
+  ASSERT_OK(ConvertSchema({GroupNode::Make("variant_shredded", Repetition::OPTIONAL,
+                                          {metadata, typed_value},
+                                          LogicalType::Variant())},
+                          /*metadata=*/nullptr, props));
+
+  auto field = result_schema_->field(0);
+  ASSERT_EQ(::arrow::Type::EXTENSION, field->type()->id());
+  auto variant_type =
+      std::dynamic_pointer_cast<::arrow::extension::VariantExtensionType>(field->type());
+  ASSERT_NE(nullptr, variant_type);
+  ASSERT_EQ(nullptr, variant_type->value());
+  ASSERT_NE(nullptr, variant_type->typed_value());
+  ASSERT_TRUE(variant_type->storage_type()->Equals(storage_type));
 }
 
 TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) {
@@ -1594,7 +1621,16 @@ TEST_F(TestConvertArrowSchema, ParquetVariantShreddedWithoutValue) {
   auto storage_type =
       ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false),
                         ::arrow::field("typed_value", ::arrow::int64())});
-  auto variant_type = ::arrow::extension::variant(storage_type);
+  auto registered_storage_type =
+      ::arrow::struct_({::arrow::field("metadata", ::arrow::binary(), false),
+                        ::arrow::field("value", ::arrow::binary(), false)});
+  ASSERT_OK_AND_ASSIGN(auto registered_variant_type,
+                       ::arrow::extension::VariantExtensionType::Make(
+                           registered_storage_type));
+  auto registered_variant =
+      ::arrow::internal::checked_pointer_cast<::arrow::extension::VariantExtensionType>(
+          registered_variant_type);
+  ASSERT_OK_AND_ASSIGN(auto variant_type, registered_variant->Deserialize(storage_type, ""));
 
   ASSERT_RAISES(Invalid, ConvertSchema({::arrow::field("variant", variant_type)}));
 }

cpp/src/parquet/arrow/reader.cc

@@ -27,6 +27,7 @@
 
 #include "arrow/array.h"
 #include "arrow/buffer.h"
+#include "arrow/extension/parquet_variant.h"
 #include "arrow/extension_type.h"
 #include "arrow/io/memory.h"
 #include "arrow/memory_pool.h"
@@ -54,6 +55,7 @@
 #include "parquet/page_index.h"
 #include "parquet/properties.h"
 #include "parquet/schema.h"
+#include "parquet/variant/validate.h"
 
 using arrow::Array;
 using arrow::ArrayData;
@@ -583,6 +585,26 @@ class ExtensionReader : public ColumnReaderImpl {
   std::unique_ptr<ColumnReaderImpl> storage_reader_;
 };
 
+class VariantReader : public ExtensionReader {
+ public:
+  VariantReader(std::shared_ptr<ReaderContext> ctx, std::shared_ptr<Field> field,
+                std::unique_ptr<ColumnReaderImpl> storage_reader)
+      : ExtensionReader(std::move(field), std::move(storage_reader)),
+        ctx_(std::move(ctx)) {}
+
+  Status BuildArray(int64_t length_upper_bound,
+                    std::shared_ptr<ChunkedArray>* out) override {
+    RETURN_NOT_OK(ExtensionReader::BuildArray(length_upper_bound, out));
+    if (ctx_->reader_properties->variant_validation_enabled()) {
+      PARQUET_CATCH_NOT_OK(variant::ValidateVariants<false>(**out, ctx_->pool));
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::shared_ptr<ReaderContext> ctx_;
+};
+
 template <typename IndexType>
 class ListReader : public ColumnReaderImpl {
  public:
@@ -892,8 +914,8 @@ Status GetReader(const SchemaField& field, const std::shared_ptr<Field>& arrow_f
   auto type_id = arrow_field->type()->id();
 
   if (type_id == ::arrow::Type::EXTENSION) {
-    auto storage_field = arrow_field->WithType(
-        checked_cast<const ExtensionType&>(*arrow_field->type()).storage_type());
+    const auto& ext_type = checked_cast<const ExtensionType&>(*arrow_field->type());
+    auto storage_field = arrow_field->WithType(ext_type.storage_type());
     RETURN_NOT_OK(GetReader(field, storage_field, ctx, out));
     if (*out) {
       auto storage_type = (*out)->field()->type();
@@ -902,7 +924,11 @@ Status GetReader(const SchemaField& field, const std::shared_ptr<Field>& arrow_f
             "Due to column pruning only part of an extension's storage type was loaded.  "
             "An extension type cannot be created without all of its fields");
       }
-      *out = std::make_unique<ExtensionReader>(arrow_field, std::move(*out));
+      if (ext_type.extension_name() == ::arrow::extension::kVariantExtensionName) {
+        *out = std::make_unique<VariantReader>(ctx, arrow_field, std::move(*out));
+      } else {
+        *out = std::make_unique<ExtensionReader>(arrow_field, std::move(*out));
+      }
     }
     return Status::OK();
   }