polars_ops/chunked_array/
scatter.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use polars_arrow::array::{
3    Array, BinaryViewArrayGeneric, BooleanArray, PrimitiveArray, View, ViewType,
4};
5use polars_buffer::Buffer;
6use polars_core::prelude::*;
7use polars_core::series::IsSorted;
8use polars_core::utils::polars_arrow::bitmap::MutableBitmap;
9use polars_core::utils::polars_arrow::types::NativeType;
10use polars_utils::index::check_bounds;
11
12pub trait ChunkedSet<T: Copy> {
13    /// Invariant for implementations: if the scatter() fails, typically because
14    /// of bad indexes, then self should remain unmodified.
15    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
16    where
17        V: IntoIterator<Item = Option<T>>;
18}
19
20trait PolarsOpsNumericType: PolarsNumericType {}
21
22impl PolarsOpsNumericType for UInt8Type {}
23impl PolarsOpsNumericType for UInt16Type {}
24impl PolarsOpsNumericType for UInt32Type {}
25impl PolarsOpsNumericType for UInt64Type {}
26#[cfg(feature = "dtype-u128")]
27impl PolarsOpsNumericType for UInt128Type {}
28impl PolarsOpsNumericType for Int8Type {}
29impl PolarsOpsNumericType for Int16Type {}
30impl PolarsOpsNumericType for Int32Type {}
31impl PolarsOpsNumericType for Int64Type {}
32#[cfg(feature = "dtype-i128")]
33impl PolarsOpsNumericType for Int128Type {}
34#[cfg(feature = "dtype-f16")]
35impl PolarsOpsNumericType for Float16Type {}
36impl PolarsOpsNumericType for Float32Type {}
37impl PolarsOpsNumericType for Float64Type {}
38
39unsafe fn scatter_primitive_impl<V, T: NativeType>(
40    set_values: V,
41    arr: &mut PrimitiveArray<T>,
42    idx: &[IdxSize],
43) where
44    V: IntoIterator<Item = Option<T>>,
45{
46    let mut values_iter = set_values.into_iter();
47
48    if let Some(validity) = arr.take_validity() {
49        let mut mut_validity = validity.make_mut();
50        arr.with_values_mut(|cur_values| {
51            for (idx, val) in idx.iter().zip(&mut values_iter) {
52                match val {
53                    Some(value) => {
54                        mut_validity.set_unchecked(*idx as usize, true);
55                        *cur_values.get_unchecked_mut(*idx as usize) = value
56                    },
57                    None => mut_validity.set_unchecked(*idx as usize, false),
58                }
59            }
60        });
61        arr.set_validity(mut_validity.into())
62    } else {
63        let mut null_idx = vec![];
64        arr.with_values_mut(|cur_values| {
65            for (idx, val) in idx.iter().zip(values_iter) {
66                match val {
67                    Some(value) => *cur_values.get_unchecked_mut(*idx as usize) = value,
68                    None => {
69                        null_idx.push(*idx);
70                    },
71                }
72            }
73        });
74
75        // Only make a validity bitmap when null values are set.
76        if !null_idx.is_empty() {
77            let mut validity = MutableBitmap::with_capacity(arr.len());
78            validity.extend_constant(arr.len(), true);
79            for idx in null_idx {
80                validity.set_unchecked(idx as usize, false)
81            }
82            arr.set_validity(Some(validity.into()))
83        }
84    }
85}
86
87unsafe fn scatter_bool_impl<V>(set_values: V, arr: &mut BooleanArray, idx: &[IdxSize])
88where
89    V: IntoIterator<Item = Option<bool>>,
90{
91    let mut values_iter = set_values.into_iter();
92
93    if let Some(validity) = arr.take_validity() {
94        let mut mut_validity = validity.make_mut();
95        arr.apply_values_mut(|cur_values| {
96            for (idx, val) in idx.iter().zip(&mut values_iter) {
97                match val {
98                    Some(value) => {
99                        mut_validity.set_unchecked(*idx as usize, true);
100                        cur_values.set_unchecked(*idx as usize, value);
101                    },
102                    None => mut_validity.set_unchecked(*idx as usize, false),
103                }
104            }
105        });
106        arr.set_validity(mut_validity.into())
107    } else {
108        let mut null_idx = vec![];
109        arr.apply_values_mut(|cur_values| {
110            for (idx, val) in idx.iter().zip(values_iter) {
111                match val {
112                    Some(value) => cur_values.set_unchecked(*idx as usize, value),
113                    None => {
114                        null_idx.push(*idx);
115                    },
116                }
117            }
118        });
119
120        // Only make a validity bitmap when null values are set.
121        if !null_idx.is_empty() {
122            let mut validity = MutableBitmap::with_capacity(arr.len());
123            validity.extend_constant(arr.len(), true);
124            for idx in null_idx {
125                validity.set_unchecked(idx as usize, false)
126            }
127            arr.set_validity(Some(validity.into()))
128        }
129    }
130}
131
132unsafe fn scatter_binview_impl<'a, V, T: ViewType + ?Sized>(
133    set_values: V,
134    arr: &mut BinaryViewArrayGeneric<T>,
135    idx: &[IdxSize],
136) where
137    V: IntoIterator<Item = Option<&'a T>>,
138{
139    let mut values_iter = set_values.into_iter();
140    let buffer_offset = arr.data_buffers().len() as u32;
141    let mut new_buffers = Vec::new();
142
143    if let Some(validity) = arr.take_validity() {
144        let mut mut_validity = validity.make_mut();
145        arr.with_views_mut(|views| {
146            for (idx, val) in idx.iter().zip(&mut values_iter) {
147                if let Some(v) = val {
148                    let view =
149                        View::new_with_buffers(v.to_bytes(), buffer_offset, &mut new_buffers);
150                    *views.get_unchecked_mut(*idx as usize) = view;
151                    mut_validity.set_unchecked(*idx as usize, true);
152                } else {
153                    mut_validity.set_unchecked(*idx as usize, false);
154                }
155            }
156        });
157        arr.set_validity(mut_validity.into())
158    } else {
159        let mut null_idx = vec![];
160        arr.with_views_mut(|views| {
161            for (idx, val) in idx.iter().zip(values_iter) {
162                if let Some(v) = val {
163                    let view =
164                        View::new_with_buffers(v.to_bytes(), buffer_offset, &mut new_buffers);
165                    *views.get_unchecked_mut(*idx as usize) = view;
166                } else {
167                    null_idx.push(*idx);
168                }
169            }
170        });
171
172        // Only make a validity bitmap when null values are set.
173        if !null_idx.is_empty() {
174            let mut validity = MutableBitmap::with_capacity(arr.len());
175            validity.extend_constant(arr.len(), true);
176            for idx in null_idx {
177                validity.set_unchecked(idx as usize, false)
178            }
179            arr.set_validity(Some(validity.into()))
180        }
181    }
182
183    let mut buffers = Buffer::to_vec(core::mem::take(arr.data_buffers_mut()));
184    buffers.extend(new_buffers.into_iter().map(Buffer::from));
185    *arr.data_buffers_mut() = Buffer::from(buffers);
186}
187
188impl<T: PolarsOpsNumericType> ChunkedSet<T::Native> for &mut ChunkedArray<T> {
189    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
190    where
191        V: IntoIterator<Item = Option<T::Native>>,
192    {
193        check_bounds(idx, self.len() as IdxSize)?;
194        let mut ca = std::mem::take(self);
195
196        // SAFETY: we will not modify the length and we unset the sorted flag,
197        // making sure to update the null count as well.
198        unsafe {
199            ca.rechunk_mut();
200            let arr = ca.downcast_iter_mut().next().unwrap();
201            scatter_primitive_impl(values, arr, idx);
202            let null_count = arr.null_count();
203            ca.set_sorted_flag(IsSorted::Not);
204            ca.set_null_count(null_count);
205        }
206
207        Ok(ca.into_series())
208    }
209}
210
211impl<'a> ChunkedSet<&'a [u8]> for &mut BinaryChunked {
212    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
213    where
214        V: IntoIterator<Item = Option<&'a [u8]>>,
215    {
216        check_bounds(idx, self.len() as IdxSize)?;
217        let mut ca = std::mem::take(self);
218
219        unsafe {
220            ca.rechunk_mut();
221            let arr = ca.downcast_iter_mut().next().unwrap();
222            scatter_binview_impl(values, arr, idx);
223            let null_count = arr.null_count();
224            ca.set_sorted_flag(IsSorted::Not);
225            ca.set_null_count(null_count);
226        }
227
228        Ok(ca.into_series())
229    }
230}
231
232impl<'a> ChunkedSet<&'a str> for &mut StringChunked {
233    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
234    where
235        V: IntoIterator<Item = Option<&'a str>>,
236    {
237        check_bounds(idx, self.len() as IdxSize)?;
238        let mut ca = std::mem::take(self);
239
240        unsafe {
241            ca.rechunk_mut();
242            let arr = ca.downcast_iter_mut().next().unwrap();
243            scatter_binview_impl(values, arr, idx);
244            let null_count = arr.null_count();
245            ca.set_sorted_flag(IsSorted::Not);
246            ca.set_null_count(null_count);
247        }
248
249        Ok(ca.into_series())
250    }
251}
252impl ChunkedSet<bool> for &mut BooleanChunked {
253    fn scatter<V>(self, idx: &[IdxSize], values: V) -> PolarsResult<Series>
254    where
255        V: IntoIterator<Item = Option<bool>>,
256    {
257        check_bounds(idx, self.len() as IdxSize)?;
258        let mut ca = std::mem::take(self);
259
260        unsafe {
261            ca.rechunk_mut();
262            let arr = ca.downcast_iter_mut().next().unwrap();
263            scatter_bool_impl(values, arr, idx);
264            let null_count = arr.null_count();
265            ca.set_sorted_flag(IsSorted::Not);
266            ca.set_null_count(null_count);
267        }
268
269        Ok(ca.into_series())
270    }
271}
polars_ops/chunked_array/scatter.rs

polars_ops/chunked_array/
scatter.rs
