Skip to main content

volas_core/
column.rs

1//! Column: a single typed, contiguous buffer.
2//!
3//! Each variant holds its buffer behind an `Arc`, so a `Column` — and the
4//! `DataFrame` / `Series` that contain it — is cheap to clone: cloning shares the
5//! buffer (an O(1) refcount bump, not an O(n) copy). Mutation (`append`) is
6//! copy-on-write via `Arc::make_mut`: it grows the `Vec` in place when the buffer
7//! is uniquely owned, and copies only when a view (another `Series`, a zero-copy
8//! export) is still alive. `F64` columns use `NaN` for missing values (matching
9//! stock-pandas / pandas semantics).
10
11use std::cmp::Ordering;
12use std::collections::HashMap;
13
14use crate::buffer::Buffer;
15use crate::strbuf::StrBuffer;
16use crate::datetime;
17use crate::dtype::DType;
18use crate::error::{Result, VolasError};
19use crate::numeric::{binary_supertype, fits, Numeric};
20use crate::stats;
21use crate::validity::Validity;
22
23/// Run a numeric kernel over a column's element type, monomorphised per dtype
24/// (`F64` / `I64`) with no f64 round-trip. `$slice` is bound to the typed slice;
25/// `$body` must produce a `Column` (via [`Numeric::into_column`]). `Bool` is
26/// handled per-op by the caller (pandas treats it inconsistently — `cumsum -> int`
27/// but `abs/cummax -> bool`); a `Bool` / `Str` / `Datetime` column here is an error.
28macro_rules! numeric_dispatch {
29    ($col:expr, $slice:ident => $body:expr) => {
30        match $col {
31            Column::F64(buf) => {
32                let $slice: &[f64] = buf.as_slice();
33                Ok($body)
34            }
35            Column::F32(buf) => {
36                let $slice: &[f32] = buf.as_slice();
37                Ok($body)
38            }
39            Column::I64(buf, _) => {
40                let $slice: &[i64] = buf.as_slice();
41                Ok($body)
42            }
43            Column::I32(buf, _) => {
44                let $slice: &[i32] = buf.as_slice();
45                Ok($body)
46            }
47            other => Err(VolasError::DType(format!(
48                "expected a numeric column, got {}",
49                other.dtype()
50            ))),
51        }
52    };
53}
54
55/// A typed, contiguous column of values. The value buffer is a [`Buffer`] —
56/// owned (`Arc`-shared, cheap clone, copy-on-write mutation) or a zero-copy borrow
57/// of foreign (Arrow / NumPy) memory.
58#[derive(Clone, Debug, PartialEq)]
59pub enum Column {
60    /// 64-bit floats; `NaN` denotes missing.
61    F64(Buffer<f64>),
62    /// 32-bit floats (narrow storage); `NaN` denotes missing.
63    F32(Buffer<f32>),
64    /// Booleans (comparison / signal results); `Validity` marks missing cells.
65    Bool(Buffer<bool>, Validity),
66    /// 64-bit signed integers; `Validity` marks missing cells.
67    I64(Buffer<i64>, Validity),
68    /// 32-bit signed integers (narrow storage); `Validity` marks missing cells.
69    I32(Buffer<i32>, Validity),
70    /// UTF-8 strings; `Validity` marks missing cells.
71    Str(StrBuffer, Validity),
72    /// Datetimes as i64 nanoseconds since the Unix epoch (UTC-naive).
73    Datetime(Buffer<i64>),
74}
75
76/// An in-place single-cell reduction op for the live tf-fold ([`Column::combine_at`]):
77/// fold one incoming value into a forming aggregate without re-reducing the period
78/// or allocating. The dual of [`crate::Agg`]'s reduce — `first` keeps, `last`
79/// replaces, the rest combine associatively.
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum CombineOp {
82    /// Keep the destination unchanged (the period's `first`).
83    Keep,
84    /// Overwrite with the source value (the period's `last`).
85    Replace,
86    /// Running maximum.
87    Max,
88    /// Running minimum.
89    Min,
90    /// Running sum.
91    Sum,
92}
93
94/// The result of a dtype-preserving scalar reduction ([`Column::sum`] etc.),
95/// carrying the value in its pandas result dtype so the binding can box it as the
96/// matching numpy scalar (`np.int64` / `np.float64` / `np.bool_`).
97#[derive(Clone, Copy, Debug, PartialEq)]
98pub enum Scalar {
99    /// A float64 result (`mean`, an f64 column's `sum`/`min`/`max`, …).
100    F64(f64),
101    /// A float32 result (an f32 column's `sum`/`min`/`max`).
102    F32(f32),
103    /// An int64 result (an i64/bool column's `sum`/`prod`, an i64 `min`/`max`).
104    I64(i64),
105    /// An int32 result (an i32 column's `sum`/`prod`/`min`/`max`).
106    I32(i32),
107    /// A boolean result (a bool column's `min`/`max`).
108    Bool(bool),
109}
110
111/// A dtype-preserving binary arithmetic op for [`Column::binary`] (pandas
112/// `+ - *`). True division is always float, so it is [`Column::div`], not here.
113#[derive(Clone, Copy, Debug, PartialEq, Eq)]
114pub enum BinOp {
115    /// Addition.
116    Add,
117    /// Subtraction.
118    Sub,
119    /// Multiplication.
120    Mul,
121}
122
123/// A three-valued logical op for [`Column::logical`] (pandas `&` / `|` / `^` on
124/// bool columns, Kleene semantics).
125#[derive(Clone, Copy, Debug, PartialEq, Eq)]
126pub enum BoolOp {
127    /// Logical AND (a present `false` short-circuits to `false`).
128    And,
129    /// Logical OR (a present `true` short-circuits to `true`).
130    Or,
131    /// Logical XOR (missing if either operand is missing).
132    Xor,
133}
134
135/// An element-wise comparison op for [`Column::compare`] (pandas `== != < <= > >=`).
136#[derive(Clone, Copy, Debug, PartialEq, Eq)]
137pub enum CmpOp {
138    Eq,
139    Ne,
140    Lt,
141    Le,
142    Gt,
143    Ge,
144}
145
146impl CmpOp {
147    /// Whether `ordering` (the result of comparing two **present** values)
148    /// satisfies this op.
149    fn matches(self, o: Ordering) -> bool {
150        match self {
151            CmpOp::Eq => o == Ordering::Equal,
152            CmpOp::Ne => o != Ordering::Equal,
153            CmpOp::Lt => o == Ordering::Less,
154            CmpOp::Le => o != Ordering::Greater,
155            CmpOp::Gt => o == Ordering::Greater,
156            CmpOp::Ge => o != Ordering::Less,
157        }
158    }
159
160    /// The op applied to two `f64`s (the numeric path); `NaN` follows IEEE — `!=`
161    /// is `true`, every other op `false`.
162    fn matches_f64(self, a: f64, b: f64) -> bool {
163        match self {
164            CmpOp::Eq => a == b,
165            CmpOp::Ne => a != b,
166            CmpOp::Lt => a < b,
167            CmpOp::Le => a <= b,
168            CmpOp::Gt => a > b,
169            CmpOp::Ge => a >= b,
170        }
171    }
172}
173
174mod cast;
175mod ops;
176mod order;
177mod transform;
178
179impl Column {
180    /// Build an `F64` column.
181    pub fn f64(v: Vec<f64>) -> Column {
182        Column::F64(Buffer::from_vec(v))
183    }
184
185    /// Build an `F32` column.
186    pub fn f32(v: Vec<f32>) -> Column {
187        Column::F32(Buffer::from_vec(v))
188    }
189
190    /// Build an `I32` column (all values present).
191    pub fn i32(v: Vec<i32>) -> Column {
192        Column::I32(Buffer::from_vec(v), Validity::dense())
193    }
194
195    /// Build a `Bool` column (all values present).
196    pub fn bool(v: Vec<bool>) -> Column {
197        Column::Bool(Buffer::from_vec(v), Validity::dense())
198    }
199
200    /// Build an `I64` column (all values present).
201    pub fn i64(v: Vec<i64>) -> Column {
202        Column::I64(Buffer::from_vec(v), Validity::dense())
203    }
204
205    /// Build an `I64` column with an explicit validity (missing-aware).
206    pub fn i64_with(v: Vec<i64>, validity: Validity) -> Column {
207        Column::I64(Buffer::from_vec(v), validity)
208    }
209
210    /// Build an `I32` column with an explicit validity (missing-aware).
211    pub fn i32_with(v: Vec<i32>, validity: Validity) -> Column {
212        Column::I32(Buffer::from_vec(v), validity)
213    }
214
215    /// Build a `Bool` column with an explicit validity (missing-aware).
216    pub fn bool_with(v: Vec<bool>, validity: Validity) -> Column {
217        Column::Bool(Buffer::from_vec(v), validity)
218    }
219
220    /// Build a `Str` column (all values present).
221    pub fn str(v: Vec<String>) -> Column {
222        Column::Str(StrBuffer::from_vec(v), Validity::dense())
223    }
224
225    /// Build a `Str` column with an explicit validity (missing-aware).
226    pub fn str_with(v: Vec<String>, validity: Validity) -> Column {
227        Column::Str(StrBuffer::from_vec(v), validity)
228    }
229
230    /// Build a `Datetime` column (epoch nanoseconds).
231    pub fn datetime(v: Vec<i64>) -> Column {
232        Column::Datetime(Buffer::from_vec(v))
233    }
234
235    /// An all-missing column of `dtype` with `len` rows — the dtype-preserving
236    /// default `other` for `where` / `mask` (so a kept value stays in its dtype
237    /// and a replaced one becomes NA, never an f64-funnel NaN).
238    pub fn na_of(dtype: DType, len: usize) -> Column {
239        let all_na = || Validity::from_valid_iter(len, std::iter::repeat_n(false, len));
240        match dtype {
241            DType::F64 => Column::f64(vec![f64::NAN; len]),
242            DType::F32 => Column::f32(vec![f32::NAN; len]),
243            DType::I64 => Column::i64_with(vec![0; len], all_na()),
244            DType::I32 => Column::i32_with(vec![0; len], all_na()),
245            DType::Bool => Column::bool_with(vec![false; len], all_na()),
246            DType::Utf8 => Column::str_with(vec![String::new(); len], all_na()),
247            DType::Datetime => Column::datetime(vec![i64::MIN; len]),
248        }
249    }
250
251    /// Number of elements.
252    pub fn len(&self) -> usize {
253        match self {
254            Column::F64(v) => v.len(),
255            Column::F32(v) => v.len(),
256            Column::Bool(v, _) => v.len(),
257            Column::I64(v, _) => v.len(),
258            Column::I32(v, _) => v.len(),
259            Column::Str(v, _) => v.len(),
260            Column::Datetime(v) => v.len(),
261        }
262    }
263
264    /// Whether the column has no elements.
265    pub fn is_empty(&self) -> bool {
266        self.len() == 0
267    }
268
269    /// The logical dtype of the column.
270    pub fn dtype(&self) -> DType {
271        match self {
272            Column::F64(_) => DType::F64,
273            Column::F32(_) => DType::F32,
274            Column::Bool(_, _) => DType::Bool,
275            Column::I64(_, _) => DType::I64,
276            Column::I32(_, _) => DType::I32,
277            Column::Str(_, _) => DType::Utf8,
278            Column::Datetime(_) => DType::Datetime,
279        }
280    }
281
282    /// Whether the value at `i` is present (not `volas.NA`). Each dtype reads its
283    /// natural missing representation: a float `NaN`, an int/bool validity mask,
284    /// or a datetime `NaT` (`i64::MIN`). `i` is assumed in bounds.
285    pub fn is_valid(&self, i: usize) -> bool {
286        match self {
287            Column::F64(v) => !v[i].is_nan(),
288            Column::F32(v) => !v[i].is_nan(),
289            Column::Bool(_, val)
290            | Column::I64(_, val)
291            | Column::I32(_, val)
292            | Column::Str(_, val) => val.is_valid(i),
293            Column::Datetime(v) => v[i] != i64::MIN,
294        }
295    }
296
297    /// Read cell `i` as `i64` (any integer / datetime / bool dtype; truncates a
298    /// float, `0` for `Str`). The integer counterpart of [`Self::get_f64`], for the
299    /// in-place tf-fold's combine on integer / datetime columns.
300    pub fn get_i64(&self, i: usize) -> i64 {
301        match self {
302            Column::I64(v, _) => v[i],
303            Column::I32(v, _) => v[i] as i64,
304            Column::Datetime(v) => v[i],
305            Column::F64(v) => v[i] as i64,
306            Column::F32(v) => v[i] as i64,
307            Column::Bool(v, _) => v[i] as i64,
308            Column::Str(_, _) => 0,
309        }
310    }
311
312    /// Count of missing (`volas.NA`) values.
313    pub fn null_count(&self) -> usize {
314        match self {
315            Column::F64(v) => v.iter().filter(|x| x.is_nan()).count(),
316            Column::F32(v) => v.iter().filter(|x| x.is_nan()).count(),
317            Column::Bool(_, val)
318            | Column::I64(_, val)
319            | Column::I32(_, val)
320            | Column::Str(_, val) => val.null_count(),
321            Column::Datetime(v) => v.iter().filter(|&&x| x == i64::MIN).count(),
322        }
323    }
324
325    /// Attach a validity to a (nullable) column produced by a value-only kernel,
326    /// so an element-wise transform carries the input's missing positions through.
327    /// A float/str/datetime column is returned unchanged (its missing lives in the
328    /// values, not a side mask).
329    fn with_validity(self, validity: Validity) -> Column {
330        match self {
331            Column::I64(v, _) => Column::I64(v, validity),
332            Column::I32(v, _) => Column::I32(v, validity),
333            // bool transforms build their column with the mask directly; float /
334            // str / datetime carry missing in the values, so leave them unchanged.
335            other => other,
336        }
337    }
338
339    /// The validity of a nullable (int / bool) column, else `None` (a float column
340    /// carries its missing in `NaN`, not a mask).
341    fn validity(&self) -> Option<&Validity> {
342        match self {
343            Column::Bool(_, val)
344            | Column::I64(_, val)
345            | Column::I32(_, val)
346            | Column::Str(_, val) => Some(val),
347            _ => None,
348        }
349    }
350
351    /// Borrow the underlying `f64` slice, if this is an `F64` column.
352    pub fn as_f64(&self) -> Option<&[f64]> {
353        if let Column::F64(v) = self {
354            Some(v.as_slice())
355        } else {
356            None
357        }
358    }
359
360    /// Borrow the underlying `bool` slice, if this is a `Bool` column.
361    pub fn as_bool(&self) -> Option<&[bool]> {
362        if let Column::Bool(v, _) = self {
363            Some(v.as_slice())
364        } else {
365            None
366        }
367    }
368
369    /// Borrow the underlying `i64` slice, if this is an `I64` column.
370    pub fn as_i64(&self) -> Option<&[i64]> {
371        if let Column::I64(v, _) = self {
372            Some(v.as_slice())
373        } else {
374            None
375        }
376    }
377
378    /// The `i`-th string of a `Str` column as `&str`, if this is a `Str` column.
379    pub fn str_at(&self, i: usize) -> Option<&str> {
380        if let Column::Str(v, _) = self {
381            Some(v.get(i))
382        } else {
383            None
384        }
385    }
386
387    /// Borrow the underlying epoch-ns slice, if this is a `Datetime` column.
388    pub fn as_datetime(&self) -> Option<&[i64]> {
389        if let Column::Datetime(v) = self {
390            Some(v.as_slice())
391        } else {
392            None
393        }
394    }
395
396    /// Materialize the values as `f64` (`bool` -> 0.0/1.0, `i64` / `datetime` ->
397    /// as f64, `str` -> NaN). Used to feed indicator kernels, which operate on
398    /// `f64`.
399    pub fn to_f64_vec(&self) -> Vec<f64> {
400        match self {
401            Column::F64(v) => v.to_vec(),
402            Column::F32(v) => v.iter().map(|&x| x as f64).collect(),
403            Column::Bool(v, val) => mask_f64(v.iter().map(|&b| if b { 1.0 } else { 0.0 }), val),
404            Column::I64(v, val) => mask_f64(v.iter().map(|&i| i as f64), val),
405            Column::I32(v, val) => mask_f64(v.iter().map(|&i| i as f64), val),
406            Column::Str(v, _) => vec![f64::NAN; v.len()],
407            Column::Datetime(v) => v
408                .iter()
409                .map(|&i| if i == i64::MIN { f64::NAN } else { i as f64 })
410                .collect(),
411        }
412    }
413
414    /// Value at position `i` coerced to `f64`, **ignoring validity** (a missing
415    /// cell reads its raw placeholder). Used only by the logical-op coercion in
416    /// `bool_at`; NumPy export goes through the validity-aware `to_f64_vec`.
417    pub fn get_f64(&self, i: usize) -> f64 {
418        match self {
419            Column::F64(v) => v[i],
420            Column::F32(v) => v[i] as f64,
421            Column::Bool(v, _) => {
422                if v[i] {
423                    1.0
424                } else {
425                    0.0
426                }
427            }
428            Column::I64(v, _) => v[i] as f64,
429            Column::I32(v, _) => v[i] as f64,
430            Column::Str(_, _) => f64::NAN,
431            Column::Datetime(v) => v[i] as f64,
432        }
433    }
434
435    /// The column as `bool` values (a `Bool` column directly; else an error).
436    fn as_bool_vec(&self) -> Result<Vec<bool>> {
437        match self {
438            Column::Bool(v, _) => Ok(v.to_vec()),
439            other => Err(VolasError::DType(format!(
440                "expected a bool column, got {}",
441                other.dtype()
442            ))),
443        }
444    }
445
446    /// The column as `i64` values: an `I64` column directly (exact), otherwise via
447    /// lossless narrowing (errors if any value is non-integral / out of range).
448    /// Used by the int paths of `select` / `binary`, where the caller has already
449    /// established the values fit.
450    fn as_i64_vec(&self) -> Result<Vec<i64>> {
451        match self {
452            Column::I64(v, _) => Ok(v.to_vec()),
453            _ => self
454                .to_f64_vec()
455                .iter()
456                .enumerate()
457                .map(|(i, &x)| {
458                    if self.is_valid(i) {
459                        i64::try_from_f64(x).ok_or_else(|| {
460                            VolasError::DType(format!("value {x} does not fit int64"))
461                        })
462                    } else {
463                        Ok(0) // placeholder for a missing value (masked by the result validity)
464                    }
465                })
466                .collect(),
467        }
468    }
469
470    /// The column's `String` values (a `Str` column directly); errors otherwise.
471    /// Used by `select` so a str `where` / `mask` keeps its values dtype-preserving.
472    fn as_str_vec(&self) -> Result<Vec<String>> {
473        match self {
474            Column::Str(v, _) => Ok(v.to_vec()),
475            other => Err(VolasError::DType(format!(
476                "cannot select a {} column as str",
477                other.dtype()
478            ))),
479        }
480    }
481
482    /// The column's epoch-ns values (a `Datetime` column directly); errors otherwise.
483    fn as_datetime_vec(&self) -> Result<Vec<i64>> {
484        match self {
485            Column::Datetime(v) => Ok(v.to_vec()),
486            other => Err(VolasError::DType(format!(
487                "cannot select a {} column as datetime",
488                other.dtype()
489            ))),
490        }
491    }
492
493    /// The column as `i32` values: an `I32` column directly, a `Bool` as 0/1,
494    /// otherwise via lossless narrowing (errors if out of range / non-integral).
495    fn as_i32_vec(&self) -> Result<Vec<i32>> {
496        match self {
497            Column::I32(v, _) => Ok(v.to_vec()),
498            Column::Bool(v, _) => Ok(v.iter().map(|&b| b as i32).collect()),
499            _ => self
500                .to_f64_vec()
501                .iter()
502                .enumerate()
503                .map(|(i, &x)| {
504                    if self.is_valid(i) {
505                        i32::try_from_f64(x).ok_or_else(|| {
506                            VolasError::DType(format!("value {x} does not fit int32"))
507                        })
508                    } else {
509                        Ok(0) // placeholder for a missing value (masked by the result validity)
510                    }
511                })
512                .collect(),
513        }
514    }
515
516    /// The column as `f32` values (an `F32` column directly; else converted).
517    pub fn to_f32_vec(&self) -> Vec<f32> {
518        match self {
519            Column::F32(v) => v.to_vec(),
520            _ => self.to_f64_vec().iter().map(|&x| x as f32).collect(),
521        }
522    }
523
524    /// Render each value as a `String` (for `astype(str)`).
525    fn to_string_vec(&self) -> Vec<String> {
526        match self {
527            Column::Str(v, _) => v.to_vec(),
528            Column::F64(v) => v.iter().map(|x| x.to_string()).collect(),
529            Column::F32(v) => v.iter().map(|x| x.to_string()).collect(),
530            Column::I64(v, _) => v.iter().map(|x| x.to_string()).collect(),
531            Column::I32(v, _) => v.iter().map(|x| x.to_string()).collect(),
532            Column::Bool(v, _) => v
533                .iter()
534                .map(|&b| if b { "True" } else { "False" }.to_string())
535                .collect(),
536            Column::Datetime(v) => v.iter().map(|&ns| datetime::format_ns(ns)).collect(),
537        }
538    }
539}
540
541/// Collect a value iterator to `f64`, mapping missing positions to `NaN` so the
542/// f64-funnel reductions (mean / std / median …) skip them. Dense ⇒ a plain
543/// collect with no per-element validity check, so the indicator feed path stays
544/// unchanged.
545fn mask_f64(vals: impl Iterator<Item = f64>, validity: &Validity) -> Vec<f64> {
546    if validity.has_nulls() {
547        vals.enumerate()
548            .map(|(i, x)| if validity.is_valid(i) { x } else { f64::NAN })
549            .collect()
550    } else {
551        vals.collect()
552    }
553}
554
555/// A bool / i32 column widened to `i64` (for the i64-accumulator reductions).
556fn widen_i64<T: Copy + Into<i64>>(v: &[T]) -> Vec<i64> {
557    v.iter().map(|&x| x.into()).collect()
558}
559
560/// Collect `Option<bool>`s into a `Bool` column: `None` -> `volas.NA`.
561fn from_option_bools(n: usize, it: impl Iterator<Item = Option<bool>>) -> Column {
562    let (mut values, mut valid) = (Vec::with_capacity(n), Vec::with_capacity(n));
563    for o in it {
564        values.push(o.unwrap_or(false));
565        valid.push(o.is_some());
566    }
567    Column::bool_with(values, Validity::from_valid_iter(n, valid))
568}
569
570/// Shift `v` by `n` (pandas `shift`), filling vacated cells with `fill`; the kept
571/// region moves as a single `memcpy`.
572fn shift_fill<T: Copy>(v: &[T], n: isize, fill: T) -> Vec<T> {
573    let len = v.len();
574    let mut out = vec![fill; len];
575    if n >= 0 {
576        let n = (n as usize).min(len);
577        out[n..].copy_from_slice(&v[..len - n]);
578    } else {
579        let n = ((-n) as usize).min(len);
580        out[..len - n].copy_from_slice(&v[n..]);
581    }
582    out
583}
584
585/// `x[i] - x[i-n]` (pandas `diff`) in one pass; the shift-gap cells are `missing`.
586fn diff_kernel<T: Copy + std::ops::Sub<Output = T>>(v: &[T], n: isize, missing: T) -> Vec<T> {
587    let len = v.len();
588    let mut out = vec![missing; len];
589    if n >= 0 {
590        let k = n as usize;
591        for i in k..len {
592            out[i] = v[i] - v[i - k];
593        }
594    } else {
595        let k = (-n) as usize;
596        for i in 0..len.saturating_sub(k) {
597            out[i] = v[i] - v[i + k];
598        }
599    }
600    out
601}
602
603/// Extend `av` (validity of an `alen`-long column) with `bv` (a `blen`-long
604/// column's validity) so it stays aligned after an `append`. Dense + dense is a
605/// no-op (the result is still fully present).
606fn append_validity(av: &mut Validity, alen: usize, bv: &Validity, blen: usize) {
607    if !av.has_nulls() && !bv.has_nulls() {
608        return;
609    }
610    let flags: Vec<bool> = (0..alen)
611        .map(|i| av.is_valid(i))
612        .chain((0..blen).map(|i| bv.is_valid(i)))
613        .collect();
614    *av = Validity::from_valid_iter(alen + blen, flags);
615}
616
617/// Sum of the present values in their element type (skip `volas.NA`). Dense ⇒
618/// the plain kernel (no per-element validity check).
619fn sum_valid<T: Numeric>(v: &[T], val: &Validity) -> T {
620    if !val.has_nulls() {
621        return stats::sum(v);
622    }
623    v.iter()
624        .enumerate()
625        .filter(|(i, _)| val.is_valid(*i))
626        .fold(T::ZERO, |a, (_, &x)| a.wrapping_add(x))
627}
628
629/// Product of the present values (skip `volas.NA`). Dense ⇒ the plain kernel.
630fn prod_valid<T: Numeric>(v: &[T], val: &Validity) -> T {
631    if !val.has_nulls() {
632        return stats::prod(v);
633    }
634    v.iter()
635        .enumerate()
636        .filter(|(i, _)| val.is_valid(*i))
637        .fold(T::ONE, |a, (_, &x)| a.wrapping_mul(x))
638}
639
640/// Min / max of the present values (skip `volas.NA`); `None` when none present.
641fn extreme_valid<T: Numeric>(v: &[T], val: &Validity, want_max: bool) -> Option<T> {
642    if !val.has_nulls() {
643        return stats::extreme(v, want_max);
644    }
645    let mut it = v
646        .iter()
647        .enumerate()
648        .filter(|(i, _)| val.is_valid(*i))
649        .map(|(_, &x)| x);
650    let first = it.next()?;
651    Some(if want_max {
652        it.fold(first, |a, x| if x > a { x } else { a })
653    } else {
654        it.fold(first, |a, x| if x < a { x } else { a })
655    })
656}
657
658/// Running cumulative `op` over present values; a missing position gets
659/// `placeholder` (masked by the carried validity, so its value is irrelevant).
660fn cum_valid<T: Copy>(v: &[T], val: &Validity, placeholder: T, op: impl Fn(T, T) -> T) -> Vec<T> {
661    let mut acc: Option<T> = None;
662    v.iter()
663        .enumerate()
664        .map(|(i, &x)| {
665            if val.is_valid(i) {
666                let next = acc.map_or(x, |a| op(a, x));
667                acc = Some(next);
668                next
669            } else {
670                placeholder
671            }
672        })
673        .collect()
674}
675
676/// Dtype-preserving cumulative for a numeric column: the dense kernel when fully
677/// present, else a skip-and-propagate fold carrying the input validity through.
678fn cum<T: Numeric>(
679    v: &[T],
680    val: &Validity,
681    dense: fn(&[T]) -> Vec<T>,
682    op: fn(T, T) -> T,
683) -> Column {
684    let out = if val.has_nulls() {
685        cum_valid(v, val, T::ZERO, op)
686    } else {
687        dense(v)
688    };
689    T::into_column(out).with_validity(val.clone())
690}
691
692/// Banker's (half-to-even) round of `x` to `decimals` places; `NaN` stays `NaN`.
693fn round_f64(x: f64, decimals: i32) -> f64 {
694    if x.is_nan() {
695        return x;
696    }
697    let f = 10f64.powi(decimals);
698    (x * f).round_ties_even() / f
699}
700
701/// Round an int to `decimals` places: identity for `decimals >= 0`; for negative
702/// `decimals`, banker's round to the nearest `10^|decimals|` multiple (integer
703/// arithmetic, exact for all i64 — no f64 round-trip).
704fn round_i64(x: i64, decimals: i32) -> i64 {
705    if decimals >= 0 {
706        return x;
707    }
708    let factor = match 10i64.checked_pow(decimals.unsigned_abs()) {
709        Some(f) => f,
710        None => return 0, // 10^k beyond i64 -> everything rounds to 0
711    };
712    let q = x.div_euclid(factor);
713    let r = x.rem_euclid(factor); // 0..factor
714    let half = factor / 2; // factor = 10^k is even, so this is exact
715    let up = r > half || (r == half && q.rem_euclid(2) != 0); // tie -> even multiple
716    if up { q + 1 } else { q }.wrapping_mul(factor)
717}
718
719/// Clamp each element to `[lo, hi]` (either optional); missing passes through.
720/// Bounds are narrowed to `T` losslessly (the caller only stays in an int dtype
721/// when the bounds fit it).
722fn clip_vec<T: Numeric>(v: &[T], lo: Option<f64>, hi: Option<f64>) -> Vec<T> {
723    let lo = lo.and_then(T::try_from_f64);
724    let hi = hi.and_then(T::try_from_f64);
725    v.iter()
726        .map(|&x| {
727            if x.is_missing() {
728                return x;
729            }
730            let mut y = x;
731            if let Some(l) = lo {
732                if y < l {
733                    y = l;
734                }
735            }
736            if let Some(h) = hi {
737                if y > h {
738                    y = h;
739                }
740            }
741            y
742        })
743        .collect()
744}
745
746/// Build a non-nullable bool mask from a per-index ordering of present pairs: a
747/// present pair (`Some(ord)`) applies `op`; a missing slot (`None`) follows IEEE,
748/// so only `!=` is `true`. Backs the typed (`str` / `datetime` / `bool`) arms of
749/// [`Column::compare`].
750fn cmp_typed(n: usize, op: CmpOp, key: impl Fn(usize) -> Option<Ordering>) -> Vec<bool> {
751    (0..n)
752        .map(|i| match key(i) {
753            Some(ord) => op.matches(ord),
754            None => op == CmpOp::Ne,
755        })
756        .collect()
757}
758
759/// Integer floor division (toward negative infinity, unlike Rust's truncating
760/// `/`), wrapping on overflow like numpy. The divisor is non-zero (callers route
761/// a zero divisor through the float path).
762fn ifloordiv(a: i64, b: i64) -> i64 {
763    let q = a.wrapping_div(b);
764    let r = a.wrapping_rem(b);
765    if r != 0 && (r < 0) != (b < 0) {
766        q - 1
767    } else {
768        q
769    }
770}
771
772/// A hashable key for a float value, or `None` for `NaN` (which groups as NA).
773/// `-0.0` is canonicalized to `+0.0` so the two zeros count as one value.
774fn float_key(x: f64) -> Option<u64> {
775    if x.is_nan() {
776        None
777    } else if x == 0.0 {
778        Some(0)
779    } else {
780        Some(x.to_bits())
781    }
782}
783
784/// Group `0..len` by `key(i)` (a `None` key is the single NA group), returning
785/// `(first_index, count, is_na)` per distinct value in order of first appearance.
786fn group_by<K: Eq + std::hash::Hash>(
787    len: usize,
788    key: impl Fn(usize) -> Option<K>,
789) -> Vec<(usize, usize, bool)> {
790    let mut order: Vec<(usize, usize, bool)> = Vec::new();
791    let mut seen: HashMap<K, usize> = HashMap::new();
792    let mut na: Option<usize> = None;
793    for i in 0..len {
794        match key(i) {
795            None => match na {
796                Some(g) => order[g].1 += 1,
797                None => {
798                    na = Some(order.len());
799                    order.push((i, 1, true));
800                }
801            },
802            Some(k) => match seen.get(&k) {
803                Some(&g) => order[g].1 += 1,
804                None => {
805                    seen.insert(k, order.len());
806                    order.push((i, 1, false));
807                }
808            },
809        }
810    }
811    order
812}
813
814/// Running OR (`or = true`, backs bool `cummax`) / running AND (`cummin`).
815fn bool_running(v: &[bool], or: bool) -> Vec<bool> {
816    let mut acc = !or; // OR seeds false; AND seeds true
817    v.iter()
818        .map(|&b| {
819            acc = if or { acc || b } else { acc && b };
820            acc
821        })
822        .collect()
823}
824
825/// Clamp a bool column (pandas `clip`): a `True` lower bound forces all true, a
826/// `False` upper bound forces all false, otherwise unchanged (bool has only the
827/// two values, so the bounds decide the whole column).
828fn clip_bool(v: &[bool], lo: Option<f64>, hi: Option<f64>) -> Vec<bool> {
829    let force_true = lo.is_some_and(|x| x != 0.0);
830    let force_false = hi == Some(0.0);
831    v.iter()
832        .map(|&b| {
833            if force_false {
834                false
835            } else if force_true {
836                true
837            } else {
838                b
839            }
840        })
841        .collect()
842}
843
844/// Element-wise `a ∘ b` for `Add` / `Sub` / `Mul` (wrapping, pandas overflow).
845fn binary_kernel<T: Numeric>(a: &[T], b: &[T], op: BinOp) -> Vec<T> {
846    a.iter()
847        .zip(b)
848        .map(|(&x, &y)| match op {
849            BinOp::Add => x.wrapping_add(y),
850            BinOp::Sub => x.wrapping_sub(y),
851            BinOp::Mul => x.wrapping_mul(y),
852        })
853        .collect()
854}
855
856#[cfg(test)]
857mod tests;
858#[cfg(test)]
859mod tests_na;