vortex_array/
macros.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
140
141
142
//! The core Vortex macro to create new encodings and array types.

use crate::Array;
/// Macro to generate all the necessary code for a new type of array encoding. Including:
/// 1. New Array type that implements `AsRef<Array>`, `GetArrayMetadata`, `ToArray`, `IntoArray`, and multiple useful `From`/`TryFrom` implementations.
/// 2. New Encoding type that implements `ArrayEncoding`.
/// 3. New metadata type that implements `ArrayMetadata`.
#[macro_export]
macro_rules! impl_encoding {
    ($id:literal, $code:expr, $Name:ident, $Metadata:ty) => {
        $crate::paste::paste! {
            #[derive(std::fmt::Debug, Clone)]
            #[repr(transparent)]
            pub struct [<$Name Array>]($crate::Array);

            impl $crate::IntoArray for [<$Name Array>] {
                fn into_array(self) -> $crate::Array {
                    self.0
                }
            }

            impl AsRef<$crate::Array> for [<$Name Array>] {
                fn as_ref(&self) -> &$crate::Array {
                    &self.0
                }
            }

            #[allow(dead_code)]
            impl [<$Name Array>] {
                fn try_from_parts(
                    dtype: vortex_dtype::DType,
                    len: usize,
                    metadata: $Metadata,
                    buffers: Option<Box<[vortex_buffer::ByteBuffer]>>,
                    children: Option<Box<[$crate::Array]>>,
                    stats: $crate::stats::StatsSet,
                ) -> VortexResult<Self> {
                    use $crate::SerializeMetadata;

                    Self::try_from($crate::Array::try_new_owned(
                            [<$Name Encoding>]::vtable(),
                            dtype,
                            len,
                            metadata.serialize()?,
                            buffers,
                            children,
                            stats
                    )?)
                }

                fn metadata(&self) -> <$Metadata as $crate::DeserializeMetadata>::Output {
                    use $crate::DeserializeMetadata;

                    // SAFETY: Metadata is validated during construction of Array.
                    unsafe { <$Metadata as DeserializeMetadata>::deserialize_unchecked(self.0.metadata_bytes()) }
                }

                /// Optionally downcast an [`Array`](crate::Array) instance to a specific encoding.
                ///
                /// Preferred in cases where a backtrace isn't needed, like when trying multiple encoding to go
                /// down different code paths.
                pub fn maybe_from(data: impl AsRef<$crate::Array>) -> Option<Self> {
                    let data = data.as_ref();
                    (data.encoding() == <[<$Name Encoding>] as $crate::Encoding>::ID).then_some(Self(data.clone()))
                }
            }

            impl std::ops::Deref for [<$Name Array>] {
                type Target = $crate::Array;

                fn deref(&self) -> &Self::Target {
                    &self.0
                }
            }

            impl TryFrom<$crate::Array> for [<$Name Array>] {
                type Error = vortex_error::VortexError;

                fn try_from(data: $crate::Array) -> vortex_error::VortexResult<Self> {
                    if data.encoding() != <[<$Name Encoding>] as $crate::Encoding>::ID {
                        vortex_error::vortex_bail!(
                            "Mismatched encoding {}, expected {}",
                            data.encoding().as_ref(),
                            <[<$Name Encoding>] as $crate::Encoding>::ID,
                        );
                    }
                    Ok(Self(data))
                }
            }

            // NOTE(ngates): this is the cheeky one.... Since we know that Arrays are structurally
            //  equal to Array, we can transmute a &Array to a &Array.
            impl<'a> TryFrom<&'a $crate::Array> for &'a [<$Name Array>] {
                type Error = vortex_error::VortexError;

                fn try_from(data: &'a $crate::Array) -> vortex_error::VortexResult<Self> {
                    if data.encoding() != <[<$Name Encoding>] as $crate::Encoding>::ID {
                        vortex_error::vortex_bail!(
                            "Mismatched encoding {}, expected {}",
                            data.encoding().as_ref(),
                            <[<$Name Encoding>] as $crate::Encoding>::ID,
                        );
                    }
                    Ok(unsafe { std::mem::transmute::<&$crate::Array, &[<$Name Array>]>(data) })
                }
            }

            /// The array encoding
            #[derive(std::fmt::Debug)]
            pub struct [<$Name Encoding>];

            impl [<$Name Encoding>] {
                pub const fn vtable() -> $crate::vtable::VTableRef {
                    $crate::vtable::VTableRef::from_static(&Self)
                }
            }

            impl $crate::Encoding for [<$Name Encoding>] {
                const ID: $crate::EncodingId = $crate::EncodingId::new($id, $code);
                type Array = [<$Name Array>];
                type Metadata = $Metadata;
            }

            impl $crate::vtable::EncodingVTable for [<$Name Encoding>] {
                #[inline]
                fn id(&self) -> $crate::EncodingId {
                    <[<$Name Encoding>] as $crate::Encoding>::ID
                }

                fn as_any(&self) -> &dyn std::any::Any {
                    self
                }
            }
        }
    };
}

impl AsRef<Array> for Array {
    fn as_ref(&self) -> &Array {
        self
    }
}