Skip to main content

polars_core/frame/
mod.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2//! DataFrame module.
3use std::borrow::Cow;
4
5use polars_arrow::datatypes::ArrowSchemaRef;
6use polars_row::ArrayRef;
7use polars_utils::UnitVec;
8use polars_utils::itertools::Itertools;
9use rayon::prelude::*;
10
11use crate::chunked_array::flags::StatisticsFlags;
12#[cfg(feature = "algorithm_group_by")]
13use crate::chunked_array::ops::unique::is_unique_helper;
14use crate::prelude::gather::check_bounds_ca;
15use crate::prelude::*;
16#[cfg(feature = "row_hash")]
17use crate::utils::split_df;
18use crate::utils::{Container, NoNull, slice_offsets, try_get_supertype};
19use crate::{HEAD_DEFAULT_LENGTH, TAIL_DEFAULT_LENGTH};
20
21#[cfg(feature = "dataframe_arithmetic")]
22mod arithmetic;
23pub mod builder;
24mod chunks;
25pub use chunks::chunk_df_for_writing;
26pub mod column;
27mod dataframe;
28mod filter;
29mod projection;
30pub use dataframe::DataFrame;
31use filter::filter_zero_width;
32use projection::{AmortizedColumnSelector, LINEAR_SEARCH_LIMIT};
33
34pub mod explode;
35mod from;
36#[cfg(feature = "algorithm_group_by")]
37pub mod group_by;
38pub(crate) mod horizontal;
39#[cfg(any(feature = "rows", feature = "object"))]
40pub mod row;
41mod top_k;
42mod upstream_traits;
43pub(crate) mod validation;
44
45use polars_arrow::record_batch::{RecordBatch, RecordBatchT};
46use polars_utils::pl_str::PlSmallStr;
47#[cfg(feature = "serde")]
48use serde::{Deserialize, Serialize};
49use strum_macros::IntoStaticStr;
50
51#[cfg(feature = "row_hash")]
52use crate::hashing::_df_rows_to_hashes_threaded_vertical;
53use crate::prelude::sort::arg_sort;
54use crate::runtime::RAYON;
55use crate::series::IsSorted;
56
57#[derive(Copy, Clone, Debug, PartialEq, Eq, Default, Hash, IntoStaticStr)]
58#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
59#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
60#[strum(serialize_all = "snake_case")]
61pub enum UniqueKeepStrategy {
62    /// Keep the first unique row.
63    First,
64    /// Keep the last unique row.
65    Last,
66    /// Keep None of the unique rows.
67    None,
68    /// Keep any of the unique rows
69    /// This allows more optimizations
70    #[default]
71    Any,
72}
73
74#[derive(Copy, Clone, Debug, PartialEq, Eq, Default, Hash, IntoStaticStr)]
75#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
76#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
77#[strum(serialize_all = "snake_case")]
78/// Naming strategy for the results of a pivot.
79pub enum PivotColumnNaming {
80    /// Always combine the values and on-column names.
81    Combine,
82    /// Prefix the values column name only if there is more than one values
83    /// column.
84    #[default]
85    Auto,
86}
87
88impl DataFrame {
89    pub fn materialized_column_iter(&self) -> impl ExactSizeIterator<Item = &Series> {
90        self.columns().iter().map(Column::as_materialized_series)
91    }
92
93    /// Returns an estimation of the total (heap) allocated size of the `DataFrame` in bytes.
94    ///
95    /// # Implementation
96    /// This estimation is the sum of the size of its buffers, validity, including nested arrays.
97    /// Multiple arrays may share buffers and bitmaps. Therefore, the size of 2 arrays is not the
98    /// sum of the sizes computed from this function. In particular, [`StructArray`]'s size is an upper bound.
99    ///
100    /// When an array is sliced, its allocated size remains constant because the buffer unchanged.
101    /// However, this function will yield a smaller number. This is because this function returns
102    /// the visible size of the buffer, not its total capacity.
103    ///
104    /// FFI buffers are included in this estimation.
105    pub fn estimated_size(&self) -> usize {
106        self.columns().iter().map(Column::estimated_size).sum()
107    }
108
109    pub fn try_apply_columns(
110        &self,
111        func: impl Fn(&Column) -> PolarsResult<Column> + Send + Sync,
112    ) -> PolarsResult<Vec<Column>> {
113        return inner(self, &func);
114
115        fn inner(
116            slf: &DataFrame,
117            func: &(dyn Fn(&Column) -> PolarsResult<Column> + Send + Sync),
118        ) -> PolarsResult<Vec<Column>> {
119            slf.columns().iter().map(func).collect()
120        }
121    }
122
123    pub fn apply_columns(&self, func: impl Fn(&Column) -> Column + Send + Sync) -> Vec<Column> {
124        return inner(self, &func);
125
126        fn inner(slf: &DataFrame, func: &(dyn Fn(&Column) -> Column + Send + Sync)) -> Vec<Column> {
127            slf.columns().iter().map(func).collect()
128        }
129    }
130
131    pub fn try_apply_columns_par(
132        &self,
133        func: impl Fn(&Column) -> PolarsResult<Column> + Send + Sync,
134    ) -> PolarsResult<Vec<Column>> {
135        return inner(self, &func);
136
137        fn inner(
138            slf: &DataFrame,
139            func: &(dyn Fn(&Column) -> PolarsResult<Column> + Send + Sync),
140        ) -> PolarsResult<Vec<Column>> {
141            RAYON.install(|| slf.columns().par_iter().map(func).collect())
142        }
143    }
144
145    pub fn apply_columns_par(&self, func: impl Fn(&Column) -> Column + Send + Sync) -> Vec<Column> {
146        return inner(self, &func);
147
148        fn inner(slf: &DataFrame, func: &(dyn Fn(&Column) -> Column + Send + Sync)) -> Vec<Column> {
149            RAYON.install(|| slf.columns().par_iter().map(func).collect())
150        }
151    }
152
153    /// Reserve additional slots into the chunks of the series.
154    pub(crate) fn reserve_chunks(&mut self, additional: usize) {
155        for s in unsafe { self.columns_mut_retain_schema() } {
156            if let Column::Series(s) = s {
157                // SAFETY:
158                // do not modify the data, simply resize.
159                unsafe { s.chunks_mut().reserve(additional) }
160            }
161        }
162    }
163    pub fn new_from_index(&self, index: usize, height: usize) -> Self {
164        let new_cols = self.apply_columns(|c| c.new_from_index(index, height));
165
166        unsafe { Self::_new_unchecked_impl(height, new_cols).with_schema_from(self) }
167    }
168
169    /// Create a new `DataFrame` with the given schema, only containing nulls.
170    pub fn full_null(schema: &Schema, height: usize) -> Self {
171        let columns = schema
172            .iter_fields()
173            .map(|f| Column::full_null(f.name().clone(), height, f.dtype()))
174            .collect();
175
176        unsafe { DataFrame::_new_unchecked_impl(height, columns) }
177    }
178
179    /// Ensure this DataFrame matches the given schema. Casts null columns to
180    /// the expected schema if necessary (but nothing else).
181    pub fn ensure_matches_schema(&mut self, schema: &Schema) -> PolarsResult<()> {
182        let mut did_cast = false;
183        let cached_schema = self.cached_schema().cloned();
184
185        for (col, (name, dt)) in unsafe { self.columns_mut() }.iter_mut().zip(schema.iter()) {
186            polars_ensure!(
187                col.name() == name,
188                SchemaMismatch: "column name mismatch: expected {:?}, found {:?}",
189                name,
190                col.name()
191            );
192
193            let needs_cast = col.dtype().matches_schema_type(dt)?;
194
195            if needs_cast {
196                *col = col.cast(dt)?;
197                did_cast = true;
198            }
199        }
200
201        if !did_cast {
202            unsafe { self.set_opt_schema(cached_schema) };
203        }
204
205        Ok(())
206    }
207
208    /// Add a new column at index 0 that counts the rows.
209    ///
210    /// # Example
211    ///
212    /// ```
213    /// # use polars_core::prelude::*;
214    /// let df1: DataFrame = df!("Name" => ["James", "Mary", "John", "Patricia"])?;
215    /// assert_eq!(df1.shape(), (4, 1));
216    ///
217    /// let df2: DataFrame = df1.with_row_index("Id".into(), None)?;
218    /// assert_eq!(df2.shape(), (4, 2));
219    /// println!("{}", df2);
220    ///
221    /// # Ok::<(), PolarsError>(())
222    /// ```
223    ///
224    /// Output:
225    ///
226    /// ```text
227    ///  shape: (4, 2)
228    ///  +-----+----------+
229    ///  | Id  | Name     |
230    ///  | --- | ---      |
231    ///  | u32 | str      |
232    ///  +=====+==========+
233    ///  | 0   | James    |
234    ///  +-----+----------+
235    ///  | 1   | Mary     |
236    ///  +-----+----------+
237    ///  | 2   | John     |
238    ///  +-----+----------+
239    ///  | 3   | Patricia |
240    ///  +-----+----------+
241    /// ```
242    pub fn with_row_index(&self, name: PlSmallStr, offset: Option<IdxSize>) -> PolarsResult<Self> {
243        let mut new_columns = Vec::with_capacity(self.width() + 1);
244        let offset = offset.unwrap_or(0);
245
246        if self.get_column_index(&name).is_some() {
247            polars_bail!(duplicate = name)
248        }
249
250        let col = Column::new_row_index(name, offset, self.height())?;
251        new_columns.push(col);
252        new_columns.extend_from_slice(self.columns());
253
254        Ok(unsafe { DataFrame::new_unchecked(self.height(), new_columns) })
255    }
256
257    /// Add a row index column in place.
258    ///
259    /// # Safety
260    /// The caller should ensure the DataFrame does not already contain a column with the given name.
261    ///
262    /// # Panics
263    /// Panics if the resulting column would reach or overflow IdxSize::MAX.
264    pub unsafe fn with_row_index_mut(
265        &mut self,
266        name: PlSmallStr,
267        offset: Option<IdxSize>,
268    ) -> &mut Self {
269        debug_assert!(
270            self.get_column_index(&name).is_none(),
271            "with_row_index_mut(): column with name {} already exists",
272            name
273        );
274
275        let offset = offset.unwrap_or(0);
276        let col = Column::new_row_index(name, offset, self.height()).unwrap();
277
278        unsafe { self.columns_mut() }.insert(0, col);
279        self
280    }
281
282    /// Shrink the capacity of this DataFrame to fit its length.
283    pub fn shrink_to_fit(&mut self) {
284        // Don't parallelize this. Memory overhead
285        for s in unsafe { self.columns_mut_retain_schema() } {
286            s.shrink_to_fit();
287        }
288    }
289
290    /// Aggregate all the chunks in the DataFrame to a single chunk in parallel.
291    /// This may lead to more peak memory consumption.
292    pub fn rechunk_mut_par(&mut self) -> &mut Self {
293        if self.columns().iter().any(|c| c.n_chunks() > 1) {
294            RAYON.install(|| {
295                unsafe { self.columns_mut_retain_schema() }
296                    .par_iter_mut()
297                    .for_each(|c| *c = c.rechunk());
298            })
299        }
300
301        self
302    }
303
304    /// Rechunks all columns to only have a single chunk.
305    pub fn rechunk_mut(&mut self) -> &mut Self {
306        // SAFETY: We never adjust the length or names of the columns.
307        let columns = unsafe { self.columns_mut() };
308
309        for col in columns.iter_mut().filter(|c| c.n_chunks() > 1) {
310            *col = col.rechunk();
311        }
312
313        self
314    }
315
316    /// Returns true if the chunks of the columns do not align and re-chunking should be done
317    pub fn should_rechunk(&self) -> bool {
318        // Fast check. It is also needed for correctness, as code below doesn't check if the number
319        // of chunks is equal.
320        if !self.columns().iter().map(Column::n_chunks).all_equal() {
321            return true;
322        }
323
324        // From here we check chunk lengths. Skipping the columns without chunks is
325        // safe because the counts are equal: either every count is 1, or there is
326        // no such column left.
327        let mut chunk_lengths = self
328            .columns()
329            .iter()
330            .filter_map(Column::lazy_as_materialized_series)
331            .map(|s| s.chunk_lengths());
332        match chunk_lengths.next() {
333            None => false,
334            Some(first_column_chunk_lengths) => {
335                // Fast Path for single Chunk Series
336                if first_column_chunk_lengths.size_hint().0 == 1 {
337                    return chunk_lengths.any(|cl| cl.size_hint().0 != 1);
338                }
339                // Always rechunk if we have more chunks than rows.
340                // except when we have an empty df containing a single chunk
341                let height = self.height();
342                let n_chunks = first_column_chunk_lengths.size_hint().0;
343                if n_chunks > height && !(height == 0 && n_chunks == 1) {
344                    return true;
345                }
346                // Slow Path for multi Chunk series
347                let v: Vec<_> = first_column_chunk_lengths.collect();
348                for cl in chunk_lengths {
349                    if cl.enumerate().any(|(idx, el)| Some(&el) != v.get(idx)) {
350                        return true;
351                    }
352                }
353                false
354            },
355        }
356    }
357
358    /// Ensure all the chunks in the [`DataFrame`] are aligned.
359    pub fn align_chunks_par(&mut self) -> &mut Self {
360        if self.should_rechunk() {
361            self.rechunk_mut_par()
362        } else {
363            self
364        }
365    }
366
367    /// Ensure all the chunks in the [`DataFrame`] are aligned.
368    pub fn align_chunks(&mut self) -> &mut Self {
369        if self.should_rechunk() {
370            self.rechunk_mut()
371        } else {
372            self
373        }
374    }
375
376    /// # Example
377    ///
378    /// ```rust
379    /// # use polars_core::prelude::*;
380    /// let df: DataFrame = df!("Language" => ["Rust", "Python"],
381    ///                         "Designer" => ["Graydon Hoare", "Guido van Rossum"])?;
382    ///
383    /// assert_eq!(df.get_column_names(), &["Language", "Designer"]);
384    /// # Ok::<(), PolarsError>(())
385    /// ```
386    pub fn get_column_names(&self) -> Vec<&PlSmallStr> {
387        self.columns().iter().map(|s| s.name()).collect()
388    }
389
390    /// Get the [`Vec<PlSmallStr>`] representing the column names.
391    pub fn get_column_names_owned(&self) -> Vec<PlSmallStr> {
392        self.columns().iter().map(|s| s.name().clone()).collect()
393    }
394
395    /// Set the column names.
396    /// # Example
397    ///
398    /// ```rust
399    /// # use polars_core::prelude::*;
400    /// let mut df: DataFrame = df!("Mathematical set" => ["ℕ", "ℤ", "𝔻", "ℚ", "ℝ", "ℂ"])?;
401    /// df.set_column_names(&["Set"])?;
402    ///
403    /// assert_eq!(df.get_column_names(), &["Set"]);
404    /// # Ok::<(), PolarsError>(())
405    /// ```
406    pub fn set_column_names<T>(&mut self, new_names: &[T]) -> PolarsResult<()>
407    where
408        T: AsRef<str>,
409    {
410        polars_ensure!(
411            new_names.len() == self.width(),
412            ShapeMismatch: "{} column names provided for a DataFrame of width {}",
413            new_names.len(), self.width()
414        );
415
416        validation::ensure_names_unique(new_names)?;
417
418        *unsafe { self.columns_mut() } = std::mem::take(unsafe { self.columns_mut() })
419            .into_iter()
420            .zip(new_names)
421            .map(|(c, name)| c.with_name(PlSmallStr::from_str(name.as_ref())))
422            .collect();
423
424        Ok(())
425    }
426
427    /// Get the data types of the columns in the [`DataFrame`].
428    ///
429    /// # Example
430    ///
431    /// ```rust
432    /// # use polars_core::prelude::*;
433    /// let venus_air: DataFrame = df!("Element" => ["Carbon dioxide", "Nitrogen"],
434    ///                                "Fraction" => [0.965, 0.035])?;
435    ///
436    /// assert_eq!(venus_air.dtypes(), &[DataType::String, DataType::Float64]);
437    /// # Ok::<(), PolarsError>(())
438    /// ```
439    pub fn dtypes(&self) -> Vec<DataType> {
440        self.columns().iter().map(|s| s.dtype().clone()).collect()
441    }
442
443    /// The number of chunks for the first column.
444    pub fn first_col_n_chunks(&self) -> usize {
445        match self.columns().iter().find_map(|col| col.as_series()) {
446            None if self.width() == 0 => 0,
447            None => 1,
448            Some(s) => s.n_chunks(),
449        }
450    }
451
452    /// The highest number of chunks for any column.
453    pub fn max_n_chunks(&self) -> usize {
454        self.columns()
455            .iter()
456            .map(|s| s.as_series().map(|s| s.n_chunks()).unwrap_or(1))
457            .max()
458            .unwrap_or(0)
459    }
460
461    /// Generate the schema fields of the [`DataFrame`].
462    ///
463    /// # Example
464    ///
465    /// ```rust
466    /// # use polars_core::prelude::*;
467    /// let earth: DataFrame = df!("Surface type" => ["Water", "Land"],
468    ///                            "Fraction" => [0.708, 0.292])?;
469    ///
470    /// let f1: Field = Field::new("Surface type".into(), DataType::String);
471    /// let f2: Field = Field::new("Fraction".into(), DataType::Float64);
472    ///
473    /// assert_eq!(earth.fields(), &[f1, f2]);
474    /// # Ok::<(), PolarsError>(())
475    /// ```
476    pub fn fields(&self) -> Vec<Field> {
477        self.columns()
478            .iter()
479            .map(|s| s.field().into_owned())
480            .collect()
481    }
482
483    /// Add multiple [`Series`] to a [`DataFrame`].
484    /// The added `Series` are required to have the same length.
485    ///
486    /// # Example
487    ///
488    /// ```rust
489    /// # use polars_core::prelude::*;
490    /// let df1: DataFrame = df!("Element" => ["Copper", "Silver", "Gold"])?;
491    /// let s1 = Column::new("Proton".into(), [29, 47, 79]);
492    /// let s2 = Column::new("Electron".into(), [29, 47, 79]);
493    ///
494    /// let df2: DataFrame = df1.hstack(&[s1, s2])?;
495    /// assert_eq!(df2.shape(), (3, 3));
496    /// println!("{}", df2);
497    /// # Ok::<(), PolarsError>(())
498    /// ```
499    ///
500    /// Output:
501    ///
502    /// ```text
503    /// shape: (3, 3)
504    /// +---------+--------+----------+
505    /// | Element | Proton | Electron |
506    /// | ---     | ---    | ---      |
507    /// | str     | i32    | i32      |
508    /// +=========+========+==========+
509    /// | Copper  | 29     | 29       |
510    /// +---------+--------+----------+
511    /// | Silver  | 47     | 47       |
512    /// +---------+--------+----------+
513    /// | Gold    | 79     | 79       |
514    /// +---------+--------+----------+
515    /// ```
516    pub fn hstack(&self, columns: &[Column]) -> PolarsResult<Self> {
517        let mut new_cols = Vec::with_capacity(self.width() + columns.len());
518
519        new_cols.extend(self.columns().iter().cloned());
520        new_cols.extend_from_slice(columns);
521
522        DataFrame::new(self.height(), new_cols)
523    }
524    /// Concatenate a [`DataFrame`] to this [`DataFrame`] and return as newly allocated [`DataFrame`].
525    ///
526    /// If many `vstack` operations are done, it is recommended to call [`DataFrame::align_chunks_par`].
527    ///
528    /// # Example
529    ///
530    /// ```rust
531    /// # use polars_core::prelude::*;
532    /// let df1: DataFrame = df!("Element" => ["Copper", "Silver", "Gold"],
533    ///                          "Melting Point (K)" => [1357.77, 1234.93, 1337.33])?;
534    /// let df2: DataFrame = df!("Element" => ["Platinum", "Palladium"],
535    ///                          "Melting Point (K)" => [2041.4, 1828.05])?;
536    ///
537    /// let df3: DataFrame = df1.vstack(&df2)?;
538    ///
539    /// assert_eq!(df3.shape(), (5, 2));
540    /// println!("{}", df3);
541    /// # Ok::<(), PolarsError>(())
542    /// ```
543    ///
544    /// Output:
545    ///
546    /// ```text
547    /// shape: (5, 2)
548    /// +-----------+-------------------+
549    /// | Element   | Melting Point (K) |
550    /// | ---       | ---               |
551    /// | str       | f64               |
552    /// +===========+===================+
553    /// | Copper    | 1357.77           |
554    /// +-----------+-------------------+
555    /// | Silver    | 1234.93           |
556    /// +-----------+-------------------+
557    /// | Gold      | 1337.33           |
558    /// +-----------+-------------------+
559    /// | Platinum  | 2041.4            |
560    /// +-----------+-------------------+
561    /// | Palladium | 1828.05           |
562    /// +-----------+-------------------+
563    /// ```
564    pub fn vstack(&self, other: &DataFrame) -> PolarsResult<Self> {
565        let mut df = self.clone();
566        df.vstack_mut(other)?;
567        Ok(df)
568    }
569
570    /// Concatenate a [`DataFrame`] to this [`DataFrame`]
571    ///
572    /// If many `vstack` operations are done, it is recommended to call [`DataFrame::align_chunks_par`].
573    ///
574    /// # Example
575    ///
576    /// ```rust
577    /// # use polars_core::prelude::*;
578    /// let mut df1: DataFrame = df!("Element" => ["Copper", "Silver", "Gold"],
579    ///                          "Melting Point (K)" => [1357.77, 1234.93, 1337.33])?;
580    /// let df2: DataFrame = df!("Element" => ["Platinum", "Palladium"],
581    ///                          "Melting Point (K)" => [2041.4, 1828.05])?;
582    ///
583    /// df1.vstack_mut(&df2)?;
584    ///
585    /// assert_eq!(df1.shape(), (5, 2));
586    /// println!("{}", df1);
587    /// # Ok::<(), PolarsError>(())
588    /// ```
589    ///
590    /// Output:
591    ///
592    /// ```text
593    /// shape: (5, 2)
594    /// +-----------+-------------------+
595    /// | Element   | Melting Point (K) |
596    /// | ---       | ---               |
597    /// | str       | f64               |
598    /// +===========+===================+
599    /// | Copper    | 1357.77           |
600    /// +-----------+-------------------+
601    /// | Silver    | 1234.93           |
602    /// +-----------+-------------------+
603    /// | Gold      | 1337.33           |
604    /// +-----------+-------------------+
605    /// | Platinum  | 2041.4            |
606    /// +-----------+-------------------+
607    /// | Palladium | 1828.05           |
608    /// +-----------+-------------------+
609    /// ```
610    pub fn vstack_mut(&mut self, other: &DataFrame) -> PolarsResult<&mut Self> {
611        if self.width() != other.width() {
612            polars_ensure!(
613                self.shape() == (0, 0),
614                ShapeMismatch:
615                "unable to append to a DataFrame of shape {:?} with a DataFrame of width {}",
616                self.shape(), other.width(),
617            );
618
619            self.clone_from(other);
620
621            return Ok(self);
622        }
623
624        let new_height = usize::checked_add(self.height(), other.height()).unwrap();
625
626        unsafe { self.columns_mut_retain_schema() }
627            .iter_mut()
628            .zip(other.columns())
629            .try_for_each::<_, PolarsResult<_>>(|(left, right)| {
630                ensure_can_extend(&*left, right)?;
631                left.append(right)
632                    .with_context(|| format!("failed to vstack column '{}'", right.name()))?;
633                Ok(())
634            })?;
635
636        unsafe { self.set_height(new_height) };
637
638        Ok(self)
639    }
640
641    pub fn vstack_mut_owned(&mut self, other: DataFrame) -> PolarsResult<&mut Self> {
642        if self.width() != other.width() {
643            polars_ensure!(
644                self.shape() == (0, 0),
645                ShapeMismatch:
646                "unable to append to a DataFrame of width {} with a DataFrame of width {}",
647                self.width(), other.width(),
648            );
649
650            *self = other;
651
652            return Ok(self);
653        }
654
655        let new_height = usize::checked_add(self.height(), other.height()).unwrap();
656
657        unsafe { self.columns_mut_retain_schema() }
658            .iter_mut()
659            .zip(other.into_columns())
660            .try_for_each::<_, PolarsResult<_>>(|(left, right)| {
661                ensure_can_extend(&*left, &right)?;
662                let right_name = right.name().clone();
663                left.append_owned(right)
664                    .with_context(|| format!("failed to vstack column '{right_name}'"))?;
665                Ok(())
666            })?;
667
668        unsafe { self.set_height(new_height) };
669
670        Ok(self)
671    }
672
673    /// Concatenate a [`DataFrame`] to this [`DataFrame`]
674    ///
675    /// If many `vstack` operations are done, it is recommended to call [`DataFrame::align_chunks_par`].
676    ///
677    /// # Panics
678    /// Panics if the schema's don't match.
679    pub fn vstack_mut_unchecked(&mut self, other: &DataFrame) -> &mut Self {
680        let new_height = usize::checked_add(self.height(), other.height()).unwrap();
681
682        unsafe { self.columns_mut_retain_schema() }
683            .iter_mut()
684            .zip(other.columns())
685            .for_each(|(left, right)| {
686                left.append(right)
687                    .with_context(|| format!("failed to vstack column '{}'", right.name()))
688                    .expect("should not fail");
689            });
690
691        unsafe { self.set_height(new_height) };
692
693        self
694    }
695
696    /// Concatenate a [`DataFrame`] to this [`DataFrame`]
697    ///
698    /// If many `vstack` operations are done, it is recommended to call [`DataFrame::align_chunks_par`].
699    ///
700    /// # Panics
701    /// Panics if the schema's don't match.
702    pub fn vstack_mut_owned_unchecked(&mut self, other: DataFrame) -> &mut Self {
703        let new_height = usize::checked_add(self.height(), other.height()).unwrap();
704
705        unsafe { self.columns_mut_retain_schema() }
706            .iter_mut()
707            .zip(other.into_columns())
708            .for_each(|(left, right)| {
709                left.append_owned(right).expect("should not fail");
710            });
711
712        unsafe { self.set_height(new_height) };
713
714        self
715    }
716
717    /// Extend the memory backed by this [`DataFrame`] with the values from `other`.
718    ///
719    /// Different from [`vstack`](Self::vstack) which adds the chunks from `other` to the chunks of this [`DataFrame`]
720    /// `extend` appends the data from `other` to the underlying memory locations and thus may cause a reallocation.
721    ///
722    /// If this does not cause a reallocation, the resulting data structure will not have any extra chunks
723    /// and thus will yield faster queries.
724    ///
725    /// Prefer `extend` over `vstack` when you want to do a query after a single append. For instance during
726    /// online operations where you add `n` rows and rerun a query.
727    ///
728    /// Prefer `vstack` over `extend` when you want to append many times before doing a query. For instance
729    /// when you read in multiple files and when to store them in a single `DataFrame`. In the latter case, finish the sequence
730    /// of `append` operations with a [`rechunk`](Self::align_chunks_par).
731    pub fn extend(&mut self, other: &DataFrame) -> PolarsResult<()> {
732        polars_ensure!(
733            self.width() == other.width(),
734            ShapeMismatch:
735            "unable to extend a DataFrame of width {} with a DataFrame of width {}",
736            self.width(), other.width(),
737        );
738
739        let new_height = usize::checked_add(self.height(), other.height()).unwrap();
740
741        unsafe { self.columns_mut_retain_schema() }
742            .iter_mut()
743            .zip(other.columns())
744            .try_for_each::<_, PolarsResult<_>>(|(left, right)| {
745                ensure_can_extend(&*left, right)?;
746                left.extend(right)
747                    .with_context(|| format!("failed to extend column '{}'", right.name()))?;
748                Ok(())
749            })?;
750
751        unsafe { self.set_height(new_height) };
752
753        Ok(())
754    }
755
756    /// Remove a column by name and return the column removed.
757    ///
758    /// # Example
759    ///
760    /// ```rust
761    /// # use polars_core::prelude::*;
762    /// let mut df: DataFrame = df!("Animal" => ["Tiger", "Lion", "Great auk"],
763    ///                             "IUCN" => ["Endangered", "Vulnerable", "Extinct"])?;
764    ///
765    /// let s1: PolarsResult<Column> = df.drop_in_place("Average weight");
766    /// assert!(s1.is_err());
767    ///
768    /// let s2: Column = df.drop_in_place("Animal")?;
769    /// assert_eq!(s2, Column::new("Animal".into(), &["Tiger", "Lion", "Great auk"]));
770    /// # Ok::<(), PolarsError>(())
771    /// ```
772    pub fn drop_in_place(&mut self, name: &str) -> PolarsResult<Column> {
773        let idx = self.try_get_column_index(name)?;
774        Ok(unsafe { self.columns_mut() }.remove(idx))
775    }
776
777    /// Return a new [`DataFrame`] where all null values are dropped.
778    ///
779    /// # Example
780    ///
781    /// ```no_run
782    /// # use polars_core::prelude::*;
783    /// let df1: DataFrame = df!("Country" => ["Malta", "Liechtenstein", "North Korea"],
784    ///                         "Tax revenue (% GDP)" => [Some(32.7), None, None])?;
785    /// assert_eq!(df1.shape(), (3, 2));
786    ///
787    /// let df2: DataFrame = df1.drop_nulls::<String>(None)?;
788    /// assert_eq!(df2.shape(), (1, 2));
789    /// println!("{}", df2);
790    /// # Ok::<(), PolarsError>(())
791    /// ```
792    ///
793    /// Output:
794    ///
795    /// ```text
796    /// shape: (1, 2)
797    /// +---------+---------------------+
798    /// | Country | Tax revenue (% GDP) |
799    /// | ---     | ---                 |
800    /// | str     | f64                 |
801    /// +=========+=====================+
802    /// | Malta   | 32.7                |
803    /// +---------+---------------------+
804    /// ```
805    pub fn drop_nulls<S>(&self, subset: Option<&[S]>) -> PolarsResult<Self>
806    where
807        for<'a> &'a S: AsRef<str>,
808    {
809        if let Some(v) = subset {
810            let v = self.select_to_vec(v)?;
811            self._drop_nulls_impl(v.as_slice())
812        } else {
813            self._drop_nulls_impl(self.columns())
814        }
815    }
816
817    fn _drop_nulls_impl(&self, subset: &[Column]) -> PolarsResult<Self> {
818        // fast path for no nulls in df
819        if subset.iter().all(|s| !s.has_nulls()) {
820            return Ok(self.clone());
821        }
822
823        let mut iter = subset.iter();
824
825        let mask = iter
826            .next()
827            .ok_or_else(|| polars_err!(NoData: "no data to drop nulls from"))?;
828        let mut mask = mask.is_not_null();
829
830        for c in iter {
831            mask = mask & c.is_not_null();
832        }
833        self.filter(&mask)
834    }
835
836    /// Drop a column by name.
837    /// This is a pure method and will return a new [`DataFrame`] instead of modifying
838    /// the current one in place.
839    ///
840    /// # Example
841    ///
842    /// ```rust
843    /// # use polars_core::prelude::*;
844    /// let df1: DataFrame = df!("Ray type" => ["α", "β", "X", "γ"])?;
845    /// let df2: DataFrame = df1.drop("Ray type")?;
846    ///
847    /// assert_eq!(df2.width(), 0);
848    /// # Ok::<(), PolarsError>(())
849    /// ```
850    pub fn drop(&self, name: &str) -> PolarsResult<Self> {
851        let idx = self.try_get_column_index(name)?;
852        let mut new_cols = Vec::with_capacity(self.width() - 1);
853
854        self.columns().iter().enumerate().for_each(|(i, s)| {
855            if i != idx {
856                new_cols.push(s.clone())
857            }
858        });
859
860        Ok(unsafe { DataFrame::_new_unchecked_impl(self.height(), new_cols) })
861    }
862
863    /// Drop columns that are in `names`.
864    pub fn drop_many<I, S>(&self, names: I) -> Self
865    where
866        I: IntoIterator<Item = S>,
867        S: Into<PlSmallStr>,
868    {
869        let names: PlHashSet<PlSmallStr> = names.into_iter().map(|s| s.into()).collect();
870        self.drop_many_amortized(&names)
871    }
872
873    /// Drop columns that are in `names` without allocating a [`HashSet`](std::collections::HashSet).
874    pub fn drop_many_amortized(&self, names: &PlHashSet<PlSmallStr>) -> DataFrame {
875        if names.is_empty() {
876            return self.clone();
877        }
878        let mut new_cols = Vec::with_capacity(self.width().saturating_sub(names.len()));
879        self.columns().iter().for_each(|s| {
880            if !names.contains(s.name()) {
881                new_cols.push(s.clone())
882            }
883        });
884
885        unsafe { DataFrame::new_unchecked(self.height(), new_cols) }
886    }
887
888    /// Insert a new column at a given index without checking for duplicates.
889    /// This can leave the [`DataFrame`] at an invalid state
890    fn insert_column_no_namecheck(
891        &mut self,
892        index: usize,
893        column: Column,
894    ) -> PolarsResult<&mut Self> {
895        polars_ensure!(
896            column.len() == self.height(),
897            ShapeMismatch:
898            "unable to add a column of length {} to a DataFrame of height {}",
899            column.len(), self.height(),
900        );
901
902        unsafe { self.columns_mut() }.insert(index, column);
903        Ok(self)
904    }
905
906    /// Insert a new column at a given index.
907    pub fn insert_column(&mut self, index: usize, column: Column) -> PolarsResult<&mut Self> {
908        let name = column.name();
909
910        polars_ensure!(
911            self.get_column_index(name).is_none(),
912            Duplicate:
913            "column with name {:?} is already present in the DataFrame", name
914        );
915
916        self.insert_column_no_namecheck(index, column)
917    }
918
919    /// Add a new column to this [`DataFrame`] or replace an existing one. Broadcasts unit-length
920    /// columns.
921    pub fn with_column(&mut self, mut column: Column) -> PolarsResult<&mut Self> {
922        column.broadcast_in_place_to(self.height())?;
923
924        if let Some(i) = self.get_column_index(column.name()) {
925            *unsafe { self.columns_mut() }.get_mut(i).unwrap() = column
926        } else {
927            unsafe { self.columns_mut() }.push(column)
928        };
929
930        Ok(self)
931    }
932
933    /// Adds a column to the [`DataFrame`] without doing any checks
934    /// on length or duplicates.
935    ///
936    /// # Safety
937    /// The caller must ensure `column.len() == self.height()` .
938    pub unsafe fn push_column_unchecked(&mut self, column: Column) -> &mut Self {
939        unsafe { self.columns_mut() }.push(column);
940        self
941    }
942
943    /// Add or replace columns to this [`DataFrame`] or replace an existing one.
944    /// Broadcasts unit-length columns, and uses an existing schema to amortize lookups.
945    pub fn with_columns_mut(
946        &mut self,
947        columns: impl IntoIterator<Item = Column>,
948        output_schema: &Schema,
949    ) -> PolarsResult<()> {
950        let columns = columns.into_iter();
951
952        unsafe {
953            self.columns_mut_retain_schema()
954                .reserve(columns.size_hint().0)
955        }
956
957        for c in columns {
958            self.with_column_and_schema_mut(c, output_schema)?;
959        }
960
961        Ok(())
962    }
963
964    fn with_column_and_schema_mut(
965        &mut self,
966        mut column: Column,
967        output_schema: &Schema,
968    ) -> PolarsResult<&mut Self> {
969        column.broadcast_in_place_to(self.height())?;
970
971        let i = output_schema
972            .index_of(column.name())
973            .or_else(|| self.get_column_index(column.name()))
974            .unwrap_or(self.width());
975
976        if i < self.width() {
977            *unsafe { self.columns_mut() }.get_mut(i).unwrap() = column
978        } else if i == self.width() {
979            unsafe { self.columns_mut() }.push(column)
980        } else {
981            // Unordered column insertion is not handled.
982            panic!("{:?}, {}", output_schema, column.name());
983        }
984
985        Ok(self)
986    }
987
988    /// Get a row in the [`DataFrame`]. Beware this is slow.
989    ///
990    /// # Example
991    ///
992    /// ```
993    /// # use polars_core::prelude::*;
994    /// fn example(df: &mut DataFrame, idx: usize) -> Option<Vec<AnyValue>> {
995    ///     df.get(idx)
996    /// }
997    /// ```
998    pub fn get(&self, idx: usize) -> Option<Vec<AnyValue<'_>>> {
999        (idx < self.height()).then(|| self.columns().iter().map(|c| c.get(idx).unwrap()).collect())
1000    }
1001
1002    /// Select a [`Series`] by index.
1003    ///
1004    /// # Example
1005    ///
1006    /// ```rust
1007    /// # use polars_core::prelude::*;
1008    /// let df: DataFrame = df!("Star" => ["Sun", "Betelgeuse", "Sirius A", "Sirius B"],
1009    ///                         "Absolute magnitude" => [4.83, -5.85, 1.42, 11.18])?;
1010    ///
1011    /// let s1: Option<&Column> = df.select_at_idx(0);
1012    /// let s2 = Column::new("Star".into(), ["Sun", "Betelgeuse", "Sirius A", "Sirius B"]);
1013    ///
1014    /// assert_eq!(s1, Some(&s2));
1015    /// # Ok::<(), PolarsError>(())
1016    /// ```
1017    pub fn select_at_idx(&self, idx: usize) -> Option<&Column> {
1018        self.columns().get(idx)
1019    }
1020
1021    /// Get column index of a [`Series`] by name.
1022    /// # Example
1023    ///
1024    /// ```rust
1025    /// # use polars_core::prelude::*;
1026    /// let df: DataFrame = df!("Name" => ["Player 1", "Player 2", "Player 3"],
1027    ///                         "Health" => [100, 200, 500],
1028    ///                         "Mana" => [250, 100, 0],
1029    ///                         "Strength" => [30, 150, 300])?;
1030    ///
1031    /// assert_eq!(df.get_column_index("Name"), Some(0));
1032    /// assert_eq!(df.get_column_index("Health"), Some(1));
1033    /// assert_eq!(df.get_column_index("Mana"), Some(2));
1034    /// assert_eq!(df.get_column_index("Strength"), Some(3));
1035    /// assert_eq!(df.get_column_index("Haste"), None);
1036    /// # Ok::<(), PolarsError>(())
1037    /// ```
1038    pub fn get_column_index(&self, name: &str) -> Option<usize> {
1039        if let Some(schema) = self.cached_schema() {
1040            schema.index_of(name)
1041        } else if self.width() <= LINEAR_SEARCH_LIMIT {
1042            self.columns().iter().position(|s| s.name() == name)
1043        } else {
1044            self.schema().index_of(name)
1045        }
1046    }
1047
1048    /// Get column index of a [`Series`] by name.
1049    pub fn try_get_column_index(&self, name: &str) -> PolarsResult<usize> {
1050        self.get_column_index(name)
1051            .ok_or_else(|| polars_err!(col_not_found = name))
1052    }
1053
1054    /// Select a single column by name.
1055    ///
1056    /// # Example
1057    ///
1058    /// ```rust
1059    /// # use polars_core::prelude::*;
1060    /// let s1 = Column::new("Password".into(), ["123456", "[]B$u$g$s$B#u#n#n#y[]{}"]);
1061    /// let s2 = Column::new("Robustness".into(), ["Weak", "Strong"]);
1062    /// let df: DataFrame = DataFrame::new_infer_height(vec![s1.clone(), s2])?;
1063    ///
1064    /// assert_eq!(df.column("Password")?, &s1);
1065    /// # Ok::<(), PolarsError>(())
1066    /// ```
1067    pub fn column(&self, name: &str) -> PolarsResult<&Column> {
1068        let idx = self.try_get_column_index(name)?;
1069        Ok(self.select_at_idx(idx).unwrap())
1070    }
1071
1072    /// Select column(s) from this [`DataFrame`] and return a new [`DataFrame`].
1073    ///
1074    /// # Examples
1075    ///
1076    /// ```
1077    /// # use polars_core::prelude::*;
1078    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
1079    ///     df.select(["foo", "bar"])
1080    /// }
1081    /// ```
1082    pub fn select<I, S>(&self, names: I) -> PolarsResult<Self>
1083    where
1084        I: IntoIterator<Item = S>,
1085        S: AsRef<str>,
1086    {
1087        DataFrame::new(self.height(), self.select_to_vec(names)?)
1088    }
1089
1090    /// Does not check for duplicates.
1091    ///
1092    /// # Safety
1093    /// `names` must not contain duplicates.
1094    pub unsafe fn select_unchecked<I, S>(&self, names: I) -> PolarsResult<Self>
1095    where
1096        I: IntoIterator<Item = S>,
1097        S: AsRef<str>,
1098    {
1099        Ok(unsafe { DataFrame::new_unchecked(self.height(), self.select_to_vec(names)?) })
1100    }
1101
1102    /// Select column(s) from this [`DataFrame`] and return them into a [`Vec`].
1103    ///
1104    /// This does not error on duplicate selections.
1105    ///
1106    /// # Example
1107    ///
1108    /// ```rust
1109    /// # use polars_core::prelude::*;
1110    /// let df: DataFrame = df!("Name" => ["Methane", "Ethane", "Propane"],
1111    ///                         "Carbon" => [1, 2, 3],
1112    ///                         "Hydrogen" => [4, 6, 8])?;
1113    /// let sv: Vec<Column> = df.select_to_vec(["Carbon", "Hydrogen"])?;
1114    ///
1115    /// assert_eq!(df["Carbon"], sv[0]);
1116    /// assert_eq!(df["Hydrogen"], sv[1]);
1117    /// # Ok::<(), PolarsError>(())
1118    /// ```
1119    pub fn select_to_vec(
1120        &self,
1121        selection: impl IntoIterator<Item = impl AsRef<str>>,
1122    ) -> PolarsResult<Vec<Column>> {
1123        AmortizedColumnSelector::new(self).select_multiple(selection)
1124    }
1125
1126    /// Take the [`DataFrame`] rows by a boolean mask.
1127    ///
1128    /// # Example
1129    ///
1130    /// ```
1131    /// # use polars_core::prelude::*;
1132    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
1133    ///     let mask = df.column("sepal_width")?.is_not_null();
1134    ///     df.filter(&mask)
1135    /// }
1136    /// ```
1137    pub fn filter(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
1138        if self.width() == 0 {
1139            filter_zero_width(self.height(), mask)
1140        } else if mask.len() == 1 && self.len() >= 1 {
1141            if mask.all() && mask.null_count() == 0 {
1142                Ok(self.clone())
1143            } else {
1144                Ok(self.clear())
1145            }
1146        } else {
1147            // Rechunk when not all chunks are aligned. This avoid O(n*m) overhead,
1148            // where n = number of chunks, and m = number of columns.
1149            let all_chunks_aligned = !self.should_rechunk()
1150                && self
1151                    .materialized_column_iter()
1152                    .next()
1153                    .is_some_and(|s| s.chunk_lengths().eq(mask.chunk_lengths()));
1154
1155            let mask = if all_chunks_aligned {
1156                Cow::Borrowed(mask)
1157            } else {
1158                mask.rechunk()
1159            };
1160
1161            let new_columns: Vec<Column> =
1162                self.try_apply_columns_par(|s| s.filter(mask.as_ref()))?;
1163            let out = unsafe {
1164                DataFrame::new_unchecked(new_columns[0].len(), new_columns).with_schema_from(self)
1165            };
1166
1167            Ok(out)
1168        }
1169    }
1170
1171    /// Same as `filter` but does not parallelize.
1172    pub fn filter_seq(&self, mask: &BooleanChunked) -> PolarsResult<Self> {
1173        if self.width() == 0 {
1174            filter_zero_width(self.height(), mask)
1175        } else if mask.len() == 1 && mask.null_count() == 0 && self.len() >= 1 {
1176            if mask.all() && mask.null_count() == 0 {
1177                Ok(self.clone())
1178            } else {
1179                Ok(self.clear())
1180            }
1181        } else {
1182            let all_chunks_aligned = !self.should_rechunk()
1183                && self
1184                    .materialized_column_iter()
1185                    .next()
1186                    .is_some_and(|s| s.chunk_lengths().eq(mask.chunk_lengths()));
1187
1188            let mask = if all_chunks_aligned {
1189                Cow::Borrowed(mask)
1190            } else {
1191                mask.rechunk()
1192            };
1193
1194            let new_columns: Vec<Column> = self.try_apply_columns(|s| s.filter(mask.as_ref()))?;
1195            let out = unsafe {
1196                DataFrame::new_unchecked(new_columns[0].len(), new_columns).with_schema_from(self)
1197            };
1198
1199            Ok(out)
1200        }
1201    }
1202
1203    /// Gather [`DataFrame`] rows by index values.
1204    ///
1205    /// # Example
1206    ///
1207    /// ```
1208    /// # use polars_core::prelude::*;
1209    /// fn example(df: &DataFrame) -> PolarsResult<DataFrame> {
1210    ///     let idx = IdxCa::new("idx".into(), [0, 1, 9]);
1211    ///     df.take(&idx)
1212    /// }
1213    /// ```
1214    pub fn take(&self, indices: &IdxCa) -> PolarsResult<Self> {
1215        check_bounds_ca(indices, self.height().try_into().unwrap_or(IdxSize::MAX))?;
1216
1217        let new_cols = self.apply_columns_par(|c| {
1218            assert_eq!(c.len(), self.height());
1219            unsafe { c.take_unchecked(indices) }
1220        });
1221
1222        Ok(unsafe { DataFrame::new_unchecked(indices.len(), new_cols).with_schema_from(self) })
1223    }
1224
1225    /// # Safety
1226    /// The indices must be in-bounds.
1227    pub unsafe fn take_unchecked(&self, idx: &IdxCa) -> Self {
1228        self.take_unchecked_impl(idx, true)
1229    }
1230
1231    /// # Safety
1232    /// The indices must be in-bounds.
1233    #[cfg(feature = "algorithm_group_by")]
1234    pub unsafe fn gather_group_unchecked(&self, group: &GroupsIndicator) -> Self {
1235        match group {
1236            GroupsIndicator::Idx((_, indices)) => unsafe {
1237                self.take_slice_unchecked_impl(indices.as_slice(), false)
1238            },
1239            GroupsIndicator::Slice([offset, len]) => self.slice(*offset as i64, *len as usize),
1240        }
1241    }
1242
1243    /// # Safety
1244    /// The indices must be in-bounds.
1245    pub unsafe fn take_unchecked_impl(&self, idx: &IdxCa, allow_threads: bool) -> Self {
1246        let cols = if allow_threads && RAYON.current_num_threads() > 1 {
1247            RAYON.install(|| {
1248                if RAYON.current_num_threads() > self.width() {
1249                    let stride = usize::max(idx.len().div_ceil(RAYON.current_num_threads()), 256);
1250                    if self.height() / stride >= 2 {
1251                        self.apply_columns_par(|c| {
1252                            // Nested types initiate a rechunk in their take_unchecked implementation.
1253                            // If we do not rechunk, it will result in rechunk storms downstream.
1254                            let c = if c.dtype().is_nested() {
1255                                &c.rechunk()
1256                            } else {
1257                                c
1258                            };
1259
1260                            (0..idx.len().div_ceil(stride))
1261                                .into_par_iter()
1262                                .map(|i| c.take_unchecked(&idx.slice((i * stride) as i64, stride)))
1263                                .reduce(
1264                                    || Column::new_empty(c.name().clone(), c.dtype()),
1265                                    |mut a, b| {
1266                                        a.append_owned(b).unwrap();
1267                                        a
1268                                    },
1269                                )
1270                        })
1271                    } else {
1272                        self.apply_columns_par(|c| c.take_unchecked(idx))
1273                    }
1274                } else {
1275                    self.apply_columns_par(|c| c.take_unchecked(idx))
1276                }
1277            })
1278        } else {
1279            self.apply_columns(|s| s.take_unchecked(idx))
1280        };
1281
1282        unsafe { DataFrame::new_unchecked(idx.len(), cols).with_schema_from(self) }
1283    }
1284
1285    /// # Safety
1286    /// The indices must be in-bounds.
1287    pub unsafe fn take_slice_unchecked(&self, idx: &[IdxSize]) -> Self {
1288        self.take_slice_unchecked_impl(idx, true)
1289    }
1290
1291    /// # Safety
1292    /// The indices must be in-bounds.
1293    pub unsafe fn take_slice_unchecked_impl(&self, idx: &[IdxSize], allow_threads: bool) -> Self {
1294        let cols = if allow_threads && RAYON.current_num_threads() > 1 {
1295            RAYON.install(|| {
1296                if RAYON.current_num_threads() > self.width() {
1297                    let stride = usize::max(idx.len().div_ceil(RAYON.current_num_threads()), 256);
1298                    if self.height() / stride >= 2 {
1299                        self.apply_columns_par(|c| {
1300                            // Nested types initiate a rechunk in their take_unchecked implementation.
1301                            // If we do not rechunk, it will result in rechunk storms downstream.
1302                            let c = if c.dtype().is_nested() {
1303                                &c.rechunk()
1304                            } else {
1305                                c
1306                            };
1307
1308                            (0..idx.len().div_ceil(stride))
1309                                .into_par_iter()
1310                                .map(|i| {
1311                                    let idx = &idx[i * stride..];
1312                                    let idx = &idx[..idx.len().min(stride)];
1313                                    c.take_slice_unchecked(idx)
1314                                })
1315                                .reduce(
1316                                    || Column::new_empty(c.name().clone(), c.dtype()),
1317                                    |mut a, b| {
1318                                        a.append_owned(b).unwrap();
1319                                        a
1320                                    },
1321                                )
1322                        })
1323                    } else {
1324                        self.apply_columns_par(|s| s.take_slice_unchecked(idx))
1325                    }
1326                } else {
1327                    self.apply_columns_par(|s| s.take_slice_unchecked(idx))
1328                }
1329            })
1330        } else {
1331            self.apply_columns(|s| s.take_slice_unchecked(idx))
1332        };
1333        unsafe { DataFrame::new_unchecked(idx.len(), cols).with_schema_from(self) }
1334    }
1335
1336    /// Rename a column in the [`DataFrame`].
1337    ///
1338    /// Should not be called in a loop as that can lead to quadratic behavior.
1339    ///
1340    /// # Example
1341    ///
1342    /// ```
1343    /// # use polars_core::prelude::*;
1344    /// fn example(df: &mut DataFrame) -> PolarsResult<&mut DataFrame> {
1345    ///     let original_name = "foo";
1346    ///     let new_name = "bar";
1347    ///     df.rename(original_name, new_name.into())
1348    /// }
1349    /// ```
1350    pub fn rename(&mut self, column: &str, name: PlSmallStr) -> PolarsResult<&mut Self> {
1351        if column == name.as_str() {
1352            return Ok(self);
1353        }
1354        polars_ensure!(
1355            !self.schema().contains(&name),
1356            Duplicate: "column rename attempted with already existing name \"{name}\""
1357        );
1358
1359        self.get_column_index(column)
1360            .and_then(|idx| unsafe { self.columns_mut() }.get_mut(idx))
1361            .ok_or_else(|| polars_err!(col_not_found = column))
1362            .map(|c| c.rename(name))?;
1363
1364        Ok(self)
1365    }
1366
1367    pub fn rename_many<'a>(
1368        mut self,
1369        renames: impl Iterator<Item = (&'a str, PlSmallStr)>,
1370    ) -> PolarsResult<Self> {
1371        let schema = self.schema().clone();
1372
1373        for (from, to) in renames {
1374            if from == to.as_str() {
1375                continue;
1376            }
1377
1378            let idx = schema
1379                .index_of(from)
1380                .ok_or_else(|| polars_err!(col_not_found = from))?;
1381
1382            unsafe { self.columns_mut() }
1383                .get_mut(idx)
1384                .unwrap()
1385                .rename(to);
1386        }
1387
1388        // Check for duplicates.
1389        let schema = Schema::from_iter_check_duplicates(
1390            self.columns()
1391                .iter()
1392                .map(|c| c.name().clone())
1393                .zip_eq(schema.iter_values().cloned()),
1394        )?;
1395
1396        unsafe { self.set_schema(Arc::new(schema)) };
1397
1398        Ok(self)
1399    }
1400
1401    /// Sort [`DataFrame`] in place.
1402    ///
1403    /// See [`DataFrame::sort`] for more instruction.
1404    pub fn sort_in_place(
1405        &mut self,
1406        by: impl IntoIterator<Item = impl AsRef<str>>,
1407        sort_options: SortMultipleOptions,
1408    ) -> PolarsResult<&mut Self> {
1409        let by: Vec<_> = by.into_iter().collect();
1410        // Several keys are sorted through a row encoding of single chunks; one
1411        // key may skip the sort by its sorted flag.
1412        if by.len() > 1 {
1413            self.rechunk_mut_par();
1414        }
1415        let by_column = self.select_to_vec(by)?;
1416
1417        let mut out = self.sort_impl(by_column, sort_options, None)?;
1418        unsafe { out.set_schema_from(self) };
1419
1420        *self = out;
1421
1422        Ok(self)
1423    }
1424
1425    #[doc(hidden)]
1426    /// This is the dispatch of Self::sort, and exists to reduce compile bloat by monomorphization.
1427    pub fn sort_impl(
1428        &self,
1429        by_column: Vec<Column>,
1430        sort_options: SortMultipleOptions,
1431        slice: Option<(i64, usize)>,
1432    ) -> PolarsResult<Self> {
1433        if by_column.is_empty() {
1434            // If no columns selected, any order (including original order) is correct.
1435            return if let Some((offset, len)) = slice {
1436                Ok(self.slice(offset, len))
1437            } else {
1438                Ok(self.clone())
1439            };
1440        }
1441
1442        for column in &by_column {
1443            if column.dtype().is_object() {
1444                polars_bail!(
1445                    InvalidOperation: "column '{}' has a dtype of '{}', which does not support sorting", column.name(), column.dtype()
1446                )
1447            }
1448        }
1449
1450        // note that the by_column argument also contains evaluated expression from
1451        // polars-lazy that may not even be present in this dataframe. therefore
1452        // when we try to set the first columns as sorted, we ignore the error as
1453        // expressions are not present (they are renamed to _POLARS_SORT_COLUMN_i.
1454        let first_descending = sort_options.descending[0];
1455        let first_by_column = by_column[0].name().to_string();
1456
1457        let set_sorted = |df: &mut DataFrame| {
1458            // Mark the first sort column as sorted; if the column does not exist it
1459            // is ok, because we sorted by an expression not present in the dataframe
1460            let _ = df.apply(&first_by_column, |s| {
1461                let mut s = s.clone();
1462                if first_descending {
1463                    s.set_sorted_flag(IsSorted::Descending)
1464                } else {
1465                    s.set_sorted_flag(IsSorted::Ascending)
1466                }
1467                s
1468            });
1469        };
1470
1471        if self.shape_has_zero() {
1472            let mut out = self.clone();
1473            set_sorted(&mut out);
1474            return Ok(out);
1475        }
1476
1477        if let Some((0, k)) = slice {
1478            if k < self.height() {
1479                return self.bottom_k_impl(k, by_column, sort_options);
1480            }
1481        }
1482        // Check if the required column is already sorted; if so we can exit early
1483        // We can do so when there is only one column to sort by, for multiple columns
1484        // it will be complicated to do so
1485        #[cfg(feature = "dtype-categorical")]
1486        let is_not_categorical_enum =
1487            !(matches!(by_column[0].dtype(), DataType::Categorical(_, _))
1488                || matches!(by_column[0].dtype(), DataType::Enum(_, _)));
1489
1490        #[cfg(not(feature = "dtype-categorical"))]
1491        #[allow(non_upper_case_globals)]
1492        const is_not_categorical_enum: bool = true;
1493
1494        if by_column.len() == 1 && is_not_categorical_enum {
1495            let required_sorting = if sort_options.descending[0] {
1496                IsSorted::Descending
1497            } else {
1498                IsSorted::Ascending
1499            };
1500            // If null count is 0 then nulls_last doesnt matter
1501            // Safe to get value at last position since the dataframe is not empty (taken care above)
1502            let no_sorting_required = (by_column[0].is_sorted_flag() == required_sorting)
1503                && ((by_column[0].null_count() == 0)
1504                    || by_column[0].get(by_column[0].len() - 1).unwrap().is_null()
1505                        == sort_options.nulls_last[0]);
1506
1507            if no_sorting_required {
1508                return if let Some((offset, len)) = slice {
1509                    Ok(self.slice(offset, len))
1510                } else {
1511                    Ok(self.clone())
1512                };
1513            }
1514        }
1515
1516        let has_nested = by_column.iter().any(|s| s.dtype().is_nested());
1517        let allow_threads = sort_options.multithreaded;
1518
1519        // a lot of indirection in both sorting and take
1520        let mut df = self.clone();
1521        let df = df.rechunk_mut_par();
1522        let mut take = match (by_column.len(), has_nested) {
1523            (1, false) => {
1524                let s = &by_column[0];
1525                let options = SortOptions {
1526                    descending: sort_options.descending[0],
1527                    nulls_last: sort_options.nulls_last[0],
1528                    multithreaded: sort_options.multithreaded,
1529                    maintain_order: sort_options.maintain_order,
1530                    limit: sort_options.limit,
1531                };
1532                // fast path for a frame with a single series
1533                // no need to compute the sort indices and then take by these indices
1534                // simply sort and return as frame
1535                if df.width() == 1 && df.try_get_column_index(s.name().as_str()).is_ok() {
1536                    let mut out = s.sort_with(options)?;
1537                    if let Some((offset, len)) = slice {
1538                        out = out.slice(offset, len);
1539                    }
1540                    return Ok(out.into_frame());
1541                }
1542                s.arg_sort(options)
1543            },
1544            _ => arg_sort(&by_column, sort_options)?,
1545        };
1546
1547        if let Some((offset, len)) = slice {
1548            take = take.slice(offset, len);
1549        }
1550
1551        // SAFETY:
1552        // the created indices are in bounds
1553        let mut df = unsafe { df.take_unchecked_impl(&take, allow_threads) };
1554        set_sorted(&mut df);
1555        Ok(df)
1556    }
1557
1558    /// Create a `DataFrame` that has fields for all the known runtime metadata for each column.
1559    ///
1560    /// This dataframe does not necessarily have a specified schema and may be changed at any
1561    /// point. It is primarily used for debugging.
1562    pub fn _to_metadata(&self) -> DataFrame {
1563        let num_columns = self.width();
1564
1565        let mut column_names =
1566            StringChunkedBuilder::new(PlSmallStr::from_static("column_name"), num_columns);
1567        let mut repr_ca = StringChunkedBuilder::new(PlSmallStr::from_static("repr"), num_columns);
1568        let mut sorted_asc_ca =
1569            BooleanChunkedBuilder::new(PlSmallStr::from_static("sorted_asc"), num_columns);
1570        let mut sorted_dsc_ca =
1571            BooleanChunkedBuilder::new(PlSmallStr::from_static("sorted_dsc"), num_columns);
1572        let mut fast_explode_list_ca =
1573            BooleanChunkedBuilder::new(PlSmallStr::from_static("fast_explode_list"), num_columns);
1574        let mut materialized_at_ca =
1575            StringChunkedBuilder::new(PlSmallStr::from_static("materialized_at"), num_columns);
1576
1577        for col in self.columns() {
1578            let flags = col.get_flags();
1579
1580            let (repr, materialized_at) = match col {
1581                Column::Series(s) => ("series", s.materialized_at()),
1582                Column::Scalar(_) => ("scalar", None),
1583            };
1584            let sorted_asc = flags.contains(StatisticsFlags::IS_SORTED_ASC);
1585            let sorted_dsc = flags.contains(StatisticsFlags::IS_SORTED_DSC);
1586            let fast_explode_list = flags.contains(StatisticsFlags::CAN_FAST_EXPLODE_LIST);
1587
1588            column_names.append_value(col.name().clone());
1589            repr_ca.append_value(repr);
1590            sorted_asc_ca.append_value(sorted_asc);
1591            sorted_dsc_ca.append_value(sorted_dsc);
1592            fast_explode_list_ca.append_value(fast_explode_list);
1593            materialized_at_ca.append_option(materialized_at.map(|v| format!("{v:#?}")));
1594        }
1595
1596        unsafe {
1597            DataFrame::new_unchecked(
1598                self.width(),
1599                vec![
1600                    column_names.finish().into_column(),
1601                    repr_ca.finish().into_column(),
1602                    sorted_asc_ca.finish().into_column(),
1603                    sorted_dsc_ca.finish().into_column(),
1604                    fast_explode_list_ca.finish().into_column(),
1605                    materialized_at_ca.finish().into_column(),
1606                ],
1607            )
1608        }
1609    }
1610    /// Return a sorted clone of this [`DataFrame`].
1611    ///
1612    /// In many cases the output chunks will be continuous in memory but this is not guaranteed
1613    /// # Example
1614    ///
1615    /// Sort by a single column with default options:
1616    /// ```
1617    /// # use polars_core::prelude::*;
1618    /// fn sort_by_sepal_width(df: &DataFrame) -> PolarsResult<DataFrame> {
1619    ///     df.sort(["sepal_width"], Default::default())
1620    /// }
1621    /// ```
1622    /// Sort by a single column with specific order:
1623    /// ```
1624    /// # use polars_core::prelude::*;
1625    /// fn sort_with_specific_order(df: &DataFrame, descending: bool) -> PolarsResult<DataFrame> {
1626    ///     df.sort(
1627    ///         ["sepal_width"],
1628    ///         SortMultipleOptions::new()
1629    ///             .with_order_descending(descending)
1630    ///     )
1631    /// }
1632    /// ```
1633    /// Sort by multiple columns with specifying order for each column:
1634    /// ```
1635    /// # use polars_core::prelude::*;
1636    /// fn sort_by_multiple_columns_with_specific_order(df: &DataFrame) -> PolarsResult<DataFrame> {
1637    ///     df.sort(
1638    ///         ["sepal_width", "sepal_length"],
1639    ///         SortMultipleOptions::new()
1640    ///             .with_order_descending_multi([false, true])
1641    ///     )
1642    /// }
1643    /// ```
1644    /// See [`SortMultipleOptions`] for more options.
1645    ///
1646    /// Also see [`DataFrame::sort_in_place`].
1647    pub fn sort(
1648        &self,
1649        by: impl IntoIterator<Item = impl AsRef<str>>,
1650        sort_options: SortMultipleOptions,
1651    ) -> PolarsResult<Self> {
1652        let mut df = self.clone();
1653        df.sort_in_place(by, sort_options)?;
1654        Ok(df)
1655    }
1656
1657    /// Replace a column with a [`Column`].
1658    ///
1659    /// # Example
1660    ///
1661    /// ```rust
1662    /// # use polars_core::prelude::*;
1663    /// let mut df: DataFrame = df!("Country" => ["United States", "China"],
1664    ///                         "Area (km²)" => [9_833_520, 9_596_961])?;
1665    /// let s: Column = Column::new("Country".into(), ["USA", "PRC"]);
1666    ///
1667    /// assert!(df.replace("Nation", s.clone()).is_err());
1668    /// assert!(df.replace("Country", s).is_ok());
1669    /// # Ok::<(), PolarsError>(())
1670    /// ```
1671    pub fn replace(&mut self, column: &str, new_col: Column) -> PolarsResult<&mut Self> {
1672        self.apply(column, |_| new_col)
1673    }
1674
1675    /// Replace column at index `idx` with a [`Series`].
1676    ///
1677    /// # Example
1678    ///
1679    /// ```ignored
1680    /// # use polars_core::prelude::*;
1681    /// let s0 = Series::new("foo".into(), ["ham", "spam", "egg"]);
1682    /// let s1 = Series::new("ascii".into(), [70, 79, 79]);
1683    /// let mut df = DataFrame::new_infer_height(vec![s0, s1])?;
1684    ///
1685    /// // Add 32 to get lowercase ascii values
1686    /// df.replace_column(1, df.select_at_idx(1).unwrap() + 32);
1687    /// # Ok::<(), PolarsError>(())
1688    /// ```
1689    pub fn replace_column(&mut self, index: usize, new_column: Column) -> PolarsResult<&mut Self> {
1690        polars_ensure!(
1691            index < self.width(),
1692            ShapeMismatch:
1693            "unable to replace at index {}, the DataFrame has only {} columns",
1694            index, self.width(),
1695        );
1696
1697        polars_ensure!(
1698            new_column.len() == self.height(),
1699            ShapeMismatch:
1700            "unable to replace a column, series length {} doesn't match the DataFrame height {}",
1701            new_column.len(), self.height(),
1702        );
1703
1704        unsafe { *self.columns_mut().get_mut(index).unwrap() = new_column };
1705
1706        Ok(self)
1707    }
1708
1709    /// Apply a closure to a column. This is the recommended way to do in place modification.
1710    ///
1711    /// # Example
1712    ///
1713    /// ```rust
1714    /// # use polars_core::prelude::*;
1715    /// let s0 = Column::new("foo".into(), ["ham", "spam", "egg"]);
1716    /// let s1 = Column::new("names".into(), ["Jean", "Claude", "van"]);
1717    /// let mut df = DataFrame::new_infer_height(vec![s0, s1])?;
1718    ///
1719    /// fn str_to_len(str_val: &Column) -> Column {
1720    ///     str_val.str()
1721    ///         .unwrap()
1722    ///         .iter()
1723    ///         .map(|opt_name: Option<&str>| {
1724    ///             opt_name.map(|name: &str| name.len() as u32)
1725    ///          })
1726    ///         .collect::<UInt32Chunked>()
1727    ///         .into_column()
1728    /// }
1729    ///
1730    /// // Replace the names column by the length of the names.
1731    /// df.apply("names", str_to_len);
1732    /// # Ok::<(), PolarsError>(())
1733    /// ```
1734    /// Results in:
1735    ///
1736    /// ```text
1737    /// +--------+-------+
1738    /// | foo    |       |
1739    /// | ---    | names |
1740    /// | str    | u32   |
1741    /// +========+=======+
1742    /// | "ham"  | 4     |
1743    /// +--------+-------+
1744    /// | "spam" | 6     |
1745    /// +--------+-------+
1746    /// | "egg"  | 3     |
1747    /// +--------+-------+
1748    /// ```
1749    pub fn apply<F, C>(&mut self, name: &str, f: F) -> PolarsResult<&mut Self>
1750    where
1751        F: FnOnce(&Column) -> C,
1752        C: IntoColumn,
1753    {
1754        let idx = self.try_get_column_index(name)?;
1755        self.apply_at_idx(idx, f)?;
1756        Ok(self)
1757    }
1758
1759    /// Apply a closure to a column at index `idx`. This is the recommended way to do in place
1760    /// modification.
1761    ///
1762    /// # Example
1763    ///
1764    /// ```rust
1765    /// # use polars_core::prelude::*;
1766    /// let s0 = Column::new("foo".into(), ["ham", "spam", "egg"]);
1767    /// let s1 = Column::new("ascii".into(), [70, 79, 79]);
1768    /// let mut df = DataFrame::new_infer_height(vec![s0, s1])?;
1769    ///
1770    /// // Add 32 to get lowercase ascii values
1771    /// df.apply_at_idx(1, |s| s + 32);
1772    /// # Ok::<(), PolarsError>(())
1773    /// ```
1774    /// Results in:
1775    ///
1776    /// ```text
1777    /// +--------+-------+
1778    /// | foo    | ascii |
1779    /// | ---    | ---   |
1780    /// | str    | i32   |
1781    /// +========+=======+
1782    /// | "ham"  | 102   |
1783    /// +--------+-------+
1784    /// | "spam" | 111   |
1785    /// +--------+-------+
1786    /// | "egg"  | 111   |
1787    /// +--------+-------+
1788    /// ```
1789    pub fn apply_at_idx<F, C>(&mut self, idx: usize, f: F) -> PolarsResult<&mut Self>
1790    where
1791        F: FnOnce(&Column) -> C,
1792        C: IntoColumn,
1793    {
1794        let df_height = self.height();
1795        let width = self.width();
1796
1797        let cached_schema = self.cached_schema().cloned();
1798
1799        let col = unsafe { self.columns_mut() }.get_mut(idx).ok_or_else(|| {
1800            polars_err!(
1801                ComputeError: "invalid column index: {} for a DataFrame with {} columns",
1802                idx, width
1803            )
1804        })?;
1805
1806        let new_col = f(col)
1807            .into_column()
1808            .with_name(col.name().clone())
1809            .broadcast_owned_to(df_height)?;
1810        let col_before = std::mem::replace(col, new_col);
1811
1812        if col.dtype() == col_before.dtype() {
1813            unsafe { self.set_opt_schema(cached_schema) };
1814        }
1815
1816        Ok(self)
1817    }
1818
1819    /// Apply a closure that may fail to a column at index `idx`. This is the recommended way to do in place
1820    /// modification.
1821    ///
1822    /// # Example
1823    ///
1824    /// This is the idiomatic way to replace some values a column of a `DataFrame` given range of indexes.
1825    ///
1826    /// ```rust
1827    /// # use polars_core::prelude::*;
1828    /// let s0 = Column::new("foo".into(), ["ham", "spam", "egg", "bacon", "quack"]);
1829    /// let s1 = Column::new("values".into(), [1, 2, 3, 4, 5]);
1830    /// let mut df = DataFrame::new_infer_height(vec![s0, s1])?;
1831    ///
1832    /// let idx = vec![0, 1, 4];
1833    ///
1834    /// df.try_apply("foo", |c| {
1835    ///     c.str()?
1836    ///     .scatter_with(idx, |opt_val| opt_val.map(|string| format!("{}-is-modified", string)))
1837    /// });
1838    /// # Ok::<(), PolarsError>(())
1839    /// ```
1840    /// Results in:
1841    ///
1842    /// ```text
1843    /// +---------------------+--------+
1844    /// | foo                 | values |
1845    /// | ---                 | ---    |
1846    /// | str                 | i32    |
1847    /// +=====================+========+
1848    /// | "ham-is-modified"   | 1      |
1849    /// +---------------------+--------+
1850    /// | "spam-is-modified"  | 2      |
1851    /// +---------------------+--------+
1852    /// | "egg"               | 3      |
1853    /// +---------------------+--------+
1854    /// | "bacon"             | 4      |
1855    /// +---------------------+--------+
1856    /// | "quack-is-modified" | 5      |
1857    /// +---------------------+--------+
1858    /// ```
1859    pub fn try_apply_at_idx<F, C>(&mut self, idx: usize, f: F) -> PolarsResult<&mut Self>
1860    where
1861        F: FnOnce(&Column) -> PolarsResult<C>,
1862        C: IntoColumn,
1863    {
1864        let df_height = self.height();
1865        let width = self.width();
1866
1867        let cached_schema = self.cached_schema().cloned();
1868
1869        let col = unsafe { self.columns_mut() }.get_mut(idx).ok_or_else(|| {
1870            polars_err!(
1871                ComputeError: "invalid column index: {} for a DataFrame with {} columns",
1872                idx, width
1873            )
1874        })?;
1875
1876        let mut new_col = f(col).map(|c| c.into_column())?;
1877
1878        polars_ensure!(
1879            new_col.len() == df_height,
1880            ShapeMismatch:
1881            "try_apply_at_idx: resulting Series has length {} while the DataFrame has height {}",
1882            new_col.len(), df_height
1883        );
1884
1885        // make sure the name remains the same after applying the closure
1886        new_col = new_col.with_name(col.name().clone());
1887        let col_before = std::mem::replace(col, new_col);
1888
1889        if col.dtype() == col_before.dtype() {
1890            unsafe { self.set_opt_schema(cached_schema) };
1891        }
1892
1893        Ok(self)
1894    }
1895
1896    /// Apply a closure that may fail to a column. This is the recommended way to do in place
1897    /// modification.
1898    ///
1899    /// # Example
1900    ///
1901    /// This is the idiomatic way to replace some values a column of a `DataFrame` given a boolean mask.
1902    ///
1903    /// ```rust
1904    /// # use polars_core::prelude::*;
1905    /// let s0 = Column::new("foo".into(), ["ham", "spam", "egg", "bacon", "quack"]);
1906    /// let s1 = Column::new("values".into(), [1, 2, 3, 4, 5]);
1907    /// let mut df = DataFrame::new_infer_height(vec![s0, s1])?;
1908    ///
1909    /// // create a mask
1910    /// let values = df.column("values")?.as_materialized_series();
1911    /// let mask = values.lt_eq(1)? | values.gt_eq(5_i32)?;
1912    ///
1913    /// df.try_apply("foo", |c| {
1914    ///     c.str()?
1915    ///     .set(&mask, Some("not_within_bounds"))
1916    /// });
1917    /// # Ok::<(), PolarsError>(())
1918    /// ```
1919    /// Results in:
1920    ///
1921    /// ```text
1922    /// +---------------------+--------+
1923    /// | foo                 | values |
1924    /// | ---                 | ---    |
1925    /// | str                 | i32    |
1926    /// +=====================+========+
1927    /// | "not_within_bounds" | 1      |
1928    /// +---------------------+--------+
1929    /// | "spam"              | 2      |
1930    /// +---------------------+--------+
1931    /// | "egg"               | 3      |
1932    /// +---------------------+--------+
1933    /// | "bacon"             | 4      |
1934    /// +---------------------+--------+
1935    /// | "not_within_bounds" | 5      |
1936    /// +---------------------+--------+
1937    /// ```
1938    pub fn try_apply<F, C>(&mut self, column: &str, f: F) -> PolarsResult<&mut Self>
1939    where
1940        F: FnOnce(&Series) -> PolarsResult<C>,
1941        C: IntoColumn,
1942    {
1943        let idx = self.try_get_column_index(column)?;
1944        self.try_apply_at_idx(idx, |c| f(c.as_materialized_series()))
1945    }
1946
1947    /// Slice the [`DataFrame`] along the rows.
1948    ///
1949    /// # Example
1950    ///
1951    /// ```rust
1952    /// # use polars_core::prelude::*;
1953    /// let df: DataFrame = df!("Fruit" => ["Apple", "Grape", "Grape", "Fig", "Fig"],
1954    ///                         "Color" => ["Green", "Red", "White", "White", "Red"])?;
1955    /// let sl: DataFrame = df.slice(2, 3);
1956    ///
1957    /// assert_eq!(sl.shape(), (3, 2));
1958    /// println!("{}", sl);
1959    /// # Ok::<(), PolarsError>(())
1960    /// ```
1961    /// Output:
1962    /// ```text
1963    /// shape: (3, 2)
1964    /// +-------+-------+
1965    /// | Fruit | Color |
1966    /// | ---   | ---   |
1967    /// | str   | str   |
1968    /// +=======+=======+
1969    /// | Grape | White |
1970    /// +-------+-------+
1971    /// | Fig   | White |
1972    /// +-------+-------+
1973    /// | Fig   | Red   |
1974    /// +-------+-------+
1975    /// ```
1976    #[must_use]
1977    pub fn slice(&self, offset: i64, length: usize) -> Self {
1978        if offset == 0 && length == self.height() {
1979            return self.clone();
1980        }
1981
1982        if length == 0 {
1983            return self.clear();
1984        }
1985
1986        let cols = self.apply_columns(|s| s.slice(offset, length));
1987
1988        let height = if let Some(fst) = cols.first() {
1989            fst.len()
1990        } else {
1991            let (_, length) = slice_offsets(offset, length, self.height());
1992            length
1993        };
1994
1995        unsafe { DataFrame::_new_unchecked_impl(height, cols).with_schema_from(self) }
1996    }
1997
1998    /// Split [`DataFrame`] at the given `offset`.
1999    pub fn split_at(&self, offset: i64) -> (Self, Self) {
2000        let (a, b) = self.columns().iter().map(|s| s.split_at(offset)).unzip();
2001
2002        let (idx, _) = slice_offsets(offset, 0, self.height());
2003
2004        let a = unsafe { DataFrame::new_unchecked(idx, a).with_schema_from(self) };
2005        let b = unsafe { DataFrame::new_unchecked(self.height() - idx, b).with_schema_from(self) };
2006        (a, b)
2007    }
2008
2009    #[must_use]
2010    pub fn clear(&self) -> Self {
2011        let cols = self.columns().iter().map(|s| s.clear()).collect::<Vec<_>>();
2012        unsafe { DataFrame::_new_unchecked_impl(0, cols).with_schema_from(self) }
2013    }
2014
2015    /// Get the head of the [`DataFrame`].
2016    ///
2017    /// # Example
2018    ///
2019    /// ```rust
2020    /// # use polars_core::prelude::*;
2021    /// let countries: DataFrame =
2022    ///     df!("Rank by GDP (2021)" => [1, 2, 3, 4, 5],
2023    ///         "Continent" => ["North America", "Asia", "Asia", "Europe", "Europe"],
2024    ///         "Country" => ["United States", "China", "Japan", "Germany", "United Kingdom"],
2025    ///         "Capital" => ["Washington", "Beijing", "Tokyo", "Berlin", "London"])?;
2026    /// assert_eq!(countries.shape(), (5, 4));
2027    ///
2028    /// println!("{}", countries.head(Some(3)));
2029    /// # Ok::<(), PolarsError>(())
2030    /// ```
2031    ///
2032    /// Output:
2033    ///
2034    /// ```text
2035    /// shape: (3, 4)
2036    /// +--------------------+---------------+---------------+------------+
2037    /// | Rank by GDP (2021) | Continent     | Country       | Capital    |
2038    /// | ---                | ---           | ---           | ---        |
2039    /// | i32                | str           | str           | str        |
2040    /// +====================+===============+===============+============+
2041    /// | 1                  | North America | United States | Washington |
2042    /// +--------------------+---------------+---------------+------------+
2043    /// | 2                  | Asia          | China         | Beijing    |
2044    /// +--------------------+---------------+---------------+------------+
2045    /// | 3                  | Asia          | Japan         | Tokyo      |
2046    /// +--------------------+---------------+---------------+------------+
2047    /// ```
2048    #[must_use]
2049    pub fn head(&self, length: Option<usize>) -> Self {
2050        let new_height = usize::min(self.height(), length.unwrap_or(HEAD_DEFAULT_LENGTH));
2051        let new_cols = self.apply_columns(|c| c.head(Some(new_height)));
2052
2053        unsafe { DataFrame::new_unchecked(new_height, new_cols).with_schema_from(self) }
2054    }
2055
2056    /// Get the tail of the [`DataFrame`].
2057    ///
2058    /// # Example
2059    ///
2060    /// ```rust
2061    /// # use polars_core::prelude::*;
2062    /// let countries: DataFrame =
2063    ///     df!("Rank (2021)" => [105, 106, 107, 108, 109],
2064    ///         "Apple Price (€/kg)" => [0.75, 0.70, 0.70, 0.65, 0.52],
2065    ///         "Country" => ["Kosovo", "Moldova", "North Macedonia", "Syria", "Turkey"])?;
2066    /// assert_eq!(countries.shape(), (5, 3));
2067    ///
2068    /// println!("{}", countries.tail(Some(2)));
2069    /// # Ok::<(), PolarsError>(())
2070    /// ```
2071    ///
2072    /// Output:
2073    ///
2074    /// ```text
2075    /// shape: (2, 3)
2076    /// +-------------+--------------------+---------+
2077    /// | Rank (2021) | Apple Price (€/kg) | Country |
2078    /// | ---         | ---                | ---     |
2079    /// | i32         | f64                | str     |
2080    /// +=============+====================+=========+
2081    /// | 108         | 0.65               | Syria   |
2082    /// +-------------+--------------------+---------+
2083    /// | 109         | 0.52               | Turkey  |
2084    /// +-------------+--------------------+---------+
2085    /// ```
2086    #[must_use]
2087    pub fn tail(&self, length: Option<usize>) -> Self {
2088        let new_height = usize::min(self.height(), length.unwrap_or(TAIL_DEFAULT_LENGTH));
2089        let new_cols = self.apply_columns(|c| c.tail(Some(new_height)));
2090
2091        unsafe { DataFrame::new_unchecked(new_height, new_cols).with_schema_from(self) }
2092    }
2093
2094    /// Iterator over the rows in this [`DataFrame`] as Arrow RecordBatches.
2095    ///
2096    /// # Panics
2097    ///
2098    /// Panics if the [`DataFrame`] that is passed is not rechunked.
2099    ///
2100    /// This responsibility is left to the caller as we don't want to take mutable references here,
2101    /// but we also don't want to rechunk here, as this operation is costly and would benefit the caller
2102    /// as well.
2103    pub fn iter_chunks(
2104        &self,
2105        compat_level: CompatLevel,
2106        parallel: bool,
2107    ) -> impl Iterator<Item = RecordBatch> + '_ {
2108        debug_assert!(!self.should_rechunk(), "expected equal chunks");
2109
2110        if self.width() == 0 {
2111            return RecordBatchIterWrap::new_zero_width(self.height());
2112        }
2113
2114        // If any of the columns is binview and we don't convert `compat_level` we allow parallelism
2115        // as we must allocate arrow strings/binaries.
2116        let must_convert = compat_level.0 == 0;
2117        let parallel = parallel
2118            && must_convert
2119            && self.width() > 1
2120            && self
2121                .columns()
2122                .iter()
2123                .any(|s| matches!(s.dtype(), DataType::String | DataType::Binary));
2124
2125        RecordBatchIterWrap::Batches(RecordBatchIter {
2126            df: self,
2127            schema: Arc::new(
2128                self.columns()
2129                    .iter()
2130                    .map(|c| c.field().to_arrow(compat_level))
2131                    .collect(),
2132            ),
2133            idx: 0,
2134            n_chunks: usize::max(1, self.first_col_n_chunks()),
2135            compat_level,
2136            parallel,
2137        })
2138    }
2139
2140    /// Iterator over the rows in this [`DataFrame`] as Arrow RecordBatches as physical values.
2141    ///
2142    /// # Panics
2143    ///
2144    /// Panics if the [`DataFrame`] that is passed is not rechunked.
2145    ///
2146    /// This responsibility is left to the caller as we don't want to take mutable references here,
2147    /// but we also don't want to rechunk here, as this operation is costly and would benefit the caller
2148    /// as well.
2149    pub fn iter_chunks_physical(&self) -> impl Iterator<Item = RecordBatch> + '_ {
2150        debug_assert!(!self.should_rechunk());
2151
2152        if self.width() == 0 {
2153            return RecordBatchIterWrap::new_zero_width(self.height());
2154        }
2155
2156        RecordBatchIterWrap::PhysicalBatches(PhysRecordBatchIter {
2157            schema: Arc::new(
2158                self.columns()
2159                    .iter()
2160                    .map(|c| c.field().to_arrow(CompatLevel::newest()))
2161                    .collect(),
2162            ),
2163            arr_iters: self
2164                .materialized_column_iter()
2165                .map(|s| s.chunks().iter())
2166                .collect(),
2167        })
2168    }
2169
2170    /// Get a [`DataFrame`] with all the columns in reversed order.
2171    #[must_use]
2172    pub fn reverse(&self) -> Self {
2173        let new_cols = self.apply_columns(Column::reverse);
2174        unsafe { DataFrame::new_unchecked(self.height(), new_cols).with_schema_from(self) }
2175    }
2176
2177    /// Shift the values by a given period and fill the parts that will be empty due to this operation
2178    /// with `Nones`.
2179    ///
2180    /// See the method on [Series](crate::series::SeriesTrait::shift) for more info on the `shift` operation.
2181    #[must_use]
2182    pub fn shift(&self, periods: i64) -> Self {
2183        let col = self.apply_columns_par(|s| s.shift(periods));
2184        unsafe { DataFrame::new_unchecked(self.height(), col).with_schema_from(self) }
2185    }
2186
2187    /// Replace None values with one of the following strategies:
2188    /// * Forward fill (replace None with the previous value)
2189    /// * Backward fill (replace None with the next value)
2190    /// * Mean fill (replace None with the mean of the whole array)
2191    /// * Min fill (replace None with the minimum of the whole array)
2192    /// * Max fill (replace None with the maximum of the whole array)
2193    ///
2194    /// See the method on [Series](crate::series::Series::fill_null) for more info on the `fill_null` operation.
2195    pub fn fill_null(&self, strategy: FillNullStrategy) -> PolarsResult<Self> {
2196        let col = self.try_apply_columns_par(|s| s.fill_null(strategy))?;
2197
2198        Ok(unsafe { DataFrame::new_unchecked(self.height(), col) })
2199    }
2200
2201    /// Drop duplicate rows from a [`DataFrame`].
2202    /// *This fails when there is a column of type List in DataFrame*
2203    ///
2204    /// Stable means that the order is maintained. This has a higher cost than an unstable distinct.
2205    ///
2206    /// # Example
2207    ///
2208    /// ```no_run
2209    /// # use polars_core::prelude::*;
2210    /// let df = df! {
2211    ///               "flt" => [1., 1., 2., 2., 3., 3.],
2212    ///               "int" => [1, 1, 2, 2, 3, 3, ],
2213    ///               "str" => ["a", "a", "b", "b", "c", "c"]
2214    ///           }?;
2215    ///
2216    /// println!("{}", df.unique_stable(None, UniqueKeepStrategy::First, None)?);
2217    /// # Ok::<(), PolarsError>(())
2218    /// ```
2219    /// Returns
2220    ///
2221    /// ```text
2222    /// +-----+-----+-----+
2223    /// | flt | int | str |
2224    /// | --- | --- | --- |
2225    /// | f64 | i32 | str |
2226    /// +=====+=====+=====+
2227    /// | 1   | 1   | "a" |
2228    /// +-----+-----+-----+
2229    /// | 2   | 2   | "b" |
2230    /// +-----+-----+-----+
2231    /// | 3   | 3   | "c" |
2232    /// +-----+-----+-----+
2233    /// ```
2234    #[cfg(feature = "algorithm_group_by")]
2235    pub fn unique_stable(
2236        &self,
2237        subset: Option<&[String]>,
2238        keep: UniqueKeepStrategy,
2239        slice: Option<(i64, usize)>,
2240    ) -> PolarsResult<DataFrame> {
2241        self.unique_impl(
2242            true,
2243            subset.map(|v| v.iter().map(|x| PlSmallStr::from_str(x.as_str())).collect()),
2244            keep,
2245            slice,
2246        )
2247    }
2248
2249    /// Unstable distinct. See [`DataFrame::unique_stable`].
2250    #[cfg(feature = "algorithm_group_by")]
2251    pub fn unique<I, S>(
2252        &self,
2253        subset: Option<&[String]>,
2254        keep: UniqueKeepStrategy,
2255        slice: Option<(i64, usize)>,
2256    ) -> PolarsResult<DataFrame> {
2257        self.unique_impl(
2258            false,
2259            subset.map(|v| v.iter().map(|x| PlSmallStr::from_str(x.as_str())).collect()),
2260            keep,
2261            slice,
2262        )
2263    }
2264
2265    #[cfg(feature = "algorithm_group_by")]
2266    pub fn unique_impl(
2267        &self,
2268        maintain_order: bool,
2269        subset: Option<Vec<PlSmallStr>>,
2270        keep: UniqueKeepStrategy,
2271        slice: Option<(i64, usize)>,
2272    ) -> PolarsResult<Self> {
2273        if self.width() == 0 {
2274            let height = usize::min(self.height(), 1);
2275            return Ok(DataFrame::empty_with_height(height));
2276        }
2277
2278        let names = subset.unwrap_or_else(|| self.get_column_names_owned());
2279        let mut df = self.clone();
2280        // take on multiple chunks is terrible
2281        df.rechunk_mut_par();
2282
2283        let columns = match (keep, maintain_order) {
2284            (UniqueKeepStrategy::First | UniqueKeepStrategy::Any, true) => {
2285                let gb = df.group_by_stable(names)?;
2286                let groups = gb.get_groups();
2287                let (offset, len) = slice.unwrap_or((0, groups.len()));
2288                let groups = groups.slice(offset, len);
2289                df.apply_columns_par(|s| unsafe { s.agg_first(&groups) })
2290            },
2291            (UniqueKeepStrategy::Last, true) => {
2292                // maintain order by last values, so the sorted groups are not correct as they
2293                // are sorted by the first value
2294                let gb = df.group_by_stable(names)?;
2295                let groups = gb.get_groups();
2296
2297                let last_idx: NoNull<IdxCa> = groups
2298                    .iter()
2299                    .map(|g| match g {
2300                        GroupsIndicator::Idx((_first, idx)) => idx[idx.len() - 1],
2301                        GroupsIndicator::Slice([first, len]) => first + len - 1,
2302                    })
2303                    .collect();
2304
2305                let mut last_idx = last_idx.into_inner().sort(false);
2306
2307                if let Some((offset, len)) = slice {
2308                    last_idx = last_idx.slice(offset, len);
2309                }
2310
2311                let last_idx = NoNull::new(last_idx);
2312                let out = unsafe { df.take_unchecked(&last_idx) };
2313                return Ok(out);
2314            },
2315            (UniqueKeepStrategy::First | UniqueKeepStrategy::Any, false) => {
2316                let gb = df.group_by(names)?;
2317                let groups = gb.get_groups();
2318                let (offset, len) = slice.unwrap_or((0, groups.len()));
2319                let groups = groups.slice(offset, len);
2320                df.apply_columns_par(|s| unsafe { s.agg_first(&groups) })
2321            },
2322            (UniqueKeepStrategy::Last, false) => {
2323                let gb = df.group_by(names)?;
2324                let groups = gb.get_groups();
2325                let (offset, len) = slice.unwrap_or((0, groups.len()));
2326                let groups = groups.slice(offset, len);
2327                df.apply_columns_par(|s| unsafe { s.agg_last(&groups) })
2328            },
2329            (UniqueKeepStrategy::None, _) => {
2330                let df_part = df.select(names)?;
2331                let mask = df_part.is_unique()?;
2332                let mut filtered = df.filter(&mask)?;
2333
2334                if let Some((offset, len)) = slice {
2335                    filtered = filtered.slice(offset, len);
2336                }
2337                return Ok(filtered);
2338            },
2339        };
2340        Ok(unsafe { DataFrame::new_unchecked_infer_height(columns).with_schema_from(self) })
2341    }
2342
2343    /// Get a mask of all the unique rows in the [`DataFrame`].
2344    ///
2345    /// # Example
2346    ///
2347    /// ```no_run
2348    /// # use polars_core::prelude::*;
2349    /// let df: DataFrame = df!("Company" => ["Apple", "Microsoft"],
2350    ///                         "ISIN" => ["US0378331005", "US5949181045"])?;
2351    /// let ca: ChunkedArray<BooleanType> = df.is_unique()?;
2352    ///
2353    /// assert!(ca.all());
2354    /// # Ok::<(), PolarsError>(())
2355    /// ```
2356    #[cfg(feature = "algorithm_group_by")]
2357    pub fn is_unique(&self) -> PolarsResult<BooleanChunked> {
2358        let gb = self.group_by(self.get_column_names_owned())?;
2359        let groups = gb.get_groups();
2360        Ok(is_unique_helper(
2361            groups,
2362            self.height() as IdxSize,
2363            true,
2364            false,
2365        ))
2366    }
2367
2368    /// Get a mask of all the duplicated rows in the [`DataFrame`].
2369    ///
2370    /// # Example
2371    ///
2372    /// ```no_run
2373    /// # use polars_core::prelude::*;
2374    /// let df: DataFrame = df!("Company" => ["Alphabet", "Alphabet"],
2375    ///                         "ISIN" => ["US02079K3059", "US02079K1079"])?;
2376    /// let ca: ChunkedArray<BooleanType> = df.is_duplicated()?;
2377    ///
2378    /// assert!(!ca.all());
2379    /// # Ok::<(), PolarsError>(())
2380    /// ```
2381    #[cfg(feature = "algorithm_group_by")]
2382    pub fn is_duplicated(&self) -> PolarsResult<BooleanChunked> {
2383        let gb = self.group_by(self.get_column_names_owned())?;
2384        let groups = gb.get_groups();
2385        Ok(is_unique_helper(
2386            groups,
2387            self.height() as IdxSize,
2388            false,
2389            true,
2390        ))
2391    }
2392
2393    /// Create a new [`DataFrame`] that shows the null counts per column.
2394    #[must_use]
2395    pub fn null_count(&self) -> Self {
2396        let cols =
2397            self.apply_columns(|c| Column::new(c.name().clone(), [c.null_count() as IdxSize]));
2398        unsafe { Self::new_unchecked(1, cols) }
2399    }
2400
2401    /// Hash and combine the row values
2402    #[cfg(feature = "row_hash")]
2403    pub fn hash_rows(
2404        &mut self,
2405        hasher_builder: Option<PlSeedableRandomStateQuality>,
2406    ) -> PolarsResult<UInt64Chunked> {
2407        let dfs = split_df(self, RAYON.current_num_threads(), false);
2408        let (cas, _) = _df_rows_to_hashes_threaded_vertical(&dfs, hasher_builder)?;
2409
2410        let mut iter = cas.into_iter();
2411        let mut acc_ca = iter.next().unwrap();
2412        for ca in iter {
2413            acc_ca.append(&ca)?;
2414        }
2415        Ok(acc_ca.rechunk().into_owned())
2416    }
2417
2418    /// Get the supertype of the columns in this DataFrame
2419    pub fn get_supertype(&self) -> Option<PolarsResult<DataType>> {
2420        self.columns()
2421            .iter()
2422            .map(|s| Ok(s.dtype().clone()))
2423            .reduce(|acc, b| try_get_supertype(&acc?, &b.unwrap()))
2424    }
2425
2426    /// Take by index values given by the slice `idx`.
2427    /// # Warning
2428    /// Be careful with allowing threads when calling this in a large hot loop
2429    /// every thread split may be on rayon stack and lead to SO
2430    #[doc(hidden)]
2431    pub unsafe fn _take_unchecked_slice(&self, idx: &[IdxSize], allow_threads: bool) -> Self {
2432        self._take_unchecked_slice_sorted(idx, allow_threads, IsSorted::Not)
2433    }
2434
2435    /// Take by index values given by the slice `idx`. Use this over `_take_unchecked_slice`
2436    /// if the index value in `idx` are sorted. This will maintain sorted flags.
2437    ///
2438    /// # Warning
2439    /// Be careful with allowing threads when calling this in a large hot loop
2440    /// every thread split may be on rayon stack and lead to SO
2441    #[doc(hidden)]
2442    pub unsafe fn _take_unchecked_slice_sorted(
2443        &self,
2444        idx: &[IdxSize],
2445        allow_threads: bool,
2446        sorted: IsSorted,
2447    ) -> Self {
2448        #[cfg(debug_assertions)]
2449        {
2450            if idx.len() > 2 {
2451                use crate::series::IsSorted;
2452
2453                match sorted {
2454                    IsSorted::Ascending => {
2455                        assert!(idx[0] <= idx[idx.len() - 1]);
2456                    },
2457                    IsSorted::Descending => {
2458                        assert!(idx[0] >= idx[idx.len() - 1]);
2459                    },
2460                    _ => {},
2461                }
2462            }
2463        }
2464        let mut ca = IdxCa::mmap_slice(PlSmallStr::EMPTY, idx);
2465        ca.set_sorted_flag(sorted);
2466        self.take_unchecked_impl(&ca, allow_threads)
2467    }
2468    #[cfg(all(feature = "partition_by", feature = "algorithm_group_by"))]
2469    #[doc(hidden)]
2470    pub fn _partition_by_impl(
2471        &self,
2472        cols: &[PlSmallStr],
2473        stable: bool,
2474        include_key: bool,
2475        parallel: bool,
2476    ) -> PolarsResult<Vec<DataFrame>> {
2477        let selected_keys = self.select_to_vec(cols.iter().cloned())?;
2478        let groups = self.group_by_with_series(selected_keys, parallel, stable)?;
2479        let groups = groups.into_groups();
2480
2481        // drop key columns prior to calculation if requested
2482        let df = if include_key {
2483            self.clone()
2484        } else {
2485            self.drop_many(cols.iter().cloned())
2486        };
2487
2488        if parallel {
2489            // don't parallelize this
2490            // there is a lot of parallelization in take and this may easily SO
2491            RAYON.install(|| {
2492                match groups.as_ref() {
2493                    GroupsType::Idx(idx) => {
2494                        // Rechunk as the gather may rechunk for every group #17562.
2495                        let mut df = df.clone();
2496                        df.rechunk_mut_par();
2497                        Ok(idx
2498                            .into_par_iter()
2499                            .map(|(_, group)| {
2500                                // groups are in bounds
2501                                unsafe {
2502                                    df._take_unchecked_slice_sorted(
2503                                        group,
2504                                        false,
2505                                        IsSorted::Ascending,
2506                                    )
2507                                }
2508                            })
2509                            .collect())
2510                    },
2511                    GroupsType::Slice { groups, .. } => Ok(groups
2512                        .into_par_iter()
2513                        .map(|[first, len]| df.slice(*first as i64, *len as usize))
2514                        .collect()),
2515                }
2516            })
2517        } else {
2518            match groups.as_ref() {
2519                GroupsType::Idx(idx) => {
2520                    // Rechunk as the gather may rechunk for every group #17562.
2521                    let mut df = df;
2522                    df.rechunk_mut();
2523                    Ok(idx
2524                        .into_iter()
2525                        .map(|(_, group)| {
2526                            // groups are in bounds
2527                            unsafe {
2528                                df._take_unchecked_slice_sorted(group, false, IsSorted::Ascending)
2529                            }
2530                        })
2531                        .collect())
2532                },
2533                GroupsType::Slice { groups, .. } => Ok(groups
2534                    .iter()
2535                    .map(|[first, len]| df.slice(*first as i64, *len as usize))
2536                    .collect()),
2537            }
2538        }
2539    }
2540
2541    /// Split into multiple DataFrames partitioned by groups
2542    #[cfg(feature = "partition_by")]
2543    pub fn partition_by<I, S>(&self, cols: I, include_key: bool) -> PolarsResult<Vec<DataFrame>>
2544    where
2545        I: IntoIterator<Item = S>,
2546        S: Into<PlSmallStr>,
2547    {
2548        let cols: UnitVec<PlSmallStr> = cols.into_iter().map(Into::into).collect();
2549        self._partition_by_impl(cols.as_slice(), false, include_key, true)
2550    }
2551
2552    /// Split into multiple DataFrames partitioned by groups
2553    /// Order of the groups are maintained.
2554    #[cfg(feature = "partition_by")]
2555    pub fn partition_by_stable<I, S>(
2556        &self,
2557        cols: I,
2558        include_key: bool,
2559    ) -> PolarsResult<Vec<DataFrame>>
2560    where
2561        I: IntoIterator<Item = S>,
2562        S: Into<PlSmallStr>,
2563    {
2564        let cols: UnitVec<PlSmallStr> = cols.into_iter().map(Into::into).collect();
2565        self._partition_by_impl(cols.as_slice(), true, include_key, true)
2566    }
2567
2568    /// Unnest the given `Struct` columns. This means that the fields of the `Struct` type will be
2569    /// inserted as columns.
2570    #[cfg(feature = "dtype-struct")]
2571    pub fn unnest(
2572        &self,
2573        cols: impl IntoIterator<Item = impl Into<PlSmallStr>>,
2574        separator: Option<&str>,
2575    ) -> PolarsResult<DataFrame> {
2576        self.unnest_impl(cols.into_iter().map(Into::into).collect(), separator)
2577    }
2578
2579    #[cfg(feature = "dtype-struct")]
2580    fn unnest_impl(
2581        &self,
2582        cols: PlHashSet<PlSmallStr>,
2583        separator: Option<&str>,
2584    ) -> PolarsResult<DataFrame> {
2585        let mut new_cols = Vec::with_capacity(std::cmp::min(self.width() * 2, self.width() + 128));
2586        let mut count = 0;
2587        for s in self.columns() {
2588            if cols.contains(s.name()) {
2589                let ca = s.struct_()?.clone();
2590                new_cols.extend(ca.fields_as_series().into_iter().map(|mut f| {
2591                    if let Some(separator) = &separator {
2592                        f.rename(polars_utils::format_pl_smallstr!(
2593                            "{}{}{}",
2594                            s.name(),
2595                            separator,
2596                            f.name()
2597                        ));
2598                    }
2599                    Column::from(f)
2600                }));
2601                count += 1;
2602            } else {
2603                new_cols.push(s.clone())
2604            }
2605        }
2606        if count != cols.len() {
2607            // one or more columns not found
2608            // the code below will return an error with the missing name
2609            let schema = self.schema();
2610            for col in cols {
2611                let _ = schema
2612                    .get(col.as_str())
2613                    .ok_or_else(|| polars_err!(col_not_found = col))?;
2614            }
2615        }
2616
2617        DataFrame::new(self.height(), new_cols)
2618    }
2619
2620    pub fn append_record_batch(&mut self, rb: RecordBatchT<ArrayRef>) -> PolarsResult<()> {
2621        // @Optimize: this does a lot of unnecessary allocations. We should probably have a
2622        // append_chunk or something like this. It is just quite difficult to make that safe.
2623        let df = DataFrame::from(rb);
2624        polars_ensure!(
2625            self.schema() == df.schema(),
2626            SchemaMismatch: "cannot append record batch with different schema\n\n
2627        Got {:?}\nexpected: {:?}", df.schema(), self.schema(),
2628        );
2629        self.vstack_mut_owned_unchecked(df);
2630        Ok(())
2631    }
2632}
2633
2634pub struct RecordBatchIter<'a> {
2635    df: &'a DataFrame,
2636    schema: ArrowSchemaRef,
2637    idx: usize,
2638    n_chunks: usize,
2639    compat_level: CompatLevel,
2640    parallel: bool,
2641}
2642
2643impl Iterator for RecordBatchIter<'_> {
2644    type Item = RecordBatch;
2645
2646    fn next(&mut self) -> Option<Self::Item> {
2647        if self.idx >= self.n_chunks {
2648            return None;
2649        }
2650
2651        // Create a batch of the columns with the same chunk no.
2652        let batch_cols: Vec<ArrayRef> = if self.parallel {
2653            let iter = self
2654                .df
2655                .columns()
2656                .par_iter()
2657                .map(Column::as_materialized_series)
2658                .map(|s| s.to_arrow(self.idx, self.compat_level));
2659            RAYON.install(|| iter.collect())
2660        } else {
2661            self.df
2662                .columns()
2663                .iter()
2664                .map(Column::as_materialized_series)
2665                .map(|s| s.to_arrow(self.idx, self.compat_level))
2666                .collect()
2667        };
2668
2669        let length = batch_cols.first().map_or(0, |arr| arr.len());
2670
2671        self.idx += 1;
2672
2673        Some(RecordBatch::new(length, self.schema.clone(), batch_cols))
2674    }
2675
2676    fn size_hint(&self) -> (usize, Option<usize>) {
2677        let n = self.n_chunks - self.idx;
2678        (n, Some(n))
2679    }
2680}
2681
2682pub struct PhysRecordBatchIter<'a> {
2683    schema: ArrowSchemaRef,
2684    arr_iters: Vec<std::slice::Iter<'a, ArrayRef>>,
2685}
2686
2687impl Iterator for PhysRecordBatchIter<'_> {
2688    type Item = RecordBatch;
2689
2690    fn next(&mut self) -> Option<Self::Item> {
2691        let arrs = self
2692            .arr_iters
2693            .iter_mut()
2694            .map(|phys_iter| phys_iter.next().cloned())
2695            .collect::<Option<Vec<_>>>()?;
2696
2697        let length = arrs.first().map_or(0, |arr| arr.len());
2698        Some(RecordBatch::new(length, self.schema.clone(), arrs))
2699    }
2700
2701    fn size_hint(&self) -> (usize, Option<usize>) {
2702        if let Some(iter) = self.arr_iters.first() {
2703            iter.size_hint()
2704        } else {
2705            (0, None)
2706        }
2707    }
2708}
2709
2710pub enum RecordBatchIterWrap<'a> {
2711    ZeroWidth {
2712        remaining_height: usize,
2713        chunk_size: usize,
2714    },
2715    Batches(RecordBatchIter<'a>),
2716    PhysicalBatches(PhysRecordBatchIter<'a>),
2717}
2718
2719impl<'a> RecordBatchIterWrap<'a> {
2720    fn new_zero_width(height: usize) -> Self {
2721        Self::ZeroWidth {
2722            remaining_height: height,
2723            chunk_size: polars_config::config().ideal_morsel_size() as usize,
2724        }
2725    }
2726}
2727
2728impl Iterator for RecordBatchIterWrap<'_> {
2729    type Item = RecordBatch;
2730
2731    fn next(&mut self) -> Option<Self::Item> {
2732        match self {
2733            Self::ZeroWidth {
2734                remaining_height,
2735                chunk_size,
2736            } => {
2737                let n = usize::min(*remaining_height, *chunk_size);
2738                *remaining_height -= n;
2739
2740                (n > 0).then(|| RecordBatch::new(n, ArrowSchemaRef::default(), vec![]))
2741            },
2742            Self::Batches(v) => v.next(),
2743            Self::PhysicalBatches(v) => v.next(),
2744        }
2745    }
2746
2747    fn size_hint(&self) -> (usize, Option<usize>) {
2748        match self {
2749            Self::ZeroWidth {
2750                remaining_height,
2751                chunk_size,
2752            } => {
2753                let n = remaining_height.div_ceil(*chunk_size);
2754                (n, Some(n))
2755            },
2756            Self::Batches(v) => v.size_hint(),
2757            Self::PhysicalBatches(v) => v.size_hint(),
2758        }
2759    }
2760}
2761
2762// utility to test if we can vstack/extend the columns
2763fn ensure_can_extend(left: &Column, right: &Column) -> PolarsResult<()> {
2764    polars_ensure!(
2765        left.name() == right.name(),
2766        ShapeMismatch: "unable to vstack, column names don't match: {:?} and {:?}",
2767        left.name(), right.name(),
2768    );
2769    Ok(())
2770}
2771
2772#[cfg(test)]
2773mod test {
2774    use super::*;
2775
2776    fn create_frame() -> DataFrame {
2777        let s0 = Column::new("days".into(), [0, 1, 2].as_ref());
2778        let s1 = Column::new("temp".into(), [22.1, 19.9, 7.].as_ref());
2779        DataFrame::new_infer_height(vec![s0, s1]).unwrap()
2780    }
2781
2782    #[test]
2783    fn sort_in_place_keeps_the_chunks_of_a_frame_sorted_by_one_key() {
2784        let mut df = df!("a" => [1, 2], "b" => [1, 1]).unwrap();
2785        df.vstack_mut(&df!("a" => [3, 4], "b" => [1, 1]).unwrap())
2786            .unwrap();
2787        df.apply("a", |c| {
2788            let mut c = c.clone();
2789            c.set_sorted_flag(IsSorted::Ascending);
2790            c
2791        })
2792        .unwrap();
2793        assert_eq!(df.first_col_n_chunks(), 2);
2794
2795        df.sort_in_place(["a"], SortMultipleOptions::default())
2796            .unwrap();
2797        assert_eq!(df.first_col_n_chunks(), 2);
2798
2799        df.sort_in_place(["a", "b"], SortMultipleOptions::default())
2800            .unwrap();
2801        assert_eq!(df.first_col_n_chunks(), 1);
2802        let a = df.column("a").unwrap().as_materialized_series();
2803        assert_eq!(
2804            a.i32().unwrap().to_vec(),
2805            [Some(1), Some(2), Some(3), Some(4)]
2806        );
2807    }
2808
2809    #[test]
2810    #[cfg_attr(miri, ignore)]
2811    fn test_recordbatch_iterator() {
2812        let df = df!(
2813            "foo" => [1, 2, 3, 4, 5]
2814        )
2815        .unwrap();
2816        let mut iter = df.iter_chunks(CompatLevel::newest(), false);
2817        assert_eq!(5, iter.next().unwrap().len());
2818        assert!(iter.next().is_none());
2819    }
2820
2821    #[test]
2822    #[cfg_attr(miri, ignore)]
2823    fn test_select() {
2824        let df = create_frame();
2825        assert_eq!(
2826            df.column("days")
2827                .unwrap()
2828                .as_series()
2829                .unwrap()
2830                .equal(1)
2831                .unwrap()
2832                .sum(),
2833            Some(1)
2834        );
2835    }
2836
2837    #[test]
2838    #[cfg_attr(miri, ignore)]
2839    fn test_filter_broadcast_on_string_col() {
2840        let col_name = "some_col";
2841        let v = vec!["test".to_string()];
2842        let s0 = Column::new(PlSmallStr::from_str(col_name), v);
2843        let mut df = DataFrame::new_infer_height(vec![s0]).unwrap();
2844
2845        df = df
2846            .filter(
2847                &df.column(col_name)
2848                    .unwrap()
2849                    .as_materialized_series()
2850                    .equal("")
2851                    .unwrap(),
2852            )
2853            .unwrap();
2854        assert_eq!(
2855            df.column(col_name)
2856                .unwrap()
2857                .as_materialized_series()
2858                .n_chunks(),
2859            1
2860        );
2861    }
2862
2863    #[test]
2864    #[cfg_attr(miri, ignore)]
2865    fn test_filter_broadcast_on_list_col() {
2866        let s1 = Series::new(PlSmallStr::EMPTY, [true, false, true]);
2867        let ll: ListChunked = [&s1].iter().copied().collect();
2868
2869        let mask = BooleanChunked::from_slice(PlSmallStr::EMPTY, &[false]);
2870        let new = ll.filter(&mask).unwrap();
2871
2872        assert_eq!(new.chunks.len(), 1);
2873        assert_eq!(new.len(), 0);
2874    }
2875
2876    #[test]
2877    fn slice() {
2878        let df = create_frame();
2879        let sliced_df = df.slice(0, 2);
2880        assert_eq!(sliced_df.shape(), (2, 2));
2881    }
2882
2883    #[test]
2884    fn rechunk_false() {
2885        let df = create_frame();
2886        assert!(!df.should_rechunk())
2887    }
2888
2889    #[test]
2890    fn rechunk_true() -> PolarsResult<()> {
2891        let mut base = df!(
2892            "a" => [1, 2, 3],
2893            "b" => [1, 2, 3]
2894        )?;
2895
2896        // Create a series with multiple chunks
2897        let mut s = Series::new("foo".into(), 0..2);
2898        let s2 = Series::new("bar".into(), 0..1);
2899        s.append(&s2)?;
2900
2901        // Append series to frame
2902        let out = base.with_column(s.into_column())?;
2903
2904        // Now we should rechunk
2905        assert!(out.should_rechunk());
2906        Ok(())
2907    }
2908
2909    #[test]
2910    fn test_duplicate_column() {
2911        let mut df = df! {
2912            "foo" => [1, 2, 3]
2913        }
2914        .unwrap();
2915        // check if column is replaced
2916        assert!(
2917            df.with_column(Column::new("foo".into(), &[1, 2, 3]))
2918                .is_ok()
2919        );
2920        assert!(
2921            df.with_column(Column::new("bar".into(), &[1, 2, 3]))
2922                .is_ok()
2923        );
2924        assert!(df.column("bar").is_ok())
2925    }
2926
2927    #[test]
2928    #[cfg_attr(miri, ignore)]
2929    fn distinct() {
2930        let df = df! {
2931            "flt" => [1., 1., 2., 2., 3., 3.],
2932            "int" => [1, 1, 2, 2, 3, 3, ],
2933            "str" => ["a", "a", "b", "b", "c", "c"]
2934        }
2935        .unwrap();
2936        let df = df
2937            .unique_stable(None, UniqueKeepStrategy::First, None)
2938            .unwrap()
2939            .sort(["flt"], SortMultipleOptions::default())
2940            .unwrap();
2941        let valid = df! {
2942            "flt" => [1., 2., 3.],
2943            "int" => [1, 2, 3],
2944            "str" => ["a", "b", "c"]
2945        }
2946        .unwrap();
2947        assert!(df.equals(&valid));
2948    }
2949
2950    #[test]
2951    fn test_vstack() {
2952        // check that it does not accidentally rechunks
2953        let mut df = df! {
2954            "flt" => [1., 1., 2., 2., 3., 3.],
2955            "int" => [1, 1, 2, 2, 3, 3, ],
2956            "str" => ["a", "a", "b", "b", "c", "c"]
2957        }
2958        .unwrap();
2959
2960        df.vstack_mut(&df.slice(0, 3)).unwrap();
2961        assert_eq!(df.first_col_n_chunks(), 2)
2962    }
2963
2964    #[test]
2965    fn test_vstack_on_empty_dataframe() {
2966        let mut df = DataFrame::empty();
2967
2968        let df_data = df! {
2969            "flt" => [1., 1., 2., 2., 3., 3.],
2970            "int" => [1, 1, 2, 2, 3, 3, ],
2971            "str" => ["a", "a", "b", "b", "c", "c"]
2972        }
2973        .unwrap();
2974
2975        df.vstack_mut(&df_data).unwrap();
2976        assert_eq!(df.height(), 6)
2977    }
2978
2979    #[test]
2980    fn test_unique_keep_none_with_slice() {
2981        let df = df! {
2982            "x" => [1, 2, 3, 2, 1]
2983        }
2984        .unwrap();
2985        let out = df
2986            .unique_stable(
2987                Some(&["x".to_string()][..]),
2988                UniqueKeepStrategy::None,
2989                Some((0, 2)),
2990            )
2991            .unwrap();
2992        let expected = df! {
2993            "x" => [3]
2994        }
2995        .unwrap();
2996        assert!(out.equals(&expected));
2997    }
2998
2999    #[test]
3000    #[cfg(feature = "dtype-i8")]
3001    fn test_apply_result_schema() {
3002        let mut df = df! {
3003            "x" => [1, 2, 3, 2, 1]
3004        }
3005        .unwrap();
3006
3007        let schema_before = df.schema().clone();
3008        df.apply("x", |f| f.cast(&DataType::Int8).unwrap()).unwrap();
3009        assert_ne!(&schema_before, df.schema());
3010    }
3011}