Skip to main content

vortex_array/scalar/
constructor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Typed constructors for [`Scalar`].
5
6use std::sync::Arc;
7
8use vortex_buffer::BufferString;
9use vortex_buffer::ByteBuffer;
10use vortex_dtype::DType;
11use vortex_dtype::DecimalDType;
12use vortex_dtype::ExtDType;
13use vortex_dtype::ExtDTypeRef;
14use vortex_dtype::NativePType;
15use vortex_dtype::Nullability;
16use vortex_dtype::PType;
17use vortex_dtype::extension::ExtDTypeVTable;
18use vortex_error::VortexExpect;
19use vortex_error::vortex_panic;
20
21use crate::scalar::DecimalValue;
22use crate::scalar::PValue;
23use crate::scalar::Scalar;
24use crate::scalar::ScalarValue;
25
26// TODO(connor): Really, we want `try_` constructors that return errors instead of just panic.
27impl Scalar {
28    /// Creates a new boolean scalar with the given value and nullability.
29    pub fn bool(value: bool, nullability: Nullability) -> Self {
30        Self::try_new(DType::Bool(nullability), Some(ScalarValue::Bool(value)))
31            .vortex_expect("unable to construct a boolean `Scalar`")
32    }
33
34    /// Creates a new primitive scalar from a native value.
35    pub fn primitive<T: NativePType + Into<PValue>>(value: T, nullability: Nullability) -> Self {
36        Self::primitive_value(value.into(), T::PTYPE, nullability)
37    }
38
39    /// Create a PrimitiveScalar from a PValue.
40    ///
41    /// Note that an explicit PType is passed since any compatible PValue may be used as the value
42    /// for a primitive type.
43    pub fn primitive_value(value: PValue, ptype: PType, nullability: Nullability) -> Self {
44        Self::try_new(
45            DType::Primitive(ptype, nullability),
46            Some(ScalarValue::Primitive(value)),
47        )
48        .vortex_expect("unable to construct a primitive `Scalar`")
49    }
50
51    /// Creates a new decimal scalar with the given value, precision, scale, and nullability.
52    pub fn decimal(
53        value: DecimalValue,
54        decimal_type: DecimalDType,
55        nullability: Nullability,
56    ) -> Self {
57        Self::try_new(
58            DType::Decimal(decimal_type, nullability),
59            Some(ScalarValue::Decimal(value)),
60        )
61        .vortex_expect("unable to construct a decimal `Scalar`")
62    }
63
64    /// Creates a new UTF-8 scalar from a string-like value.
65    ///
66    /// # Panics
67    ///
68    /// Panics if the input cannot be converted to a valid UTF-8 string.
69    pub fn utf8<B>(str: B, nullability: Nullability) -> Self
70    where
71        B: Into<BufferString>,
72    {
73        Self::try_utf8(str, nullability).unwrap()
74    }
75
76    /// Tries to create a new UTF-8 scalar from a string-like value.
77    ///
78    /// # Errors
79    ///
80    /// Returns an error if the input cannot be converted to a valid UTF-8 string.
81    pub fn try_utf8<B>(
82        str: B,
83        nullability: Nullability,
84    ) -> Result<Self, <B as TryInto<BufferString>>::Error>
85    where
86        B: TryInto<BufferString>,
87    {
88        Ok(Self::try_new(
89            DType::Utf8(nullability),
90            Some(ScalarValue::Utf8(str.try_into()?)),
91        )
92        .vortex_expect("unable to construct a UTF-8 `Scalar`"))
93    }
94
95    /// Creates a new binary scalar from a byte buffer.
96    pub fn binary(buffer: impl Into<ByteBuffer>, nullability: Nullability) -> Self {
97        Self::try_new(
98            DType::Binary(nullability),
99            Some(ScalarValue::Binary(buffer.into())),
100        )
101        .vortex_expect("unable to construct a binary `Scalar`")
102    }
103
104    /// Creates a new list scalar with the given element type and children.
105    ///
106    /// # Panics
107    ///
108    /// Panics if any child scalar has a different type than the element type, or if there are too
109    /// many children.
110    pub fn list(
111        element_dtype: impl Into<Arc<DType>>,
112        children: Vec<Scalar>,
113        nullability: Nullability,
114    ) -> Self {
115        Self::create_list(element_dtype, children, nullability, ListKind::Variable)
116    }
117
118    /// Creates a new empty list scalar with the given element type.
119    pub fn list_empty(element_dtype: Arc<DType>, nullability: Nullability) -> Self {
120        Self::create_list(element_dtype, vec![], nullability, ListKind::Variable)
121    }
122
123    /// Creates a new fixed-size list scalar with the given element type and children.
124    ///
125    /// # Panics
126    ///
127    /// Panics if any child scalar has a different type than the element type, or if there are too
128    /// many children.
129    pub fn fixed_size_list(
130        element_dtype: impl Into<Arc<DType>>,
131        children: Vec<Scalar>,
132        nullability: Nullability,
133    ) -> Self {
134        Self::create_list(element_dtype, children, nullability, ListKind::FixedSize)
135    }
136
137    /// Creates a list [`Scalar`] from an element dtype, children, nullability, and list kind.
138    fn create_list(
139        element_dtype: impl Into<Arc<DType>>,
140        children: Vec<Scalar>,
141        nullability: Nullability,
142        list_kind: ListKind,
143    ) -> Self {
144        let element_dtype = element_dtype.into();
145
146        let children: Vec<Option<ScalarValue>> = children
147            .into_iter()
148            .map(|child| {
149                if child.dtype() != &*element_dtype {
150                    vortex_panic!(
151                        "tried to create list of {} with values of type {}",
152                        element_dtype,
153                        child.dtype()
154                    );
155                }
156                child.into_value()
157            })
158            .collect();
159        let size: u32 = children
160            .len()
161            .try_into()
162            .vortex_expect("tried to create a list that was too large");
163
164        let dtype = match list_kind {
165            ListKind::Variable => DType::List(element_dtype, nullability),
166            ListKind::FixedSize => DType::FixedSizeList(element_dtype, size, nullability),
167        };
168
169        Self::try_new(dtype, Some(ScalarValue::List(children)))
170            .vortex_expect("unable to construct a list `Scalar`")
171    }
172
173    /// Creates a new extension scalar wrapping the given storage value.
174    pub fn extension<V: ExtDTypeVTable + Default>(options: V::Metadata, value: Scalar) -> Self {
175        let ext_dtype = ExtDType::<V>::try_new(options, value.dtype().clone())
176            .vortex_expect("Failed to create extension dtype");
177        Self::try_new(DType::Extension(ext_dtype.erased()), value.into_value())
178            .vortex_expect("unable to construct an extension `Scalar`")
179    }
180
181    /// Creates a new extension scalar wrapping the given storage value.
182    ///
183    /// # Panics
184    ///
185    /// Panics if the storage dtype of `ext_dtype` does not match `value`'s dtype.
186    pub fn extension_ref(ext_dtype: ExtDTypeRef, value: Scalar) -> Self {
187        assert_eq!(ext_dtype.storage_dtype(), value.dtype());
188        Self::try_new(DType::Extension(ext_dtype), value.into_value())
189            .vortex_expect("unable to construct an extension `Scalar`")
190    }
191}
192
193/// A helper enum for creating a [`ListScalar`].
194enum ListKind {
195    /// Variable-length list.
196    Variable,
197    /// Fixed-size list.
198    FixedSize,
199}