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