polars_core/chunked_array/ops/
append.rs

1use polars_error::constants::LENGTH_LIMIT_MSG;
2
3use crate::prelude::*;
4use crate::series::IsSorted;
5
6pub(crate) fn new_chunks(chunks: &mut Vec<ArrayRef>, other: &[ArrayRef], len: usize) {
7    // Replace an empty array.
8    if chunks.len() == 1 && len == 0 {
9        other.clone_into(chunks);
10    } else {
11        for chunk in other {
12            if !chunk.is_empty() {
13                chunks.push(chunk.clone());
14            }
15        }
16    }
17}
18
19pub(crate) fn new_chunks_owned(chunks: &mut Vec<ArrayRef>, other: Vec<ArrayRef>, len: usize) {
20    // Replace an empty array.
21    if chunks.len() == 1 && len == 0 {
22        *chunks = other;
23    } else {
24        chunks.reserve(other.len());
25        chunks.extend(other.into_iter().filter(|c| !c.is_empty()));
26    }
27}
28
29/// Unit-length arrays are sorted but don't necessarily carry the flag.
30fn unit_or_sorted_flag<T: PolarsDataType>(ca: &ChunkedArray<T>) -> IsSorted {
31    if ca.len() == 1 {
32        IsSorted::Ascending
33    } else {
34        ca.is_sorted_flag()
35    }
36}
37
38pub(super) fn update_sorted_flag_before_append<T>(ca: &mut ChunkedArray<T>, other: &ChunkedArray<T>)
39where
40    T: PolarsDataType,
41    for<'a> T::Physical<'a>: TotalOrd,
42{
43    // Note: Do not call (first|last)_non_null on an array here before checking
44    // it is sorted, otherwise it will lead to quadratic behavior.
45    let sorted_flag = match (
46        ca.null_count() != ca.len(),
47        other.null_count() != other.len(),
48    ) {
49        (false, false) => IsSorted::Ascending,
50        (false, true) => {
51            if
52            // lhs is empty, just take sorted flag from rhs
53            ca.is_empty()
54                || (
55                    // lhs is non-empty and all-null, so rhs must have nulls ordered first
56                    unit_or_sorted_flag(other) != IsSorted::Not
57                        && 1 + other.last_non_null().unwrap() == other.len()
58                )
59            {
60                unit_or_sorted_flag(other)
61            } else {
62                IsSorted::Not
63            }
64        },
65        (true, false) => {
66            if
67            // rhs is empty, just take sorted flag from lhs
68            other.is_empty()
69                || (
70                    // rhs is non-empty and all-null, so lhs must have nulls ordered last
71                    unit_or_sorted_flag(ca) != IsSorted::Not && ca.first_non_null().unwrap() == 0
72                )
73            {
74                unit_or_sorted_flag(ca)
75            } else {
76                IsSorted::Not
77            }
78        },
79        (true, true) => {
80            // both arrays have non-null values.
81            // for arrays of unit length we can ignore the sorted flag, as it is
82            // not necessarily set.
83            let both_sorted = (ca.is_sorted_any() || ca.len() == 1)
84                && (other.is_sorted_any() || other.len() == 1);
85
86            // We will coerce for a single distinct value, which is sorted in both directions.
87            // Only compare values when the flags differ, as that is relatively expensive.
88            let flags_differ = ca.is_sorted_flag() != other.is_sorted_flag();
89            let is_single_value = |arr: &ChunkedArray<T>| {
90                arr.len() - arr.null_count() == 1
91                    || (both_sorted
92                        && flags_differ
93                        && unsafe {
94                            arr.value_unchecked(arr.first_non_null().unwrap())
95                                .tot_eq(&arr.value_unchecked(arr.last_non_null().unwrap()))
96                        })
97            };
98            let l_single = is_single_value(ca);
99            let r_single = is_single_value(other);
100
101            if !both_sorted || (flags_differ && !(l_single || r_single)) {
102                IsSorted::Not
103            } else {
104                let l_idx = ca.last_non_null().unwrap();
105                let r_idx = other.first_non_null().unwrap();
106
107                let null_pos_check =
108                    // check null positions
109                    // lhs does not end in nulls
110                    (1 + l_idx == ca.len())
111                    // rhs does not start with nulls
112                    && (r_idx == 0)
113                    // if there are nulls, they are all on one end
114                    && !(ca.first_non_null().unwrap() != 0 && 1 + other.last_non_null().unwrap() != other.len());
115
116                if !null_pos_check {
117                    IsSorted::Not
118                } else {
119                    #[allow(unused_assignments)]
120                    let mut out = IsSorted::Not;
121
122                    // This can be relatively expensive because of chunks, so delay as much as possible.
123                    let l_val = unsafe { ca.value_unchecked(l_idx) };
124                    let r_val = unsafe { other.value_unchecked(r_idx) };
125
126                    match (l_single, r_single) {
127                        (true, true) => {
128                            out = [IsSorted::Descending, IsSorted::Ascending]
129                                [l_val.tot_le(&r_val) as usize];
130                            drop(l_val);
131                            drop(r_val);
132                            ca.set_sorted_flag(out);
133                            return;
134                        },
135                        (true, false) => out = other.is_sorted_flag(),
136                        _ => out = ca.is_sorted_flag(),
137                    }
138
139                    debug_assert!(!matches!(out, IsSorted::Not));
140
141                    let check = if matches!(out, IsSorted::Ascending) {
142                        l_val.tot_le(&r_val)
143                    } else {
144                        l_val.tot_ge(&r_val)
145                    };
146
147                    if !check {
148                        out = IsSorted::Not
149                    }
150
151                    out
152                }
153            }
154        },
155    };
156
157    ca.set_sorted_flag(sorted_flag);
158}
159
160impl<T> ChunkedArray<T>
161where
162    T: PolarsDataType<IsNested = FalseT, IsObject = FalseT>,
163    for<'a> T::Physical<'a>: TotalOrd,
164{
165    /// Append in place. This is done by adding the chunks of `other` to this [`ChunkedArray`].
166    ///
167    /// See also [`extend`](Self::extend) for appends to the underlying memory
168    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
169        update_sorted_flag_before_append::<T>(self, other);
170        let len = self.len();
171        self.length = self
172            .length
173            .checked_add(other.length)
174            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
175        self.null_count += other.null_count;
176        new_chunks(&mut self.chunks, &other.chunks, len);
177        Ok(())
178    }
179
180    /// Append in place. This is done by adding the chunks of `other` to this [`ChunkedArray`].
181    ///
182    /// See also [`extend`](Self::extend) for appends to the underlying memory
183    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
184        update_sorted_flag_before_append::<T>(self, &other);
185        let len = self.len();
186        self.length = self
187            .length
188            .checked_add(other.length)
189            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
190        self.null_count += other.null_count;
191        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
192        Ok(())
193    }
194}
195
196#[doc(hidden)]
197impl ListChunked {
198    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
199        self.append_owned(other.clone())
200    }
201
202    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
203        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
204        self.field = Arc::new(Field::new(self.name().clone(), dtype));
205
206        let len = self.len();
207        self.length = self
208            .length
209            .checked_add(other.length)
210            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
211        self.null_count += other.null_count;
212        self.set_sorted_flag(IsSorted::Not);
213        if !other.get_fast_explode_list() {
214            self.unset_fast_explode_list()
215        }
216
217        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
218        Ok(())
219    }
220}
221
222#[cfg(feature = "dtype-array")]
223#[doc(hidden)]
224impl ArrayChunked {
225    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
226        self.append_owned(other.clone())
227    }
228
229    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
230        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
231        self.field = Arc::new(Field::new(self.name().clone(), dtype));
232
233        let len = self.len();
234
235        self.length = self
236            .length
237            .checked_add(other.length)
238            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
239        self.null_count += other.null_count;
240
241        self.set_sorted_flag(IsSorted::Not);
242
243        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
244        Ok(())
245    }
246}
247
248#[cfg(feature = "dtype-struct")]
249#[doc(hidden)]
250impl StructChunked {
251    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
252        self.append_owned(other.clone())
253    }
254
255    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
256        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
257        self.field = Arc::new(Field::new(self.name().clone(), dtype));
258
259        let len = self.len();
260
261        self.length = self
262            .length
263            .checked_add(other.length)
264            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
265        self.null_count += other.null_count;
266
267        self.set_sorted_flag(IsSorted::Not);
268
269        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
270        Ok(())
271    }
272}
273
274#[cfg(feature = "dtype-categorical")]
275#[doc(hidden)]
276impl<T: PolarsCategoricalType> CategoricalChunked<T> {
277    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
278        assert!(self.dtype() == other.dtype());
279        self.phys.append(&other.phys)
280    }
281
282    pub fn append_owned(&mut self, other: Self) -> PolarsResult<()> {
283        assert!(self.dtype() == other.dtype());
284        self.phys.append_owned(other.phys)
285    }
286}
287
288#[cfg(feature = "object")]
289#[doc(hidden)]
290impl<T: PolarsObject> ObjectChunked<T> {
291    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
292        self.append_owned(other.clone())
293    }
294
295    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
296        let len = self.len();
297        self.length = self
298            .length
299            .checked_add(other.length)
300            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
301        self.null_count += other.null_count;
302        self.set_sorted_flag(IsSorted::Not);
303
304        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
305        Ok(())
306    }
307}
polars_core/chunked_array/ops/append.rs

polars_core/chunked_array/ops/
append.rs
