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

use crate::prelude::*;
use arrow::bitmap::MutableBitmap;
use num::{FromPrimitive, Zero};
use std::ops::{Add, Div, Mul, Sub};

fn linear_itp<T>(low: T, step: T, diff: T, steps_n: T) -> T
where
    T: Sub<Output = T> + Mul<Output = T> + Add<Output = T> + Div<Output = T>,
{
    low + step * diff / steps_n
}

impl<T> Interpolate for ChunkedArray<T>
where
    T: PolarsNumericType,
{
    fn interpolate(&self) -> Self {
        // This implementation differs from pandas as that boundary None's are not removed
        // this prevents a lot of errors due to expressions leading to different lengths
        if self.null_count() == 0 || self.null_count() == self.len() {
            return self.clone();
        }

        // we first find the first and last so that we can set the null buffer
        let mut first = 0;
        let mut last = self.len();
        // find first non None
        for i in 0..self.len() {
            // Safety: we just bound checked
            if unsafe { self.get_unchecked(i).is_some() } {
                first = i;
                break;
            }
        }

        // find last non None
        for i in (0..self.len()).rev() {
            if unsafe { self.get_unchecked(i).is_some() } {
                last = i + 1;
                break;
            }
        }

        // fill av with first
        let mut av = AlignedVec::with_capacity(self.len());
        let mut iter = self.into_iter();
        for _ in 0..first {
            av.push(Zero::zero())
        }

        let mut low_val = None;
        loop {
            let next = iter.next();
            match next {
                Some(Some(v)) => {
                    av.push(v);
                    low_val = Some(v);
                }
                Some(None) => {
                    match low_val {
                        // not a non-null value encountered yet
                        // so we skip
                        None => continue,
                        Some(low) => {
                            let mut steps = 1u32;
                            loop {
                                steps += 1;
                                match iter.next() {
                                    // end of iterator, break
                                    None => break,
                                    // another null
                                    Some(None) => {}
                                    Some(Some(high)) => {
                                        let diff = high - low;
                                        let steps_n = T::Native::from_u32(steps).unwrap();
                                        for step_i in 1..steps {
                                            let step_i = T::Native::from_u32(step_i).unwrap();
                                            let v = linear_itp(low, step_i, diff, steps_n);
                                            av.push(v)
                                        }
                                        av.push(high);
                                        break;
                                    }
                                }
                            }
                        }
                    }
                }
                None => {
                    break;
                }
            }
        }
        if first != 0 || last != self.len() {
            let mut validity = MutableBitmap::with_capacity(self.len());
            validity.extend_constant(self.len(), true);

            for i in 0..first {
                validity.set(i, false);
            }

            for i in last..self.len() {
                validity.set(i, false);
                av.push(Zero::zero())
            }

            let array = PrimitiveArray::from_data(
                T::get_dtype().to_arrow(),
                av.into(),
                Some(validity.into()),
            );
            Self::new_from_chunks(self.name(), vec![Arc::new(array)])
        } else {
            Self::new_from_aligned_vec(self.name(), av)
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_interpolate() {
        let ca = UInt32Chunked::new_from_opt_slice("", &[Some(1), None, None, Some(4), Some(5)]);
        let out = ca.interpolate();
        assert_eq!(
            Vec::from(&out),
            &[Some(1), Some(2), Some(3), Some(4), Some(5)]
        );

        let ca =
            UInt32Chunked::new_from_opt_slice("", &[None, Some(1), None, None, Some(4), Some(5)]);
        let out = ca.interpolate();
        assert_eq!(
            Vec::from(&out),
            &[None, Some(1), Some(2), Some(3), Some(4), Some(5)]
        );

        let ca = UInt32Chunked::new_from_opt_slice(
            "",
            &[None, Some(1), None, None, Some(4), Some(5), None],
        );
        let out = ca.interpolate();
        assert_eq!(
            Vec::from(&out),
            &[None, Some(1), Some(2), Some(3), Some(4), Some(5), None]
        );
    }
}