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)]
80pub enum Scalar {
81 F64(f64),
83 F32(f32),
85 I64(i64),
87 I32(i32),
89 Bool(bool),
91}
92
93#[derive(Clone, Copy, Debug, PartialEq, Eq)]
96pub enum BinOp {
97 Add,
99 Sub,
101 Mul,
103}
104
105#[derive(Clone, Copy, Debug, PartialEq, Eq)]
108pub enum BoolOp {
109 And,
111 Or,
113 Xor,
115}
116
117#[derive(Clone, Copy, Debug, PartialEq, Eq)]
119pub enum CmpOp {
120 Eq,
121 Ne,
122 Lt,
123 Le,
124 Gt,
125 Ge,
126}
127
128impl CmpOp {
129 fn matches(self, o: Ordering) -> bool {
132 match self {
133 CmpOp::Eq => o == Ordering::Equal,
134 CmpOp::Ne => o != Ordering::Equal,
135 CmpOp::Lt => o == Ordering::Less,
136 CmpOp::Le => o != Ordering::Greater,
137 CmpOp::Gt => o == Ordering::Greater,
138 CmpOp::Ge => o != Ordering::Less,
139 }
140 }
141
142 fn matches_f64(self, a: f64, b: f64) -> bool {
145 match self {
146 CmpOp::Eq => a == b,
147 CmpOp::Ne => a != b,
148 CmpOp::Lt => a < b,
149 CmpOp::Le => a <= b,
150 CmpOp::Gt => a > b,
151 CmpOp::Ge => a >= b,
152 }
153 }
154}
155
156mod cast;
157mod ops;
158mod order;
159mod transform;
160
161impl Column {
162 pub fn f64(v: Vec<f64>) -> Column {
164 Column::F64(Buffer::from_vec(v))
165 }
166
167 pub fn f32(v: Vec<f32>) -> Column {
169 Column::F32(Buffer::from_vec(v))
170 }
171
172 pub fn i32(v: Vec<i32>) -> Column {
174 Column::I32(Buffer::from_vec(v), Validity::dense())
175 }
176
177 pub fn bool(v: Vec<bool>) -> Column {
179 Column::Bool(Buffer::from_vec(v), Validity::dense())
180 }
181
182 pub fn i64(v: Vec<i64>) -> Column {
184 Column::I64(Buffer::from_vec(v), Validity::dense())
185 }
186
187 pub fn i64_with(v: Vec<i64>, validity: Validity) -> Column {
189 Column::I64(Buffer::from_vec(v), validity)
190 }
191
192 pub fn i32_with(v: Vec<i32>, validity: Validity) -> Column {
194 Column::I32(Buffer::from_vec(v), validity)
195 }
196
197 pub fn bool_with(v: Vec<bool>, validity: Validity) -> Column {
199 Column::Bool(Buffer::from_vec(v), validity)
200 }
201
202 pub fn str(v: Vec<String>) -> Column {
204 Column::Str(StrBuffer::from_vec(v), Validity::dense())
205 }
206
207 pub fn str_with(v: Vec<String>, validity: Validity) -> Column {
209 Column::Str(StrBuffer::from_vec(v), validity)
210 }
211
212 pub fn datetime(v: Vec<i64>) -> Column {
214 Column::Datetime(Buffer::from_vec(v))
215 }
216
217 pub fn na_of(dtype: DType, len: usize) -> Column {
221 let all_na = || Validity::from_valid_iter(len, std::iter::repeat_n(false, len));
222 match dtype {
223 DType::F64 => Column::f64(vec![f64::NAN; len]),
224 DType::F32 => Column::f32(vec![f32::NAN; len]),
225 DType::I64 => Column::i64_with(vec![0; len], all_na()),
226 DType::I32 => Column::i32_with(vec![0; len], all_na()),
227 DType::Bool => Column::bool_with(vec![false; len], all_na()),
228 DType::Utf8 => Column::str_with(vec![String::new(); len], all_na()),
229 DType::Datetime => Column::datetime(vec![i64::MIN; len]),
230 }
231 }
232
233 pub fn len(&self) -> usize {
235 match self {
236 Column::F64(v) => v.len(),
237 Column::F32(v) => v.len(),
238 Column::Bool(v, _) => v.len(),
239 Column::I64(v, _) => v.len(),
240 Column::I32(v, _) => v.len(),
241 Column::Str(v, _) => v.len(),
242 Column::Datetime(v) => v.len(),
243 }
244 }
245
246 pub fn is_empty(&self) -> bool {
248 self.len() == 0
249 }
250
251 pub fn dtype(&self) -> DType {
253 match self {
254 Column::F64(_) => DType::F64,
255 Column::F32(_) => DType::F32,
256 Column::Bool(_, _) => DType::Bool,
257 Column::I64(_, _) => DType::I64,
258 Column::I32(_, _) => DType::I32,
259 Column::Str(_, _) => DType::Utf8,
260 Column::Datetime(_) => DType::Datetime,
261 }
262 }
263
264 pub fn is_valid(&self, i: usize) -> bool {
268 match self {
269 Column::F64(v) => !v[i].is_nan(),
270 Column::F32(v) => !v[i].is_nan(),
271 Column::Bool(_, val)
272 | Column::I64(_, val)
273 | Column::I32(_, val)
274 | Column::Str(_, val) => val.is_valid(i),
275 Column::Datetime(v) => v[i] != i64::MIN,
276 }
277 }
278
279 pub fn null_count(&self) -> usize {
281 match self {
282 Column::F64(v) => v.iter().filter(|x| x.is_nan()).count(),
283 Column::F32(v) => v.iter().filter(|x| x.is_nan()).count(),
284 Column::Bool(_, val)
285 | Column::I64(_, val)
286 | Column::I32(_, val)
287 | Column::Str(_, val) => val.null_count(),
288 Column::Datetime(v) => v.iter().filter(|&&x| x == i64::MIN).count(),
289 }
290 }
291
292 fn with_validity(self, validity: Validity) -> Column {
297 match self {
298 Column::I64(v, _) => Column::I64(v, validity),
299 Column::I32(v, _) => Column::I32(v, validity),
300 other => other,
303 }
304 }
305
306 fn validity(&self) -> Option<&Validity> {
309 match self {
310 Column::Bool(_, val)
311 | Column::I64(_, val)
312 | Column::I32(_, val)
313 | Column::Str(_, val) => Some(val),
314 _ => None,
315 }
316 }
317
318 pub fn as_f64(&self) -> Option<&[f64]> {
320 if let Column::F64(v) = self {
321 Some(v.as_slice())
322 } else {
323 None
324 }
325 }
326
327 pub fn as_bool(&self) -> Option<&[bool]> {
329 if let Column::Bool(v, _) = self {
330 Some(v.as_slice())
331 } else {
332 None
333 }
334 }
335
336 pub fn as_i64(&self) -> Option<&[i64]> {
338 if let Column::I64(v, _) = self {
339 Some(v.as_slice())
340 } else {
341 None
342 }
343 }
344
345 pub fn str_at(&self, i: usize) -> Option<&str> {
347 if let Column::Str(v, _) = self {
348 Some(v.get(i))
349 } else {
350 None
351 }
352 }
353
354 pub fn as_datetime(&self) -> Option<&[i64]> {
356 if let Column::Datetime(v) = self {
357 Some(v.as_slice())
358 } else {
359 None
360 }
361 }
362
363 pub fn to_f64_vec(&self) -> Vec<f64> {
367 match self {
368 Column::F64(v) => v.to_vec(),
369 Column::F32(v) => v.iter().map(|&x| x as f64).collect(),
370 Column::Bool(v, val) => mask_f64(v.iter().map(|&b| if b { 1.0 } else { 0.0 }), val),
371 Column::I64(v, val) => mask_f64(v.iter().map(|&i| i as f64), val),
372 Column::I32(v, val) => mask_f64(v.iter().map(|&i| i as f64), val),
373 Column::Str(v, _) => vec![f64::NAN; v.len()],
374 Column::Datetime(v) => v
375 .iter()
376 .map(|&i| if i == i64::MIN { f64::NAN } else { i as f64 })
377 .collect(),
378 }
379 }
380
381 pub fn get_f64(&self, i: usize) -> f64 {
385 match self {
386 Column::F64(v) => v[i],
387 Column::F32(v) => v[i] as f64,
388 Column::Bool(v, _) => {
389 if v[i] {
390 1.0
391 } else {
392 0.0
393 }
394 }
395 Column::I64(v, _) => v[i] as f64,
396 Column::I32(v, _) => v[i] as f64,
397 Column::Str(_, _) => f64::NAN,
398 Column::Datetime(v) => v[i] as f64,
399 }
400 }
401
402 fn as_bool_vec(&self) -> Result<Vec<bool>> {
404 match self {
405 Column::Bool(v, _) => Ok(v.to_vec()),
406 other => Err(VolasError::DType(format!(
407 "expected a bool column, got {}",
408 other.dtype()
409 ))),
410 }
411 }
412
413 fn as_i64_vec(&self) -> Result<Vec<i64>> {
418 match self {
419 Column::I64(v, _) => Ok(v.to_vec()),
420 _ => self
421 .to_f64_vec()
422 .iter()
423 .enumerate()
424 .map(|(i, &x)| {
425 if self.is_valid(i) {
426 i64::try_from_f64(x).ok_or_else(|| {
427 VolasError::DType(format!("value {x} does not fit int64"))
428 })
429 } else {
430 Ok(0) }
432 })
433 .collect(),
434 }
435 }
436
437 fn as_str_vec(&self) -> Result<Vec<String>> {
440 match self {
441 Column::Str(v, _) => Ok(v.to_vec()),
442 other => Err(VolasError::DType(format!(
443 "cannot select a {} column as str",
444 other.dtype()
445 ))),
446 }
447 }
448
449 fn as_datetime_vec(&self) -> Result<Vec<i64>> {
451 match self {
452 Column::Datetime(v) => Ok(v.to_vec()),
453 other => Err(VolasError::DType(format!(
454 "cannot select a {} column as datetime",
455 other.dtype()
456 ))),
457 }
458 }
459
460 fn as_i32_vec(&self) -> Result<Vec<i32>> {
463 match self {
464 Column::I32(v, _) => Ok(v.to_vec()),
465 Column::Bool(v, _) => Ok(v.iter().map(|&b| b as i32).collect()),
466 _ => self
467 .to_f64_vec()
468 .iter()
469 .enumerate()
470 .map(|(i, &x)| {
471 if self.is_valid(i) {
472 i32::try_from_f64(x).ok_or_else(|| {
473 VolasError::DType(format!("value {x} does not fit int32"))
474 })
475 } else {
476 Ok(0) }
478 })
479 .collect(),
480 }
481 }
482
483 pub fn to_f32_vec(&self) -> Vec<f32> {
485 match self {
486 Column::F32(v) => v.to_vec(),
487 _ => self.to_f64_vec().iter().map(|&x| x as f32).collect(),
488 }
489 }
490
491 fn to_string_vec(&self) -> Vec<String> {
493 match self {
494 Column::Str(v, _) => v.to_vec(),
495 Column::F64(v) => v.iter().map(|x| x.to_string()).collect(),
496 Column::F32(v) => v.iter().map(|x| x.to_string()).collect(),
497 Column::I64(v, _) => v.iter().map(|x| x.to_string()).collect(),
498 Column::I32(v, _) => v.iter().map(|x| x.to_string()).collect(),
499 Column::Bool(v, _) => v
500 .iter()
501 .map(|&b| if b { "True" } else { "False" }.to_string())
502 .collect(),
503 Column::Datetime(v) => v.iter().map(|&ns| datetime::format_ns(ns)).collect(),
504 }
505 }
506}
507
508fn mask_f64(vals: impl Iterator<Item = f64>, validity: &Validity) -> Vec<f64> {
513 if validity.has_nulls() {
514 vals.enumerate()
515 .map(|(i, x)| if validity.is_valid(i) { x } else { f64::NAN })
516 .collect()
517 } else {
518 vals.collect()
519 }
520}
521
522fn widen_i64<T: Copy + Into<i64>>(v: &[T]) -> Vec<i64> {
524 v.iter().map(|&x| x.into()).collect()
525}
526
527fn from_option_bools(n: usize, it: impl Iterator<Item = Option<bool>>) -> Column {
529 let (mut values, mut valid) = (Vec::with_capacity(n), Vec::with_capacity(n));
530 for o in it {
531 values.push(o.unwrap_or(false));
532 valid.push(o.is_some());
533 }
534 Column::bool_with(values, Validity::from_valid_iter(n, valid))
535}
536
537fn shift_fill<T: Copy>(v: &[T], n: isize, fill: T) -> Vec<T> {
540 let len = v.len();
541 let mut out = vec![fill; len];
542 if n >= 0 {
543 let n = (n as usize).min(len);
544 out[n..].copy_from_slice(&v[..len - n]);
545 } else {
546 let n = ((-n) as usize).min(len);
547 out[..len - n].copy_from_slice(&v[n..]);
548 }
549 out
550}
551
552fn diff_kernel<T: Copy + std::ops::Sub<Output = T>>(v: &[T], n: isize, missing: T) -> Vec<T> {
554 let len = v.len();
555 let mut out = vec![missing; len];
556 if n >= 0 {
557 let k = n as usize;
558 for i in k..len {
559 out[i] = v[i] - v[i - k];
560 }
561 } else {
562 let k = (-n) as usize;
563 for i in 0..len.saturating_sub(k) {
564 out[i] = v[i] - v[i + k];
565 }
566 }
567 out
568}
569
570fn append_validity(av: &mut Validity, alen: usize, bv: &Validity, blen: usize) {
574 if !av.has_nulls() && !bv.has_nulls() {
575 return;
576 }
577 let flags: Vec<bool> = (0..alen)
578 .map(|i| av.is_valid(i))
579 .chain((0..blen).map(|i| bv.is_valid(i)))
580 .collect();
581 *av = Validity::from_valid_iter(alen + blen, flags);
582}
583
584fn sum_valid<T: Numeric>(v: &[T], val: &Validity) -> T {
587 if !val.has_nulls() {
588 return stats::sum(v);
589 }
590 v.iter()
591 .enumerate()
592 .filter(|(i, _)| val.is_valid(*i))
593 .fold(T::ZERO, |a, (_, &x)| a.wrapping_add(x))
594}
595
596fn prod_valid<T: Numeric>(v: &[T], val: &Validity) -> T {
598 if !val.has_nulls() {
599 return stats::prod(v);
600 }
601 v.iter()
602 .enumerate()
603 .filter(|(i, _)| val.is_valid(*i))
604 .fold(T::ONE, |a, (_, &x)| a.wrapping_mul(x))
605}
606
607fn extreme_valid<T: Numeric>(v: &[T], val: &Validity, want_max: bool) -> Option<T> {
609 if !val.has_nulls() {
610 return stats::extreme(v, want_max);
611 }
612 let mut it = v
613 .iter()
614 .enumerate()
615 .filter(|(i, _)| val.is_valid(*i))
616 .map(|(_, &x)| x);
617 let first = it.next()?;
618 Some(if want_max {
619 it.fold(first, |a, x| if x > a { x } else { a })
620 } else {
621 it.fold(first, |a, x| if x < a { x } else { a })
622 })
623}
624
625fn cum_valid<T: Copy>(v: &[T], val: &Validity, placeholder: T, op: impl Fn(T, T) -> T) -> Vec<T> {
628 let mut acc: Option<T> = None;
629 v.iter()
630 .enumerate()
631 .map(|(i, &x)| {
632 if val.is_valid(i) {
633 let next = acc.map_or(x, |a| op(a, x));
634 acc = Some(next);
635 next
636 } else {
637 placeholder
638 }
639 })
640 .collect()
641}
642
643fn cum<T: Numeric>(
646 v: &[T],
647 val: &Validity,
648 dense: fn(&[T]) -> Vec<T>,
649 op: fn(T, T) -> T,
650) -> Column {
651 let out = if val.has_nulls() {
652 cum_valid(v, val, T::ZERO, op)
653 } else {
654 dense(v)
655 };
656 T::into_column(out).with_validity(val.clone())
657}
658
659fn round_f64(x: f64, decimals: i32) -> f64 {
661 if x.is_nan() {
662 return x;
663 }
664 let f = 10f64.powi(decimals);
665 (x * f).round_ties_even() / f
666}
667
668fn round_i64(x: i64, decimals: i32) -> i64 {
672 if decimals >= 0 {
673 return x;
674 }
675 let factor = match 10i64.checked_pow(decimals.unsigned_abs()) {
676 Some(f) => f,
677 None => return 0, };
679 let q = x.div_euclid(factor);
680 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)
684}
685
686fn clip_vec<T: Numeric>(v: &[T], lo: Option<f64>, hi: Option<f64>) -> Vec<T> {
690 let lo = lo.and_then(T::try_from_f64);
691 let hi = hi.and_then(T::try_from_f64);
692 v.iter()
693 .map(|&x| {
694 if x.is_missing() {
695 return x;
696 }
697 let mut y = x;
698 if let Some(l) = lo {
699 if y < l {
700 y = l;
701 }
702 }
703 if let Some(h) = hi {
704 if y > h {
705 y = h;
706 }
707 }
708 y
709 })
710 .collect()
711}
712
713fn cmp_typed(n: usize, op: CmpOp, key: impl Fn(usize) -> Option<Ordering>) -> Vec<bool> {
718 (0..n)
719 .map(|i| match key(i) {
720 Some(ord) => op.matches(ord),
721 None => op == CmpOp::Ne,
722 })
723 .collect()
724}
725
726fn ifloordiv(a: i64, b: i64) -> i64 {
730 let q = a.wrapping_div(b);
731 let r = a.wrapping_rem(b);
732 if r != 0 && (r < 0) != (b < 0) {
733 q - 1
734 } else {
735 q
736 }
737}
738
739fn float_key(x: f64) -> Option<u64> {
742 if x.is_nan() {
743 None
744 } else if x == 0.0 {
745 Some(0)
746 } else {
747 Some(x.to_bits())
748 }
749}
750
751fn group_by<K: Eq + std::hash::Hash>(
754 len: usize,
755 key: impl Fn(usize) -> Option<K>,
756) -> Vec<(usize, usize, bool)> {
757 let mut order: Vec<(usize, usize, bool)> = Vec::new();
758 let mut seen: HashMap<K, usize> = HashMap::new();
759 let mut na: Option<usize> = None;
760 for i in 0..len {
761 match key(i) {
762 None => match na {
763 Some(g) => order[g].1 += 1,
764 None => {
765 na = Some(order.len());
766 order.push((i, 1, true));
767 }
768 },
769 Some(k) => match seen.get(&k) {
770 Some(&g) => order[g].1 += 1,
771 None => {
772 seen.insert(k, order.len());
773 order.push((i, 1, false));
774 }
775 },
776 }
777 }
778 order
779}
780
781fn bool_running(v: &[bool], or: bool) -> Vec<bool> {
783 let mut acc = !or; v.iter()
785 .map(|&b| {
786 acc = if or { acc || b } else { acc && b };
787 acc
788 })
789 .collect()
790}
791
792fn clip_bool(v: &[bool], lo: Option<f64>, hi: Option<f64>) -> Vec<bool> {
796 let force_true = lo.is_some_and(|x| x != 0.0);
797 let force_false = hi == Some(0.0);
798 v.iter()
799 .map(|&b| {
800 if force_false {
801 false
802 } else if force_true {
803 true
804 } else {
805 b
806 }
807 })
808 .collect()
809}
810
811fn binary_kernel<T: Numeric>(a: &[T], b: &[T], op: BinOp) -> Vec<T> {
813 a.iter()
814 .zip(b)
815 .map(|(&x, &y)| match op {
816 BinOp::Add => x.wrapping_add(y),
817 BinOp::Sub => x.wrapping_sub(y),
818 BinOp::Mul => x.wrapping_mul(y),
819 })
820 .collect()
821}
822
823#[cfg(test)]
824mod tests;
825#[cfg(test)]
826mod tests_na;