volas_core/
numeric.rs

1//! The numeric element types a column can compute over (`f64`, `i64`), behind a
2//! single trait so the cumulative / arithmetic / clip / select kernels are
3//! written once and monomorphised per dtype (zero-cost, no f64 round-trip).
4//!
5//! Missing values use the f64 `NaN` sentinel — `i64` has none, so `is_missing`
6//! is always false there — matching pandas 3.0's default (non-nullable) dtypes.
7//! `wrapping_*` matches pandas/numpy int64 overflow (it wraps, e.g.
8//! `abs(i64::MIN) == i64::MIN`) and also keeps volas's overflow-checked test
9//! build from panicking.
10
11use crate::column::Column;
12use crate::dtype::DType;
13
14/// A numeric column element type. Implemented for `f64` and `i64`; adding `i32` /
15/// `f32` / `u64` later is just another `impl` plus a dispatch arm.
16pub trait Numeric: Copy + PartialOrd {
17    /// The column dtype backed by this element type.
18    const DTYPE: DType;
19    /// Additive identity (cumsum seed).
20    const ZERO: Self;
21    /// Multiplicative identity (cumprod seed).
22    const ONE: Self;
23
24    /// Wrap a computed buffer back into a `Column` of this element's dtype.
25    fn into_column(values: Vec<Self>) -> Column;
26
27    /// Whether this element is the missing sentinel (`NaN` for f64; never i64).
28    fn is_missing(self) -> bool;
29    /// Widen to f64 (lossless for both supported types within range).
30    fn to_f64(self) -> f64;
31    /// Narrow from f64 losslessly, or `None` when the value would not fit
32    /// (non-integral / out of range / `NaN` for i64). Powers the keep-vs-promote
33    /// decision in clip / where / mask / assignment.
34    fn try_from_f64(x: f64) -> Option<Self>;
35
36    /// Wrapping add (matches pandas/numpy int64 overflow; plain `+` for f64).
37    fn wrapping_add(self, other: Self) -> Self;
38    /// Wrapping subtract.
39    fn wrapping_sub(self, other: Self) -> Self;
40    /// Wrapping multiply.
41    fn wrapping_mul(self, other: Self) -> Self;
42    /// Wrapping absolute value (`abs(i64::MIN) == i64::MIN`, like pandas).
43    fn wrapping_abs(self) -> Self;
44}
45
46impl Numeric for f64 {
47    const DTYPE: DType = DType::F64;
48    const ZERO: Self = 0.0;
49    const ONE: Self = 1.0;
50
51    #[inline]
52    fn into_column(values: Vec<Self>) -> Column {
53        Column::f64(values)
54    }
55    #[inline]
56    fn is_missing(self) -> bool {
57        self.is_nan()
58    }
59    #[inline]
60    fn to_f64(self) -> f64 {
61        self
62    }
63    #[inline]
64    fn try_from_f64(x: f64) -> Option<Self> {
65        Some(x)
66    }
67    #[inline]
68    fn wrapping_add(self, other: Self) -> Self {
69        self + other
70    }
71    #[inline]
72    fn wrapping_sub(self, other: Self) -> Self {
73        self - other
74    }
75    #[inline]
76    fn wrapping_mul(self, other: Self) -> Self {
77        self * other
78    }
79    #[inline]
80    fn wrapping_abs(self) -> Self {
81        self.abs()
82    }
83}
84
85impl Numeric for f32 {
86    const DTYPE: DType = DType::F32;
87    const ZERO: Self = 0.0;
88    const ONE: Self = 1.0;
89
90    #[inline]
91    fn into_column(values: Vec<Self>) -> Column {
92        Column::f32(values)
93    }
94    #[inline]
95    fn is_missing(self) -> bool {
96        self.is_nan()
97    }
98    #[inline]
99    fn to_f64(self) -> f64 {
100        self as f64
101    }
102    #[inline]
103    fn try_from_f64(x: f64) -> Option<Self> {
104        Some(x as f32) // f32 absorbs any float (with rounding), like f64
105    }
106    #[inline]
107    fn wrapping_add(self, other: Self) -> Self {
108        self + other
109    }
110    #[inline]
111    fn wrapping_sub(self, other: Self) -> Self {
112        self - other
113    }
114    #[inline]
115    fn wrapping_mul(self, other: Self) -> Self {
116        self * other
117    }
118    #[inline]
119    fn wrapping_abs(self) -> Self {
120        self.abs()
121    }
122}
123
124impl Numeric for i32 {
125    const DTYPE: DType = DType::I32;
126    const ZERO: Self = 0;
127    const ONE: Self = 1;
128
129    #[inline]
130    fn into_column(values: Vec<Self>) -> Column {
131        Column::i32(values)
132    }
133    #[inline]
134    fn is_missing(self) -> bool {
135        false
136    }
137    #[inline]
138    fn to_f64(self) -> f64 {
139        self as f64
140    }
141    #[inline]
142    fn try_from_f64(x: f64) -> Option<Self> {
143        if x.is_finite() && x.fract() == 0.0 && x >= i32::MIN as f64 && x <= i32::MAX as f64 {
144            Some(x as i32)
145        } else {
146            None
147        }
148    }
149    #[inline]
150    fn wrapping_add(self, other: Self) -> Self {
151        i32::wrapping_add(self, other)
152    }
153    #[inline]
154    fn wrapping_sub(self, other: Self) -> Self {
155        i32::wrapping_sub(self, other)
156    }
157    #[inline]
158    fn wrapping_mul(self, other: Self) -> Self {
159        i32::wrapping_mul(self, other)
160    }
161    #[inline]
162    fn wrapping_abs(self) -> Self {
163        i32::wrapping_abs(self)
164    }
165}
166
167impl Numeric for i64 {
168    const DTYPE: DType = DType::I64;
169    const ZERO: Self = 0;
170    const ONE: Self = 1;
171
172    #[inline]
173    fn into_column(values: Vec<Self>) -> Column {
174        Column::i64(values)
175    }
176    #[inline]
177    fn is_missing(self) -> bool {
178        false
179    }
180    #[inline]
181    fn to_f64(self) -> f64 {
182        self as f64
183    }
184    #[inline]
185    fn try_from_f64(x: f64) -> Option<Self> {
186        // Lossless only: finite, integral, and within [-2^63, 2^63). `i64::MAX as
187        // f64` rounds up to 2^63, so the upper bound must be exclusive.
188        if x.is_finite() && x.fract() == 0.0 && (-9_223_372_036_854_775_808.0..9_223_372_036_854_775_808.0).contains(&x) {
189            Some(x as i64)
190        } else {
191            None
192        }
193    }
194    #[inline]
195    fn wrapping_add(self, other: Self) -> Self {
196        i64::wrapping_add(self, other)
197    }
198    #[inline]
199    fn wrapping_sub(self, other: Self) -> Self {
200        i64::wrapping_sub(self, other)
201    }
202    #[inline]
203    fn wrapping_mul(self, other: Self) -> Self {
204        i64::wrapping_mul(self, other)
205    }
206    #[inline]
207    fn wrapping_abs(self) -> Self {
208        i64::wrapping_abs(self)
209    }
210}
211
212// --- dtype promotion (single source of truth for pandas 3.0 result dtypes) ---
213
214/// Result dtype of a binary arithmetic op (`+` `-` `*`) by operand *type*
215/// (numpy-style promotion): same numeric dtype stays (`f32+f32→f32`,
216/// `i32+i32→i32`); same kind widens (`f32+f64→f64`, `i32+i64→i64`); mixed
217/// int+float → `f64` (safe); bool counts as `i32`. Type-based, matching pandas
218/// (`int * 2.0 → float64`); the value-based rule is [`fits`] (clip / where /
219/// assign). (`/` is always f64.)
220pub fn binary_supertype(a: DType, b: DType) -> DType {
221    use DType::{Bool, F32, F64, I32, I64};
222    if a == b && matches!(a, F64 | F32 | I64 | I32) {
223        return a;
224    }
225    let int_like = |d| matches!(d, I64 | I32 | Bool);
226    if a.is_float() || b.is_float() {
227        if a.is_float() && b.is_float() {
228            // both float -> the wider one
229            if a == F64 || b == F64 { F64 } else { F32 }
230        } else {
231            F64 // int + float -> f64
232        }
233    } else if int_like(a) && int_like(b) {
234        // both int (bool counts) -> the wider one (bool -> i32)
235        if a == I64 || b == I64 { I64 } else { I32 }
236    } else {
237        F64
238    }
239}
240
241/// Whether the scalar `x` fits `target` losslessly — the keep-vs-promote test for
242/// clip / where / mask / assignment (e.g. `0.0` fits `int64`, `2.5` does not).
243pub fn fits(target: DType, x: f64) -> bool {
244    match target {
245        DType::F64 | DType::F32 => true, // any float fits a float dtype (rounding)
246        DType::I64 => i64::try_from_f64(x).is_some(),
247        DType::I32 => i32::try_from_f64(x).is_some(),
248        _ => false,
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255
256    #[test]
257    fn f64_element_semantics() {
258        assert_eq!(f64::DTYPE, DType::F64);
259        assert!(f64::NAN.is_missing() && !1.0_f64.is_missing());
260        assert_eq!(2.5_f64.to_f64(), 2.5);
261        assert_eq!(f64::try_from_f64(2.5), Some(2.5)); // f64 always fits
262        assert_eq!(3.0_f64.wrapping_add(4.0), 7.0);
263        assert_eq!((-2.0_f64).wrapping_abs(), 2.0);
264    }
265
266    #[test]
267    fn i64_element_semantics() {
268        assert_eq!(i64::DTYPE, DType::I64);
269        assert!(!5_i64.is_missing()); // i64 has no missing
270        assert_eq!(5_i64.to_f64(), 5.0);
271        assert_eq!(3_i32.to_f64(), 3.0);
272        // lossless narrowing
273        assert_eq!(i64::try_from_f64(5.0), Some(5));
274        assert_eq!(i64::try_from_f64(2.5), None); // non-integral
275        assert_eq!(i64::try_from_f64(f64::NAN), None); // NaN
276        assert_eq!(i64::try_from_f64(1e30), None); // out of range
277        assert_eq!(i64::try_from_f64(9.223372036854776e18), None); // 2^63, exclusive
278        // wrapping matches pandas int64 overflow
279        assert_eq!(i64::MAX.wrapping_add(1), i64::MIN);
280        assert_eq!(i64::MIN.wrapping_abs(), i64::MIN);
281        assert_eq!(3_i64.wrapping_mul(4), 12);
282        assert_eq!(7_i64.wrapping_sub(10), -3);
283    }
284
285    #[test]
286    fn f32_i32_element_semantics() {
287        // f32: NaN sentinel, lossless to_f64, always-fits try_from (rounding)
288        assert_eq!(f32::DTYPE, DType::F32);
289        assert!(f32::NAN.is_missing() && !1.5_f32.is_missing());
290        assert_eq!(2.5_f32.to_f64(), 2.5);
291        assert_eq!(f32::try_from_f64(2.5), Some(2.5_f32));
292        assert_eq!(f32::try_from_f64(0.1).map(|x| x.is_finite()), Some(true)); // rounds, still fits
293        assert_eq!(3.0_f32.wrapping_add(4.0), 7.0);
294        assert_eq!((-2.0_f32).wrapping_abs(), 2.0);
295        assert!(matches!(f32::into_column(vec![1.0_f32]), Column::F32(_)));
296        // i32: no missing, wrapping, range-checked try_from
297        assert_eq!(i32::DTYPE, DType::I32);
298        assert!(!5_i32.is_missing());
299        assert_eq!(i32::try_from_f64(5.0), Some(5));
300        assert_eq!(i32::try_from_f64(2.5), None); // non-integral
301        assert_eq!(i32::try_from_f64(3e9), None); // out of i32 range
302        assert_eq!(i32::MAX.wrapping_add(1), i32::MIN); // wraps
303        assert_eq!(i32::MIN.wrapping_abs(), i32::MIN);
304        assert_eq!(7_i32.wrapping_sub(10), -3);
305        assert_eq!(3_i32.wrapping_mul(4), 12);
306        assert!(matches!(i32::into_column(vec![1_i32]), Column::I32(_, _)));
307    }
308
309    #[test]
310    fn promotion_rules() {
311        use DType::{Bool, F32, F64, I32, I64, Utf8};
312        // same numeric dtype stays
313        assert_eq!(binary_supertype(I64, I64), I64);
314        assert_eq!(binary_supertype(F32, F32), F32);
315        assert_eq!(binary_supertype(I32, I32), I32);
316        // same kind widens
317        assert_eq!(binary_supertype(F32, F64), F64);
318        assert_eq!(binary_supertype(I32, I64), I64);
319        // mixed int + float -> f64
320        assert_eq!(binary_supertype(I64, F64), F64);
321        assert_eq!(binary_supertype(I32, F32), F64);
322        // bool counts as i32 (its arithmetic is special-cased elsewhere)
323        assert_eq!(binary_supertype(Bool, I32), I32);
324        assert_eq!(binary_supertype(Bool, I64), I64);
325        assert_eq!(binary_supertype(I64, Utf8), F64); // non-numeric operand -> f64
326        // value-based fits (clip / where / assign)
327        assert!(fits(I64, 3.0) && !fits(I64, 2.5));
328        assert!(fits(I32, 3.0) && !fits(I32, 1e20)); // out of i32 range
329        assert!(fits(F64, 2.5) && fits(F32, 2.5)); // any float fits a float dtype
330        assert!(!fits(Bool, 1.0));
331    }
332}
volas_core/numeric.rs

volas_core/
numeric.rs
