Skip to main content

polars_core/frame/
mod.rs

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