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

polars_core/chunked_array/ops/
append.rs
