vortex_dtype/
extension.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
use std::fmt::{Display, Formatter};
use std::sync::Arc;

use crate::{DType, Nullability};

/// A unique identifier for an extension type
#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(::serde::Serialize, ::serde::Deserialize))]
pub struct ExtID(Arc<str>);

impl ExtID {
    /// Constructs a new `ExtID` from a string
    pub fn new(value: Arc<str>) -> Self {
        Self(value)
    }
}

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

impl AsRef<str> for ExtID {
    fn as_ref(&self) -> &str {
        self.0.as_ref()
    }
}

impl From<&str> for ExtID {
    fn from(value: &str) -> Self {
        Self(value.into())
    }
}

/// Opaque metadata for an extension type
#[derive(Debug, Clone, PartialOrd, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ExtMetadata(Arc<[u8]>);

impl ExtMetadata {
    /// Constructs a new `ExtMetadata` from a byte slice
    pub fn new(value: Arc<[u8]>) -> Self {
        Self(value)
    }
}

impl AsRef<[u8]> for ExtMetadata {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl From<&[u8]> for ExtMetadata {
    fn from(value: &[u8]) -> Self {
        Self(value.into())
    }
}

/// A type descriptor for an extension type
#[derive(Debug, Clone, PartialOrd, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ExtDType {
    id: ExtID,
    storage_dtype: Arc<DType>,
    metadata: Option<ExtMetadata>,
}

impl ExtDType {
    /// Creates a new `ExtDType`.
    ///
    /// Extension data types in Vortex allows library users to express additional semantic meaning
    /// on top of a set of scalar values. Metadata can optionally be provided for the extension type
    /// to allow for parameterized types.
    ///
    /// A simple example would be if one wanted to create a `vortex.temperature` extension type. The
    /// canonical encoding for such values would be `f64`, and the metadata can contain an optional
    /// temperature unit, allowing downstream users to be sure they properly account for Celsius
    /// and Fahrenheit conversions.
    ///
    /// ```
    /// use std::sync::Arc;
    /// use vortex_dtype::{DType, ExtDType, ExtID, ExtMetadata, Nullability, PType};
    ///
    /// #[repr(u8)]
    /// enum TemperatureUnit {
    ///     C = 0u8,
    ///     F = 1u8,
    /// }
    ///
    /// // Make a new extension type that encodes the unit for a set of nullable `f64`.
    /// pub fn create_temperature_type(unit: TemperatureUnit) -> ExtDType {
    ///     ExtDType::new(
    ///         ExtID::new("vortex.temperature".into()),
    ///         Arc::new(DType::Primitive(PType::F64, Nullability::Nullable)),
    ///         Some(ExtMetadata::new([unit as u8].into()))
    ///     )
    /// }
    /// ```
    pub fn new(id: ExtID, storage_dtype: Arc<DType>, metadata: Option<ExtMetadata>) -> Self {
        assert!(
            !matches!(storage_dtype.as_ref(), &DType::Extension(_)),
            "ExtDType cannot have Extension storage_dtype"
        );

        Self {
            id,
            storage_dtype,
            metadata,
        }
    }

    /// Returns the `ExtID` for this extension type
    #[inline]
    pub fn id(&self) -> &ExtID {
        &self.id
    }

    /// Returns the `ExtMetadata` for this extension type, if it exists
    #[inline]
    pub fn storage_dtype(&self) -> &DType {
        self.storage_dtype.as_ref()
    }

    /// Returns a new `ExtDType` with the given nullability
    pub fn with_nullability(&self, nullability: Nullability) -> Self {
        Self::new(
            self.id.clone(),
            Arc::new(self.storage_dtype.with_nullability(nullability)),
            self.metadata.clone(),
        )
    }

    /// Returns the `ExtMetadata` for this extension type, if it exists
    #[inline]
    pub fn metadata(&self) -> Option<&ExtMetadata> {
        self.metadata.as_ref()
    }
}