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_proto::scalar as pb;
18use vortex_proto::scalar::ListValue;
19use vortex_proto::scalar::UnionValue as PbUnionValue;
20use vortex_proto::scalar::scalar_value::Kind;
21use vortex_session::VortexSession;
22
23use crate::dtype::DType;
24use crate::dtype::PType;
25use crate::dtype::half::f16;
26use crate::dtype::i256;
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(), fields.nfields(),
463                Serde: "expected {} struct fields in ListValue, got {}",
464                fields.nfields(),
465                v.values.len()
466            );
467
468            v.values
469                .iter()
470                .zip(fields.fields())
471                .map(|(value, field_dtype)| ScalarValue::from_proto(value, &field_dtype, session))
472                .collect::<VortexResult<Vec<_>>>()?
473        }
474        DType::Map(map, _) => {
475            let entry_dtype = map.entries_dtype();
476            v.values
477                .iter()
478                .map(|entry| ScalarValue::from_proto(entry, &entry_dtype, session))
479                .collect::<VortexResult<Vec<_>>>()?
480        }
481        _ => vortex_bail!(
482            Serde: "expected a tuple-backed dtype for ListValue, got {dtype}"
483        ),
484    };
485    Ok(ScalarValue::Tuple(values))
486}
487
488/// Deserialize a present union scalar value.
489fn union_from_proto(
490    value: &PbUnionValue,
491    dtype: &DType,
492    session: &VortexSession,
493) -> VortexResult<ScalarValue> {
494    let DType::Union(variants, _) = dtype else {
495        vortex_bail!(Serde: "expected Union dtype for UnionValue, got {dtype}");
496    };
497
498    let type_id = u8::try_from(value.type_id).map_err(
499        |_| vortex_err!(Serde: "union type ID {} is outside the u8 range", value.type_id),
500    )?;
501
502    let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| {
503        vortex_err!(
504            Serde: "union type ID {type_id} is not present in {:?}",
505            variants.type_ids()
506        )
507    })?;
508
509    let child_dtype = variants
510        .variant_by_index(child_index)
511        .ok_or_else(|| vortex_err!(Serde: "union type ID {type_id} resolved out of bounds"))?;
512
513    let child_proto = value
514        .value
515        .as_deref()
516        .ok_or_else(|| vortex_err!(Serde: "UnionValue missing child value"))?;
517
518    let child_value = ScalarValue::from_proto(child_proto, &child_dtype, session)?;
519    Scalar::validate(&child_dtype, child_value.as_ref()).map_err(|error| {
520        vortex_err!(
521            Serde: "union type ID {type_id} has invalid child for dtype {child_dtype}: {error}"
522        )
523    })?;
524
525    Ok(ScalarValue::Union(UnionValue::new(type_id, child_value)))
526}
527
528#[cfg(test)]
529mod tests {
530    use std::f32;
531    use std::f64;
532    use std::sync::Arc;
533
534    use vortex_buffer::BufferString;
535    use vortex_error::VortexError;
536    use vortex_error::vortex_panic;
537    use vortex_proto::scalar as pb;
538    use vortex_session::VortexSession;
539
540    use super::*;
541    use crate::dtype::DType;
542    use crate::dtype::DecimalDType;
543    use crate::dtype::Nullability;
544    use crate::dtype::PType;
545    use crate::dtype::UnionVariants;
546    use crate::dtype::half::f16;
547    use crate::scalar::DecimalValue;
548    use crate::scalar::Scalar;
549    use crate::scalar::ScalarValue;
550
551    fn session() -> VortexSession {
552        VortexSession::empty()
553    }
554
555    fn round_trip(scalar: Scalar) {
556        assert_eq!(
557            scalar,
558            Scalar::from_proto(&pb::Scalar::from(&scalar), &session()).unwrap(),
559        );
560    }
561
562    #[test]
563    fn test_null() {
564        round_trip(Scalar::null(DType::Null));
565    }
566
567    #[test]
568    fn test_bool() {
569        round_trip(Scalar::new(
570            DType::Bool(Nullability::Nullable),
571            Some(ScalarValue::Bool(true)),
572        ));
573    }
574
575    #[test]
576    fn test_primitive() {
577        round_trip(Scalar::new(
578            DType::Primitive(PType::I32, Nullability::Nullable),
579            Some(ScalarValue::Primitive(42i32.into())),
580        ));
581    }
582
583    #[test]
584    fn test_buffer() {
585        round_trip(Scalar::new(
586            DType::Binary(Nullability::Nullable),
587            Some(ScalarValue::Binary(vec![1, 2, 3].into())),
588        ));
589    }
590
591    #[test]
592    fn test_buffer_string() {
593        round_trip(Scalar::new(
594            DType::Utf8(Nullability::Nullable),
595            Some(ScalarValue::Utf8(BufferString::from("hello".to_string()))),
596        ));
597    }
598
599    #[test]
600    fn test_list() {
601        round_trip(Scalar::new(
602            DType::List(
603                Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
604                Nullability::Nullable,
605            ),
606            Some(ScalarValue::Tuple(vec![
607                Some(ScalarValue::Primitive(42i32.into())),
608                Some(ScalarValue::Primitive(43i32.into())),
609            ])),
610        ));
611    }
612
613    #[test]
614    fn test_map() {
615        let dtype = DType::map(
616            DType::Primitive(PType::I32, Nullability::NonNullable),
617            DType::Utf8(Nullability::Nullable),
618            true,
619            Nullability::Nullable,
620        )
621        .unwrap();
622        round_trip(
623            Scalar::try_map(
624                dtype.clone(),
625                [
626                    (
627                        Scalar::primitive(1i32, Nullability::NonNullable),
628                        Scalar::utf8("one", Nullability::Nullable),
629                    ),
630                    (
631                        Scalar::primitive(2i32, Nullability::NonNullable),
632                        Scalar::null(DType::Utf8(Nullability::Nullable)),
633                    ),
634                ],
635            )
636            .unwrap(),
637        );
638        round_trip(Scalar::null(dtype));
639    }
640
641    #[test]
642    fn test_f16() {
643        round_trip(Scalar::primitive(
644            f16::from_f32(0.42),
645            Nullability::Nullable,
646        ));
647    }
648
649    #[test]
650    fn test_i8() {
651        round_trip(Scalar::new(
652            DType::Primitive(PType::I8, Nullability::Nullable),
653            Some(ScalarValue::Primitive(i8::MIN.into())),
654        ));
655
656        round_trip(Scalar::new(
657            DType::Primitive(PType::I8, Nullability::Nullable),
658            Some(ScalarValue::Primitive(0i8.into())),
659        ));
660
661        round_trip(Scalar::new(
662            DType::Primitive(PType::I8, Nullability::Nullable),
663            Some(ScalarValue::Primitive(i8::MAX.into())),
664        ));
665    }
666
667    #[test]
668    fn test_decimal_i32_roundtrip() {
669        // A typical decimal with moderate precision and scale.
670        round_trip(Scalar::decimal(
671            DecimalValue::I32(123_456),
672            DecimalDType::new(10, 2),
673            Nullability::NonNullable,
674        ));
675    }
676
677    #[test]
678    fn test_decimal_i128_roundtrip() {
679        // A large decimal value that requires i128 storage.
680        round_trip(Scalar::decimal(
681            DecimalValue::I128(99_999_999_999_999_999_999),
682            DecimalDType::new(38, 6),
683            Nullability::Nullable,
684        ));
685    }
686
687    #[test]
688    fn test_decimal_null_roundtrip() {
689        round_trip(Scalar::null(DType::Decimal(
690            DecimalDType::new(10, 2),
691            Nullability::Nullable,
692        )));
693    }
694
695    #[test]
696    fn test_scalar_value_serde_roundtrip_binary() {
697        round_trip(Scalar::binary(
698            ByteBuffer::copy_from(b"hello"),
699            Nullability::NonNullable,
700        ));
701    }
702
703    #[test]
704    fn test_scalar_value_serde_roundtrip_utf8() {
705        round_trip(Scalar::utf8("hello", Nullability::NonNullable));
706    }
707
708    #[test]
709    fn test_variant_scalar_roundtrip() {
710        let nums = Scalar::list(
711            Arc::new(DType::Variant(Nullability::NonNullable)),
712            vec![
713                Scalar::variant(Scalar::primitive(-7_i16, Nullability::NonNullable)),
714                Scalar::variant(Scalar::primitive(42_u32, Nullability::NonNullable)),
715                Scalar::variant(Scalar::decimal(
716                    DecimalValue::I128(123_456_789),
717                    DecimalDType::new(18, 0),
718                    Nullability::NonNullable,
719                )),
720            ],
721            Nullability::NonNullable,
722        );
723
724        let nested = Scalar::list(
725            Arc::new(DType::Variant(Nullability::NonNullable)),
726            vec![
727                Scalar::variant(Scalar::from(true)),
728                Scalar::variant(nums),
729                Scalar::variant(Scalar::binary(
730                    ByteBuffer::copy_from(b"abc"),
731                    Nullability::NonNullable,
732                )),
733                Scalar::variant(Scalar::null(DType::Null)),
734            ],
735            Nullability::NonNullable,
736        );
737
738        round_trip(Scalar::variant(nested));
739    }
740
741    #[test]
742    fn test_variant_scalar_proto_preserves_scalar_null_vs_variant_null() {
743        let scalar_null = Scalar::null(DType::Variant(Nullability::Nullable));
744        let variant_null = Scalar::variant(Scalar::null(DType::Null));
745
746        let scalar_null_pb = pb::Scalar::from(&scalar_null);
747        let variant_null_pb = pb::Scalar::from(&variant_null);
748
749        assert_ne!(scalar_null_pb, variant_null_pb);
750        assert_eq!(
751            Scalar::from_proto(&scalar_null_pb, &session()).unwrap(),
752            scalar_null,
753        );
754        assert_eq!(
755            Scalar::from_proto(&variant_null_pb, &session()).unwrap(),
756            variant_null,
757        );
758    }
759
760    #[test]
761    fn test_union_scalar_roundtrip() -> VortexResult<()> {
762        let variants = UnionVariants::try_new(
763            ["int", "string"].into(),
764            vec![
765                DType::Primitive(PType::I32, Nullability::Nullable),
766                DType::Utf8(Nullability::NonNullable),
767            ],
768            vec![5, 9],
769        )?;
770
771        round_trip(Scalar::union(
772            variants.clone(),
773            5,
774            Scalar::primitive(42_i32, Nullability::Nullable),
775            Nullability::Nullable,
776        )?);
777
778        let inner_null = Scalar::union(
779            variants.clone(),
780            5,
781            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
782            Nullability::Nullable,
783        )?;
784        let inner_null_proto = pb::Scalar::from(&inner_null);
785
786        assert!(matches!(
787            inner_null_proto
788                .value
789                .as_deref()
790                .and_then(|value| value.kind.as_ref()),
791            Some(Kind::UnionValue(union_value))
792                if matches!(
793                    union_value
794                        .value
795                        .as_deref()
796                        .and_then(|value| value.kind.as_ref()),
797                    Some(Kind::NullValue(_))
798                )
799        ));
800        let outer_null = Scalar::null(DType::Union(variants, Nullability::Nullable));
801        let outer_null_proto = pb::Scalar::from(&outer_null);
802        assert!(matches!(
803            outer_null_proto
804                .value
805                .as_deref()
806                .and_then(|value| value.kind.as_ref()),
807            Some(Kind::NullValue(_))
808        ));
809
810        assert_ne!(inner_null_proto, outer_null_proto);
811        round_trip(inner_null);
812        round_trip(outer_null);
813
814        let struct_dtype = DType::struct_(
815            [
816                (
817                    "number",
818                    DType::Primitive(PType::I32, Nullability::NonNullable),
819                ),
820                ("label", DType::Utf8(Nullability::Nullable)),
821            ],
822            Nullability::NonNullable,
823        );
824        let struct_scalar = Scalar::struct_(
825            struct_dtype.clone(),
826            [
827                Scalar::primitive(42_i32, Nullability::NonNullable),
828                Scalar::utf8("answer", Nullability::Nullable),
829            ],
830        );
831        let struct_variants =
832            UnionVariants::try_new(["record"].into(), vec![struct_dtype], vec![13])?;
833        round_trip(Scalar::union(
834            struct_variants,
835            13,
836            struct_scalar,
837            Nullability::NonNullable,
838        )?);
839
840        Ok(())
841    }
842
843    #[test]
844    fn test_union_proto_rejects_malformed_values() -> VortexResult<()> {
845        let variants = UnionVariants::try_new(
846            ["int"].into(),
847            vec![DType::Primitive(PType::I32, Nullability::Nullable)],
848            vec![5],
849        )?;
850        let dtype = DType::Union(variants, Nullability::NonNullable);
851
852        let unknown_tag = pb::ScalarValue {
853            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
854                type_id: 7,
855                value: Some(Box::new(ScalarValue::to_proto(
856                    Scalar::primitive(42_i32, Nullability::Nullable).value(),
857                ))),
858            }))),
859        };
860
861        assert!(ScalarValue::from_proto(&unknown_tag, &dtype, &session()).is_err());
862
863        let missing_child = pb::ScalarValue {
864            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
865                type_id: 5,
866                value: None,
867            }))),
868        };
869
870        assert!(matches!(
871            ScalarValue::from_proto(&missing_child, &dtype, &session()),
872            Err(VortexError::Serde(..))
873        ));
874
875        let wrong_child_value = pb::ScalarValue {
876            kind: Some(Kind::UnionValue(Box::new(PbUnionValue {
877                type_id: 5,
878                value: Some(Box::new(ScalarValue::to_proto(
879                    Scalar::utf8("wrong", Nullability::NonNullable).value(),
880                ))),
881            }))),
882        };
883
884        assert!(matches!(
885            ScalarValue::from_proto(&wrong_child_value, &dtype, &session()),
886            Err(VortexError::Serde(..))
887        ));
888
889        Ok(())
890    }
891
892    #[test]
893    fn test_backcompat_f16_serialized_as_u64() {
894        // Backwards compatibility test for the legacy f16 serialization format.
895        //
896        // Previously, f16 ScalarValues were serialized as `Uint64Value(v.to_bits() as u64)` because
897        // the proto schema only had 64-bit integer types, and f16's underlying representation is
898        // u16 which got widened to u64.
899        //
900        // The current implementation uses a dedicated `F16Value` proto field, but we must still be
901        // able to deserialize the old format. This test verifies that:
902        //
903        // 1. A `Uint64Value` containing f16 bits can be read as a U64 primitive (the raw bits).
904        // 2. When wrapped in a Scalar with F16 dtype, the value is correctly interpreted as f16.
905        //
906        // This ensures data written with the old serialization format remains readable.
907
908        // Simulate the old serialization: f16(0.42) stored as Uint64Value with its bit pattern.
909        let f16_value = f16::from_f32(0.42);
910        let f16_bits_as_u64 = f16_value.to_bits() as u64; // 14008
911
912        let pb_scalar_value = pb::ScalarValue {
913            kind: Some(Kind::Uint64Value(f16_bits_as_u64)),
914        };
915
916        // Step 1: Verify the normal U64 scalar.
917        let scalar_value = ScalarValue::from_proto(
918            &pb_scalar_value,
919            &DType::Primitive(PType::U64, Nullability::NonNullable),
920            &session(),
921        )
922        .unwrap();
923        assert_eq!(
924            scalar_value.as_ref().map(|v| v.as_primitive()),
925            Some(&PValue::U64(14008u64)),
926        );
927
928        // Step 2: Verify that when we use F16 dtype, the Uint64Value is correctly interpreted.
929        let scalar_value_f16 = ScalarValue::from_proto(
930            &pb_scalar_value,
931            &DType::Primitive(PType::F16, Nullability::Nullable),
932            &session(),
933        )
934        .unwrap();
935
936        let scalar = Scalar::new(
937            DType::Primitive(PType::F16, Nullability::Nullable),
938            scalar_value_f16,
939        );
940
941        assert_eq!(
942            scalar.as_primitive().pvalue().unwrap(),
943            PValue::F16(f16::from_f32(0.42)),
944            "Uint64Value should be correctly interpreted as f16 when dtype is F16"
945        );
946    }
947
948    #[test]
949    fn test_scalar_value_direct_roundtrip_f16() {
950        // Test that ScalarValue with f16 roundtrips correctly without going through Scalar.
951        let f16_values = vec![
952            f16::from_f32(0.0),
953            f16::from_f32(1.0),
954            f16::from_f32(-1.0),
955            f16::from_f32(0.42),
956            f16::from_f32(5.722046e-6),
957            f16::from_f32(f32::consts::PI),
958            f16::INFINITY,
959            f16::NEG_INFINITY,
960            f16::NAN,
961        ];
962
963        for f16_val in f16_values {
964            let scalar_value = ScalarValue::Primitive(PValue::F16(f16_val));
965            let pb_value = ScalarValue::to_proto(Some(&scalar_value));
966            let read_back = ScalarValue::from_proto(
967                &pb_value,
968                &DType::Primitive(PType::F16, Nullability::NonNullable),
969                &session(),
970            )
971            .unwrap();
972
973            match (&scalar_value, read_back.as_ref()) {
974                (
975                    ScalarValue::Primitive(PValue::F16(original)),
976                    Some(ScalarValue::Primitive(PValue::F16(roundtripped))),
977                ) => {
978                    if original.is_nan() && roundtripped.is_nan() {
979                        // NaN values are equal for our purposes.
980                        continue;
981                    }
982                    assert_eq!(
983                        original, roundtripped,
984                        "F16 value {original:?} did not roundtrip correctly"
985                    );
986                }
987                _ => {
988                    vortex_panic!(
989                        "Expected f16 primitive values, got {scalar_value:?} and {read_back:?}"
990                    )
991                }
992            }
993        }
994    }
995
996    #[test]
997    fn test_scalar_value_direct_roundtrip_preserves_values() {
998        // Test that ScalarValue roundtripping preserves values (but not necessarily exact types).
999        // Note: Proto encoding consolidates integer types (u8/u16/u32 → u64, i8/i16/i32 → i64).
1000
1001        // Test cases that should roundtrip exactly.
1002        let exact_roundtrip_cases: Vec<(&str, Option<ScalarValue>, DType)> = vec![
1003            ("null", None, DType::Null),
1004            (
1005                "bool_true",
1006                Some(ScalarValue::Bool(true)),
1007                DType::Bool(Nullability::Nullable),
1008            ),
1009            (
1010                "bool_false",
1011                Some(ScalarValue::Bool(false)),
1012                DType::Bool(Nullability::Nullable),
1013            ),
1014            (
1015                "u64",
1016                Some(ScalarValue::Primitive(PValue::U64(18446744073709551615))),
1017                DType::Primitive(PType::U64, Nullability::Nullable),
1018            ),
1019            (
1020                "i64",
1021                Some(ScalarValue::Primitive(PValue::I64(-9223372036854775808))),
1022                DType::Primitive(PType::I64, Nullability::Nullable),
1023            ),
1024            (
1025                "f32",
1026                Some(ScalarValue::Primitive(PValue::F32(f32::consts::E))),
1027                DType::Primitive(PType::F32, Nullability::Nullable),
1028            ),
1029            (
1030                "f64",
1031                Some(ScalarValue::Primitive(PValue::F64(f64::consts::PI))),
1032                DType::Primitive(PType::F64, Nullability::Nullable),
1033            ),
1034            (
1035                "string",
1036                Some(ScalarValue::Utf8(BufferString::from("test"))),
1037                DType::Utf8(Nullability::Nullable),
1038            ),
1039            (
1040                "bytes",
1041                Some(ScalarValue::Binary(vec![1, 2, 3, 4, 5].into())),
1042                DType::Binary(Nullability::Nullable),
1043            ),
1044        ];
1045
1046        for (name, value, dtype) in exact_roundtrip_cases {
1047            let pb_value = ScalarValue::to_proto(value.as_ref());
1048            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1049
1050            let original_debug = format!("{value:?}");
1051            let roundtrip_debug = format!("{read_back:?}");
1052            assert_eq!(
1053                original_debug, roundtrip_debug,
1054                "ScalarValue {name} did not roundtrip exactly"
1055            );
1056        }
1057
1058        // Test cases where type changes but value is preserved.
1059        // Unsigned integers consolidate to U64.
1060        let unsigned_cases = vec![
1061            (
1062                "u8",
1063                ScalarValue::Primitive(PValue::U8(255)),
1064                DType::Primitive(PType::U8, Nullability::Nullable),
1065                255u64,
1066            ),
1067            (
1068                "u16",
1069                ScalarValue::Primitive(PValue::U16(65535)),
1070                DType::Primitive(PType::U16, Nullability::Nullable),
1071                65535u64,
1072            ),
1073            (
1074                "u32",
1075                ScalarValue::Primitive(PValue::U32(4294967295)),
1076                DType::Primitive(PType::U32, Nullability::Nullable),
1077                4294967295u64,
1078            ),
1079        ];
1080
1081        for (name, value, dtype, expected) in unsigned_cases {
1082            let pb_value = ScalarValue::to_proto(Some(&value));
1083            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1084
1085            match read_back.as_ref() {
1086                Some(ScalarValue::Primitive(pv)) => {
1087                    let v = match pv {
1088                        PValue::U8(v) => *v as u64,
1089                        PValue::U16(v) => *v as u64,
1090                        PValue::U32(v) => *v as u64,
1091                        PValue::U64(v) => *v,
1092                        _ => vortex_panic!("Unexpected primitive type for {name}: {pv:?}"),
1093                    };
1094                    assert_eq!(
1095                        v, expected,
1096                        "ScalarValue {name} value not preserved: expected {expected}, got {v}"
1097                    );
1098                }
1099                _ => vortex_panic!("Unexpected type after roundtrip for {name}: {read_back:?}"),
1100            }
1101        }
1102
1103        // Signed integers consolidate to I64.
1104        let signed_cases = vec![
1105            (
1106                "i8",
1107                ScalarValue::Primitive(PValue::I8(-128)),
1108                DType::Primitive(PType::I8, Nullability::Nullable),
1109                -128i64,
1110            ),
1111            (
1112                "i16",
1113                ScalarValue::Primitive(PValue::I16(-32768)),
1114                DType::Primitive(PType::I16, Nullability::Nullable),
1115                -32768i64,
1116            ),
1117            (
1118                "i32",
1119                ScalarValue::Primitive(PValue::I32(-2147483648)),
1120                DType::Primitive(PType::I32, Nullability::Nullable),
1121                -2147483648i64,
1122            ),
1123        ];
1124
1125        for (name, value, dtype, expected) in signed_cases {
1126            let pb_value = ScalarValue::to_proto(Some(&value));
1127            let read_back = ScalarValue::from_proto(&pb_value, &dtype, &session()).unwrap();
1128
1129            match read_back.as_ref() {
1130                Some(ScalarValue::Primitive(pv)) => {
1131                    let v = match pv {
1132                        PValue::I8(v) => *v as i64,
1133                        PValue::I16(v) => *v as i64,
1134                        PValue::I32(v) => *v as i64,
1135                        PValue::I64(v) => *v,
1136                        _ => vortex_panic!("Unexpected primitive type for {name}: {pv:?}"),
1137                    };
1138                    assert_eq!(
1139                        v, expected,
1140                        "ScalarValue {name} value not preserved: expected {expected}, got {v}"
1141                    );
1142                }
1143                _ => vortex_panic!("Unexpected type after roundtrip for {name}: {read_back:?}"),
1144            }
1145        }
1146    }
1147
1148    // Backwards compatibility: signed integer stats could previously be serialized as unsigned.
1149    // Therefore, we allow casting between signed and unsigned integers of the same bit width.
1150    #[test]
1151    fn test_backcompat_signed_integer_deserialized_as_unsigned() {
1152        let v = ScalarValue::Primitive(PValue::I64(0));
1153        assert_eq!(
1154            Scalar::from_proto_value(
1155                &pb::ScalarValue::from(&v),
1156                &DType::Primitive(PType::U64, Nullability::Nullable),
1157                &session()
1158            )
1159            .unwrap(),
1160            Scalar::primitive(0u64, Nullability::Nullable)
1161        );
1162    }
1163
1164    // Backwards compatibility: unsigned integer stats could previously be serialized as signed.
1165    // Therefore, we allow casting between signed and unsigned integers of the same bit width.
1166    #[test]
1167    fn test_backcompat_unsigned_integer_deserialized_as_signed() {
1168        let v = ScalarValue::Primitive(PValue::U64(0));
1169        assert_eq!(
1170            Scalar::from_proto_value(
1171                &pb::ScalarValue::from(&v),
1172                &DType::Primitive(PType::I64, Nullability::Nullable),
1173                &session()
1174            )
1175            .unwrap(),
1176            Scalar::primitive(0i64, Nullability::Nullable)
1177        );
1178    }
1179}
vortex_array/scalar/proto.rs

vortex_array/scalar/
proto.rs
