feat: impl subtract for primitive_scalar

vortex-scalar/src/primitive.rs

@@ -2,8 +2,13 @@ use std::any::type_name;
 
 use num_traits::{FromPrimitive, NumCast};
 use vortex_dtype::half::f16;
-use vortex_dtype::{match_each_native_ptype, DType, NativePType, Nullability, PType};
-use vortex_error::{vortex_err, vortex_panic, VortexError, VortexResult, VortexUnwrap};
+use vortex_dtype::{
+    match_each_float_ptype, match_each_integer_ptype, match_each_native_ptype, DType, NativePType,
+    Nullability, PType,
+};
+use vortex_error::{
+    vortex_bail, vortex_err, vortex_panic, VortexError, VortexResult, VortexUnwrap,
+};
 
 use crate::pvalue::PValue;
 use crate::value::ScalarValue;
@@ -103,6 +108,46 @@ impl<'a> PrimitiveScalar<'a> {
             }?)),
         }
     }
+
+    pub fn checked_sub(&self, other: &PrimitiveScalar) -> VortexResult<Self> {
+        if self.ptype != other.ptype {
+            vortex_bail!("Failed to subtract {} and {}", self.ptype, other.ptype)
+        }
+        let result_pvalue: Option<PValue> = if !self.ptype.is_float() {
+            // checked_sub only implemented by integer type
+            match_each_integer_ptype!(self.ptype, |$T| {
+                let lhs = self.as_::<$T>()?;
+                let rhs = other.as_::<$T>()?;
+                match (lhs, rhs) {
+                    (Some(lv), Some(rv)) => {
+                        lv.checked_sub(rv).map(PValue::from)
+                    },
+                    _ => None
+                }
+            })
+        } else {
+            match_each_float_ptype!(self.ptype, |$T| {
+                let lhs = self.as_::<$T>()?;
+                let rhs = other.as_::<$T>()?;
+                match (lhs, rhs) {
+                    (Some(lv), Some(rv)) => {
+                        let f = lv - rv;
+                        if f.is_infinite() || f.is_nan() {
+                            None
+                        } else {
+                            Some(f.into())
+                        }
+                    },
+                    _ => None
+                }
+            })
+        };
+        Ok(Self {
+            dtype: self.dtype,
+            ptype: self.ptype,
+            pvalue: result_pvalue,
+        })
+    }
 }
 
 pub trait FromPrimitiveOrF16: FromPrimitive {
@@ -154,6 +199,14 @@ impl<'a> TryFrom<&'a Scalar> for PrimitiveScalar<'a> {
     }
 }
 
+impl std::ops::Sub for PrimitiveScalar<'_> {
+    type Output = VortexResult<Self>;
+
+    fn sub(self, rhs: Self) -> Self::Output {
+        self.checked_sub(&rhs)
+    }
+}
+
 impl Scalar {
     pub fn primitive<T: NativePType + Into<PValue>>(value: T, nullability: Nullability) -> Self {
         Self {
@@ -270,3 +323,97 @@ impl From<usize> for Scalar {
         Scalar::primitive(value as u64, Nullability::NonNullable)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::{DType, Nullability, PType};
+
+    use crate::value::InnerScalarValue;
+    use crate::{PValue, PrimitiveScalar, ScalarValue};
+
+    #[test]
+    fn test_integer_subtract() {
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
+        let p_scalar1 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::I32(5))),
+        )
+        .unwrap();
+        let p_scalar2 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::I32(4))),
+        )
+        .unwrap();
+        let res = p_scalar1.checked_sub(&p_scalar2).unwrap();
+        assert_eq!(res.as_::<i32>().unwrap().unwrap(), 1);
+
+        assert_eq!(
+            (p_scalar1 - p_scalar2)
+                .unwrap()
+                .as_::<i32>()
+                .unwrap()
+                .unwrap(),
+            1
+        );
+    }
+
+    #[test]
+    fn test_integer_subtract_overflow() {
+        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
+        let p_scalar1 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::I32(i32::MIN))),
+        )
+        .unwrap();
+        let p_scalar2 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::I32(i32::MAX))),
+        )
+        .unwrap();
+        let res = p_scalar1 - p_scalar2;
+        assert!(res.unwrap().pvalue.is_none());
+    }
+
+    #[test]
+    fn test_float_subtract() {
+        let dtype = DType::Primitive(PType::F32, Nullability::NonNullable);
+        let p_scalar1 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::F32(1.99f32))),
+        )
+        .unwrap();
+        let p_scalar2 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::F32(1.0f32))),
+        )
+        .unwrap();
+        let res = p_scalar1.checked_sub(&p_scalar2).unwrap();
+        assert_eq!(res.as_::<f32>().unwrap().unwrap(), 0.99f32);
+
+        assert_eq!(
+            (p_scalar1 - p_scalar2)
+                .unwrap()
+                .as_::<f32>()
+                .unwrap()
+                .unwrap(),
+            0.99f32
+        );
+    }
+
+    #[test]
+    fn test_float_subtract_overflow() {
+        let dtype = DType::Primitive(PType::F32, Nullability::NonNullable);
+        let p_scalar1 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::F32(f32::MIN))),
+        )
+        .unwrap();
+        let p_scalar2 = PrimitiveScalar::try_new(
+            &dtype,
+            &ScalarValue(InnerScalarValue::Primitive(PValue::F32(f32::MAX))),
+        )
+        .unwrap();
+        let res = p_scalar1 - p_scalar2;
+        assert!(res.unwrap().pvalue.is_none());
+    }
+}