polars_core/chunked_array/array/
mod.rs

1//! Special fixed-size-list utility methods
2
3mod iterator;
4
5use std::borrow::Cow;
6
7use either::Either;
8
9use super::align_inner_chunks;
10use crate::prelude::*;
11
12impl ArrayChunked {
13    /// Get the inner data type of the fixed size list.
14    pub fn inner_dtype(&self) -> &DataType {
15        match self.dtype() {
16            DataType::Array(dt, _size) => dt.as_ref(),
17            _ => unreachable!(),
18        }
19    }
20
21    /// Relabel the inner dtype, checking its physical representation.
22    ///
23    /// # Safety
24    /// The values must be valid for `dtype`, see [`Self::to_logical`].
25    ///
26    /// # Panics
27    /// Panics if the physical representation of `dtype` differs the physical
28    /// representation of the existing inner `dtype`.
29    pub unsafe fn set_inner_dtype(&mut self, dtype: DataType) {
30        assert_eq!(dtype.to_physical(), self.inner_dtype().to_physical());
31        unsafe { self.to_logical(dtype) }
32    }
33
34    pub fn width(&self) -> usize {
35        match self.dtype() {
36            DataType::Array(_dt, size) => *size,
37            _ => unreachable!(),
38        }
39    }
40
41    /// Relabel the inner dtype without changing values.
42    ///
43    /// # Safety
44    /// Same requirements as [`ListChunked::to_logical`].
45    pub unsafe fn to_logical(&mut self, inner_dtype: DataType) {
46        debug_assert_eq!(inner_dtype.to_physical(), self.inner_dtype().to_physical());
47        let width = self.width();
48        let fld = Arc::make_mut(&mut self.field);
49        fld.set_dtype(DataType::Array(Box::new(inner_dtype), width))
50    }
51
52    /// Convert the datatype of the array into the physical datatype.
53    pub fn to_physical_repr(&self) -> Cow<'_, ArrayChunked> {
54        let Cow::Owned(physical_repr) = self.get_inner().to_physical_repr() else {
55            return Cow::Borrowed(self);
56        };
57
58        let chunk_len_validity_iter =
59            if physical_repr.chunks().len() == 1 && self.chunks().len() > 1 {
60                // Physical repr got rechunked, rechunk our validity as well.
61                Either::Left(std::iter::once((self.len(), self.rechunk_validity())))
62            } else {
63                // No rechunking, expect the same number of chunks.
64                assert_eq!(self.chunks().len(), physical_repr.chunks().len());
65                Either::Right(
66                    self.chunks()
67                        .iter()
68                        .map(|c| (c.len(), c.validity().cloned())),
69                )
70            };
71
72        let width = self.width();
73        let chunks: Vec<_> = chunk_len_validity_iter
74            .zip(physical_repr.into_chunks())
75            .map(|((len, validity), values)| {
76                FixedSizeListArray::new(
77                    ArrowDataType::FixedSizeList(
78                        Box::new(ArrowField::new(
79                            LIST_VALUES_NAME,
80                            values.dtype().clone(),
81                            true,
82                        )),
83                        width,
84                    ),
85                    len,
86                    values,
87                    validity,
88                )
89                .to_boxed()
90            })
91            .collect();
92
93        let name = self.name().clone();
94        let dtype = DataType::Array(Box::new(self.inner_dtype().to_physical()), width);
95        Cow::Owned(unsafe { ArrayChunked::from_chunks_and_dtype_unchecked(name, chunks, dtype) })
96    }
97
98    /// Convert a non-logical [`ArrayChunked`] back into a logical [`ArrayChunked`] without casting.
99    ///
100    /// # Safety
101    ///
102    /// This can lead to invalid memory access in downstream code.
103    pub unsafe fn from_physical_unchecked(&self, to_inner_dtype: DataType) -> PolarsResult<Self> {
104        debug_assert!(!self.inner_dtype().is_logical());
105
106        let chunks = self
107            .downcast_iter()
108            .map(|chunk| chunk.values())
109            .cloned()
110            .collect();
111
112        let inner = unsafe {
113            Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, chunks, self.inner_dtype())
114        };
115        let inner = unsafe { inner.from_physical_unchecked(&to_inner_dtype) }?;
116
117        let chunks: Vec<_> = self
118            .downcast_iter()
119            .zip(inner.into_chunks())
120            .map(|(chunk, values)| {
121                FixedSizeListArray::new(
122                    ArrowDataType::FixedSizeList(
123                        Box::new(ArrowField::new(
124                            LIST_VALUES_NAME,
125                            values.dtype().clone(),
126                            true,
127                        )),
128                        self.width(),
129                    ),
130                    chunk.len(),
131                    values,
132                    chunk.validity().cloned(),
133                )
134                .to_boxed()
135            })
136            .collect();
137
138        let name = self.name().clone();
139        let dtype = DataType::Array(Box::new(to_inner_dtype), self.width());
140        Ok(unsafe { Self::from_chunks_and_dtype_unchecked(name, chunks, dtype) })
141    }
142
143    /// Get the inner values as `Series`
144    pub fn get_inner(&self) -> Series {
145        let chunks: Vec<_> = self.downcast_iter().map(|c| c.values().clone()).collect();
146
147        // SAFETY: Data type of arrays matches because they are chunks from the same array.
148        unsafe {
149            Series::from_chunks_and_dtype_unchecked(self.name().clone(), chunks, self.inner_dtype())
150        }
151    }
152
153    /// The total number of inner values across all chunks, i.e. `len() * width()`
154    /// discounting sliced-away chunks.
155    pub fn inner_length(&self) -> usize {
156        self.downcast_iter().map(|c| c.values().len()).sum()
157    }
158
159    /// Rebuild the arrays around new inner values, reusing the widths and outer validity.
160    ///
161    /// `values` must have `inner_length()` elements; its chunks need not line up with
162    /// this array's, but nothing is copied when they do.
163    pub fn with_inner_values(&self, values: &Series) -> ArrayChunked {
164        if cfg!(debug_assertions) {
165            assert_eq!(values.len(), self.inner_length());
166        }
167
168        // Align the chunks of the array's inner values and the values series.
169        let values = align_inner_chunks(self.downcast_iter().map(|arr| arr.values().len()), values);
170        let values_dtype = values.dtype().clone();
171        let width = self.width();
172
173        let chunks = self
174            .downcast_iter()
175            .zip(values.into_chunks())
176            .map(|(ca_arr, v_arr)| {
177                debug_assert_eq!(ca_arr.values().len(), v_arr.len());
178                FixedSizeListArray::new(
179                    FixedSizeListArray::default_datatype(v_arr.dtype().clone(), width),
180                    ca_arr.len(),
181                    v_arr,
182                    ca_arr.validity().cloned(),
183                )
184                .to_boxed()
185            })
186            .collect::<Vec<_>>();
187
188        // SAFETY: the chunks' inner dtype is derived from `values`' own chunks.
189        unsafe {
190            ArrayChunked::from_chunks_and_dtype_unchecked(
191                self.name().clone(),
192                chunks,
193                DataType::Array(Box::new(values_dtype), width),
194            )
195        }
196    }
197
198    /// Ignore the list indices and apply `func` to the inner type as [`Series`].
199    pub fn apply_to_inner(
200        &self,
201        func: &dyn Fn(Series) -> PolarsResult<Series>,
202    ) -> PolarsResult<ArrayChunked> {
203        // Rechunk or the generated Series will have wrong length.
204        let ca = self.rechunk();
205        let arr = ca.downcast_as_array();
206
207        // SAFETY:
208        // Inner dtype is passed correctly
209        let elements = unsafe {
210            Series::from_chunks_and_dtype_unchecked(
211                self.name().clone(),
212                vec![arr.values().clone()],
213                ca.inner_dtype(),
214            )
215        };
216
217        let expected_len = elements.len();
218        let out: Series = func(elements)?;
219        polars_ensure!(
220            out.len() == expected_len,
221            ComputeError: "the function should apply element-wise, it removed elements instead"
222        );
223        let out = out.rechunk();
224        let values = out.chunks()[0].clone();
225
226        let inner_dtype = FixedSizeListArray::default_datatype(values.dtype().clone(), ca.width());
227        let arr = FixedSizeListArray::new(inner_dtype, arr.len(), values, arr.validity().cloned());
228
229        Ok(unsafe {
230            ArrayChunked::from_chunks_and_dtype_unchecked(
231                self.name().clone(),
232                vec![arr.into_boxed()],
233                DataType::Array(Box::new(out.dtype().clone()), self.width()),
234            )
235        })
236    }
237}
polars_core/chunked_array/array/mod.rs

polars_core/chunked_array/array/
mod.rs
