vortex_vector/
datum.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use crate::Scalar;
5use crate::ScalarOps;
6use crate::Vector;
7use crate::VectorMutOps;
8use crate::VectorOps;
9use crate::binaryview::BinaryType;
10use crate::binaryview::BinaryViewScalar;
11use crate::binaryview::BinaryViewType;
12use crate::binaryview::BinaryViewVector;
13use crate::binaryview::StringType;
14use crate::bool::BoolScalar;
15use crate::bool::BoolVector;
16use crate::decimal::DScalar;
17use crate::decimal::DVector;
18use crate::decimal::DecimalScalar;
19use crate::decimal::DecimalVector;
20use crate::fixed_size_list::FixedSizeListScalar;
21use crate::fixed_size_list::FixedSizeListVector;
22use crate::listview::ListViewScalar;
23use crate::listview::ListViewVector;
24use crate::null::NullScalar;
25use crate::null::NullVector;
26use crate::primitive::PScalar;
27use crate::primitive::PVector;
28use crate::primitive::PrimitiveScalar;
29use crate::primitive::PrimitiveVector;
30use crate::struct_::StructScalar;
31use crate::struct_::StructVector;
32use vortex_dtype::NativeDecimalType;
33use vortex_dtype::NativePType;
34use vortex_dtype::PType;
35use vortex_error::vortex_panic;
36
37/// Represents either a scalar or vector value.
38#[derive(Clone, Debug)]
39pub enum Datum {
40    /// A scalar value.
41    Scalar(Scalar),
42    /// A vector value.
43    Vector(Vector),
44}
45
46impl From<Scalar> for Datum {
47    fn from(value: Scalar) -> Self {
48        Datum::Scalar(value)
49    }
50}
51
52impl From<Vector> for Datum {
53    fn from(value: Vector) -> Self {
54        Datum::Vector(value)
55    }
56}
57
58impl Datum {
59    /// Ensure the datum is a vector by repeating the scalar value if necessary.
60    pub fn unwrap_into_vector(self, len: usize) -> Vector {
61        match self {
62            Datum::Scalar(scalar) => scalar.repeat(len).freeze(),
63            Datum::Vector(vector) => {
64                if vector.len() != len {
65                    vortex_panic!(
66                        "Vector length mismatch: expected length {}, got {}",
67                        len,
68                        vector.len()
69                    );
70                }
71                vector
72            }
73        }
74    }
75
76    /// Returns the scalar value if this `Datum` is a `Scalar`, otherwise returns `None`.
77    pub fn into_scalar(self) -> Option<Scalar> {
78        match self {
79            Datum::Scalar(scalar) => Some(scalar),
80            Datum::Vector(_) => None,
81        }
82    }
83
84    /// Returns the vector value if this `Datum` is a `Vector`, otherwise returns `None`.
85    pub fn into_vector(self) -> Option<Vector> {
86        match self {
87            Datum::Scalar(_) => None,
88            Datum::Vector(vector) => Some(vector),
89        }
90    }
91
92    /// Returns a reference to the scalar value if this `Datum` is a `Scalar`, otherwise returns `None`.
93    pub fn as_scalar(&self) -> Option<&Scalar> {
94        match self {
95            Datum::Scalar(scalar) => Some(scalar),
96            Datum::Vector(_) => None,
97        }
98    }
99
100    /// Returns a reference to the vector value if this `Datum` is a `Vector`, otherwise returns `None`.
101    pub fn as_vector(&self) -> Option<&Vector> {
102        match self {
103            Datum::Scalar(_) => None,
104            Datum::Vector(vector) => Some(vector),
105        }
106    }
107}
108
109impl Datum {
110    /// Converts the `Datum` into a `NullDatum`.
111    pub fn into_null(self) -> NullDatum {
112        match self {
113            Datum::Scalar(scalar) => NullDatum::Scalar(scalar.into_null()),
114            Datum::Vector(vector) => NullDatum::Vector(vector.into_null()),
115        }
116    }
117
118    /// Converts the `Datum` into a `BoolDatum`.
119    pub fn into_bool(self) -> BoolDatum {
120        match self {
121            Datum::Scalar(scalar) => BoolDatum::Scalar(scalar.into_bool()),
122            Datum::Vector(vector) => BoolDatum::Vector(vector.into_bool()),
123        }
124    }
125
126    /// Converts the `Datum` into a `PrimitiveDatum`.
127    pub fn into_primitive(self) -> PrimitiveDatum {
128        match self {
129            Datum::Scalar(scalar) => PrimitiveDatum::Scalar(scalar.into_primitive()),
130            Datum::Vector(vector) => PrimitiveDatum::Vector(vector.into_primitive()),
131        }
132    }
133
134    /// Converts the `Datum` into a `DecimalDatum`.
135    pub fn into_decimal(self) -> DecimalDatum {
136        match self {
137            Datum::Scalar(scalar) => DecimalDatum::Scalar(scalar.into_decimal()),
138            Datum::Vector(vector) => DecimalDatum::Vector(vector.into_decimal()),
139        }
140    }
141
142    /// Converts the `Datum` into a `ListViewDatum`.
143    pub fn into_list(self) -> ListViewDatum {
144        match self {
145            Datum::Scalar(scalar) => ListViewDatum::Scalar(scalar.into_list()),
146            Datum::Vector(vector) => ListViewDatum::Vector(vector.into_list()),
147        }
148    }
149
150    /// Converts the `Datum` into a `FixedSizeListDatum`.
151    pub fn into_fixed_size_list(self) -> FixedSizeListDatum {
152        match self {
153            Datum::Scalar(scalar) => FixedSizeListDatum::Scalar(scalar.into_fixed_size_list()),
154            Datum::Vector(vector) => FixedSizeListDatum::Vector(vector.into_fixed_size_list()),
155        }
156    }
157
158    /// Converts the `Datum` into a `StructDatum`.
159    pub fn into_struct(self) -> StructDatum {
160        match self {
161            Datum::Scalar(scalar) => StructDatum::Scalar(scalar.into_struct()),
162            Datum::Vector(vector) => StructDatum::Vector(vector.into_struct()),
163        }
164    }
165
166    /// Converts the `Datum` into a `TypedDatum`.
167    pub fn into_typed(self) -> TypedDatum {
168        match self {
169            Datum::Scalar(sc) => match sc {
170                Scalar::Null(s) => TypedDatum::Null(NullDatum::Scalar(s)),
171                Scalar::Bool(s) => TypedDatum::Bool(BoolDatum::Scalar(s)),
172                Scalar::Decimal(s) => TypedDatum::Decimal(DecimalDatum::Scalar(s)),
173                Scalar::Primitive(s) => TypedDatum::Primitive(PrimitiveDatum::Scalar(s)),
174                Scalar::String(s) => TypedDatum::String(BinaryViewDatum::Scalar(s)),
175                Scalar::Binary(s) => TypedDatum::Binary(BinaryViewDatum::Scalar(s)),
176                Scalar::List(s) => TypedDatum::List(ListViewDatum::Scalar(s)),
177                Scalar::FixedSizeList(s) => {
178                    TypedDatum::FixedSizeList(FixedSizeListDatum::Scalar(s))
179                }
180                Scalar::Struct(s) => TypedDatum::Struct(StructDatum::Scalar(s)),
181            },
182            Datum::Vector(vec) => match vec {
183                Vector::Null(v) => TypedDatum::Null(NullDatum::Vector(v)),
184                Vector::Bool(v) => TypedDatum::Bool(BoolDatum::Vector(v)),
185                Vector::Decimal(v) => TypedDatum::Decimal(DecimalDatum::Vector(v)),
186                Vector::Primitive(v) => TypedDatum::Primitive(PrimitiveDatum::Vector(v)),
187                Vector::String(v) => TypedDatum::String(BinaryViewDatum::Vector(v)),
188                Vector::Binary(v) => TypedDatum::Binary(BinaryViewDatum::Vector(v)),
189                Vector::List(v) => TypedDatum::List(ListViewDatum::Vector(v)),
190                Vector::FixedSizeList(v) => {
191                    TypedDatum::FixedSizeList(FixedSizeListDatum::Vector(v))
192                }
193                Vector::Struct(v) => TypedDatum::Struct(StructDatum::Vector(v)),
194            },
195        }
196    }
197}
198
199macro_rules! datum {
200    // Non-generic version
201    ($Name:ident) => {
202        paste::paste! {
203            #[derive(Clone, Debug)]
204            #[doc = concat!("Datum enum for `", stringify!($Name), "`.")]
205            pub enum [<$Name Datum>] {
206                /// Scalar variant
207                Scalar([<$Name Scalar>]),
208                /// Vector variant
209                Vector([<$Name Vector>]),
210            }
211
212            impl From<[<$Name Datum>]> for Datum {
213                fn from(val: [<$Name Datum>]) -> Self {
214                    match val {
215                        [<$Name Datum>]::Scalar(scalar) => Datum::Scalar(Scalar::from(scalar)),
216                        [<$Name Datum>]::Vector(vector) => Datum::Vector(Vector::from(vector)),
217                    }
218                }
219            }
220
221            impl From<[<$Name Scalar>]> for Datum {
222                fn from(val: [<$Name Scalar>]) -> Self {
223                    Datum::Scalar(Scalar::from(val))
224                }
225            }
226
227            impl From<[<$Name Scalar>]> for [<$Name Datum>] {
228                fn from(val: [<$Name Scalar>]) -> Self {
229                    [<$Name Datum>]::Scalar(val)
230                }
231            }
232
233            impl From<[<$Name Vector>]> for [<$Name Datum>] {
234                fn from(val: [<$Name Vector>]) -> Self {
235                    [<$Name Datum>]::Vector(val)
236                }
237            }
238        }
239    };
240
241    // Generic version with trait bound
242    ($Name:ident < $T:ident : $Bound:path >) => {
243        paste::paste! {
244            #[derive(Clone, Debug)]
245            #[doc = concat!("Datum enum for `", stringify!([<$Name Datum>]<$T: $Bound>), "`.")]
246            pub enum [<$Name Datum>]<$T: $Bound> {
247                /// Scalar variant
248                Scalar([<$Name Scalar>]<$T>),
249                /// Vector variant
250                Vector([<$Name Vector>]<$T>),
251            }
252
253            impl<$T: $Bound> From<[<$Name Datum>]<$T>> for Datum {
254                fn from(val: [<$Name Datum>]<$T>) -> Self {
255                    match val {
256                        [<$Name Datum>]::Scalar(scalar) => Datum::Scalar(Scalar::from(scalar)),
257                        [<$Name Datum>]::Vector(vector) => Datum::Vector(Vector::from(vector)),
258                    }
259                }
260            }
261
262            impl<$T: $Bound> From<[<$Name Scalar>]<$T>> for Datum {
263                fn from(val: [<$Name Scalar>]<$T>) -> Self {
264                    Datum::Scalar(Scalar::from(val))
265                }
266            }
267
268            impl<$T: $Bound> From<[<$Name Scalar>]<$T>> for [<$Name Datum>]<$T> {
269                fn from(val: [<$Name Scalar>]<$T>) -> Self {
270                    [<$Name Datum>]::Scalar(val)
271                }
272            }
273
274            impl<$T: $Bound> From<[<$Name Vector>]<$T>> for [<$Name Datum>]<$T> {
275                fn from(val: [<$Name Vector>]<$T>) -> Self {
276                    [<$Name Datum>]::Vector(val)
277                }
278            }
279        }
280    };
281}
282
283datum!(Null);
284datum!(Bool);
285datum!(Primitive);
286datum!(P<T: NativePType>);
287datum!(Decimal);
288datum!(D<D: NativeDecimalType>);
289datum!(BinaryView<T: BinaryViewType>);
290datum!(ListView);
291datum!(FixedSizeList);
292datum!(Struct);
293
294impl PrimitiveDatum {
295    /// Returns the primitive type of this datum.
296    pub fn ptype(&self) -> PType {
297        match self {
298            PrimitiveDatum::Scalar(sc) => sc.ptype(),
299            PrimitiveDatum::Vector(sc) => sc.ptype(),
300        }
301    }
302}
303
304/// A variant of [`Datum`] that is typed.
305#[derive(Debug)]
306pub enum TypedDatum {
307    /// Null datum.
308    Null(NullDatum),
309    /// Boolean datum.
310    Bool(BoolDatum),
311    /// Decimal datum.
312    Decimal(DecimalDatum),
313    /// Primitive datum.
314    Primitive(PrimitiveDatum),
315    /// String datum
316    String(BinaryViewDatum<StringType>),
317    /// Binary datum
318    Binary(BinaryViewDatum<BinaryType>),
319    /// List datum.
320    List(ListViewDatum),
321    /// fsl datum.
322    FixedSizeList(FixedSizeListDatum),
323    /// struct datum.
324    Struct(StructDatum),
325}