polars_arrow/compute/
concatenate.rs

1//r Contains the concatenate kernel
2use polars_error::{polars_bail, PolarsResult};
3
4use crate::array::growable::make_growable;
5use crate::array::Array;
6use crate::bitmap::{Bitmap, BitmapBuilder};
7
8/// Concatenate multiple [`Array`] of the same type into a single [`Array`].
9pub fn concatenate(arrays: &[&dyn Array]) -> PolarsResult<Box<dyn Array>> {
10    if arrays.is_empty() {
11        polars_bail!(InvalidOperation: "concat requires input of at least one array")
12    }
13
14    if arrays
15        .iter()
16        .any(|array| array.dtype() != arrays[0].dtype())
17    {
18        polars_bail!(InvalidOperation: "It is not possible to concatenate arrays of different data types.")
19    }
20
21    let lengths = arrays.iter().map(|array| array.len()).collect::<Vec<_>>();
22    let capacity = lengths.iter().sum();
23
24    let mut mutable = make_growable(arrays, false, capacity);
25
26    for (i, len) in lengths.iter().enumerate() {
27        // SAFETY: len is correct
28        unsafe { mutable.extend(i, 0, *len) }
29    }
30
31    Ok(mutable.as_box())
32}
33
34/// Concatenate the validities of multiple [Array]s into a single Bitmap.
35pub fn concatenate_validities(arrays: &[&dyn Array]) -> Option<Bitmap> {
36    let null_count: usize = arrays.iter().map(|a| a.null_count()).sum();
37    if null_count == 0 {
38        return None;
39    }
40
41    let total_size: usize = arrays.iter().map(|a| a.len()).sum();
42    let mut bitmap = BitmapBuilder::with_capacity(total_size);
43    for arr in arrays {
44        if arr.null_count() == arr.len() {
45            bitmap.extend_constant(arr.len(), false);
46        } else if arr.null_count() == 0 {
47            bitmap.extend_constant(arr.len(), true);
48        } else {
49            bitmap.extend_from_bitmap(arr.validity().unwrap());
50        }
51    }
52    bitmap.into_opt_validity()
53}
polars_arrow/compute/concatenate.rs

polars_arrow/compute/
concatenate.rs
