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_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure_eq;
13use vortex_error::vortex_err;
14use vortex_error::vortex_panic;
15
16use crate::dtype::DType;
17use crate::dtype::DecimalDType;
18use crate::dtype::NativePType;
19use crate::dtype::Nullability;
20use crate::dtype::PType;
21use crate::dtype::UnionVariants;
22use crate::dtype::extension::ExtDType;
23use crate::dtype::extension::ExtDTypeRef;
24use crate::dtype::extension::ExtVTable;
25use crate::scalar::DecimalValue;
26use crate::scalar::PValue;
27use crate::scalar::Scalar;
28use crate::scalar::ScalarValue;
29use crate::scalar::UnionValue;
30
31// TODO(connor): Really, we want `try_` constructors that return errors instead of just panic.
32impl Scalar {
33    /// Creates a new boolean scalar with the given value and nullability.
34    pub fn bool(value: bool, nullability: Nullability) -> Self {
35        Self::try_new(DType::Bool(nullability), Some(ScalarValue::Bool(value)))
36            .vortex_expect("unable to construct a boolean `Scalar`")
37    }
38
39    /// Creates a new primitive scalar from a native value.
40    pub fn primitive<T: NativePType + Into<PValue>>(value: T, nullability: Nullability) -> Self {
41        Self::primitive_value(value.into(), T::PTYPE, nullability)
42    }
43
44    /// Create a PrimitiveScalar from a PValue.
45    ///
46    /// Note that an explicit PType is passed since any compatible PValue may be used as the value
47    /// for a primitive type.
48    pub fn primitive_value(value: PValue, ptype: PType, nullability: Nullability) -> Self {
49        Self::try_new(
50            DType::Primitive(ptype, nullability),
51            Some(ScalarValue::Primitive(value)),
52        )
53        .vortex_expect("unable to construct a primitive `Scalar`")
54    }
55
56    /// Creates a new decimal scalar with the given value, precision, scale, and nullability.
57    pub fn decimal(
58        value: DecimalValue,
59        decimal_type: DecimalDType,
60        nullability: Nullability,
61    ) -> Self {
62        Self::try_new(
63            DType::Decimal(decimal_type, nullability),
64            Some(ScalarValue::Decimal(value)),
65        )
66        .vortex_expect("unable to construct a decimal `Scalar`")
67    }
68
69    /// Creates a new UTF-8 scalar from a string-like value.
70    ///
71    /// # Panics
72    ///
73    /// Panics if the input cannot be converted to a valid UTF-8 string.
74    pub fn utf8<B>(str: B, nullability: Nullability) -> Self
75    where
76        B: Into<BufferString>,
77    {
78        Self::try_utf8(str, nullability).unwrap()
79    }
80
81    /// Tries to create a new UTF-8 scalar from a string-like value.
82    ///
83    /// # Errors
84    ///
85    /// Returns an error if the input cannot be converted to a valid UTF-8 string.
86    pub fn try_utf8<B>(
87        str: B,
88        nullability: Nullability,
89    ) -> Result<Self, <B as TryInto<BufferString>>::Error>
90    where
91        B: TryInto<BufferString>,
92    {
93        Ok(Self::try_new(
94            DType::Utf8(nullability),
95            Some(ScalarValue::Utf8(str.try_into()?)),
96        )
97        .vortex_expect("unable to construct a UTF-8 `Scalar`"))
98    }
99
100    /// Creates a new binary scalar from a byte buffer.
101    pub fn binary(buffer: impl Into<ByteBuffer>, nullability: Nullability) -> Self {
102        Self::try_new(
103            DType::Binary(nullability),
104            Some(ScalarValue::Binary(buffer.into())),
105        )
106        .vortex_expect("unable to construct a binary `Scalar`")
107    }
108
109    /// Creates a new list scalar with the given element type and children.
110    ///
111    /// # Panics
112    ///
113    /// Panics if any child scalar has a different type than the element type, or if there are too
114    /// many children.
115    pub fn list(
116        element_dtype: impl Into<Arc<DType>>,
117        children: Vec<Scalar>,
118        nullability: Nullability,
119    ) -> Self {
120        Self::create_list(element_dtype, children, nullability, ListKind::Variable)
121    }
122
123    /// Creates a new empty list scalar with the given element type.
124    pub fn list_empty(element_dtype: Arc<DType>, nullability: Nullability) -> Self {
125        Self::create_list(element_dtype, vec![], nullability, ListKind::Variable)
126    }
127
128    /// Creates a new fixed-size list scalar with the given element type and children.
129    ///
130    /// # Panics
131    ///
132    /// Panics if any child scalar has a different type than the element type, or if there are too
133    /// many children.
134    pub fn fixed_size_list(
135        element_dtype: impl Into<Arc<DType>>,
136        children: Vec<Scalar>,
137        nullability: Nullability,
138    ) -> Self {
139        Self::create_list(element_dtype, children, nullability, ListKind::FixedSize)
140    }
141
142    /// Creates a list [`Scalar`] from an element dtype, children, nullability, and list kind.
143    fn create_list(
144        element_dtype: impl Into<Arc<DType>>,
145        children: Vec<Scalar>,
146        nullability: Nullability,
147        list_kind: ListKind,
148    ) -> Self {
149        let element_dtype = element_dtype.into();
150
151        let children: Vec<Option<ScalarValue>> = children
152            .into_iter()
153            .map(|child| {
154                if child.dtype() != &*element_dtype {
155                    vortex_panic!(
156                        "tried to create list of {} with values of type {}",
157                        element_dtype,
158                        child.dtype()
159                    );
160                }
161                child.into_value()
162            })
163            .collect();
164        let size: u32 = children
165            .len()
166            .try_into()
167            .vortex_expect("tried to create a list that was too large");
168
169        let dtype = match list_kind {
170            ListKind::Variable => DType::List(element_dtype, nullability),
171            ListKind::FixedSize => DType::FixedSizeList(element_dtype, size, nullability),
172        };
173
174        Self::try_new(dtype, Some(ScalarValue::Tuple(children)))
175            .vortex_expect("unable to construct a list `Scalar`")
176    }
177
178    /// Creates a new extension scalar wrapping the given storage value.
179    pub fn extension<V: ExtVTable + Default>(options: V::Metadata, storage_scalar: Scalar) -> Self {
180        let ext_dtype = ExtDType::<V>::try_new(options, storage_scalar.dtype().clone())
181            .vortex_expect("Failed to create extension dtype");
182
183        Self::extension_ref(ext_dtype.erased(), storage_scalar)
184    }
185
186    /// Creates a new extension scalar wrapping the given storage value.
187    ///
188    /// # Panics
189    ///
190    /// Panics if the storage dtype of `ext_dtype` does not match `value`'s dtype.
191    pub fn extension_ref(ext_dtype: ExtDTypeRef, storage_scalar: Scalar) -> Self {
192        assert_eq!(ext_dtype.storage_dtype(), storage_scalar.dtype());
193
194        Self::try_new(DType::Extension(ext_dtype), storage_scalar.into_value())
195            .vortex_expect("unable to construct an extension `Scalar`")
196    }
197
198    /// Creates a union scalar from a type ID and child scalar.
199    ///
200    /// The selected child's dtype is verified and then discarded. Its raw value is stored so an
201    /// inner null child remains distinct from a null at the outer union level. Passing a null child
202    /// creates a non-null union with a selected type ID; use [`Scalar::null`] with a nullable
203    /// [`DType::Union`] to create an outer null union with no selected type ID.
204    ///
205    /// # Errors
206    ///
207    /// Returns an error if the type ID is not present in `variants` or if the child scalar's dtype
208    /// does not exactly match the selected variant dtype.
209    pub fn union(
210        variants: UnionVariants,
211        type_id: u8,
212        child: Scalar,
213        nullability: Nullability,
214    ) -> VortexResult<Self> {
215        let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| {
216            vortex_err!(
217                "union type ID {type_id} is not present in {:?}",
218                variants.type_ids()
219            )
220        })?;
221
222        let expected_dtype = variants
223            .variant_by_index(child_index)
224            .vortex_expect("type ID resolved to a valid child index");
225
226        vortex_ensure_eq!(
227            child.dtype(),
228            &expected_dtype,
229            "union type ID {type_id} selects child dtype {expected_dtype}, got {}",
230            child.dtype()
231        );
232
233        Self::try_new(
234            DType::Union(variants, nullability),
235            Some(ScalarValue::Union(UnionValue::new(
236                type_id,
237                child.into_value(),
238            ))),
239        )
240    }
241
242    /// Creates a new variant scalar from a row-specific nested scalar.
243    ///
244    /// Use [`Scalar::null(DType::Variant(Nullability::Nullable))`][Scalar::null] for a top-level
245    /// null variant value, and
246    /// `Scalar::variant(Scalar::null(DType::Null))` for a defined variant-null.
247    pub fn variant(value: Scalar) -> Self {
248        Self::try_new(
249            DType::Variant(Nullability::NonNullable),
250            Some(ScalarValue::Variant(Box::new(value))),
251        )
252        .vortex_expect("unable to construct a variant `Scalar`")
253    }
254}
255
256/// A helper enum for creating a `ListScalar`.
257enum ListKind {
258    /// Variable-length list.
259    Variable,
260    /// Fixed-size list.
261    FixedSize,
262}