vortex_scalar/
binary.rs

1use std::fmt::{Display, Formatter};
2use std::sync::Arc;
3
4use itertools::Itertools;
5use vortex_buffer::ByteBuffer;
6use vortex_dtype::{DType, Nullability};
7use vortex_error::{VortexError, VortexExpect as _, VortexResult, vortex_bail, vortex_err};
8
9use crate::{InnerScalarValue, Scalar, ScalarValue};
10
11#[derive(Debug, Hash)]
12pub struct BinaryScalar<'a> {
13    dtype: &'a DType,
14    value: Option<Arc<ByteBuffer>>,
15}
16
17impl Display for BinaryScalar<'_> {
18    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
19        match &self.value {
20            None => write!(f, "null"),
21            Some(v) => write!(
22                f,
23                "\"{}\"",
24                v.as_slice().iter().map(|b| format!("{b:x}")).format(" ")
25            ),
26        }
27    }
28}
29
30impl PartialEq for BinaryScalar<'_> {
31    fn eq(&self, other: &Self) -> bool {
32        self.dtype.eq_ignore_nullability(other.dtype) && self.value == other.value
33    }
34}
35
36impl Eq for BinaryScalar<'_> {}
37
38impl PartialOrd for BinaryScalar<'_> {
39    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
40        Some(self.value.cmp(&other.value))
41    }
42}
43
44impl Ord for BinaryScalar<'_> {
45    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
46        self.value.cmp(&other.value)
47    }
48}
49
50impl<'a> BinaryScalar<'a> {
51    pub fn from_scalar_value(dtype: &'a DType, value: ScalarValue) -> VortexResult<Self> {
52        if !matches!(dtype, DType::Binary(..)) {
53            vortex_bail!("Can only construct binary scalar from binary dtype, found {dtype}")
54        }
55        Ok(Self {
56            dtype,
57            value: value.as_buffer()?,
58        })
59    }
60
61    #[inline]
62    pub fn dtype(&self) -> &'a DType {
63        self.dtype
64    }
65
66    pub fn value(&self) -> Option<ByteBuffer> {
67        self.value.as_ref().map(|v| v.as_ref().clone())
68    }
69
70    /// Construct a value at most `max_length` in size that's greater than ourselves.
71    ///
72    /// Will return None if constructing greater value overflows
73    pub fn upper_bound(self, max_length: usize) -> Option<Self> {
74        if let Some(value) = self.value {
75            if value.len() > max_length {
76                let sliced = value.slice(0..max_length);
77                drop(value);
78                let mut sliced_mut = sliced.into_mut();
79                for b in sliced_mut.iter_mut().rev() {
80                    let (incr, overflow) = b.overflowing_add(1);
81                    *b = incr;
82                    if !overflow {
83                        return Some(Self {
84                            dtype: self.dtype,
85                            value: Some(Arc::new(sliced_mut.freeze())),
86                        });
87                    }
88                }
89                None
90            } else {
91                Some(Self {
92                    dtype: self.dtype,
93                    value: Some(value),
94                })
95            }
96        } else {
97            Some(self)
98        }
99    }
100
101    /// Construct a value at most `max_length` in size that's less than ourselves.
102    pub fn lower_bound(self, max_length: usize) -> Self {
103        if let Some(value) = self.value {
104            if value.len() > max_length {
105                Self {
106                    dtype: self.dtype,
107                    value: Some(Arc::new(value.slice(0..max_length))),
108                }
109            } else {
110                Self {
111                    dtype: self.dtype,
112                    value: Some(value),
113                }
114            }
115        } else {
116            self
117        }
118    }
119
120    pub(crate) fn cast(&self, dtype: &DType) -> VortexResult<Scalar> {
121        if !matches!(dtype, DType::Binary(..)) {
122            vortex_bail!("Can't cast binary to {}", dtype)
123        }
124        Ok(Scalar::new(
125            dtype.clone(),
126            ScalarValue(InnerScalarValue::Buffer(
127                self.value
128                    .as_ref()
129                    .vortex_expect("nullness handled in Scalar::cast")
130                    .clone(),
131            )),
132        ))
133    }
134
135    /// Length of the scalar value or None if value is null
136    pub fn len(&self) -> Option<usize> {
137        self.value.as_ref().map(|v| v.len())
138    }
139
140    /// Returns whether its value is non-null and empty, otherwise `None`.
141    pub fn is_empty(&self) -> Option<bool> {
142        self.value.as_ref().map(|v| v.is_empty())
143    }
144
145    /// Extract value as a ScalarValue
146    pub fn into_value(self) -> ScalarValue {
147        ScalarValue(
148            self.value
149                .map(InnerScalarValue::Buffer)
150                .unwrap_or_else(|| InnerScalarValue::Null),
151        )
152    }
153}
154
155impl Scalar {
156    pub fn binary(buffer: impl Into<Arc<ByteBuffer>>, nullability: Nullability) -> Self {
157        Self {
158            dtype: DType::Binary(nullability),
159            value: ScalarValue(InnerScalarValue::Buffer(buffer.into())),
160        }
161    }
162}
163
164impl<'a> TryFrom<&'a Scalar> for BinaryScalar<'a> {
165    type Error = VortexError;
166
167    fn try_from(value: &'a Scalar) -> Result<Self, Self::Error> {
168        if !matches!(value.dtype(), DType::Binary(_)) {
169            vortex_bail!("Expected binary scalar, found {}", value.dtype())
170        }
171        Ok(Self {
172            dtype: value.dtype(),
173            value: value.value.as_buffer()?,
174        })
175    }
176}
177
178impl<'a> TryFrom<&'a Scalar> for ByteBuffer {
179    type Error = VortexError;
180
181    fn try_from(scalar: &'a Scalar) -> VortexResult<Self> {
182        let binary = scalar
183            .as_binary_opt()
184            .ok_or_else(|| vortex_err!("Cannot extract buffer from non-buffer scalar"))?;
185
186        binary
187            .value()
188            .ok_or_else(|| vortex_err!("Cannot extract present value from null scalar"))
189    }
190}
191
192impl<'a> TryFrom<&'a Scalar> for Option<ByteBuffer> {
193    type Error = VortexError;
194
195    fn try_from(scalar: &'a Scalar) -> VortexResult<Self> {
196        Ok(scalar
197            .as_binary_opt()
198            .ok_or_else(|| vortex_err!("Cannot extract buffer from non-buffer scalar"))?
199            .value())
200    }
201}
202
203impl TryFrom<Scalar> for ByteBuffer {
204    type Error = VortexError;
205
206    fn try_from(scalar: Scalar) -> VortexResult<Self> {
207        Self::try_from(&scalar)
208    }
209}
210
211impl TryFrom<Scalar> for Option<ByteBuffer> {
212    type Error = VortexError;
213
214    fn try_from(scalar: Scalar) -> VortexResult<Self> {
215        Self::try_from(&scalar)
216    }
217}
218
219impl From<&[u8]> for Scalar {
220    fn from(value: &[u8]) -> Self {
221        Scalar::from(ByteBuffer::from(value.to_vec()))
222    }
223}
224
225impl From<ByteBuffer> for Scalar {
226    fn from(value: ByteBuffer) -> Self {
227        Self {
228            dtype: DType::Binary(Nullability::NonNullable),
229            value: ScalarValue(InnerScalarValue::Buffer(Arc::new(value))),
230        }
231    }
232}
233
234impl From<Arc<ByteBuffer>> for Scalar {
235    fn from(value: Arc<ByteBuffer>) -> Self {
236        Self {
237            dtype: DType::Binary(Nullability::NonNullable),
238            value: ScalarValue(InnerScalarValue::Buffer(value)),
239        }
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use vortex_buffer::buffer;
246    use vortex_dtype::Nullability;
247    use vortex_error::{VortexExpect, VortexUnwrap};
248
249    use crate::{BinaryScalar, Scalar};
250
251    #[test]
252    fn lower_bound() {
253        let binary = Scalar::binary(buffer![0u8, 5, 47, 33, 129], Nullability::NonNullable);
254        let expected = Scalar::binary(buffer![0u8, 5], Nullability::NonNullable);
255        assert_eq!(
256            BinaryScalar::try_from(&binary)
257                .vortex_unwrap()
258                .lower_bound(2),
259            BinaryScalar::try_from(&expected).vortex_unwrap()
260        );
261    }
262
263    #[test]
264    fn upper_bound() {
265        let binary = Scalar::binary(buffer![0u8, 5, 255, 234, 23], Nullability::NonNullable);
266        let expected = Scalar::binary(buffer![0u8, 6, 0], Nullability::NonNullable);
267        assert_eq!(
268            BinaryScalar::try_from(&binary)
269                .vortex_unwrap()
270                .upper_bound(3)
271                .vortex_expect("must have upper bound"),
272            BinaryScalar::try_from(&expected).vortex_unwrap()
273        );
274    }
275
276    #[test]
277    fn upper_bound_overflow() {
278        let binary = Scalar::binary(buffer![255u8, 255, 255], Nullability::NonNullable);
279        assert!(
280            BinaryScalar::try_from(&binary)
281                .vortex_unwrap()
282                .upper_bound(2)
283                .is_none()
284        );
285    }
286}
vortex_scalar/binary.rs

vortex_scalar/
binary.rs
