GH-45937: [C++][Parquet] support to read, write and validate variant

cpp/src/arrow/extension/CMakeLists.txt

@@ -15,7 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc uuid_test.cc)
+set(CANONICAL_EXTENSION_TESTS bool8_test.cc json_test.cc parquet_variant_test.cc
+                              uuid_test.cc)
 
 if(ARROW_JSON)
   list(APPEND CANONICAL_EXTENSION_TESTS tensor_extension_array_test.cc opaque_test.cc)

cpp/src/arrow/extension/meson.build

@@ -15,7 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 
-canonical_extension_tests = ['bool8_test.cc', 'json_test.cc', 'uuid_test.cc']
+canonical_extension_tests = [
+    'bool8_test.cc',
+    'json_test.cc',
+    'parquet_variant_test.cc',
+    'uuid_test.cc',
+]
 
 if needs_json
     canonical_extension_tests += [

cpp/src/arrow/extension/parquet_variant.cc

@@ -17,28 +17,239 @@
 
 #include "arrow/extension/parquet_variant.h"
 
+#include <algorithm>
 #include <string>
 
+#include "arrow/array/array_nested.h"
 #include "arrow/extension_type.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
+#include "arrow/type.h"
+#include "arrow/type_traits.h"
 #include "arrow/util/logging_internal.h"
 
 namespace arrow::extension {
 
+namespace {
+
+// Variant typed_value decimals use disjoint precision bands, unlike ordinary
+// Arrow decimals, which may use a wider type than their precision requires.
+bool IsSupportedVariantDecimal(const DecimalType& type) {
+  const auto precision = type.precision();
+  if (precision < 1 || type.scale() < 0 || type.scale() > precision) {
+    return false;
+  }
+  switch (type.id()) {
+    case Type::DECIMAL32:
+      return precision <= Decimal32Type::kMaxPrecision;
+    case Type::DECIMAL64:
+      return precision > Decimal32Type::kMaxPrecision &&
+             precision <= Decimal64Type::kMaxPrecision;
+    case Type::DECIMAL128:
+      return precision > Decimal64Type::kMaxPrecision &&
+             precision <= Decimal128Type::kMaxPrecision;
+    default:
+      return false;
+  }
+}
+
+bool IsSupportedMetadataType(const DataType& type) {
+  switch (type.id()) {
+    case Type::DICTIONARY: {
+      const auto& dictionary = internal::checked_cast<const DictionaryType&>(type);
+      return is_integer(dictionary.index_type()->id()) &&
+             is_binary_or_binary_view(dictionary.value_type()->id());
+    }
+    case Type::RUN_END_ENCODED: {
+      const auto& ree = internal::checked_cast<const RunEndEncodedType&>(type);
+      return RunEndEncodedType::RunEndTypeValid(*ree.run_end_type()) &&
+             is_binary_or_binary_view(ree.value_type()->id());
+    }
+    default:
+      return is_binary_or_binary_view(type.storage_id());
+  }
+}
+
+bool IsSupportedPrimitiveTypedValue(const std::shared_ptr<DataType>& type) {
+  switch (type->id()) {
+    case Type::BOOL:
+    case Type::INT8:
+    case Type::INT16:
+    case Type::INT32:
+    case Type::INT64:
+    case Type::FLOAT:
+    case Type::DOUBLE:
+    case Type::DATE32:
+    case Type::BINARY:
+    case Type::LARGE_BINARY:
+    case Type::BINARY_VIEW:
+    case Type::STRING:
+    case Type::LARGE_STRING:
+    case Type::STRING_VIEW:
+      return true;
+    case Type::DECIMAL32:
+    case Type::DECIMAL64:
+    case Type::DECIMAL128: {
+      const auto& decimal = internal::checked_cast<const DecimalType&>(*type);
+      return IsSupportedVariantDecimal(decimal);
+    }
+    case Type::TIME64:
+      return internal::checked_cast<const Time64Type&>(*type).unit() == TimeUnit::MICRO;
+    case Type::TIMESTAMP: {
+      const auto unit = internal::checked_cast<const TimestampType&>(*type).unit();
+      return unit == TimeUnit::MICRO || unit == TimeUnit::NANO;
+    }
+    case Type::EXTENSION: {
+      const auto& ext_type = internal::checked_cast<const ExtensionType&>(*type);
+      return ext_type.extension_name() == "arrow.uuid";
+    }
+    default:
+      return false;
+  }
+}
+
+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;
+  }
+
+  std::shared_ptr<Field> value;
+  std::shared_ptr<Field> typed_value;
+  for (const auto& field : type->fields()) {
+    if (field->name() == "value") {
+      if (value != nullptr || !field->nullable() ||
+          !is_binary_or_binary_view(field->type()->storage_id())) {
+        return false;
+      }
+      value = field;
+    } else if (field->name() == "typed_value") {
+      if (typed_value != nullptr || !IsSupportedTypedValue<strict>(field)) {
+        return false;
+      }
+      typed_value = field;
+    } else {
+      return false;
+    }
+  }
+
+  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;
+  }
+
+  switch (field->type()->id()) {
+    case Type::STRUCT:
+      return field->type()->num_fields() > 0 &&
+             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: {
+      auto& inner_field = field->type()->field(0);
+      return !inner_field->nullable() &&
+             IsVariantFieldGroup<strict, /*in_object=*/false>(inner_field->type());
+    }
+    default:
+      return IsSupportedPrimitiveTypedValue(field->type());
+  }
+}
+
+template <bool strict>
+bool IsSupportedStorageTypeImpl(const std::shared_ptr<DataType>& storage_type) {
+  if (storage_type->id() != Type::STRUCT) {
+    return false;
+  }
+
+  std::shared_ptr<Field> metadata;
+  std::shared_ptr<Field> value;
+  std::shared_ptr<Field> typed_value;
+
+  for (const auto& field : storage_type->fields()) {
+    if (field->name() == "metadata") {
+      if (metadata != nullptr || !IsSupportedMetadataType(*field->type()) ||
+          field->nullable()) {
+        return false;
+      }
+      metadata = field;
+    } else if (field->name() == "value") {
+      if (value != nullptr || !is_binary_or_binary_view(field->type()->storage_id())) {
+        return false;
+      }
+      value = field;
+    } else if (field->name() == "typed_value") {
+      if (typed_value != nullptr || !IsSupportedTypedValue<strict>(field)) {
+        return false;
+      }
+      typed_value = field;
+    } else {
+      return false;
+    }
+  }
+
+  if (metadata == nullptr) {
+    return false;
+  }
+
+  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;
+  }
+}
+
+}  // namespace
+
+std::shared_ptr<Array> VariantArray::metadata() const {
+  return internal::checked_cast<const StructArray&>(*storage())
+      .GetFieldByName("metadata");
+}
+
+std::shared_ptr<Array> VariantArray::value() const {
+  return internal::checked_cast<const StructArray&>(*storage()).GetFieldByName("value");
+}
+
+std::shared_ptr<Array> VariantArray::typed_value() const {
+  return internal::checked_cast<const StructArray&>(*storage())
+      .GetFieldByName("typed_value");
+}
+
+bool VariantArray::is_shredded() const {
+  return internal::checked_cast<const VariantExtensionType&>(*type()).typed_value() !=
+         nullptr;
+}
+
 VariantExtensionType::VariantExtensionType(const std::shared_ptr<DataType>& storage_type)
     : ExtensionType(storage_type) {
-  // GH-45948: Shredded variants will need to handle an optional shredded_value as
-  // well as value_ becoming optional.
-
-  // IsSupportedStorageType should have been called already, asserting that both
-  // metadata and value are present.
-  if (storage_type->field(0)->name() == "metadata") {
-    metadata_ = storage_type->field(0);
-    value_ = storage_type->field(1);
-  } else {
-    value_ = storage_type->field(0);
-    metadata_ = storage_type->field(1);
+  // 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;
+    }
   }
 }
 
@@ -49,58 +260,29 @@ bool VariantExtensionType::ExtensionEquals(const ExtensionType& other) const {
 
 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));
+  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 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("arrow.parquet.variant",
+  DCHECK_EQ(kVariantExtensionName,
             internal::checked_cast<const ExtensionType&>(*data->type).extension_name());
   return std::make_shared<VariantArray>(data);
 }
 
-namespace {
-bool IsBinaryField(const std::shared_ptr<Field> field) {
-  return field->type()->storage_id() == Type::BINARY ||
-         field->type()->storage_id() == Type::LARGE_BINARY;
-}
-}  // namespace
-
 bool VariantExtensionType::IsSupportedStorageType(
     const std::shared_ptr<DataType>& storage_type) {
-  // For now we only supported unshredded variants. Unshredded variant storage
-  // type should be a struct with a binary metadata and binary value.
-  //
-  // GH-45948: In shredded variants, the binary value field can be replaced
-  // with one or more of the following: object, array, typed_value, and
-  // variant_value.
-  if (storage_type->id() == Type::STRUCT) {
-    if (storage_type->num_fields() == 2) {
-      // Ordering of metadata and value fields does not matter, as we will assign
-      // these to the VariantExtensionType's member shared_ptrs in the constructor.
-      // Here we just need to check that they are both present.
-
-      const auto& field0 = storage_type->field(0);
-      const auto& field1 = storage_type->field(1);
-
-      bool metadata_and_value_present =
-          (field0->name() == "metadata" && field1->name() == "value") ||
-          (field1->name() == "metadata" && field0->name() == "value");
-
-      if (metadata_and_value_present) {
-        // Both metadata and value must be non-nullable binary types for unshredded
-        // variants. This will change in GH-46948, when we will require a Visitor
-        // to traverse the structure of the variant.
-        return IsBinaryField(field0) && IsBinaryField(field1) && !field0->nullable() &&
-               !field1->nullable();
-      }
-    }
-  }
-
-  return false;
+  return IsSupportedStorageTypeImpl</*strict=*/true>(storage_type);
 }
 
 Result<std::shared_ptr<DataType>> VariantExtensionType::Make(
@@ -113,9 +295,6 @@ Result<std::shared_ptr<DataType>> VariantExtensionType::Make(
   return std::make_shared<VariantExtensionType>(std::move(storage_type));
 }
 
-/// NOTE: this is still experimental. GH-45948 will add shredding support, at which point
-/// we need to separate this into unshredded_variant and shredded_variant helper
-/// functions.
 std::shared_ptr<DataType> variant(std::shared_ptr<DataType> storage_type) {
   return VariantExtensionType::Make(std::move(storage_type)).ValueOrDie();
 }