fix: add FromPrimitiveOrF16 to NativePType replacing FromPrimitive

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

encodings/fastlanes/src/for/compress.rs

@@ -14,7 +14,6 @@ use vortex_dtype::{
     match_each_unsigned_integer_ptype,
 };
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
-use vortex_scalar::FromPrimitiveOrF16;
 
 use crate::unpack_iter::{UnpackStrategy, UnpackedChunks};
 use crate::{BitPackedArray, BitPackedVTable, FoRArray, bitpack_compress};
@@ -102,9 +101,7 @@ pub fn decompress(array: &FoRArray) -> PrimitiveArray {
     })
 }
 
-fn fused_decompress<
-    T: PhysicalPType<Physical = T> + UnsignedPType + FoR + FromPrimitiveOrF16 + WrappingAdd,
->(
+fn fused_decompress<T: PhysicalPType<Physical = T> + UnsignedPType + FoR + WrappingAdd>(
     for_: &FoRArray,
     bp: &BitPackedArray,
 ) -> PrimitiveArray {

vortex-array/src/arrays/chunked/compute/sum.rs

@@ -5,7 +5,7 @@ use num_traits::PrimInt;
 use vortex_dtype::Nullability::Nullable;
 use vortex_dtype::{DType, DecimalDType, NativePType, i256, match_each_native_ptype};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
-use vortex_scalar::{DecimalScalar, DecimalValue, FromPrimitiveOrF16, Scalar};
+use vortex_scalar::{DecimalScalar, DecimalValue, Scalar};
 
 use crate::arrays::{ChunkedArray, ChunkedVTable};
 use crate::compute::{SumKernel, SumKernelAdapter, sum};
@@ -39,9 +39,7 @@ impl SumKernel for ChunkedVTable {
 
 register_kernel!(SumKernelAdapter(ChunkedVTable).lift());
 
-fn sum_int<T: NativePType + PrimInt + FromPrimitiveOrF16>(
-    chunks: &[ArrayRef],
-) -> VortexResult<Option<T>> {
+fn sum_int<T: NativePType + PrimInt>(chunks: &[ArrayRef]) -> VortexResult<Option<T>> {
     let mut result: T = T::zero();
     for chunk in chunks {
         let chunk_sum = sum(chunk)?;

vortex-array/src/arrays/constant/compute/sum.rs

@@ -4,9 +4,7 @@
 use num_traits::{CheckedMul, ToPrimitive};
 use vortex_dtype::{DType, DecimalDType, NativePType, Nullability, i256, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
-use vortex_scalar::{
-    DecimalScalar, DecimalValue, FromPrimitiveOrF16, PrimitiveScalar, Scalar, ScalarValue,
-};
+use vortex_scalar::{DecimalScalar, DecimalValue, PrimitiveScalar, Scalar, ScalarValue};
 
 use crate::arrays::{ConstantArray, ConstantVTable};
 use crate::compute::{SumKernel, SumKernelAdapter};
@@ -84,7 +82,7 @@ fn sum_integral<T>(
     array_len: usize,
 ) -> VortexResult<Option<T>>
 where
-    T: FromPrimitiveOrF16 + NativePType + CheckedMul,
+    T: NativePType + CheckedMul,
     Scalar: From<Option<T>>,
 {
     let v = primitive_scalar.as_::<T>();

vortex-dtype/src/f16.rs

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use half::f16;
+use num_traits::FromPrimitive;
+
+/// A trait for types that can be created from primitive values, including f16.
+///
+/// This extends the `FromPrimitive` trait to also support conversion from f16 values.
+pub trait FromPrimitiveOrF16: FromPrimitive {
+    /// Converts an f16 value to this type, returning None if the conversion fails.
+    fn from_f16(v: f16) -> Option<Self>;
+}
+
+macro_rules! from_primitive_or_f16_for_non_floating_point {
+    ($T:ty) => {
+        impl FromPrimitiveOrF16 for $T {
+            fn from_f16(_: f16) -> Option<Self> {
+                None
+            }
+        }
+    };
+}
+
+from_primitive_or_f16_for_non_floating_point!(usize);
+from_primitive_or_f16_for_non_floating_point!(u8);
+from_primitive_or_f16_for_non_floating_point!(u16);
+from_primitive_or_f16_for_non_floating_point!(u32);
+from_primitive_or_f16_for_non_floating_point!(u64);
+from_primitive_or_f16_for_non_floating_point!(i8);
+from_primitive_or_f16_for_non_floating_point!(i16);
+from_primitive_or_f16_for_non_floating_point!(i32);
+from_primitive_or_f16_for_non_floating_point!(i64);
+
+impl FromPrimitiveOrF16 for f16 {
+    fn from_f16(v: f16) -> Option<Self> {
+        Some(v)
+    }
+}
+
+impl FromPrimitiveOrF16 for f32 {
+    fn from_f16(v: f16) -> Option<Self> {
+        Some(v.to_f32())
+    }
+}
+
+impl FromPrimitiveOrF16 for f64 {
+    fn from_f16(v: f16) -> Option<Self> {
+        Some(v.to_f64())
+    }
+}

vortex-dtype/src/lib.rs

@@ -18,6 +18,7 @@ pub mod datetime;
 mod decimal;
 mod dtype;
 mod extension;
+mod f16;
 mod field;
 mod field_mask;
 mod field_names;
@@ -31,6 +32,7 @@ pub use bigint::*;
 pub use decimal::*;
 pub use dtype::{DType, NativeDType};
 pub use extension::*;
+pub use f16::*;
 pub use field::*;
 pub use field_mask::*;
 pub use field_names::*;

vortex-dtype/src/ptype.rs

@@ -10,14 +10,12 @@ use std::ops::AddAssign;
 use std::panic::RefUnwindSafe;
 
 use num_traits::bounds::UpperBounded;
-use num_traits::{
-    AsPrimitive, Bounded, FromPrimitive, Num, NumCast, PrimInt, ToPrimitive, Unsigned,
-};
+use num_traits::{AsPrimitive, Bounded, Num, NumCast, PrimInt, ToPrimitive, Unsigned};
 use vortex_error::{VortexError, VortexResult, vortex_err};
 
-use crate::DType;
 use crate::half::f16;
 use crate::nullability::Nullability::NonNullable;
+use crate::{DType, FromPrimitiveOrF16};
 
 /// Physical type enum, represents the in-memory physical layout but might represent a different logical type.
 #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Hash, prost::Enumeration)]
@@ -96,7 +94,7 @@ pub trait NativePType:
     + RefUnwindSafe
     + Num
     + NumCast
-    + FromPrimitive
+    + FromPrimitiveOrF16
     + ToBytes
     + TryFromBytes
     + private::Sealed

vortex-scalar/src/primitive.rs

@@ -6,9 +6,11 @@ use std::cmp::Ordering;
 use std::fmt::{Debug, Display, Formatter};
 use std::ops::{Add, Sub};
 
-use num_traits::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, FromPrimitive};
+use num_traits::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub};
 use vortex_dtype::half::f16;
-use vortex_dtype::{DType, NativePType, Nullability, PType, match_each_native_ptype};
+use vortex_dtype::{
+    DType, FromPrimitiveOrF16, NativePType, Nullability, PType, match_each_native_ptype,
+};
 use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_err, vortex_panic};
 
 use crate::pvalue::{CoercePValue, PValue};
@@ -104,7 +106,7 @@ impl<'a> PrimitiveScalar<'a> {
     /// # Panics
     ///
     /// Panics if the primitive type of this scalar does not match the requested type.
-    pub fn typed_value<T: NativePType + TryFrom<PValue, Error = VortexError>>(&self) -> Option<T> {
+    pub fn typed_value<T: NativePType>(&self) -> Option<T> {
         assert_eq!(
             self.ptype,
             T::PTYPE,
@@ -215,52 +217,6 @@ impl<'a> PrimitiveScalar<'a> {
     }
 }
 
-/// A trait for types that can be created from primitive values, including f16.
-///
-/// This extends the `FromPrimitive` trait to also support conversion from f16 values.
-pub trait FromPrimitiveOrF16: FromPrimitive {
-    /// Converts an f16 value to this type, returning None if the conversion fails.
-    fn from_f16(v: f16) -> Option<Self>;
-}
-
-macro_rules! from_primitive_or_f16_for_non_floating_point {
-    ($T:ty) => {
-        impl FromPrimitiveOrF16 for $T {
-            fn from_f16(_: f16) -> Option<Self> {
-                None
-            }
-        }
-    };
-}
-
-from_primitive_or_f16_for_non_floating_point!(usize);
-from_primitive_or_f16_for_non_floating_point!(u8);
-from_primitive_or_f16_for_non_floating_point!(u16);
-from_primitive_or_f16_for_non_floating_point!(u32);
-from_primitive_or_f16_for_non_floating_point!(u64);
-from_primitive_or_f16_for_non_floating_point!(i8);
-from_primitive_or_f16_for_non_floating_point!(i16);
-from_primitive_or_f16_for_non_floating_point!(i32);
-from_primitive_or_f16_for_non_floating_point!(i64);
-
-impl FromPrimitiveOrF16 for f16 {
-    fn from_f16(v: f16) -> Option<Self> {
-        Some(v)
-    }
-}
-
-impl FromPrimitiveOrF16 for f32 {
-    fn from_f16(v: f16) -> Option<Self> {
-        Some(v.to_f32())
-    }
-}
-
-impl FromPrimitiveOrF16 for f64 {
-    fn from_f16(v: f16) -> Option<Self> {
-        Some(v.to_f64())
-    }
-}
-
 impl<'a> TryFrom<&'a Scalar> for PrimitiveScalar<'a> {
     type Error = VortexError;
 

vortex-scalar/src/pvalue.rs

@@ -202,7 +202,7 @@ impl PValue {
             PValue::I16(i) => T::from_i16(i),
             PValue::I32(i) => T::from_i32(i),
             PValue::I64(i) => T::from_i64(i),
-            PValue::F16(f) => <T as NumCast>::from(f),
+            PValue::F16(f) => T::from_f16(f),
             PValue::F32(f) => T::from_f32(f),
             PValue::F64(f) => T::from_f64(f),
         }
@@ -533,7 +533,7 @@ mod test {
     use std::collections::HashSet;
 
     use vortex_dtype::half::f16;
-    use vortex_dtype::{PType, ToBytes};
+    use vortex_dtype::{FromPrimitiveOrF16, NativePType, PType, ToBytes};
 
     use crate::PValue;
     use crate::pvalue::CoercePValue;
@@ -912,4 +912,13 @@ mod test {
             1.0f64 // 0x3ff0000000000000 is 1.0 in f64
         );
     }
+
+    #[test]
+    fn test_f16_nans_equal() {
+        let nan = f16::NAN;
+        let nan2 = f16::from_le_bytes([154, 253]);
+        assert!(nan2.is_nan());
+        let nan3 = f16::from_f16(nan2).unwrap();
+        assert_eq!(nan2.to_bits(), nan3.to_bits(),);
+    }
 }