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            if !(ca.is_sorted_any() || ca.len() == 1)
74                || !(other.is_sorted_any() || other.len() == 1)
75                || !(
76                    // We will coerce for single values
77                    ca.len() - ca.null_count() == 1
78                        || other.len() - other.null_count() == 1
79                        || ca.is_sorted_flag() == other.is_sorted_flag()
80                )
81            {
82                IsSorted::Not
83            } else {
84                let l_idx = ca.last_non_null().unwrap();
85                let r_idx = other.first_non_null().unwrap();
86
87                let null_pos_check =
88                    // check null positions
89                    // lhs does not end in nulls
90                    (1 + l_idx == ca.len())
91                    // rhs does not start with nulls
92                    && (r_idx == 0)
93                    // if there are nulls, they are all on one end
94                    && !(ca.first_non_null().unwrap() != 0 && 1 + other.last_non_null().unwrap() != other.len());
95
96                if !null_pos_check {
97                    IsSorted::Not
98                } else {
99                    #[allow(unused_assignments)]
100                    let mut out = IsSorted::Not;
101
102                    // This can be relatively expensive because of chunks, so delay as much as possible.
103                    let l_val = unsafe { ca.value_unchecked(l_idx) };
104                    let r_val = unsafe { other.value_unchecked(r_idx) };
105
106                    match (
107                        ca.len() - ca.null_count() == 1,
108                        other.len() - other.null_count() == 1,
109                    ) {
110                        (true, true) => {
111                            out = [IsSorted::Descending, IsSorted::Ascending]
112                                [l_val.tot_le(&r_val) as usize];
113                            drop(l_val);
114                            drop(r_val);
115                            ca.set_sorted_flag(out);
116                            return;
117                        },
118                        (true, false) => out = other.is_sorted_flag(),
119                        _ => out = ca.is_sorted_flag(),
120                    }
121
122                    debug_assert!(!matches!(out, IsSorted::Not));
123
124                    let check = if matches!(out, IsSorted::Ascending) {
125                        l_val.tot_le(&r_val)
126                    } else {
127                        l_val.tot_ge(&r_val)
128                    };
129
130                    if !check {
131                        out = IsSorted::Not
132                    }
133
134                    out
135                }
136            }
137        },
138    };
139
140    ca.set_sorted_flag(sorted_flag);
141}
142
143impl<T> ChunkedArray<T>
144where
145    T: PolarsDataType<IsNested = FalseT, IsObject = FalseT>,
146    for<'a> T::Physical<'a>: TotalOrd,
147{
148    /// Append in place. This is done by adding the chunks of `other` to this [`ChunkedArray`].
149    ///
150    /// See also [`extend`](Self::extend) for appends to the underlying memory
151    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
152        self.append_owned(other.clone())
153    }
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_owned(&mut self, mut 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_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
167        Ok(())
168    }
169}
170
171#[doc(hidden)]
172impl ListChunked {
173    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
174        self.append_owned(other.clone())
175    }
176
177    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
178        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
179        self.field = Arc::new(Field::new(self.name().clone(), dtype));
180
181        let len = self.len();
182        self.length = self
183            .length
184            .checked_add(other.length)
185            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
186        self.null_count += other.null_count;
187        self.set_sorted_flag(IsSorted::Not);
188        if !other.get_fast_explode_list() {
189            self.unset_fast_explode_list()
190        }
191
192        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
193        Ok(())
194    }
195}
196
197#[cfg(feature = "dtype-array")]
198#[doc(hidden)]
199impl ArrayChunked {
200    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
201        self.append_owned(other.clone())
202    }
203
204    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
205        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
206        self.field = Arc::new(Field::new(self.name().clone(), dtype));
207
208        let len = self.len();
209
210        self.length = self
211            .length
212            .checked_add(other.length)
213            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
214        self.null_count += other.null_count;
215
216        self.set_sorted_flag(IsSorted::Not);
217
218        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
219        Ok(())
220    }
221}
222
223#[cfg(feature = "dtype-struct")]
224#[doc(hidden)]
225impl StructChunked {
226    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
227        self.append_owned(other.clone())
228    }
229
230    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
231        let dtype = merge_dtypes(self.dtype(), other.dtype())?;
232        self.field = Arc::new(Field::new(self.name().clone(), dtype));
233
234        let len = self.len();
235
236        self.length = self
237            .length
238            .checked_add(other.length)
239            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
240        self.null_count += other.null_count;
241
242        self.set_sorted_flag(IsSorted::Not);
243
244        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
245        Ok(())
246    }
247}
248
249#[cfg(feature = "object")]
250#[doc(hidden)]
251impl<T: PolarsObject> ObjectChunked<T> {
252    pub fn append(&mut self, other: &Self) -> PolarsResult<()> {
253        self.append_owned(other.clone())
254    }
255
256    pub fn append_owned(&mut self, mut other: Self) -> PolarsResult<()> {
257        let len = self.len();
258        self.length = self
259            .length
260            .checked_add(other.length)
261            .ok_or_else(|| polars_err!(ComputeError: LENGTH_LIMIT_MSG))?;
262        self.null_count += other.null_count;
263        self.set_sorted_flag(IsSorted::Not);
264
265        new_chunks_owned(&mut self.chunks, std::mem::take(&mut other.chunks), len);
266        Ok(())
267    }
268}
polars_core/chunked_array/ops/append.rs

polars_core/chunked_array/ops/
append.rs
