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

vortex_array/scalar/
proto.rs
