vortex_array/
encoding.rs

1//! Traits and types to define shared unique encoding identifiers.
2
3use std::any::Any;
4use std::fmt::{Debug, Display, Formatter};
5use std::sync::Arc;
6
7use arcref::ArcRef;
8use vortex_buffer::ByteBuffer;
9use vortex_dtype::DType;
10use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
11
12use crate::serde::ArrayChildren;
13use crate::vtable::{EncodeVTable, SerdeVTable, VTable};
14use crate::{Array, ArrayRef, Canonical, DeserializeMetadata};
15
16/// EncodingId is a globally unique name of the array's encoding.
17pub type EncodingId = ArcRef<str>;
18
19pub type EncodingRef = ArcRef<dyn Encoding>;
20
21/// Marker trait for array encodings with their associated Array type.
22pub trait Encoding: 'static + private::Sealed + Send + Sync + Debug {
23    /// Downcast the encoding to [`Any`].
24    fn as_any(&self) -> &dyn Any;
25
26    fn to_encoding(&self) -> EncodingRef;
27
28    fn into_encoding(self) -> EncodingRef
29    where
30        Self: Sized;
31
32    /// Returns the ID of the encoding.
33    fn id(&self) -> EncodingId;
34
35    /// Build an array from its parts.
36    fn build(
37        &self,
38        dtype: &DType,
39        len: usize,
40        metadata: &[u8],
41        buffers: &[ByteBuffer],
42        children: &dyn ArrayChildren,
43    ) -> VortexResult<ArrayRef>;
44
45    /// Encode the canonical array into this encoding implementation.
46    /// Returns `None` if this encoding does not support the given canonical array, for example
47    /// if the data type is incompatible.
48    ///
49    /// Panics if `like` is encoded with a different encoding.
50    fn encode(&self, input: &Canonical, like: Option<&dyn Array>)
51    -> VortexResult<Option<ArrayRef>>;
52}
53
54/// Adapter struct used to lift the [`VTable`] trait into an object-safe [`Encoding`]
55/// implementation.
56///
57/// Since this is a unit struct with `repr(transparent)`, we are able to turn un-adapted array
58/// structs into [`dyn Encoding`] using some cheeky casting inside [`std::ops::Deref`] and
59/// [`AsRef`]. See the `vtable!` macro for more details.
60#[repr(transparent)]
61pub struct EncodingAdapter<V: VTable>(V::Encoding);
62
63impl<V: VTable> Encoding for EncodingAdapter<V> {
64    fn as_any(&self) -> &dyn Any {
65        self
66    }
67
68    fn to_encoding(&self) -> EncodingRef {
69        ArcRef::new_arc(Arc::new(EncodingAdapter::<V>(self.0.clone())))
70    }
71
72    fn into_encoding(self) -> EncodingRef
73    where
74        Self: Sized,
75    {
76        todo!()
77    }
78
79    fn id(&self) -> EncodingId {
80        V::id(&self.0)
81    }
82
83    fn build(
84        &self,
85        dtype: &DType,
86        len: usize,
87        metadata: &[u8],
88        buffers: &[ByteBuffer],
89        children: &dyn ArrayChildren,
90    ) -> VortexResult<ArrayRef> {
91        let metadata =
92            <<V::SerdeVTable as SerdeVTable<V>>::Metadata as DeserializeMetadata>::deserialize(
93                metadata,
94            )?;
95        let array = <V::SerdeVTable as SerdeVTable<V>>::build(
96            &self.0, dtype, len, &metadata, buffers, children,
97        )?;
98        assert_eq!(array.len(), len, "Array length mismatch after building");
99        assert_eq!(array.dtype(), dtype, "Array dtype mismatch after building");
100        Ok(array.to_array())
101    }
102
103    fn encode(
104        &self,
105        input: &Canonical,
106        like: Option<&dyn Array>,
107    ) -> VortexResult<Option<ArrayRef>> {
108        let downcast_like = like
109            .map(|like| {
110                like.as_opt::<V>().ok_or_else(|| {
111                    vortex_err!(
112                        "Like array {} does not match requested encoding {}",
113                        like.encoding_id(),
114                        self.id()
115                    )
116                })
117            })
118            .transpose()?;
119
120        let Some(array) =
121            <V::EncodeVTable as EncodeVTable<V>>::encode(&self.0, input, downcast_like)?
122        else {
123            return Ok(None);
124        };
125
126        let input = input.as_ref();
127        if array.len() != input.len() {
128            vortex_bail!(
129                "Array length mismatch after encoding: {} != {}",
130                array.len(),
131                input.len()
132            );
133        }
134        if array.dtype() != input.dtype() {
135            vortex_bail!(
136                "Array dtype mismatch after encoding: {} != {}",
137                array.dtype(),
138                input.dtype()
139            );
140        }
141
142        Ok(Some(array.to_array()))
143    }
144}
145
146impl<V: VTable> Debug for EncodingAdapter<V> {
147    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148        f.debug_struct("Encoding").field("id", &self.id()).finish()
149    }
150}
151
152impl Display for dyn Encoding + '_ {
153    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
154        write!(f, "{}", self.id())
155    }
156}
157
158impl PartialEq for dyn Encoding + '_ {
159    fn eq(&self, other: &Self) -> bool {
160        self.id() == other.id()
161    }
162}
163
164impl Eq for dyn Encoding + '_ {}
165
166impl dyn Encoding + '_ {
167    pub fn as_<V: VTable>(&self) -> &V::Encoding {
168        self.as_any()
169            .downcast_ref::<EncodingAdapter<V>>()
170            .map(|e| &e.0)
171            .vortex_expect("Encoding is not of the expected type")
172    }
173}
174
175mod private {
176    use super::*;
177
178    pub trait Sealed {}
179
180    impl<V: VTable> Sealed for EncodingAdapter<V> {}
181}
vortex_array/encoding.rs

vortex_array/
encoding.rs
