polars_core/chunked_array/list/
mod.rs

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

polars_core/chunked_array/list/
mod.rs
