Store PValue in FoR metadata

spiraldb · Jan 1, 2025 · 4e1fb1d · 4e1fb1d
1 parent 9532cbf
commit 4e1fb1d
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 10 deletions.
diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs
@@ -12,7 +12,7 @@ use vortex_array::{
 };
 use vortex_dtype::DType;
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
-use vortex_scalar::{Scalar, ScalarValue};
+use vortex_scalar::{PValue, Scalar};
 
 mod compress;
 mod compute;
@@ -21,7 +21,7 @@ impl_encoding!("fastlanes.for", ids::FL_FOR, FoR);
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FoRMetadata {
-    reference: ScalarValue,
+    reference: PValue,
     shift: u8,
 }
 
@@ -43,13 +43,18 @@ impl FoRArray {
                 .with_nullability(child.dtype().nullability()),
         )?;
 
+        let dtype = reference.dtype().clone();
+
+        // Convert the reference into a PValue which is smaller to store.
+        let reference = reference
+            .as_primitive()
+            .pvalue()
+            .vortex_expect("Reference value is non-null");
+
         Self::try_from_parts(
-            reference.dtype().clone(),
+            dtype,
             child.len(),
-            FoRMetadata {
-                reference: reference.into_value(),
-                shift,
-            },
+            FoRMetadata { reference, shift },
             [child].into(),
             StatsSet::default(),
         )
@@ -69,7 +74,11 @@ impl FoRArray {
 
     #[inline]
     pub fn reference_scalar(&self) -> Scalar {
-        Scalar::new(self.dtype().clone(), self.metadata().reference.clone())
+        Scalar::primitive_value(
+            self.metadata().reference,
+            self.ptype(),
+            self.dtype().nullability(),
+        )
     }
 
     #[inline]

diff --git a/vortex-scalar/src/primitive.rs b/vortex-scalar/src/primitive.rs
@@ -176,9 +176,17 @@ impl std::ops::Sub for PrimitiveScalar<'_> {
 
 impl Scalar {
     pub fn primitive<T: NativePType + Into<PValue>>(value: T, nullability: Nullability) -> Self {
+        Self::primitive_value(value.into(), T::PTYPE, nullability)
+    }
+
+    /// Create a PrimitiveScalar from a PValue.
+    ///
+    /// Note that an explicit PType is passed since any compatible PValue may be used as the value
+    /// for a primitive type.
+    pub fn primitive_value(value: PValue, ptype: PType, nullability: Nullability) -> Self {
         Self {
-            dtype: DType::Primitive(T::PTYPE, nullability),
-            value: ScalarValue(InnerScalarValue::Primitive(value.into())),
+            dtype: DType::Primitive(ptype, nullability),
+            value: ScalarValue(InnerScalarValue::Primitive(value)),
         }
     }
 

diff --git a/vortex-scalar/src/serde/serde.rs b/vortex-scalar/src/serde/serde.rs
@@ -184,3 +184,16 @@ impl Serialize for PValue {
         }
     }
 }
+
+impl<'de> Deserialize<'de> for PValue {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        ScalarValue::deserialize(deserializer)
+            .and_then(|scalar| scalar.0.as_pvalue().map_err(|e| Error::custom(e)))
+            .and_then(|pvalue| {
+                pvalue.ok_or_else(|| Error::custom("Expected a non-null primitive scalar value"))
+            })
+    }
+}