polars_core/chunked_array/ops/
nesting_utils.rs

1use polars_arrow::array::{Array, IntoBoxedArray};
2use polars_compute::find_validity_mismatch::find_validity_mismatch;
3use polars_utils::IdxSize;
4
5use super::ListChunked;
6use crate::chunked_array::flags::StatisticsFlags;
7use crate::prelude::{ChunkedArray, FalseT, PolarsDataType};
8use crate::series::Series;
9use crate::series::implementations::null::NullChunked;
10use crate::utils::align_chunks_binary_ca_series;
11
12/// Utility methods for dealing with nested chunked arrays.
13pub trait ChunkNestingUtils: Sized {
14    /// Propagate nulls of nested datatype to all levels of nesting.
15    fn propagate_nulls(&self) -> Option<Self>;
16
17    /// Trim all lists of unused start and end elements recursively.
18    fn trim_lists_to_normalized_offsets(&self) -> Option<Self>;
19
20    /// Find the indices of the values where the validity mismatches.
21    ///
22    /// This is done recursively.
23    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>);
24}
25
26impl ChunkNestingUtils for ListChunked {
27    fn propagate_nulls(&self) -> Option<Self> {
28        use polars_compute::propagate_nulls::propagate_nulls_list;
29
30        let flags = self.get_flags();
31
32        if flags.has_propagated_nulls() {
33            return None;
34        }
35
36        if !self.inner_dtype().is_nested() && !self.has_nulls() {
37            self.flags
38                .set(flags | StatisticsFlags::HAS_PROPAGATED_NULLS);
39            return None;
40        }
41
42        let mut chunks = Vec::new();
43        for (i, chunk) in self.downcast_iter().enumerate() {
44            if let Some(propagated_chunk) = propagate_nulls_list(chunk) {
45                chunks.reserve(self.chunks.len());
46                chunks.extend(self.chunks[..i].iter().cloned());
47                chunks.push(propagated_chunk.into_boxed());
48                break;
49            }
50        }
51
52        // If we found a chunk that needs propagating, create a new ListChunked
53        let out = if chunks.is_empty() {
54            None
55        } else {
56            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
57                match propagate_nulls_list(chunk) {
58                    None => chunk.to_boxed(),
59                    Some(chunk) => chunk.into_boxed(),
60                }
61            }));
62
63            // SAFETY: The length and null_count should remain the same.
64            Some(unsafe {
65                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
66            })
67        };
68
69        finish_propagate_nulls(out, self, flags)
70    }
71
72    fn trim_lists_to_normalized_offsets(&self) -> Option<Self> {
73        use polars_compute::trim_lists_to_normalized_offsets::trim_lists_to_normalized_offsets_list;
74
75        let flags = self.get_flags();
76
77        if flags.has_trimmed_lists_to_normalized_offsets() {
78            return None;
79        }
80
81        let mut chunks = Vec::new();
82        for (i, chunk) in self.downcast_iter().enumerate() {
83            if let Some(trimmed) = trim_lists_to_normalized_offsets_list(chunk) {
84                chunks.reserve(self.chunks.len());
85                chunks.extend(self.chunks[..i].iter().cloned());
86                chunks.push(trimmed.into_boxed());
87                break;
88            }
89        }
90
91        // If we found a chunk that needs compacting, create a new ArrayChunked
92        if !chunks.is_empty() {
93            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
94                match trim_lists_to_normalized_offsets_list(chunk) {
95                    Some(chunk) => chunk.into_boxed(),
96                    None => chunk.to_boxed(),
97                }
98            }));
99
100            // SAFETY: The length and null_count should remain the same.
101            let mut ca = unsafe {
102                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
103            };
104
105            ca.set_flags(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
106            return Some(ca);
107        }
108
109        self.flags
110            .set(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
111        None
112    }
113
114    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
115        let (slf, other) = align_chunks_binary_ca_series(self, other);
116        let mut offset: IdxSize = 0;
117        for (l, r) in slf.downcast_iter().zip(other.chunks()) {
118            let start_length = idxs.len();
119            find_validity_mismatch(l, r.as_ref(), idxs);
120            for idx in idxs[start_length..].iter_mut() {
121                *idx += offset;
122            }
123            offset += l.len() as IdxSize;
124        }
125    }
126}
127
128#[cfg(feature = "dtype-array")]
129impl ChunkNestingUtils for super::ArrayChunked {
130    fn propagate_nulls(&self) -> Option<Self> {
131        use polars_compute::propagate_nulls::propagate_nulls_fsl;
132
133        let flags = self.get_flags();
134
135        if flags.has_propagated_nulls() {
136            return None;
137        }
138
139        if !self.inner_dtype().is_nested() && !self.has_nulls() {
140            self.flags
141                .set(flags | StatisticsFlags::HAS_PROPAGATED_NULLS);
142            return None;
143        }
144
145        let mut chunks = Vec::new();
146        for (i, chunk) in self.downcast_iter().enumerate() {
147            if let Some(propagated_chunk) = propagate_nulls_fsl(chunk) {
148                chunks.reserve(self.chunks.len());
149                chunks.extend(self.chunks[..i].iter().cloned());
150                chunks.push(propagated_chunk.into_boxed());
151                break;
152            }
153        }
154
155        let out = if chunks.is_empty() {
156            None
157        } else {
158            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
159                match propagate_nulls_fsl(chunk) {
160                    None => chunk.to_boxed(),
161                    Some(chunk) => chunk.into_boxed(),
162                }
163            }));
164
165            // SAFETY: The length and null_count should remain the same.
166            Some(unsafe {
167                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
168            })
169        };
170
171        finish_propagate_nulls(out, self, flags)
172    }
173
174    fn trim_lists_to_normalized_offsets(&self) -> Option<Self> {
175        use polars_compute::trim_lists_to_normalized_offsets::trim_lists_to_normalized_offsets_fsl;
176
177        let flags = self.get_flags();
178
179        if flags.has_trimmed_lists_to_normalized_offsets()
180            || !self.inner_dtype().contains_list_recursive()
181        {
182            return None;
183        }
184
185        let mut chunks = Vec::new();
186        for (i, chunk) in self.downcast_iter().enumerate() {
187            if let Some(trimmed) = trim_lists_to_normalized_offsets_fsl(chunk) {
188                chunks.reserve(self.chunks.len());
189                chunks.extend(self.chunks[..i].iter().cloned());
190                chunks.push(trimmed.into_boxed());
191                break;
192            }
193        }
194
195        // If we found a chunk that needs compacting, create a new ArrayChunked
196        if !chunks.is_empty() {
197            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
198                match trim_lists_to_normalized_offsets_fsl(chunk) {
199                    Some(chunk) => chunk.into_boxed(),
200                    None => chunk.to_boxed(),
201                }
202            }));
203
204            // SAFETY: The length and null_count should remain the same.
205            let mut ca = unsafe {
206                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
207            };
208            ca.set_flags(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
209            return Some(ca);
210        }
211
212        self.flags
213            .set(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
214        None
215    }
216
217    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
218        let (slf, other) = align_chunks_binary_ca_series(self, other);
219        let mut offset: IdxSize = 0;
220        for (l, r) in slf.downcast_iter().zip(other.chunks()) {
221            let start_length = idxs.len();
222            find_validity_mismatch(l, r.as_ref(), idxs);
223            for idx in idxs[start_length..].iter_mut() {
224                *idx += offset;
225            }
226            offset += l.len() as IdxSize;
227        }
228    }
229}
230
231#[cfg(feature = "dtype-struct")]
232impl ChunkNestingUtils for super::StructChunked {
233    fn propagate_nulls(&self) -> Option<Self> {
234        use polars_compute::propagate_nulls::propagate_nulls_struct;
235
236        let flags = self.get_flags();
237
238        if flags.has_propagated_nulls() {
239            return None;
240        }
241
242        if self.struct_fields().iter().all(|f| !f.dtype().is_nested()) && !self.has_nulls() {
243            self.flags
244                .set(flags | StatisticsFlags::HAS_PROPAGATED_NULLS);
245            return None;
246        }
247
248        let mut chunks = Vec::new();
249        for (i, chunk) in self.downcast_iter().enumerate() {
250            if let Some(propagated_chunk) = propagate_nulls_struct(chunk) {
251                chunks.reserve(self.chunks.len());
252                chunks.extend(self.chunks[..i].iter().cloned());
253                chunks.push(propagated_chunk.into_boxed());
254                break;
255            }
256        }
257
258        let out = if chunks.is_empty() {
259            None
260        } else {
261            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
262                match propagate_nulls_struct(chunk) {
263                    None => chunk.to_boxed(),
264                    Some(chunk) => chunk.into_boxed(),
265                }
266            }));
267
268            // SAFETY: The length and null_count should remain the same.
269            Some(unsafe {
270                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
271            })
272        };
273
274        finish_propagate_nulls(out, self, flags)
275    }
276
277    fn trim_lists_to_normalized_offsets(&self) -> Option<Self> {
278        use polars_compute::trim_lists_to_normalized_offsets::trim_lists_to_normalized_offsets_struct;
279
280        let flags = self.get_flags();
281
282        if flags.has_trimmed_lists_to_normalized_offsets()
283            || !self
284                .struct_fields()
285                .iter()
286                .any(|f| f.dtype().contains_list_recursive())
287        {
288            return None;
289        }
290
291        let mut chunks = Vec::new();
292        for (i, chunk) in self.downcast_iter().enumerate() {
293            if let Some(trimmed) = trim_lists_to_normalized_offsets_struct(chunk) {
294                chunks.reserve(self.chunks.len());
295                chunks.extend(self.chunks[..i].iter().cloned());
296                chunks.push(trimmed.into_boxed());
297                break;
298            }
299        }
300
301        // If we found a chunk that needs compacting, create a new ArrayChunked
302        if !chunks.is_empty() {
303            chunks.extend(self.downcast_iter().skip(chunks.len()).map(|chunk| {
304                match trim_lists_to_normalized_offsets_struct(chunk) {
305                    Some(chunk) => chunk.into_boxed(),
306                    None => chunk.to_boxed(),
307                }
308            }));
309
310            // SAFETY: The length and null_count should remain the same.
311            let mut ca = unsafe {
312                Self::new_with_dims(self.field.clone(), chunks, self.length, self.null_count)
313            };
314            ca.set_flags(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
315            return Some(ca);
316        }
317
318        self.flags
319            .set(flags | StatisticsFlags::HAS_TRIMMED_LISTS_TO_NORMALIZED_OFFSETS);
320        None
321    }
322
323    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
324        let (slf, other) = align_chunks_binary_ca_series(self, other);
325        let mut offset: IdxSize = 0;
326        for (l, r) in slf.downcast_iter().zip(other.chunks()) {
327            let start_length = idxs.len();
328            find_validity_mismatch(l, r.as_ref(), idxs);
329            for idx in idxs[start_length..].iter_mut() {
330                *idx += offset;
331            }
332            offset += l.len() as IdxSize;
333        }
334    }
335}
336
337/// Mark `out` or `orig` as having propagated nulls.
338fn finish_propagate_nulls<T: PolarsDataType>(
339    out: Option<ChunkedArray<T>>,
340    orig: &ChunkedArray<T>,
341    flags: StatisticsFlags,
342) -> Option<ChunkedArray<T>> {
343    match out {
344        Some(mut ca) => {
345            ca.set_flags(flags | StatisticsFlags::HAS_PROPAGATED_NULLS);
346            Some(ca)
347        },
348        None => {
349            orig.flags
350                .set(flags | StatisticsFlags::HAS_PROPAGATED_NULLS);
351            None
352        },
353    }
354}
355
356impl<T: PolarsDataType<IsNested = FalseT>> ChunkNestingUtils for ChunkedArray<T> {
357    fn propagate_nulls(&self) -> Option<Self> {
358        None
359    }
360
361    fn trim_lists_to_normalized_offsets(&self) -> Option<Self> {
362        None
363    }
364
365    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
366        let slf_nc = self.null_count();
367        let other_nc = other.null_count();
368
369        // Fast path for non-nested datatypes.
370        if slf_nc == other_nc && (slf_nc == 0 || slf_nc == self.len()) {
371            return;
372        }
373
374        let (slf, other) = align_chunks_binary_ca_series(self, other);
375        let mut offset: IdxSize = 0;
376        for (l, r) in slf.downcast_iter().zip(other.chunks()) {
377            let start_length = idxs.len();
378            find_validity_mismatch(l, r.as_ref(), idxs);
379            for idx in idxs[start_length..].iter_mut() {
380                *idx += offset;
381            }
382            offset += l.len() as IdxSize;
383        }
384    }
385}
386
387impl ChunkNestingUtils for NullChunked {
388    fn propagate_nulls(&self) -> Option<Self> {
389        None
390    }
391
392    fn trim_lists_to_normalized_offsets(&self) -> Option<Self> {
393        None
394    }
395
396    fn find_validity_mismatch(&self, other: &Series, idxs: &mut Vec<IdxSize>) {
397        let other_nc = other.null_count();
398
399        // Fast path for non-nested datatypes.
400        if other_nc == self.len() {
401            return;
402        }
403
404        match other.rechunk_validity() {
405            None => idxs.extend(0..self.len() as IdxSize),
406            Some(v) => idxs.extend(v.true_idx_iter().map(|v| v as IdxSize)),
407        }
408    }
409}
polars_core/chunked_array/ops/nesting_utils.rs

polars_core/chunked_array/ops/
nesting_utils.rs
