1use 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
23macro_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#[derive(Clone, Debug, PartialEq)]
59pub enum Column {
60 F64(Buffer<f64>),
62 F32(Buffer<f32>),
64 Bool(Buffer<bool>, Validity),
66 I64(Buffer<i64>, Validity),
68 I32(Buffer<i32>, Validity),
70 Str(StrBuffer, Validity),
72 Datetime(Buffer<i64>),
74}
75
76#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum CombineOp {
82 Keep,
84 Replace,
86 Max,
88 Min,
90 Sum,
92}
93
94#[derive(Clone, Copy, Debug, PartialEq)]
98pub enum Scalar {
99 F64(f64),
101 F32(f32),
103 I64(i64),
105 I32(i32),
107 Bool(bool),
109}
110
111#[derive(Clone, Copy, Debug, PartialEq, Eq)]
114pub enum BinOp {
115 Add,
117 Sub,
119 Mul,
121}
122
123#[derive(Clone, Copy, Debug, PartialEq, Eq)]
126pub enum BoolOp {
127 And,
129 Or,
131 Xor,
133}
134
135#[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 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 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 pub fn f64(v: Vec<f64>) -> Column {
182 Column::F64(Buffer::from_vec(v))
183 }
184
185 pub fn f32(v: Vec<f32>) -> Column {
187 Column::F32(Buffer::from_vec(v))
188 }
189
190 pub fn i32(v: Vec<i32>) -> Column {
192 Column::I32(Buffer::from_vec(v), Validity::dense())
193 }
194
195 pub fn bool(v: Vec<bool>) -> Column {
197 Column::Bool(Buffer::from_vec(v), Validity::dense())
198 }
199
200 pub fn i64(v: Vec<i64>) -> Column {
202 Column::I64(Buffer::from_vec(v), Validity::dense())
203 }
204
205 pub fn i64_with(v: Vec<i64>, validity: Validity) -> Column {
207 Column::I64(Buffer::from_vec(v), validity)
208 }
209
210 pub fn i32_with(v: Vec<i32>, validity: Validity) -> Column {
212 Column::I32(Buffer::from_vec(v), validity)
213 }
214
215 pub fn bool_with(v: Vec<bool>, validity: Validity) -> Column {
217 Column::Bool(Buffer::from_vec(v), validity)
218 }
219
220 pub fn str(v: Vec<String>) -> Column {
222 Column::Str(StrBuffer::from_vec(v), Validity::dense())
223 }
224
225 pub fn str_with(v: Vec<String>, validity: Validity) -> Column {
227 Column::Str(StrBuffer::from_vec(v), validity)
228 }
229
230 pub fn datetime(v: Vec<i64>) -> Column {
232 Column::Datetime(Buffer::from_vec(v))
233 }
234
235 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 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 pub fn is_empty(&self) -> bool {
266 self.len() == 0
267 }
268
269 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 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 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 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 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 other => other,
336 }
337 }
338
339 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 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 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 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 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 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 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 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 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 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) }
465 })
466 .collect(),
467 }
468 }
469
470 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 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 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) }
511 })
512 .collect(),
513 }
514 }
515
516 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 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
541fn 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
555fn widen_i64<T: Copy + Into<i64>>(v: &[T]) -> Vec<i64> {
557 v.iter().map(|&x| x.into()).collect()
558}
559
560fn 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
570fn 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
585fn 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
603fn 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
617fn 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
629fn 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
640fn 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
658fn 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
676fn 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
692fn 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
701fn 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, };
712 let q = x.div_euclid(factor);
713 let r = x.rem_euclid(factor); let half = factor / 2; let up = r > half || (r == half && q.rem_euclid(2) != 0); if up { q + 1 } else { q }.wrapping_mul(factor)
717}
718
719fn 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
746fn 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
759fn 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
772fn 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
784fn 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
814fn bool_running(v: &[bool], or: bool) -> Vec<bool> {
816 let mut acc = !or; v.iter()
818 .map(|&b| {
819 acc = if or { acc || b } else { acc && b };
820 acc
821 })
822 .collect()
823}
824
825fn 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
844fn 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;