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

#[cfg(feature = "object")]
use crate::chunked_array::object::ObjectArray;
use crate::prelude::*;
use crate::utils::index_to_chunked_index;
use arrow::array::*;
use std::marker::PhantomData;

pub struct Chunks<'a, T> {
    chunks: &'a [ArrayRef],
    phantom: PhantomData<T>,
}

impl<'a, T> Chunks<'a, T> {
    fn new(chunks: &'a [ArrayRef]) -> Self {
        Chunks {
            chunks,
            phantom: PhantomData,
        }
    }

    pub fn get(&self, index: usize) -> Option<&'a T> {
        self.chunks.get(index).map(|arr| {
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const T) }
        })
    }

    pub unsafe fn get_unchecked(&self, index: usize) -> &'a T {
        let arr = self.chunks.get_unchecked(index);
        let arr = &**arr;
        &*(arr as *const dyn Array as *const T)
    }

    pub fn len(&self) -> usize {
        self.chunks.len()
    }
}

impl<T> ChunkedArray<T>
where
    T: PolarsNumericType,
{
    pub fn downcast_iter(
        &self,
    ) -> impl Iterator<Item = &PrimitiveArray<T::Native>> + DoubleEndedIterator {
        self.chunks.iter().map(|arr| {
            // Safety:
            // This should be the array type in PolarsNumericType
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const PrimitiveArray<T::Native>) }
        })
    }
    pub fn downcast_chunks(&self) -> Chunks<'_, PrimitiveArray<T::Native>> {
        Chunks::new(&self.chunks)
    }

    /// Get the index of the chunk and the index of the value in that chunk
    #[inline]
    pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
        if self.chunks.len() == 1 {
            return (0, index);
        }
        index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
    }
}

impl BooleanChunked {
    pub fn downcast_iter(&self) -> impl Iterator<Item = &BooleanArray> + DoubleEndedIterator {
        self.chunks.iter().map(|arr| {
            // Safety:
            // This should be the array type in BooleanChunked
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const BooleanArray) }
        })
    }
    pub fn downcast_chunks(&self) -> Chunks<'_, BooleanArray> {
        Chunks::new(&self.chunks)
    }

    #[inline]
    pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
        if self.chunks.len() == 1 {
            return (0, index);
        }
        index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
    }
}

impl Utf8Chunked {
    pub fn downcast_iter(&self) -> impl Iterator<Item = &Utf8Array<i64>> + DoubleEndedIterator {
        // Safety:
        // This is the array type that must be in a Utf8Chunked
        self.chunks.iter().map(|arr| {
            // Safety:
            // This should be the array type in Utf8Chunked
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const Utf8Array<i64>) }
        })
    }
    pub fn downcast_chunks(&self) -> Chunks<'_, Utf8Array<i64>> {
        Chunks::new(&self.chunks)
    }

    #[inline]
    pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
        if self.chunks.len() == 1 {
            return (0, index);
        }
        index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
    }
}

impl ListChunked {
    pub fn downcast_iter(&self) -> impl Iterator<Item = &ListArray<i64>> + DoubleEndedIterator {
        // Safety:
        // This is the array type that must be in a ListChunked
        self.chunks.iter().map(|arr| {
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const ListArray<i64>) }
        })
    }
    pub fn downcast_chunks(&self) -> Chunks<'_, ListArray<i64>> {
        Chunks::new(&self.chunks)
    }

    #[inline]
    pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
        if self.chunks.len() == 1 {
            return (0, index);
        }
        index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
    }
}

#[cfg(feature = "object")]
impl<'a, T> ObjectChunked<T>
where
    T: PolarsObject,
{
    pub fn downcast_iter(&self) -> impl Iterator<Item = &ObjectArray<T>> + DoubleEndedIterator {
        self.chunks.iter().map(|arr| {
            let arr = &**arr;
            unsafe { &*(arr as *const dyn Array as *const ObjectArray<T>) }
        })
    }
    pub fn downcast_chunks(&self) -> Chunks<'_, ObjectArray<T>> {
        Chunks::new(&self.chunks)
    }

    #[inline]
    pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
        if self.chunks.len() == 1 {
            return (0, index);
        }
        index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
    }
}