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