polars_ops/chunked_array/
scatter.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use arrow::array::{Array, PrimitiveArray};
3use polars_core::prelude::*;
4use polars_core::series::IsSorted;
5use polars_core::utils::arrow::bitmap::MutableBitmap;
6use polars_core::utils::arrow::types::NativeType;
7use polars_utils::index::check_bounds;
8
9pub trait ChunkedSet<T: Copy> {
10    /// Invariant for implementations: if the scatter() fails, typically because
11    /// of bad indexes, then self should remain unmodified.
12    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
13    where
14        V: IntoIterator<Item = Option<T>>;
15}
16fn check_sorted(idx: &[IdxSize]) -> PolarsResult<()> {
17    if idx.is_empty() {
18        return Ok(());
19    }
20    let mut sorted = true;
21    let mut previous = idx[0];
22    for &i in &idx[1..] {
23        if i < previous {
24            // we will not break here as that prevents SIMD
25            sorted = false;
26        }
27        previous = i;
28    }
29    polars_ensure!(sorted, ComputeError: "set indices must be sorted");
30    Ok(())
31}
32
33trait PolarsOpsNumericType: PolarsNumericType {}
34
35impl PolarsOpsNumericType for UInt8Type {}
36impl PolarsOpsNumericType for UInt16Type {}
37impl PolarsOpsNumericType for UInt32Type {}
38impl PolarsOpsNumericType for UInt64Type {}
39impl PolarsOpsNumericType for Int8Type {}
40impl PolarsOpsNumericType for Int16Type {}
41impl PolarsOpsNumericType for Int32Type {}
42impl PolarsOpsNumericType for Int64Type {}
43#[cfg(feature = "dtype-i128")]
44impl PolarsOpsNumericType for Int128Type {}
45impl PolarsOpsNumericType for Float32Type {}
46impl PolarsOpsNumericType for Float64Type {}
47
48unsafe fn scatter_impl<V, T: NativeType>(
49    new_values_slice: &mut [T],
50    set_values: V,
51    arr: &mut PrimitiveArray<T>,
52    idx: &[IdxSize],
53    len: usize,
54) where
55    V: IntoIterator<Item = Option<T>>,
56{
57    let mut values_iter = set_values.into_iter();
58
59    if arr.null_count() > 0 {
60        arr.apply_validity(|v| {
61            let mut mut_validity = v.make_mut();
62
63            for (idx, val) in idx.iter().zip(&mut values_iter) {
64                match val {
65                    Some(value) => {
66                        mut_validity.set_unchecked(*idx as usize, true);
67                        *new_values_slice.get_unchecked_mut(*idx as usize) = value
68                    },
69                    None => mut_validity.set_unchecked(*idx as usize, false),
70                }
71            }
72            mut_validity.into()
73        })
74    } else {
75        let mut null_idx = vec![];
76        for (idx, val) in idx.iter().zip(values_iter) {
77            match val {
78                Some(value) => *new_values_slice.get_unchecked_mut(*idx as usize) = value,
79                None => {
80                    null_idx.push(*idx);
81                },
82            }
83        }
84        // only make a validity bitmap when null values are set
85        if !null_idx.is_empty() {
86            let mut validity = MutableBitmap::with_capacity(len);
87            validity.extend_constant(len, true);
88            for idx in null_idx {
89                validity.set_unchecked(idx as usize, false)
90            }
91            arr.set_validity(Some(validity.into()))
92        }
93    }
94}
95
96impl<T: PolarsOpsNumericType> ChunkedSet<T::Native> for &mut ChunkedArray<T>
97where
98    ChunkedArray<T>: IntoSeries,
99{
100    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
101    where
102        V: IntoIterator<Item = Option<T::Native>>,
103    {
104        check_bounds(idx, self.len() as IdxSize)?;
105        let mut ca = std::mem::take(self);
106        ca.rechunk_mut();
107
108        // SAFETY:
109        // we will not modify the length
110        // and we unset the sorted flag.
111        ca.set_sorted_flag(IsSorted::Not);
112        let arr = unsafe { ca.downcast_iter_mut() }.next().unwrap();
113        let len = arr.len();
114
115        match arr.get_mut_values() {
116            Some(current_values) => {
117                let ptr = current_values.as_mut_ptr();
118
119                // reborrow because the bck does not allow it
120                let current_values = unsafe { &mut *std::slice::from_raw_parts_mut(ptr, len) };
121                // SAFETY:
122                // we checked bounds
123                unsafe { scatter_impl(current_values, values, arr, idx, len) };
124            },
125            None => {
126                let mut new_values = arr.values().as_slice().to_vec();
127                // SAFETY:
128                // we checked bounds
129                unsafe { scatter_impl(&mut new_values, values, arr, idx, len) };
130                arr.set_values(new_values.into());
131            },
132        };
133
134        // The null count may have changed - make sure to update the ChunkedArray
135        let new_null_count = arr.null_count();
136        unsafe { ca.set_null_count(new_null_count) };
137
138        Ok(ca.into_series())
139    }
140}
141
142impl<'a> ChunkedSet<&'a str> for &'a StringChunked {
143    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
144    where
145        V: IntoIterator<Item = Option<&'a str>>,
146    {
147        check_bounds(idx, self.len() as IdxSize)?;
148        check_sorted(idx)?;
149        let mut ca_iter = self.into_iter().enumerate();
150        let mut builder = StringChunkedBuilder::new(self.name().clone(), self.len());
151
152        for (current_idx, current_value) in idx.iter().zip(values) {
153            for (cnt_idx, opt_val_self) in &mut ca_iter {
154                if cnt_idx == *current_idx as usize {
155                    builder.append_option(current_value);
156                    break;
157                } else {
158                    builder.append_option(opt_val_self);
159                }
160            }
161        }
162        // the last idx is probably not the last value so we finish the iterator
163        for (_, opt_val_self) in ca_iter {
164            builder.append_option(opt_val_self);
165        }
166
167        let ca = builder.finish();
168        Ok(ca.into_series())
169    }
170}
171impl ChunkedSet<bool> for &BooleanChunked {
172    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
173    where
174        V: IntoIterator<Item = Option<bool>>,
175    {
176        check_bounds(idx, self.len() as IdxSize)?;
177        check_sorted(idx)?;
178        let mut ca_iter = self.into_iter().enumerate();
179        let mut builder = BooleanChunkedBuilder::new(self.name().clone(), self.len());
180
181        for (current_idx, current_value) in idx.iter().zip(values) {
182            for (cnt_idx, opt_val_self) in &mut ca_iter {
183                if cnt_idx == *current_idx as usize {
184                    builder.append_option(current_value);
185                    break;
186                } else {
187                    builder.append_option(opt_val_self);
188                }
189            }
190        }
191        // the last idx is probably not the last value so we finish the iterator
192        for (_, opt_val_self) in ca_iter {
193            builder.append_option(opt_val_self);
194        }
195
196        let ca = builder.finish();
197        Ok(ca.into_series())
198    }
199}
polars_ops/chunked_array/scatter.rs

polars_ops/chunked_array/
scatter.rs
