vortex_array/scalar/
proto.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Protobuf serialization and deserialization for scalars.
5
6use num_traits::ToBytes;
7use num_traits::ToPrimitive;
8use prost::Message;
9use vortex_buffer::BufferString;
10use vortex_buffer::ByteBuffer;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_ensure;
15use vortex_error::vortex_ensure_eq;
16use vortex_error::vortex_err;
17use vortex_session::VortexSession;
18
19use crate::dtype::DType;
20use crate::dtype::PType;
21use crate::dtype::half::f16;
22use crate::dtype::i256;
23use crate::proto::scalar as pb;
24use crate::proto::scalar::ListValue;
25use crate::proto::scalar::UnionValue as PbUnionValue;
26use crate::proto::scalar::scalar_value::Kind;
27use crate::scalar::DecimalValue;
28use crate::scalar::PValue;
29use crate::scalar::Scalar;
30use crate::scalar::ScalarValue;
31use crate::scalar::UnionValue;
32
33////////////////////////////////////////////////////////////////////////////////////////////////////
34// Serialize INTO proto.
35////////////////////////////////////////////////////////////////////////////////////////////////////
36
37impl From<&Scalar> for pb::Scalar {
38    fn from(value: &Scalar) -> Self {
39        pb::Scalar {
40            dtype: Some(
41                (value.dtype())
42                    .try_into()
43                    .vortex_expect("Failed to convert DType to proto"),
44            ),
45            value: Some(Box::new(ScalarValue::to_proto(value.value()))),
46        }
47    }
48}
49
50impl ScalarValue {
51    /// Ideally, we would not have this function and instead implement this `From` implementation:
52    ///
53    /// ```ignore
54    /// impl From<Option<&ScalarValue>> for pb::ScalarValue { ... }
55    /// ```
56    ///
57    /// However, we are not allowed to do this because of the Orphan rule (`Option` and
58    /// `pb::ScalarValue` are not types defined in this crate). So we must make this a method on
59    /// `vortex_array::scalar::ScalarValue` directly.
60    pub fn to_proto(this: Option<&Self>) -> pb::ScalarValue {
61        match this {
62            None => pb::ScalarValue {
63                kind: Some(Kind::NullValue(0)),
64            },
65            Some(this) => pb::ScalarValue::from(this),
66        }
67    }
68
69    /// Serialize an optional [`ScalarValue`] to protobuf bytes (handles null values).
70    pub fn to_proto_bytes<B: Default + bytes::BufMut>(value: Option<&ScalarValue>) -> B {
71        let proto = Self::to_proto(value);
72        let mut buf = B::default();
73        proto
74            .encode(&mut buf)
75            .vortex_expect("Failed to encode scalar value");
76        buf
77    }
78}
79
80impl From<&ScalarValue> for pb::ScalarValue {
81    fn from(value: &ScalarValue) -> Self {
82        match value {
83            ScalarValue::Bool(v) => pb::ScalarValue {
84                kind: Some(Kind::BoolValue(*v)),
85            },
86            ScalarValue::Primitive(v) => pb::ScalarValue::from(v),
87            ScalarValue::Decimal(v) => {
88                let inner_value = match v {
89                    DecimalValue::I8(v) => v.to_le_bytes().to_vec(),
90                    DecimalValue::I16(v) => v.to_le_bytes().to_vec(),
91                    DecimalValue::I32(v) => v.to_le_bytes().to_vec(),
92                    DecimalValue::I64(v) => v.to_le_bytes().to_vec(),
93                    DecimalValue::I128(v128) => v128.to_le_bytes().to_vec(),
94                    DecimalValue::I256(v256) => v256.to_le_bytes().to_vec(),
95                };
96
97                pb::ScalarValue {
98                    kind: Some(Kind::BytesValue(inner_value)),
99                }
100            }
101            ScalarValue::Utf8(v) => pb::ScalarValue {
102                kind: Some(Kind::StringValue(v.to_string())),
103            },
104            ScalarValue::Binary(v) => pb::ScalarValue {
105                kind: Some(Kind::BytesValue(v.to_vec())),
106            },
107            ScalarValue::Tuple(v) => {
108                let mut values = Vec::with_capacity(v.len());
109                for elem in v.iter() {
110                    values.push(ScalarValue::to_proto(elem.as_ref()));
111                }
112                pb::ScalarValue {
113                    kind: Some(Kind::ListValue(ListValue { values })),
114                }
115            }
116            ScalarValue::Union(v) => pb::ScalarValue {
117                kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
118                    type_id: u32::from(v.type_id()),
119                    value: Some(Box::new(ScalarValue::to_proto(v.child_value()))),
120                }))),
121            },
122            ScalarValue::Variant(v) => pb::ScalarValue {
123                kind: Some(Kind::VariantValue(Box::new(pb::Scalar::from(v.as_ref())))),
124            },
125        }
126    }
127}
128
129impl From<&PValue> for pb::ScalarValue {
130    fn from(value: &PValue) -> Self {
131        match value {
132            PValue::I8(v) => pb::ScalarValue {
133                kind: Some(Kind::Int64Value(*v as i64)),
134            },
135            PValue::I16(v) => pb::ScalarValue {
136                kind: Some(Kind::Int64Value(*v as i64)),
137            },
138            PValue::I32(v) => pb::ScalarValue {
139                kind: Some(Kind::Int64Value(*v as i64)),
140            },
141            PValue::I64(v) => pb::ScalarValue {
142                kind: Some(Kind::Int64Value(*v)),
143            },
144            PValue::U8(v) => pb::ScalarValue {
145                kind: Some(Kind::Uint64Value(*v as u64)),
146            },
147            PValue::U16(v) => pb::ScalarValue {
148                kind: Some(Kind::Uint64Value(*v as u64)),
149            },
150            PValue::U32(v) => pb::ScalarValue {
151                kind: Some(Kind::Uint64Value(*v as u64)),
152            },
153            PValue::U64(v) => pb::ScalarValue {
154                kind: Some(Kind::Uint64Value(*v)),
155            },
156            PValue::F16(v) => pb::ScalarValue {
157                kind: Some(Kind::F16Value(v.to_bits() as u64)),
158            },
159            PValue::F32(v) => pb::ScalarValue {
160                kind: Some(Kind::F32Value(*v)),
161            },
162            PValue::F64(v) => pb::ScalarValue {
163                kind: Some(Kind::F64Value(*v)),
164            },
165        }
166    }
167}
168
169////////////////////////////////////////////////////////////////////////////////////////////////////
170// Serialize FROM proto.
171////////////////////////////////////////////////////////////////////////////////////////////////////
172
173impl Scalar {
174    /// Creates a [`Scalar`] from a [protobuf `ScalarValue`](pb::ScalarValue) representation.
175    ///
176    /// Note that we need to provide a [`DType`] since protobuf serialization only supports 64-bit
177    /// integers, and serializing _into_ protobuf loses that type information.
178    ///
179    /// # Errors
180    ///
181    /// Returns an error if type validation fails.
182    pub fn from_proto_value(
183        value: &pb::ScalarValue,
184        dtype: &DType,
185        session: &VortexSession,
186    ) -> VortexResult<Self> {
187        let scalar_value = ScalarValue::from_proto(value, dtype, session)?;
188
189        Scalar::try_new(dtype.clone(), scalar_value)
190    }
191
192    /// Creates a [`Scalar`] from its [protobuf](pb::Scalar) representation.
193    ///
194    /// # Errors
195    ///
196    /// Returns an error if the protobuf is missing required fields or if type validation fails.
197    pub fn from_proto(value: &pb::Scalar, session: &VortexSession) -> VortexResult<Self> {
198        let dtype = DType::from_proto(
199            value
200                .dtype
201                .as_ref()
202                .ok_or_else(|| vortex_err!(Serde: "Scalar missing dtype"))?,
203            session,
204        )?;
205
206        let pb_scalar_value: &pb::ScalarValue = value
207            .value
208            .as_ref()
209            .ok_or_else(|| vortex_err!(Serde: "Scalar missing value"))?;
210
211        let value: Option<ScalarValue> = ScalarValue::from_proto(pb_scalar_value, &dtype, session)?;
212
213        Scalar::try_new(dtype, value)
214    }
215}
216
217impl ScalarValue {
218    /// Deserialize a [`ScalarValue`] from protobuf bytes.
219    ///
220    /// Note that we need to provide a [`DType`] since protobuf serialization only supports 64-bit
221    /// integers, and serializing _into_ protobuf loses that type information.
222    ///
223    /// # Errors
224    ///
225    /// Returns an error if decoding or type validation fails.
226    pub fn from_proto_bytes(
227        bytes: &[u8],
228        dtype: &DType,
229        session: &VortexSession,
230    ) -> VortexResult<Option<Self>> {
231        let proto = pb::ScalarValue::decode(bytes)?;
232        Self::from_proto(&proto, dtype, session)
233    }
234
235    /// Creates a [`ScalarValue`] from its [protobuf](pb::ScalarValue) representation.
236    ///
237    /// Note that we need to provide a [`DType`] since protobuf serialization only supports 64-bit
238    /// integers, and serializing _into_ protobuf loses that type information.
239    ///
240    /// # Errors
241    ///
242    /// Returns an error if the protobuf value cannot be converted to the given [`DType`].
243    pub fn from_proto(
244        value: &pb::ScalarValue,
245        dtype: &DType,
246        session: &VortexSession,
247    ) -> VortexResult<Option<Self>> {
248        let kind = value
249            .kind
250            .as_ref()
251            .ok_or_else(|| vortex_err!(Serde: "Scalar value missing kind"))?;
252
253        // `DType::Extension` store their serialized values using the storage `DType`.
254        let dtype = match dtype {
255            DType::Extension(ext) => ext.storage_dtype(),
256            _ => dtype,
257        };
258
259        Ok(match kind {
260            Kind::NullValue(_) => None,
261            Kind::BoolValue(v) => Some(bool_from_proto(*v, dtype)?),
262            Kind::Int64Value(v) => Some(int64_from_proto(*v, dtype)?),
263            Kind::Uint64Value(v) => Some(uint64_from_proto(*v, dtype)?),
264            Kind::F16Value(v) => Some(f16_from_proto(*v, dtype)?),
265            Kind::F32Value(v) => Some(f32_from_proto(*v, dtype)?),
266            Kind::F64Value(v) => Some(f64_from_proto(*v, dtype)?),
267            Kind::StringValue(s) => Some(string_from_proto(s, dtype)?),
268            Kind::BytesValue(b) => Some(bytes_from_proto(b, dtype)?),
269            Kind::ListValue(v) => Some(list_from_proto(v, dtype, session)?),
270            Kind::UnionValue(v) => Some(union_from_proto(v, dtype, session)?),
271            Kind::VariantValue(v) => match dtype {
272                DType::Variant(_) => Some(ScalarValue::Variant(Box::new(Scalar::from_proto(
273                    v, session,
274                )?))),
275                _ => vortex_bail!(Serde: "expected non-Variant scalar proto for dtype {dtype}"),
276            },
277        })
278    }
279}
280
281/// Deserialize a [`ScalarValue::Bool`] from a protobuf `BoolValue`.
282fn bool_from_proto(v: bool, dtype: &DType) -> VortexResult<ScalarValue> {
283    vortex_ensure!(
284        dtype.is_boolean(),
285        Serde: "expected Bool dtype for BoolValue, got {dtype}"
286    );
287
288    Ok(ScalarValue::Bool(v))
289}
290
291/// Deserialize a [`ScalarValue::Primitive`] from a protobuf `Int64Value`.
292///
293/// Protobuf consolidates all signed integers into `i64`, so we narrow back to the original
294/// type using the provided [`DType`].
295fn int64_from_proto(v: i64, dtype: &DType) -> VortexResult<ScalarValue> {
296    vortex_ensure!(
297        dtype.is_primitive(),
298        Serde: "expected Primitive dtype for Int64Value, got {dtype}"
299    );
300
301    let pvalue = match dtype.as_ptype() {
302        PType::I8 => v.to_i8().map(PValue::I8),
303        PType::I16 => v.to_i16().map(PValue::I16),
304        PType::I32 => v.to_i32().map(PValue::I32),
305        PType::I64 => Some(PValue::I64(v)),
306        // It was previously possible for unsigned types to get their stats serialised as signed,
307        // so we allow casting back to unsigned for backwards compatibility.
308        PType::U8 => v.to_u8().map(PValue::U8),
309        PType::U16 => v.to_u16().map(PValue::U16),
310        PType::U32 => v.to_u32().map(PValue::U32),
311        PType::U64 => v.to_u64().map(PValue::U64),
312        ftype @ (PType::F16 | PType::F32 | PType::F64) => vortex_bail!(
313            Serde: "expected signed integer ptype for serialized Int64Value, got float {ftype}"
314        ),
315    }
316    .ok_or_else(|| vortex_err!(Serde: "Int64 value {v} out of range for dtype {dtype}"))?;
317
318    Ok(ScalarValue::Primitive(pvalue))
319}
320
321/// Deserialize a [`ScalarValue::Primitive`] from a protobuf `Uint64Value`.
322///
323/// Protobuf consolidates all unsigned integers into `u64`, so we narrow back to the original
324/// type using the provided [`DType`]. Also handles the backwards-compatible case where `f16`
325/// values were serialized as `u64` (via `f16::to_bits() as u64`).
326fn uint64_from_proto(v: u64, dtype: &DType) -> VortexResult<ScalarValue> {
327    vortex_ensure!(
328        dtype.is_primitive(),
329        Serde: "expected Primitive dtype for Uint64Value, got {dtype}"
330    );
331
332    let pvalue = match dtype.as_ptype() {
333        PType::U8 => v.to_u8().map(PValue::U8),
334        PType::U16 => v.to_u16().map(PValue::U16),
335        PType::U32 => v.to_u32().map(PValue::U32),
336        PType::U64 => Some(PValue::U64(v)),
337        // It was previously possible for signed types to get their stats serialised as unsigned,
338        // so we allow casting back to signed for backwards compatibility.
339        PType::I8 => v.to_i8().map(PValue::I8),
340        PType::I16 => v.to_i16().map(PValue::I16),
341        PType::I32 => v.to_i32().map(PValue::I32),
342        PType::I64 => v.to_i64().map(PValue::I64),
343        // f16 values used to be serialized as u64, so we need to be able to read an f16 from a u64.
344        PType::F16 => v.to_u16().map(f16::from_bits).map(PValue::F16),
345        ftype @ (PType::F32 | PType::F64) => vortex_bail!(
346            Serde: "expected unsigned integer ptype for serialized Uint64Value, got {ftype}"
347        ),
348    }
349    .ok_or_else(|| vortex_err!(Serde: "Uint64 value {v} out of range for dtype {dtype}"))?;
350
351    Ok(ScalarValue::Primitive(pvalue))
352}
353
354/// Deserialize a [`ScalarValue::Primitive`] from a protobuf `F16Value`.
355fn f16_from_proto(v: u64, dtype: &DType) -> VortexResult<ScalarValue> {
356    vortex_ensure!(
357        matches!(dtype, DType::Primitive(PType::F16, _)),
358        Serde: "expected F16 dtype for F16Value, got {dtype}"
359    );
360
361    let bits = u16::try_from(v)
362        .map_err(|_| vortex_err!(Serde: "f16 bitwise representation has more than 16 bits: {v}"))?;
363
364    Ok(ScalarValue::Primitive(PValue::F16(f16::from_bits(bits))))
365}
366
367/// Deserialize a [`ScalarValue::Primitive`] from a protobuf `F32Value`.
368fn f32_from_proto(v: f32, dtype: &DType) -> VortexResult<ScalarValue> {
369    vortex_ensure!(
370        matches!(dtype, DType::Primitive(PType::F32, _)),
371        Serde: "expected F32 dtype for F32Value, got {dtype}"
372    );
373
374    Ok(ScalarValue::Primitive(PValue::F32(v)))
375}
376
377/// Deserialize a [`ScalarValue::Primitive`] from a protobuf `F64Value`.
378fn f64_from_proto(v: f64, dtype: &DType) -> VortexResult<ScalarValue> {
379    vortex_ensure!(
380        matches!(dtype, DType::Primitive(PType::F64, _)),
381        Serde: "expected F64 dtype for F64Value, got {dtype}"
382    );
383
384    Ok(ScalarValue::Primitive(PValue::F64(v)))
385}
386
387/// Deserialize a [`ScalarValue::Utf8`] or [`ScalarValue::Binary`] from a protobuf
388/// `StringValue`.
389fn string_from_proto(s: &str, dtype: &DType) -> VortexResult<ScalarValue> {
390    match dtype {
391        DType::Utf8(_) => Ok(ScalarValue::Utf8(BufferString::from(s))),
392        DType::Binary(_) => Ok(ScalarValue::Binary(ByteBuffer::copy_from(s.as_bytes()))),
393        _ => vortex_bail!(
394            Serde: "expected Utf8 or Binary dtype for StringValue, got {dtype}"
395        ),
396    }
397}
398
399/// Deserialize a [`ScalarValue`] from a protobuf bytes and a `DType`.
400///
401/// Handles [`Utf8`](ScalarValue::Utf8), [`Binary`](ScalarValue::Binary), and
402/// [`Decimal`](ScalarValue::Decimal) dtypes.
403fn bytes_from_proto(bytes: &[u8], dtype: &DType) -> VortexResult<ScalarValue> {
404    match dtype {
405        DType::Utf8(_) => Ok(ScalarValue::Utf8(BufferString::try_from(bytes)?)),
406        DType::Binary(_) => Ok(ScalarValue::Binary(ByteBuffer::copy_from(bytes))),
407        // TODO(connor): This is incorrect, we need to verify this matches the inner decimal_dtype.
408        DType::Decimal(..) => Ok(ScalarValue::Decimal(match bytes.len() {
409            1 => DecimalValue::I8(bytes[0] as i8),
410            2 => DecimalValue::I16(i16::from_le_bytes(
411                bytes
412                    .try_into()
413                    .ok()
414                    .vortex_expect("Buffer has invalid number of bytes"),
415            )),
416            4 => DecimalValue::I32(i32::from_le_bytes(
417                bytes
418                    .try_into()
419                    .ok()
420                    .vortex_expect("Buffer has invalid number of bytes"),
421            )),
422            8 => DecimalValue::I64(i64::from_le_bytes(
423                bytes
424                    .try_into()
425                    .ok()
426                    .vortex_expect("Buffer has invalid number of bytes"),
427            )),
428            16 => DecimalValue::I128(i128::from_le_bytes(
429                bytes
430                    .try_into()
431                    .ok()
432                    .vortex_expect("Buffer has invalid number of bytes"),
433            )),
434            32 => DecimalValue::I256(i256::from_le_bytes(
435                bytes
436                    .try_into()
437                    .ok()
438                    .vortex_expect("Buffer has invalid number of bytes"),
439            )),
440            l => vortex_bail!(Serde: "invalid decimal byte length: {l}"),
441        })),
442        _ => vortex_bail!(
443            Serde: "expected Utf8, Binary, or Decimal dtype for BytesValue, got {dtype}"
444        ),
445    }
446}
447
448/// Deserialize a [`ScalarValue::Tuple`] from a protobuf `ListValue`.
449fn list_from_proto(
450    v: &ListValue,
451    dtype: &DType,
452    session: &VortexSession,
453) -> VortexResult<ScalarValue> {
454    let values = match dtype {
455        DType::List(element_dtype, _) | DType::FixedSizeList(element_dtype, ..) => v
456            .values
457            .iter()
458            .map(|elem| ScalarValue::from_proto(elem, element_dtype.as_ref(), session))
459            .collect::<VortexResult<Vec<_>>>()?,
460        DType::Struct(fields, _) => {
461            vortex_ensure_eq!(
462                v.values.len(),
463                fields.nfields(),
464                Serde: "ListValue has the wrong number of struct fields"
465            );
466
467            v.values
468                .iter()
469                .zip(fields.fields())
470                .map(|(value, field_dtype)| ScalarValue::from_proto(value, &field_dtype, session))
471                .collect::<VortexResult<Vec<_>>>()?
472        }
473        DType::Map(map, _) => {
474            let entry_dtype = map.entries_dtype();
475            v.values
476                .iter()
477                .map(|entry| ScalarValue::from_proto(entry, &entry_dtype, session))
478                .collect::<VortexResult<Vec<_>>>()?
479        }
480        _ => vortex_bail!(
481            Serde: "expected a tuple-backed dtype for ListValue, got {dtype}"
482        ),
483    };
484    Ok(ScalarValue::Tuple(values))
485}
486
487/// Deserialize a present union scalar value.
488fn union_from_proto(
489    value: &PbUnionValue,
490    dtype: &DType,
491    session: &VortexSession,
492) -> VortexResult<ScalarValue> {
493    let DType::Union(variants, _) = dtype else {
494        vortex_bail!(Serde: "expected Union dtype for UnionValue, got {dtype}");
495    };
496
497    let type_id = u8::try_from(value.type_id).map_err(
498        |_| vortex_err!(Serde: "union type ID {} is outside the u8 range", value.type_id),
499    )?;
500
501    let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| {
502        vortex_err!(
503            Serde: "union type ID {type_id} is not present in {:?}",
504            variants.type_ids()
505        )
506    })?;
507
508    let child_dtype = variants
509        .variant_by_index(child_index)
510        .ok_or_else(|| vortex_err!(Serde: "union type ID {type_id} resolved out of bounds"))?;
511
512    let child_proto = value
513        .value
514        .as_deref()
515        .ok_or_else(|| vortex_err!(Serde: "UnionValue missing child value"))?;
516
517    let child_value = ScalarValue::from_proto(child_proto, &child_dtype, session)?;
518    Scalar::validate(&child_dtype, child_value.as_ref()).map_err(|error| {
519        vortex_err!(
520            Serde: "union type ID {type_id} has invalid child for dtype {child_dtype}: {error}"
521        )
522    })?;
523
524    Ok(ScalarValue::Union(UnionValue::new(type_id, child_value)))
525}
526
527#[cfg(test)]
528mod tests {
529    use std::f32;
530    use std::f64;
531    use std::sync::Arc;
532
533    use vortex_buffer::BufferString;
534    use vortex_error::VortexError;
535    use vortex_error::vortex_panic;
536    use vortex_session::VortexSession;
537
538    use super::*;
539    use crate::dtype::DType;
540    use crate::dtype::DecimalDType;
541    use crate::dtype::Nullability;
542    use crate::dtype::PType;
543    use crate::dtype::UnionVariants;
544    use crate::dtype::half::f16;
545    use crate::proto::scalar as pb;
546    use crate::scalar::DecimalValue;
547    use crate::scalar::Scalar;
548    use crate::scalar::ScalarValue;
549
550    fn session() -> VortexSession {
551        VortexSession::empty()
552    }
553
554    fn round_trip(scalar: Scalar) {
555        assert_eq!(
556            scalar,
557            Scalar::from_proto(&pb::Scalar::from(&scalar), &session()).unwrap(),
558        );
559    }
560
561    #[test]
562    fn test_null() {
563        round_trip(Scalar::null(DType::Null));
564    }
565
566    #[test]
567    fn test_bool() {
568        round_trip(Scalar::new(
569            DType::Bool(Nullability::Nullable),
570            Some(ScalarValue::Bool(true)),
571        ));
572    }
573
574    #[test]
575    fn test_primitive() {
576        round_trip(Scalar::new(
577            DType::Primitive(PType::I32, Nullability::Nullable),
578            Some(ScalarValue::Primitive(42i32.into())),
579        ));
580    }
581
582    #[test]
583    fn test_buffer() {
584        round_trip(Scalar::new(
585            DType::Binary(Nullability::Nullable),
586            Some(ScalarValue::Binary(vec![1, 2, 3].into())),
587        ));
588    }
589
590    #[test]
591    fn test_buffer_string() {
592        round_trip(Scalar::new(
593            DType::Utf8(Nullability::Nullable),
594            Some(ScalarValue::Utf8(BufferString::from("hello".to_string()))),
595        ));
596    }
597
598    #[test]
599    fn test_list() {
600        round_trip(Scalar::new(
601            DType::List(
602                Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
603                Nullability::Nullable,
604            ),
605            Some(ScalarValue::Tuple(vec![
606                Some(ScalarValue::Primitive(42i32.into())),
607                Some(ScalarValue::Primitive(43i32.into())),
608            ])),
609        ));
610    }
611
612    #[test]
613    fn test_map() {
614        let dtype = DType::map(
615            DType::Primitive(PType::I32, Nullability::NonNullable),
616            DType::Utf8(Nullability::Nullable),
617            true,
618            Nullability::Nullable,
619        )
620        .unwrap();
621        round_trip(
622            Scalar::try_map(
623                dtype.clone(),
624                [
625                    (
626                        Scalar::primitive(1i32, Nullability::NonNullable),
627                        Scalar::utf8("one", Nullability::Nullable),
628                    ),
629                    (
630                        Scalar::primitive(2i32, Nullability::NonNullable),
631                        Scalar::null(DType::Utf8(Nullability::Nullable)),
632                    ),
633                ],
634            )
635            .unwrap(),
636        );
637        round_trip(Scalar::null(dtype));
638    }
639
640    #[test]
641    fn test_f16() {
642        round_trip(Scalar::primitive(
643            f16::from_f32(0.42),
644            Nullability::Nullable,
645        ));
646    }
647
648    #[test]
649    fn test_i8() {
650        round_trip(Scalar::new(
651            DType::Primitive(PType::I8, Nullability::Nullable),
652            Some(ScalarValue::Primitive(i8::MIN.into())),
653        ));
654
655        round_trip(Scalar::new(
656            DType::Primitive(PType::I8, Nullability::Nullable),
657            Some(ScalarValue::Primitive(0i8.into())),
658        ));
659
660        round_trip(Scalar::new(
661            DType::Primitive(PType::I8, Nullability::Nullable),
662            Some(ScalarValue::Primitive(i8::MAX.into())),
663        ));
664    }
665
666    #[test]
667    fn test_decimal_i32_roundtrip() {
668        // A typical decimal with moderate precision and scale.
669        round_trip(Scalar::decimal(
670            DecimalValue::I32(123_456),
671            DecimalDType::new(10, 2),
672            Nullability::NonNullable,
673        ));
674    }
675
676    #[test]
677    fn test_decimal_i128_roundtrip() {
678        // A large decimal value that requires i128 storage.
679        round_trip(Scalar::decimal(
680            DecimalValue::I128(99_999_999_999_999_999_999),
681            DecimalDType::new(38, 6),
682            Nullability::Nullable,
683        ));
684    }
685
686    #[test]
687    fn test_decimal_null_roundtrip() {
688        round_trip(Scalar::null(DType::Decimal(
689            DecimalDType::new(10, 2),
690            Nullability::Nullable,
691        )));
692    }
693
694    #[test]
695    fn test_scalar_value_serde_roundtrip_binary() {
696        round_trip(Scalar::binary(
697            ByteBuffer::copy_from(b"hello"),
698            Nullability::NonNullable,
699        ));
700    }
701
702    #[test]
703    fn test_scalar_value_serde_roundtrip_utf8() {
704        round_trip(Scalar::utf8("hello", Nullability::NonNullable));
705    }
706
707    #[test]
708    fn test_variant_scalar_roundtrip() {
709        let nums = Scalar::list(
710            Arc::new(DType::Variant(Nullability::NonNullable)),
711            vec![
712                Scalar::variant(Scalar::primitive(-7_i16, Nullability::NonNullable)),
713                Scalar::variant(Scalar::primitive(42_u32, Nullability::NonNullable)),
714                Scalar::variant(Scalar::decimal(
715                    DecimalValue::I128(123_456_789),
716                    DecimalDType::new(18, 0),
717                    Nullability::NonNullable,
718                )),
719            ],
720            Nullability::NonNullable,
721        );
722
723        let nested = Scalar::list(
724            Arc::new(DType::Variant(Nullability::NonNullable)),
725            vec![
726                Scalar::variant(Scalar::from(true)),
727                Scalar::variant(nums),
728                Scalar::variant(Scalar::binary(
729                    ByteBuffer::copy_from(b"abc"),
730                    Nullability::NonNullable,
731                )),
732                Scalar::variant(Scalar::null(DType::Null)),
733            ],
734            Nullability::NonNullable,
735        );
736
737        round_trip(Scalar::variant(nested));
738    }
739
740    #[test]
741    fn test_variant_scalar_proto_preserves_scalar_null_vs_variant_null() {
742        let scalar_null = Scalar::null(DType::Variant(Nullability::Nullable));
743        let variant_null = Scalar::variant(Scalar::null(DType::Null));
744
745        let scalar_null_pb = pb::Scalar::from(&scalar_null);
746        let variant_null_pb = pb::Scalar::from(&variant_null);
747
748        assert_ne!(scalar_null_pb, variant_null_pb);
749        assert_eq!(
750            Scalar::from_proto(&scalar_null_pb, &session()).unwrap(),
751            scalar_null,
752        );
753        assert_eq!(
754            Scalar::from_proto(&variant_null_pb, &session()).unwrap(),
755            variant_null,
756        );
757    }
758
759    #[test]
760    fn test_union_scalar_roundtrip() -> VortexResult<()> {
761        let variants = UnionVariants::try_new(
762            ["int", "string"].into(),
763            vec![
764                DType::Primitive(PType::I32, Nullability::Nullable),
765                DType::Utf8(Nullability::NonNullable),
766            ],
767            vec![5, 9],
768        )?;
769
770        round_trip(Scalar::union(
771            variants.clone(),
772            5,
773            Scalar::primitive(42_i32, Nullability::Nullable),
774            Nullability::Nullable,
775        )?);
776
777        let inner_null = Scalar::union(
778            variants.clone(),
779            5,
780            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
781            Nullability::Nullable,
782        )?;
783        let inner_null_proto = pb::Scalar::from(&inner_null);
784
785        assert!(matches!(
786            inner_null_proto
787                .value
788                .as_deref()
789                .and_then(|value| value.kind.as_ref()),
790            Some(Kind::UnionValue(union_value))
791                if matches!(
792                    union_value
793                        .value
794                        .as_deref()
795                        .and_then(|value| value.kind.as_ref()),
796                    Some(Kind::NullValue(_))
797                )
798        ));
799        let outer_null = Scalar::null(DType::Union(variants, Nullability::Nullable));
800        let outer_null_proto = pb::Scalar::from(&outer_null);
801        assert!(matches!(
802            outer_null_proto
803                .value
804                .as_deref()
805                .and_then(|value| value.kind.as_ref()),
806            Some(Kind::NullValue(_))
807        ));
808
809        assert_ne!(inner_null_proto, outer_null_proto);
810        round_trip(inner_null);
811        round_trip(outer_null);
812
813        let struct_dtype = DType::struct_(
814            [
815                (
816                    "number",
817                    DType::Primitive(PType::I32, Nullability::NonNullable),
818                ),
819                ("label", DType::Utf8(Nullability::Nullable)),
820            ],
821            Nullability::NonNullable,
822        );
823        let struct_scalar = Scalar::struct_(
824            struct_dtype.clone(),
825            [
826                Scalar::primitive(42_i32, Nullability::NonNullable),
827                Scalar::utf8("answer", Nullability::Nullable),
828            ],
829        );
830        let struct_variants =
831            UnionVariants::try_new(["record"].into(), vec![struct_dtype], vec![13])?;
832        round_trip(Scalar::union(
833            struct_variants,
834            13,
835            struct_scalar,
836            Nullability::NonNullable,
837        )?);
838
839        Ok(())
840    }
841
842    #[test]
843    fn test_union_proto_rejects_malformed_values() -> VortexResult<()> {
844        let variants = UnionVariants::try_new(
845            ["int"].into(),
846            vec![DType::Primitive(PType::I32, Nullability::Nullable)],
847            vec![5],
848        )?;
849        let dtype = DType::Union(variants, Nullability::NonNullable);
850
851        let unknown_tag = pb::ScalarValue {
852            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
853                type_id: 7,
854                value: Some(Box::new(ScalarValue::to_proto(
855                    Scalar::primitive(42_i32, Nullability::Nullable).value(),
856                ))),
857            }))),
858        };
859
860        assert!(ScalarValue::from_proto(&unknown_tag, &dtype, &session()).is_err());
861
862        let missing_child = pb::ScalarValue {
863            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
864                type_id: 5,
865                value: None,
866            }))),
867        };
868
869        assert!(matches!(
870            ScalarValue::from_proto(&missing_child, &dtype, &session()),
871            Err(VortexError::Serde(..))
872        ));
873
874        let wrong_child_value = pb::ScalarValue {
875            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
876                type_id: 5,
877                value: Some(Box::new(ScalarValue::to_proto(
878                    Scalar::utf8("wrong", Nullability::NonNullable).value(),
879                ))),
880            }))),
881        };
882
883        assert!(matches!(
884            ScalarValue::from_proto(&wrong_child_value, &dtype, &session()),
885            Err(VortexError::Serde(..))
886        ));
887
888        Ok(())
889    }
890
891    #[test]
892    fn test_backcompat_f16_serialized_as_u64() {
893        // Backwards compatibility test for the legacy f16 serialization format.
894        //
895        // Previously, f16 ScalarValues were serialized as `Uint64Value(v.to_bits() as u64)` because
896        // the proto schema only had 64-bit integer types, and f16's underlying representation is
897        // u16 which got widened to u64.
898        //
899        // The current implementation uses a dedicated `F16Value` proto field, but we must still be
900        // able to deserialize the old format. This test verifies that:
901        //
902        // 1. A `Uint64Value` containing f16 bits can be read as a U64 primitive (the raw bits).
903        // 2. When wrapped in a Scalar with F16 dtype, the value is correctly interpreted as f16.
904        //
905        // This ensures data written with the old serialization format remains readable.
906
907        // Simulate the old serialization: f16(0.42) stored as Uint64Value with its bit pattern.
908        let f16_value = f16::from_f32(0.42);
909        let f16_bits_as_u64 = f16_value.to_bits() as u64; // 14008
910
911        let pb_scalar_value = pb::ScalarValue {
912            kind: Some(Kind::Uint64Value(f16_bits_as_u64)),
913        };
914
915        // Step 1: Verify the normal U64 scalar.
916        let scalar_value = ScalarValue::from_proto(
917            &pb_scalar_value,
918            &DType::Primitive(PType::U64, Nullability::NonNullable),
919            &session(),
920        )
921        .unwrap();
922        assert_eq!(
923            scalar_value.as_ref().map(|v| v.as_primitive()),
924            Some(&PValue::U64(14008u64)),
925        );
926
927        // Step 2: Verify that when we use F16 dtype, the Uint64Value is correctly interpreted.
928        let scalar_value_f16 = ScalarValue::from_proto(
929            &pb_scalar_value,
930            &DType::Primitive(PType::F16, Nullability::Nullable),
931            &session(),
932        )
933        .unwrap();
934
935        let scalar = Scalar::new(
936            DType::Primitive(PType::F16, Nullability::Nullable),
937            scalar_value_f16,
938        );
939
940        assert_eq!(
941            scalar.as_primitive().pvalue().unwrap(),
942            PValue::F16(f16::from_f32(0.42)),
943            "Uint64Value should be correctly interpreted as f16 when dtype is F16"
944        );
945    }
946
947    #[test]
948    fn test_scalar_value_direct_roundtrip_f16() {
949        // Test that ScalarValue with f16 roundtrips correctly without going through Scalar.
950        let f16_values = vec![
951            f16::from_f32(0.0),
952            f16::from_f32(1.0),
953            f16::from_f32(-1.0),
954            f16::from_f32(0.42),
955            f16::from_f32(5.722046e-6),
956            f16::from_f32(f32::consts::PI),
957            f16::INFINITY,
958            f16::NEG_INFINITY,
959            f16::NAN,
960        ];
961
962        for f16_val in f16_values {
963            let scalar_value = ScalarValue::Primitive(PValue::F16(f16_val));
964            let pb_value = ScalarValue::to_proto(Some(&scalar_value));
965            let read_back = ScalarValue::from_proto(
966                &pb_value,
967                &DType::Primitive(PType::F16, Nullability::NonNullable),
968                &session(),
969            )
970            .unwrap();
971
972            match (&scalar_value, read_back.as_ref()) {
973                (
974                    ScalarValue::Primitive(PValue::F16(original)),
975                    Some(ScalarValue::Primitive(PValue::F16(roundtripped))),
976                ) => {
977                    if original.is_nan() && roundtripped.is_nan() {
978                        // NaN values are equal for our purposes.
979                        continue;
980                    }
981                    assert_eq!(
982                        original, roundtripped,
983                        "F16 value {original:?} did not roundtrip correctly"
984                    );
985                }
986                _ => {
987                    vortex_panic!(
988                        "Expected f16 primitive values, got {scalar_value:?} and {read_back:?}"
989                    )
990                }
991            }
992        }
993    }
994
995    #[test]
996    fn test_scalar_value_direct_roundtrip_preserves_values() {
997        // Test that ScalarValue roundtripping preserves values (but not necessarily exact types).
998        // Note: Proto encoding consolidates integer types (u8/u16/u32 → u64, i8/i16/i32 → i64).
999
1000        // Test cases that should roundtrip exactly.
1001        let exact_roundtrip_cases: Vec<(&str, Option<ScalarValue>, DType)> = vec![
1002            ("null", None, DType::Null),
1003            (
1004                "bool_true",
1005                Some(ScalarValue::Bool(true)),
1006                DType::Bool(Nullability::Nullable),
1007            ),
1008            (
1009                "bool_false",
1010                Some(ScalarValue::Bool(false)),
1011                DType::Bool(Nullability::Nullable),
1012            ),
1013            (
1014                "u64",
1015                Some(ScalarValue::Primitive(PValue::U64(18446744073709551615))),
1016                DType::Primitive(PType::U64, Nullability::Nullable),
1017            ),
1018            (
1019                "i64",
1020                Some(ScalarValue::Primitive(PValue::I64(-9223372036854775808))),
1021                DType::Primitive(PType::I64, Nullability::Nullable),
1022            ),
1023            (
1024                "f32",
1025                Some(ScalarValue::Primitive(PValue::F32(f32::consts::E))),
1026                DType::Primitive(PType::F32, Nullability::Nullable),
1027            ),
1028            (
1029                "f64",
1030                Some(ScalarValue::Primitive(PValue::F64(f64::consts::PI))),
1031                DType::Primitive(PType::F64, Nullability::Nullable),
1032            ),
1033            (
1034                "string",
1035                Some(ScalarValue::Utf8(BufferString::from("test"))),
1036                DType::Utf8(Nullability::Nullable),
1037            ),
1038            (
1039                "bytes",
1040                Some(ScalarValue::Binary(vec![1, 2, 3, 4, 5].into())),
1041                DType::Binary(Nullability::Nullable),
1042            ),
1043        ];
1044
1045        for (name, value, dtype) in exact_roundtrip_cases {
1046            let pb_value = ScalarValue::to_proto(value.as_ref());
1047            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1048
1049            let original_debug = format!("{value:?}");
1050            let roundtrip_debug = format!("{read_back:?}");
1051            assert_eq!(
1052                original_debug, roundtrip_debug,
1053                "ScalarValue {name} did not roundtrip exactly"
1054            );
1055        }
1056
1057        // Test cases where type changes but value is preserved.
1058        // Unsigned integers consolidate to U64.
1059        let unsigned_cases = vec![
1060            (
1061                "u8",
1062                ScalarValue::Primitive(PValue::U8(255)),
1063                DType::Primitive(PType::U8, Nullability::Nullable),
1064                255u64,
1065            ),
1066            (
1067                "u16",
1068                ScalarValue::Primitive(PValue::U16(65535)),
1069                DType::Primitive(PType::U16, Nullability::Nullable),
1070                65535u64,
1071            ),
1072            (
1073                "u32",
1074                ScalarValue::Primitive(PValue::U32(4294967295)),
1075                DType::Primitive(PType::U32, Nullability::Nullable),
1076                4294967295u64,
1077            ),
1078        ];
1079
1080        for (name, value, dtype, expected) in unsigned_cases {
1081            let pb_value = ScalarValue::to_proto(Some(&value));
1082            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1083
1084            match read_back.as_ref() {
1085                Some(ScalarValue::Primitive(pv)) => {
1086                    let v = match pv {
1087                        PValue::U8(v) => *v as u64,
1088                        PValue::U16(v) => *v as u64,
1089                        PValue::U32(v) => *v as u64,
1090                        PValue::U64(v) => *v,
1091                        _ => vortex_panic!("Unexpected primitive type for {name}: {pv:?}"),
1092                    };
1093                    assert_eq!(
1094                        v, expected,
1095                        "ScalarValue {name} value not preserved: expected {expected}, got {v}"
1096                    );
1097                }
1098                _ => vortex_panic!("Unexpected type after roundtrip for {name}: {read_back:?}"),
1099            }
1100        }
1101
1102        // Signed integers consolidate to I64.
1103        let signed_cases = vec![
1104            (
1105                "i8",
1106                ScalarValue::Primitive(PValue::I8(-128)),
1107                DType::Primitive(PType::I8, Nullability::Nullable),
1108                -128i64,
1109            ),
1110            (
1111                "i16",
1112                ScalarValue::Primitive(PValue::I16(-32768)),
1113                DType::Primitive(PType::I16, Nullability::Nullable),
1114                -32768i64,
1115            ),
1116            (
1117                "i32",
1118                ScalarValue::Primitive(PValue::I32(-2147483648)),
1119                DType::Primitive(PType::I32, Nullability::Nullable),
1120                -2147483648i64,
1121            ),
1122        ];
1123
1124        for (name, value, dtype, expected) in signed_cases {
1125            let pb_value = ScalarValue::to_proto(Some(&value));
1126            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1127
1128            match read_back.as_ref() {
1129                Some(ScalarValue::Primitive(pv)) => {
1130                    let v = match pv {
1131                        PValue::I8(v) => *v as i64,
1132                        PValue::I16(v) => *v as i64,
1133                        PValue::I32(v) => *v as i64,
1134                        PValue::I64(v) => *v,
1135                        _ => vortex_panic!("Unexpected primitive type for {name}: {pv:?}"),
1136                    };
1137                    assert_eq!(
1138                        v, expected,
1139                        "ScalarValue {name} value not preserved: expected {expected}, got {v}"
1140                    );
1141                }
1142                _ => vortex_panic!("Unexpected type after roundtrip for {name}: {read_back:?}"),
1143            }
1144        }
1145    }
1146
1147    // Backwards compatibility: signed integer stats could previously be serialized as unsigned.
1148    // Therefore, we allow casting between signed and unsigned integers of the same bit width.
1149    #[test]
1150    fn test_backcompat_signed_integer_deserialized_as_unsigned() {
1151        let v = ScalarValue::Primitive(PValue::I64(0));
1152        assert_eq!(
1153            Scalar::from_proto_value(
1154                &pb::ScalarValue::from(&v),
1155                &DType::Primitive(PType::U64, Nullability::Nullable),
1156                &session()
1157            )
1158            .unwrap(),
1159            Scalar::primitive(0u64, Nullability::Nullable)
1160        );
1161    }
1162
1163    // Backwards compatibility: unsigned integer stats could previously be serialized as signed.
1164    // Therefore, we allow casting between signed and unsigned integers of the same bit width.
1165    #[test]
1166    fn test_backcompat_unsigned_integer_deserialized_as_signed() {
1167        let v = ScalarValue::Primitive(PValue::U64(0));
1168        assert_eq!(
1169            Scalar::from_proto_value(
1170                &pb::ScalarValue::from(&v),
1171                &DType::Primitive(PType::I64, Nullability::Nullable),
1172                &session()
1173            )
1174            .unwrap(),
1175            Scalar::primitive(0i64, Nullability::Nullable)
1176        );
1177    }
1178}
vortex_array/scalar/proto.rs

vortex_array/scalar/
proto.rs
