polars-ops 0.54.4

More operations on Polars data structures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

use arrow::array::PrimitiveArray;
use arrow::compute::utils::combine_validities_and3;
use polars_core::prelude::*;
use polars_core::utils::align_chunks_ternary;
use polars_core::with_match_physical_numeric_polars_type;

// (a * b) + c
fn fma_arr<T: NumericNative>(
    a: &PrimitiveArray<T>,
    b: &PrimitiveArray<T>,
    c: &PrimitiveArray<T>,
) -> PrimitiveArray<T> {
    assert_eq!(a.len(), b.len());
    let validity = combine_validities_and3(a.validity(), b.validity(), c.validity());
    let a = a.values().as_slice();
    let b = b.values().as_slice();
    let c = c.values().as_slice();

    assert_eq!(a.len(), b.len());
    assert_eq!(b.len(), c.len());
    let out = a
        .iter()
        .zip(b.iter())
        .zip(c.iter())
        .map(|((a, b), c)| *a * *b + *c)
        .collect::<Vec<_>>();
    PrimitiveArray::from_data_default(out.into(), validity)
}

fn fma_ca<T: PolarsNumericType>(
    a: &ChunkedArray<T>,
    b: &ChunkedArray<T>,
    c: &ChunkedArray<T>,
) -> ChunkedArray<T> {
    let (a, b, c) = align_chunks_ternary(a, b, c);
    let chunks = a
        .downcast_iter()
        .zip(b.downcast_iter())
        .zip(c.downcast_iter())
        .map(|((a, b), c)| fma_arr(a, b, c));
    ChunkedArray::from_chunk_iter(a.name().clone(), chunks)
}

pub fn fma_columns(a: &Column, b: &Column, c: &Column) -> Column {
    if a.len() == b.len() && a.len() == c.len() {
        with_match_physical_numeric_polars_type!(a.dtype(), |$T| {
            let a: &ChunkedArray<$T> = a.as_materialized_series().as_ref().as_ref().as_ref();
            let b: &ChunkedArray<$T> = b.as_materialized_series().as_ref().as_ref().as_ref();
            let c: &ChunkedArray<$T> = c.as_materialized_series().as_ref().as_ref().as_ref();

            fma_ca(a, b, c).into_column()
        })
    } else {
        (&(a * b).unwrap() + c).unwrap()
    }
}

// a - (b * c)
fn fsm_arr<T: NumericNative>(
    a: &PrimitiveArray<T>,
    b: &PrimitiveArray<T>,
    c: &PrimitiveArray<T>,
) -> PrimitiveArray<T> {
    assert_eq!(a.len(), b.len());
    let validity = combine_validities_and3(a.validity(), b.validity(), c.validity());
    let a = a.values().as_slice();
    let b = b.values().as_slice();
    let c = c.values().as_slice();

    assert_eq!(a.len(), b.len());
    assert_eq!(b.len(), c.len());
    let out = a
        .iter()
        .zip(b.iter())
        .zip(c.iter())
        .map(|((a, b), c)| *a - (*b * *c))
        .collect::<Vec<_>>();
    PrimitiveArray::from_data_default(out.into(), validity)
}

fn fsm_ca<T: PolarsNumericType>(
    a: &ChunkedArray<T>,
    b: &ChunkedArray<T>,
    c: &ChunkedArray<T>,
) -> ChunkedArray<T> {
    let (a, b, c) = align_chunks_ternary(a, b, c);
    let chunks = a
        .downcast_iter()
        .zip(b.downcast_iter())
        .zip(c.downcast_iter())
        .map(|((a, b), c)| fsm_arr(a, b, c));
    ChunkedArray::from_chunk_iter(a.name().clone(), chunks)
}

pub fn fsm_columns(a: &Column, b: &Column, c: &Column) -> Column {
    if a.len() == b.len() && a.len() == c.len() {
        with_match_physical_numeric_polars_type!(a.dtype(), |$T| {
            let a: &ChunkedArray<$T> = a.as_materialized_series().as_ref().as_ref().as_ref();
            let b: &ChunkedArray<$T> = b.as_materialized_series().as_ref().as_ref().as_ref();
            let c: &ChunkedArray<$T> = c.as_materialized_series().as_ref().as_ref().as_ref();

            fsm_ca(a, b, c).into_column()
        })
    } else {
        (a - &(b * c).unwrap()).unwrap()
    }
}

fn fms_arr<T: NumericNative>(
    a: &PrimitiveArray<T>,
    b: &PrimitiveArray<T>,
    c: &PrimitiveArray<T>,
) -> PrimitiveArray<T> {
    assert_eq!(a.len(), b.len());
    let validity = combine_validities_and3(a.validity(), b.validity(), c.validity());
    let a = a.values().as_slice();
    let b = b.values().as_slice();
    let c = c.values().as_slice();

    assert_eq!(a.len(), b.len());
    assert_eq!(b.len(), c.len());
    let out = a
        .iter()
        .zip(b.iter())
        .zip(c.iter())
        .map(|((a, b), c)| (*a * *b) - *c)
        .collect::<Vec<_>>();
    PrimitiveArray::from_data_default(out.into(), validity)
}

fn fms_ca<T: PolarsNumericType>(
    a: &ChunkedArray<T>,
    b: &ChunkedArray<T>,
    c: &ChunkedArray<T>,
) -> ChunkedArray<T> {
    let (a, b, c) = align_chunks_ternary(a, b, c);
    let chunks = a
        .downcast_iter()
        .zip(b.downcast_iter())
        .zip(c.downcast_iter())
        .map(|((a, b), c)| fms_arr(a, b, c));
    ChunkedArray::from_chunk_iter(a.name().clone(), chunks)
}

pub fn fms_columns(a: &Column, b: &Column, c: &Column) -> Column {
    if a.len() == b.len() && a.len() == c.len() {
        with_match_physical_numeric_polars_type!(a.dtype(), |$T| {
            let a: &ChunkedArray<$T> = a.as_materialized_series().as_ref().as_ref().as_ref();
            let b: &ChunkedArray<$T> = b.as_materialized_series().as_ref().as_ref().as_ref();
            let c: &ChunkedArray<$T> = c.as_materialized_series().as_ref().as_ref().as_ref();

            fms_ca(a, b, c).into_column()
        })
    } else {
        (&(a * b).unwrap() - c).unwrap()
    }
}