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
158
159
160
161
162
163
164
165
166
167
168
169
170

//! Traits that indicate the allowed arguments in a ChunkedArray::take operation.
use crate::prelude::*;
use arrow::array::{Array, UInt32Array};

// Utility traits
pub trait TakeIterator: Iterator<Item = usize> {
    fn check_bounds(&self, bound: usize) -> Result<()>;
}
pub trait TakeIteratorNulls: Iterator<Item = Option<usize>> {
    fn check_bounds(&self, bound: usize) -> Result<()>;
}

// Implement for the ref as well
impl TakeIterator for &mut dyn TakeIterator {
    fn check_bounds(&self, bound: usize) -> Result<()> {
        (**self).check_bounds(bound)
    }
}
impl TakeIteratorNulls for &mut dyn TakeIteratorNulls {
    fn check_bounds(&self, bound: usize) -> Result<()> {
        (**self).check_bounds(bound)
    }
}

// Clonable iterators may implement the traits above
impl<I> TakeIterator for I
where
    I: Iterator<Item = usize> + Clone + Sized,
{
    fn check_bounds(&self, bound: usize) -> Result<()> {
        // clone so that the iterator can be used again.
        let iter = self.clone();
        let mut inbounds = true;

        for i in iter {
            if i >= bound {
                inbounds = false;
                break;
            }
        }
        if inbounds {
            Ok(())
        } else {
            Err(PolarsError::OutOfBounds(
                "take indices are out of bounds".into(),
            ))
        }
    }
}
impl<I> TakeIteratorNulls for I
where
    I: Iterator<Item = Option<usize>> + Clone + Sized,
{
    fn check_bounds(&self, bound: usize) -> Result<()> {
        // clone so that the iterator can be used again.
        let iter = self.clone();
        let mut inbounds = true;

        for i in iter.flatten() {
            if i >= bound {
                inbounds = false;
                break;
            }
        }
        if inbounds {
            Ok(())
        } else {
            Err(PolarsError::OutOfBounds(
                "take indices are out of bounds".into(),
            ))
        }
    }
}

/// One of the three arguments allowed in unchecked_take
pub enum TakeIdx<'a, I, INulls>
where
    I: TakeIterator,
    INulls: TakeIteratorNulls,
{
    Array(&'a UInt32Array),
    Iter(I),
    // will return a null where None
    IterNulls(INulls),
}

impl<'a, I, INulls> TakeIdx<'a, I, INulls>
where
    I: TakeIterator,
    INulls: TakeIteratorNulls,
{
    pub(crate) fn check_bounds(&self, bound: usize) -> Result<()> {
        match self {
            TakeIdx::Iter(i) => i.check_bounds(bound),
            TakeIdx::IterNulls(i) => i.check_bounds(bound),
            TakeIdx::Array(arr) => {
                let mut inbounds = true;
                let len = bound as u32;
                if arr.null_count() == 0 {
                    for &i in arr.values().as_slice() {
                        if i >= len {
                            inbounds = false;
                            break;
                        }
                    }
                } else {
                    for opt_v in *arr {
                        match opt_v {
                            Some(&v) if v >= len => {
                                inbounds = false;
                                break;
                            }
                            _ => {}
                        }
                    }
                }
                if inbounds {
                    Ok(())
                } else {
                    Err(PolarsError::OutOfBounds(
                        "take indices are out of bounds".into(),
                    ))
                }
            }
        }
    }
}

/// Dummy type, we need to instantiate all generic types, so we fill one with a dummy.
pub type Dummy<T> = std::iter::Once<T>;

// Below the conversions from
// * UInt32Chunked
// * Iterator<Item=usize>
// * Iterator<Item=Option<usize>>
//
// To the checked and unchecked TakeIdx enums

// Unchecked conversions

/// Conversion from UInt32Chunked to Unchecked TakeIdx
impl<'a> From<&'a UInt32Chunked> for TakeIdx<'a, Dummy<usize>, Dummy<Option<usize>>> {
    fn from(ca: &'a UInt32Chunked) -> Self {
        if ca.chunks.len() == 1 {
            TakeIdx::Array(ca.downcast_iter().next().unwrap())
        } else {
            panic!("implementation error, should be transformed to an iterator by the caller")
        }
    }
}

/// Conversion from Iterator<Item=usize> to Unchecked TakeIdx
impl<'a, I> From<I> for TakeIdx<'a, I, Dummy<Option<usize>>>
where
    I: TakeIterator,
{
    fn from(iter: I) -> Self {
        TakeIdx::Iter(iter)
    }
}

/// Conversion from Iterator<Item=Option<usize>> to Unchecked TakeIdx
impl<'a, I> From<I> for TakeIdx<'a, Dummy<usize>, I>
where
    I: TakeIteratorNulls,
{
    fn from(iter: I) -> Self {
        TakeIdx::IterNulls(iter)
    }
}