polars_core/chunked_array/
from_iterator_par.rs

1//! Parallel iterator collection into [`ChunkedArray<T>`]
2//!
3//! Two strategies:
4//!
5//! - `collect_into_linked_list*` — for iterators of unknown length. Folds into
6//!   one accumulator per rayon task and concatenates, so the resulting chunk
7//!   count equals the task count and varies with thread count and stealing.
8//!   `optional_rechunk` is the only backstop.
9//!
10//! - `collect_*_par` — for callers with O(1) random access and a known output
11//!   length. Writes into a single preallocated buffer, so the result is always
12//!   one chunk regardless of how rayon split the work. Preferable.
13
14use std::collections::LinkedList;
15use std::sync::Mutex;
16
17use arrow::bitmap::Bitmap;
18use arrow::pushable::{NoOption, Pushable};
19use rayon::prelude::*;
20
21use super::from_iterator::PolarsAsRef;
22use crate::chunked_array::builder::get_list_builder;
23use crate::datatypes::BooleanChunked;
24use crate::prelude::*;
25use crate::utils::NoNull;
26use crate::utils::flatten::flatten_par;
27
28/// FromParallelIterator trait
29// Code taken from https://docs.rs/rayon/1.3.1/src/rayon/iter/extend.rs.html#356-366
30fn vec_push<T>(mut vec: Vec<T>, elem: T) -> Vec<T> {
31    vec.push(elem);
32    vec
33}
34
35fn as_list<T>(item: T) -> LinkedList<T> {
36    let mut list = LinkedList::new();
37    list.push_back(item);
38    list
39}
40
41fn list_append<T>(mut list1: LinkedList<T>, mut list2: LinkedList<T>) -> LinkedList<T> {
42    list1.append(&mut list2);
43    list1
44}
45
46fn collect_into_linked_list_vec<I>(par_iter: I) -> LinkedList<Vec<I::Item>>
47where
48    I: IntoParallelIterator,
49{
50    let it = par_iter.into_par_iter();
51    // be careful optimizing allocations. Its hard to figure out the size
52    // needed
53    // https://github.com/pola-rs/polars/issues/1562
54    it.fold(Vec::new, vec_push)
55        .map(as_list)
56        .reduce(LinkedList::new, list_append)
57}
58
59fn collect_into_linked_list<I, P, F>(par_iter: I, identity: F) -> LinkedList<P::Freeze>
60where
61    I: IntoParallelIterator,
62    P: Pushable<I::Item> + Send + Sync,
63    F: Fn() -> P + Sync + Send,
64    P::Freeze: Send,
65{
66    let it = par_iter.into_par_iter();
67    it.fold(identity, |mut v, item| {
68        v.push(item);
69        v
70    })
71    // The freeze on this line, ensures the null count is done in parallel
72    .map(|p| as_list(p.freeze()))
73    .reduce(LinkedList::new, list_append)
74}
75
76fn get_capacity_from_par_results<T>(ll: &LinkedList<Vec<T>>) -> usize {
77    ll.iter().map(|list| list.len()).sum()
78}
79
80impl<T> FromParallelIterator<T::Native> for NoNull<ChunkedArray<T>>
81where
82    T: PolarsNumericType,
83{
84    fn from_par_iter<I: IntoParallelIterator<Item = T::Native>>(iter: I) -> Self {
85        // Get linkedlist filled with different vec result from different threads
86        let vectors = collect_into_linked_list_vec(iter);
87        let vectors = vectors.into_iter().collect::<Vec<_>>();
88        let values = flatten_par(&vectors);
89        NoNull::new(ChunkedArray::new_vec(PlSmallStr::EMPTY, values))
90    }
91}
92
93impl<T> FromParallelIterator<Option<T::Native>> for ChunkedArray<T>
94where
95    T: PolarsNumericType,
96{
97    fn from_par_iter<I: IntoParallelIterator<Item = Option<T::Native>>>(iter: I) -> Self {
98        let chunks = collect_into_linked_list(iter, MutablePrimitiveArray::new);
99        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
100    }
101}
102
103impl FromParallelIterator<bool> for BooleanChunked {
104    fn from_par_iter<I: IntoParallelIterator<Item = bool>>(iter: I) -> Self {
105        let chunks = collect_into_linked_list(iter, MutableBooleanArray::new);
106        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
107    }
108}
109
110impl FromParallelIterator<Option<bool>> for BooleanChunked {
111    fn from_par_iter<I: IntoParallelIterator<Item = Option<bool>>>(iter: I) -> Self {
112        let chunks = collect_into_linked_list(iter, MutableBooleanArray::new);
113        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
114    }
115}
116
117impl<Ptr> FromParallelIterator<Ptr> for StringChunked
118where
119    Ptr: PolarsAsRef<str> + Send + Sync + NoOption,
120{
121    fn from_par_iter<I: IntoParallelIterator<Item = Ptr>>(iter: I) -> Self {
122        let chunks = collect_into_linked_list(iter, MutableBinaryViewArray::new);
123        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
124    }
125}
126
127impl<Ptr> FromParallelIterator<Ptr> for BinaryChunked
128where
129    Ptr: PolarsAsRef<[u8]> + Send + Sync + NoOption,
130{
131    fn from_par_iter<I: IntoParallelIterator<Item = Ptr>>(iter: I) -> Self {
132        let chunks = collect_into_linked_list(iter, MutableBinaryViewArray::new);
133        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
134    }
135}
136
137impl<Ptr> FromParallelIterator<Option<Ptr>> for StringChunked
138where
139    Ptr: AsRef<str> + Send + Sync,
140{
141    fn from_par_iter<I: IntoParallelIterator<Item = Option<Ptr>>>(iter: I) -> Self {
142        let chunks = collect_into_linked_list(iter, MutableBinaryViewArray::new);
143        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
144    }
145}
146
147impl<Ptr> FromParallelIterator<Option<Ptr>> for BinaryChunked
148where
149    Ptr: AsRef<[u8]> + Send + Sync,
150{
151    fn from_par_iter<I: IntoParallelIterator<Item = Option<Ptr>>>(iter: I) -> Self {
152        let chunks = collect_into_linked_list(iter, MutableBinaryViewArray::new);
153        Self::from_chunk_iter(PlSmallStr::EMPTY, chunks).optional_rechunk()
154    }
155}
156
157pub trait FromParIterWithDtype<K> {
158    fn from_par_iter_with_dtype<I>(iter: I, name: PlSmallStr, dtype: DataType) -> Self
159    where
160        I: IntoParallelIterator<Item = K>,
161        Self: Sized;
162}
163
164fn get_value_cap(vectors: &LinkedList<Vec<Option<Series>>>) -> usize {
165    vectors
166        .iter()
167        .map(|list| {
168            list.iter()
169                .map(|opt_s| opt_s.as_ref().map(|s| s.len()).unwrap_or(0))
170                .sum::<usize>()
171        })
172        .sum::<usize>()
173}
174
175fn get_dtype(vectors: &LinkedList<Vec<Option<Series>>>) -> DataType {
176    for v in vectors {
177        for s in v.iter().flatten() {
178            let dtype = s.dtype();
179            if !matches!(dtype, DataType::Null) {
180                return dtype.clone();
181            }
182        }
183    }
184    DataType::Null
185}
186
187fn materialize_list(
188    name: PlSmallStr,
189    vectors: &LinkedList<Vec<Option<Series>>>,
190    dtype: DataType,
191    value_capacity: usize,
192    list_capacity: usize,
193) -> PolarsResult<ListChunked> {
194    let mut builder = get_list_builder(&dtype, value_capacity, list_capacity, name);
195    for v in vectors {
196        for val in v {
197            builder.append_opt_series(val.as_ref())?;
198        }
199    }
200    Ok(builder.finish())
201}
202
203impl FromParallelIterator<Option<Series>> for ListChunked {
204    fn from_par_iter<I>(par_iter: I) -> Self
205    where
206        I: IntoParallelIterator<Item = Option<Series>>,
207    {
208        list_from_par_iter(par_iter, PlSmallStr::EMPTY).unwrap()
209    }
210}
211
212pub fn list_from_par_iter<I>(par_iter: I, name: PlSmallStr) -> PolarsResult<ListChunked>
213where
214    I: IntoParallelIterator<Item = Option<Series>>,
215{
216    let vectors = collect_into_linked_list_vec(par_iter);
217
218    let list_capacity: usize = get_capacity_from_par_results(&vectors);
219    let value_capacity = get_value_cap(&vectors);
220    let dtype = get_dtype(&vectors);
221    if let DataType::Null = dtype {
222        Ok(ListChunked::full_null_with_dtype(
223            name,
224            list_capacity,
225            &DataType::Null,
226        ))
227    } else {
228        materialize_list(name, &vectors, dtype, value_capacity, list_capacity)
229    }
230}
231
232pub fn try_list_from_par_iter<I>(par_iter: I, name: PlSmallStr) -> PolarsResult<ListChunked>
233where
234    I: IntoParallelIterator<Item = PolarsResult<Option<Series>>>,
235{
236    fn ok<T, E>(saved: &Mutex<Option<E>>) -> impl Fn(Result<T, E>) -> Option<T> + '_ {
237        move |item| match item {
238            Ok(item) => Some(item),
239            Err(error) => {
240                // We don't need a blocking `lock()`, as anybody
241                // else holding the lock will also be writing
242                // `Some(error)`, and then ours is irrelevant.
243                if let Ok(mut guard) = saved.try_lock() {
244                    if guard.is_none() {
245                        *guard = Some(error);
246                    }
247                }
248                None
249            },
250        }
251    }
252
253    let saved_error = Mutex::new(None);
254    let iter = par_iter.into_par_iter().map(ok(&saved_error)).while_some();
255
256    let collection = list_from_par_iter(iter, name)?;
257
258    match saved_error.into_inner().unwrap() {
259        Some(error) => Err(error),
260        None => Ok(collection),
261    }
262}
263
264impl FromParIterWithDtype<Option<Series>> for ListChunked {
265    fn from_par_iter_with_dtype<I>(iter: I, name: PlSmallStr, dtype: DataType) -> Self
266    where
267        I: IntoParallelIterator<Item = Option<Series>>,
268        Self: Sized,
269    {
270        let vectors = collect_into_linked_list_vec(iter);
271
272        let list_capacity: usize = get_capacity_from_par_results(&vectors);
273        let value_capacity = get_value_cap(&vectors);
274        if let DataType::List(dtype) = dtype {
275            materialize_list(name, &vectors, *dtype, value_capacity, list_capacity).unwrap()
276        } else {
277            panic!("expected list dtype")
278        }
279    }
280}
281
282pub trait ChunkedCollectParIterExt: ParallelIterator {
283    fn collect_ca_with_dtype<B: FromParIterWithDtype<Self::Item>>(
284        self,
285        name: PlSmallStr,
286        dtype: DataType,
287    ) -> B
288    where
289        Self: Sized,
290    {
291        B::from_par_iter_with_dtype(self, name, dtype)
292    }
293}
294
295impl<I: ParallelIterator> ChunkedCollectParIterExt for I {}
296
297// Adapted from rayon
298impl<C, T, E> FromParIterWithDtype<Result<T, E>> for Result<C, E>
299where
300    C: FromParIterWithDtype<T>,
301    T: Send,
302    E: Send,
303{
304    fn from_par_iter_with_dtype<I>(par_iter: I, name: PlSmallStr, dtype: DataType) -> Self
305    where
306        I: IntoParallelIterator<Item = Result<T, E>>,
307    {
308        fn ok<T, E>(saved: &Mutex<Option<E>>) -> impl Fn(Result<T, E>) -> Option<T> + '_ {
309            move |item| match item {
310                Ok(item) => Some(item),
311                Err(error) => {
312                    // We don't need a blocking `lock()`, as anybody
313                    // else holding the lock will also be writing
314                    // `Some(error)`, and then ours is irrelevant.
315                    if let Ok(mut guard) = saved.try_lock() {
316                        if guard.is_none() {
317                            *guard = Some(error);
318                        }
319                    }
320                    None
321                },
322            }
323        }
324
325        let saved_error = Mutex::new(None);
326        let iter = par_iter.into_par_iter().map(ok(&saved_error)).while_some();
327
328        let collection = C::from_par_iter_with_dtype(iter, name, dtype);
329
330        match saved_error.into_inner().unwrap() {
331            Some(error) => Err(error),
332            None => Ok(collection),
333        }
334    }
335}
336
337// Collect directly into a single chunk. For directly addressable fixed-size types only.
338pub(crate) fn collect_bool_par<F>(len: usize, f: F) -> BooleanChunked
339where
340    F: Fn(usize) -> bool + Send + Sync,
341{
342    let n_bytes = len.div_ceil(8);
343    let mut values: Vec<u8> = Vec::with_capacity(n_bytes);
344
345    (0..n_bytes)
346        .into_par_iter()
347        .map(|b| {
348            let lo = b * 8;
349            let hi = (lo + 8).min(len);
350            let mut v = 0u8;
351            for (bit, g) in (lo..hi).enumerate() {
352                v |= (f(g) as u8) << bit;
353            }
354            v
355        })
356        .collect_into_vec(&mut values);
357
358    BooleanChunked::from_bitmap(PlSmallStr::EMPTY, Bitmap::from_u8_vec(values, len))
359}
360
361pub(crate) fn collect_bool_opt_par<F>(len: usize, f: F) -> BooleanChunked
362where
363    F: Fn(usize) -> Option<bool> + Send + Sync,
364{
365    let n_bytes = len.div_ceil(8);
366    let mut values: Vec<u8> = Vec::with_capacity(n_bytes);
367    let mut validity: Vec<u8> = Vec::with_capacity(n_bytes);
368
369    (0..n_bytes)
370        .into_par_iter()
371        .map(|b| {
372            let lo = b * 8;
373            let hi = (lo + 8).min(len);
374            let (mut v, mut m) = (0u8, 0u8);
375            for (bit, g) in (lo..hi).enumerate() {
376                if let Some(x) = f(g) {
377                    m |= 1 << bit;
378                    v |= (x as u8) << bit;
379                }
380            }
381            (v, m)
382        })
383        .unzip_into_vecs(&mut values, &mut validity);
384
385    let values = Bitmap::from_u8_vec(values, len);
386    let validity = Bitmap::from_u8_vec(validity, len);
387    let validity = (validity.unset_bits() > 0).then_some(validity);
388    BooleanChunked::with_chunk(
389        PlSmallStr::EMPTY,
390        BooleanArray::new(ArrowDataType::Boolean, values, validity),
391    )
392}
393
394// TODO: Placeholder for the NoNull helpers, future PR.
395#[allow(unused)]
396pub(crate) fn collect_primitive_par<T, F>(len: usize, f: F) -> ChunkedArray<T>
397where
398    T: PolarsNumericType,
399    F: Fn(usize) -> T::Native + Send + Sync,
400{
401    let mut values: Vec<T::Native> = Vec::new();
402    (0..len)
403        .into_par_iter()
404        .map(f)
405        .collect_into_vec(&mut values);
406    ChunkedArray::from_vec(PlSmallStr::EMPTY, values)
407}
408
409pub(crate) fn collect_primitive_opt_par<T, F>(len: usize, f: F) -> ChunkedArray<T>
410where
411    T: PolarsNumericType,
412    F: Fn(usize) -> Option<T::Native> + Send + Sync,
413{
414    let mut values: Vec<T::Native> = vec![T::Native::default(); len];
415    let mut validity: Vec<u8> = Vec::with_capacity(len.div_ceil(8));
416
417    values
418        .par_chunks_mut(8)
419        .enumerate()
420        .map(|(b, out)| {
421            let lo = b * 8;
422            let mut m = 0u8;
423            for (bit, slot) in out.iter_mut().enumerate() {
424                if let Some(x) = f(lo + bit) {
425                    *slot = x;
426                    m |= 1 << bit;
427                }
428            }
429            m
430        })
431        .collect_into_vec(&mut validity);
432
433    let validity = Bitmap::from_u8_vec(validity, len);
434    let validity = (validity.unset_bits() > 0).then_some(validity);
435    ChunkedArray::from_vec_validity(PlSmallStr::EMPTY, values, validity)
436}
polars_core/chunked_array/from_iterator_par.rs

polars_core/chunked_array/
from_iterator_par.rs
