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