polars_python/series/
scatter.rs

1use arrow::array::Array;
2use polars::prelude::*;
3use pyo3::prelude::*;
4
5use super::PySeries;
6use crate::utils::EnterPolarsExt;
7
8#[pymethods]
9impl PySeries {
10    fn scatter(&mut self, py: Python<'_>, idx: PySeries, values: PySeries) -> PyResult<()> {
11        // we take the value because we want a ref count of 1 so that we can
12        // have mutable access cheaply via _get_inner_mut().
13        let s = std::mem::take(&mut self.series);
14        py.enter_polars(|| {
15            let result = scatter(s, &idx.series, &values.series);
16            match result {
17                Ok(out) => {
18                    self.series = out;
19                    Ok(())
20                },
21                Err((s, e)) => {
22                    // Restore original series:
23                    self.series = s;
24                    Err(e)
25                },
26            }
27        })
28    }
29}
30
31fn scatter(mut s: Series, idx: &Series, values: &Series) -> Result<Series, (Series, PolarsError)> {
32    let logical_dtype = s.dtype().clone();
33
34    let idx = match polars_ops::prelude::convert_to_unsigned_index(idx, s.len()) {
35        Ok(idx) => idx,
36        Err(err) => return Err((s, err)),
37    };
38    let idx = idx.rechunk();
39    let idx = idx.downcast_as_array();
40
41    if idx.null_count() > 0 {
42        return Err((
43            s,
44            PolarsError::ComputeError("index values should not be null".into()),
45        ));
46    }
47
48    let idx = idx.values().as_slice();
49
50    let mut values = match values.to_physical_repr().cast(&s.dtype().to_physical()) {
51        Ok(values) => values,
52        Err(err) => return Err((s, err)),
53    };
54
55    // Broadcast values input
56    if values.len() == 1 && idx.len() > 1 {
57        values = values.new_from_index(0, idx.len());
58    }
59
60    // do not shadow, otherwise s is not dropped immediately
61    // and we want to have mutable access
62    s = s.to_physical_repr().into_owned();
63    let s_mut_ref = &mut s;
64    scatter_impl(s_mut_ref, logical_dtype, idx, &values).map_err(|err| (s, err))
65}
66
67fn scatter_impl(
68    s: &mut Series,
69    logical_dtype: DataType,
70    idx: &[IdxSize],
71    values: &Series,
72) -> PolarsResult<Series> {
73    let mutable_s = s._get_inner_mut();
74
75    let s = match logical_dtype.to_physical() {
76        DataType::Int8 => {
77            let ca: &mut ChunkedArray<Int8Type> = mutable_s.as_mut();
78            let values = values.i8()?;
79            ca.scatter(idx, values)
80        },
81        DataType::Int16 => {
82            let ca: &mut ChunkedArray<Int16Type> = mutable_s.as_mut();
83            let values = values.i16()?;
84            ca.scatter(idx, values)
85        },
86        DataType::Int32 => {
87            let ca: &mut ChunkedArray<Int32Type> = mutable_s.as_mut();
88            let values = values.i32()?;
89            ca.scatter(idx, values)
90        },
91        DataType::Int64 => {
92            let ca: &mut ChunkedArray<Int64Type> = mutable_s.as_mut();
93            let values = values.i64()?;
94            ca.scatter(idx, values)
95        },
96        DataType::UInt8 => {
97            let ca: &mut ChunkedArray<UInt8Type> = mutable_s.as_mut();
98            let values = values.u8()?;
99            ca.scatter(idx, values)
100        },
101        DataType::UInt16 => {
102            let ca: &mut ChunkedArray<UInt16Type> = mutable_s.as_mut();
103            let values = values.u16()?;
104            ca.scatter(idx, values)
105        },
106        DataType::UInt32 => {
107            let ca: &mut ChunkedArray<UInt32Type> = mutable_s.as_mut();
108            let values = values.u32()?;
109            ca.scatter(idx, values)
110        },
111        DataType::UInt64 => {
112            let ca: &mut ChunkedArray<UInt64Type> = mutable_s.as_mut();
113            let values = values.u64()?;
114            ca.scatter(idx, values)
115        },
116        DataType::Float32 => {
117            let ca: &mut ChunkedArray<Float32Type> = mutable_s.as_mut();
118            let values = values.f32()?;
119            ca.scatter(idx, values)
120        },
121        DataType::Float64 => {
122            let ca: &mut ChunkedArray<Float64Type> = mutable_s.as_mut();
123            let values = values.f64()?;
124            ca.scatter(idx, values)
125        },
126        DataType::Boolean => {
127            let ca = s.bool()?;
128            let values = values.bool()?;
129            ca.scatter(idx, values)
130        },
131        DataType::String => {
132            let ca = s.str()?;
133            let values = values.str()?;
134            ca.scatter(idx, values)
135        },
136        _ => {
137            return Err(PolarsError::ComputeError(
138                format!("not yet implemented for dtype: {logical_dtype}").into(),
139            ));
140        },
141    };
142
143    s.and_then(|s| s.cast(&logical_dtype))
144}
polars_python/series/scatter.rs

polars_python/series/
scatter.rs
