Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ExtensionType for uuid and map to parquet logical type #5822

Draft
wants to merge 7 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions arrow-schema/src/datatype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -672,6 +672,82 @@ impl DataType {
}
}

/// Canonical extension types.
///
/// The Arrow columnar format allows defining extension types so as to extend
/// standard Arrow data types with custom semantics. Often these semantics will
/// be specific to a system or application. However, it is beneficial to share
/// the definitions of well-known extension types so as to improve
/// interoperability between different systems integrating Arrow columnar data.
///
/// <https://arrow.apache.org/docs/format/CanonicalExtensions.html>
#[non_exhaustive]
#[derive(Debug, Clone)]
pub enum ExtensionType {
/// Extension name: `arrow.uuid`.
///
/// The storage type of the extension is `FixedSizeBinary` with a length of
/// 16 bytes.
///
/// Note:
/// A specific UUID version is not required or guaranteed. This extension
/// represents UUIDs as FixedSizeBinary(16) with big-endian notation and
/// does not interpret the bytes in any way.
Uuid,
}

impl fmt::Display for ExtensionType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.name())
}
}

impl ExtensionType {
/// The metadata key for the string name identifying the custom data type.
pub const NAME_KEY: &'static str = "ARROW:extension:name";

/// The metadata key for a serialized representation of the ExtensionType
/// necessary to reconstruct the custom type.
pub const METADATA_KEY: &'static str = "ARROW:extension:metadata";

/// Returns the name of this extension type.
pub fn name(&self) -> &'static str {
match self {
ExtensionType::Uuid => "arrow.uuid",
}
}

/// Returns the metadata of this extension type.
pub fn metadata(&self) -> Option<String> {
match self {
ExtensionType::Uuid => None,
}
}

/// Returns `true` iff the given [`DataType`] can be used as storage type
/// for this extension type.
pub(crate) fn supports_storage_type(&self, data_type: &DataType) -> bool {
match self {
ExtensionType::Uuid => matches!(data_type, DataType::FixedSizeBinary(16)),
}
}

/// Extract an [`ExtensionType`] from the given [`Field`].
///
/// This function returns `None` if the extension type is not supported or
/// recognized.
pub(crate) fn try_from_field(field: &Field) -> Option<Self> {
let metadata = field.metadata().get(ExtensionType::METADATA_KEY);
field
.metadata()
.get(ExtensionType::NAME_KEY)
.and_then(|name| match name.as_str() {
"arrow.uuid" if metadata.is_none() => Some(ExtensionType::Uuid),
_ => None,
})
}
}

/// The maximum precision for [DataType::Decimal128] values
pub const DECIMAL128_MAX_PRECISION: u8 = 38;

Expand Down
33 changes: 32 additions & 1 deletion arrow-schema/src/field.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use std::sync::Arc;

use crate::datatype::DataType;
use crate::schema::SchemaBuilder;
use crate::{Fields, UnionFields, UnionMode};
use crate::{ExtensionType, Fields, UnionFields, UnionMode};

/// A reference counted [`Field`]
pub type FieldRef = Arc<Field>;
Expand Down Expand Up @@ -337,6 +337,37 @@ impl Field {
self
}

/// Returns the canonical [`ExtensionType`] of this [`Field`], if set.
pub fn extension_type(&self) -> Option<ExtensionType> {
ExtensionType::try_from_field(self)
}

/// Updates the metadata of this [`Field`] with the [`ExtensionType::name`]
/// and [`ExtensionType::metadata`] of the given [`ExtensionType`].
///
/// # Panics
///
/// This function panics when the datatype of this field is not a valid
/// storage type for the given extension type.
pub fn with_extension_type(mut self, extension_type: ExtensionType) -> Self {
if extension_type.supports_storage_type(&self.data_type) {
self.metadata.insert(
ExtensionType::NAME_KEY.to_owned(),
extension_type.name().to_owned(),
);
if let Some(metadata) = extension_type.metadata() {
self.metadata
.insert(ExtensionType::METADATA_KEY.to_owned(), metadata);
}
self
} else {
panic!(
"{extension_type} does not support {} as storage type",
self.data_type
);
}
}

/// Indicates whether this [`Field`] supports null values.
#[inline]
pub const fn is_nullable(&self) -> bool {
Expand Down
6 changes: 5 additions & 1 deletion parquet/src/arrow/schema/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ use std::collections::HashMap;
use std::sync::Arc;

use arrow_ipc::writer;
use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
use arrow_schema::{DataType, ExtensionType, Field, Fields, Schema, TimeUnit};

use crate::basic::{
ConvertedType, LogicalType, Repetition, TimeUnit as ParquetTimeUnit, Type as PhysicalType,
Expand Down Expand Up @@ -468,6 +468,10 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
.with_repetition(repetition)
.with_id(id)
.with_length(*length)
.with_logical_type(match field.extension_type() {
Some(ExtensionType::Uuid) => Some(LogicalType::Uuid),
_ => None,
})
.build()
}
DataType::BinaryView => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)
Expand Down
Loading