Skip to main content

arrow_schema/extension/canonical/
variant.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! [`VariantType`] Arrow extension type
6
7use crate::extension::ExtensionType;
8use crate::{ArrowError, DataType};
9
10/// Arrow Variant [`ExtensionType`].
11///
12/// Represents the canonical Arrow Extension Type for storing variants.
13/// See [`VariantArray`] in the `parquet-variant-compute` crate for more examples
14/// of using this extension type.
15///
16/// [`VariantArray`]: https://docs.rs/parquet-variant-compute/latest/parquet_variant_compute/struct.VariantArray.html
17#[derive(Debug, Clone, PartialEq)]
18pub struct VariantType;
19
20impl ExtensionType for VariantType {
21    const NAME: &'static str = "arrow.parquet.variant";
22
23    // Variants extension metadata is an empty string
24    // <https://github.com/apache/arrow/blob/d803afcc43f5d132506318fd9e162d33b2c3d4cd/docs/source/format/CanonicalExtensions.rst?plain=1#L473>
25    type Metadata = &'static str;
26
27    fn metadata(&self) -> &Self::Metadata {
28        &""
29    }
30
31    fn serialize_metadata(&self) -> Option<String> {
32        Some(String::new())
33    }
34
35    fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
36        Ok("")
37    }
38
39    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
40        if matches!(data_type, DataType::Struct(_)) {
41            Ok(())
42        } else {
43            Err(ArrowError::InvalidArgumentError(format!(
44                "VariantType only supports StructArray, got {data_type}"
45            )))
46        }
47    }
48
49    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
50        Self.supports_data_type(data_type)?;
51        Ok(Self)
52    }
53
54    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> {
55        Self.supports_data_type(data_type)
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use crate::extension::CanonicalExtensionType;
63    use crate::{Field, Fields};
64
65    #[test]
66    fn variant_type_requires_struct() {
67        let field = Field::new("v", DataType::Struct(Fields::empty()), true)
68            .with_extension_type(VariantType);
69        assert!(field.has_valid_extension_type::<VariantType>());
70        assert_eq!(
71            field.try_canonical_extension_type().unwrap(),
72            CanonicalExtensionType::Variant(VariantType),
73        );
74
75        let err = Field::new("v", DataType::Binary, true)
76            .try_with_extension_type(VariantType)
77            .unwrap_err();
78        assert_eq!(
79            err.to_string(),
80            "Invalid argument error: VariantType only supports StructArray, got Binary"
81        );
82    }
83}