Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(python,rust!): Infer values columns in DataFrame.pivot when values is None #14477

Merged
merged 8 commits into from
Feb 25, 2024
Merged

feat(python,rust!): Infer values columns in DataFrame.pivot when values is None #14477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
928854d
Infer `values` columns when absent
mcrumiller Feb 12, 2024
ab6c80e
Remove unwanted test
mcrumiller Feb 21, 2024
1c1f1bb
Rename variables to clarify test
mcrumiller Feb 21, 2024
910ba21
Use set logic to compute columns
mcrumiller Feb 21, 2024
024e21d
Revert python arg order and fix test
mcrumiller Feb 21, 2024
ffaed2c
Preserve column order
mcrumiller Feb 22, 2024
85d9834
Re-add depr test
stinodego Feb 25, 2024
53e5293
Revert to earlier implementation of values logic
stinodego Feb 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions crates/polars-lazy/src/frame/pivot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,22 +31,22 @@ impl PhysicalAggExpr for PivotExpr {
}
}

pub fn pivot<I0, S0, I1, S1, I2, S2>(
pub fn pivot<I0, I1, I2, S0, S1, S2>(
df: &DataFrame,
values: I0,
index: I1,
columns: I2,
index: I0,
columns: I1,
values: Option<I2>,
sort_columns: bool,
agg_expr: Option<Expr>,
// used as separator/delimiter in generated column names.
separator: Option<&str>,
) -> PolarsResult<DataFrame>
where
I0: IntoIterator<Item = S0>,
S0: AsRef<str>,
I1: IntoIterator<Item = S1>,
S1: AsRef<str>,
I2: IntoIterator<Item = S2>,
S0: AsRef<str>,
S1: AsRef<str>,
S2: AsRef<str>,
{
// make sure that the root column is replaced
Expand All @@ -56,31 +56,31 @@ where
});
polars_ops::pivot::pivot(
df,
values,
index,
columns,
values,
sort_columns,
agg_expr,
separator,
)
}

pub fn pivot_stable<I0, S0, I1, S1, I2, S2>(
pub fn pivot_stable<I0, I1, I2, S0, S1, S2>(
df: &DataFrame,
values: I0,
index: I1,
columns: I2,
index: I0,
columns: I1,
values: Option<I2>,
sort_columns: bool,
agg_expr: Option<Expr>,
// used as separator/delimiter in generated column names.
separator: Option<&str>,
) -> PolarsResult<DataFrame>
where
I0: IntoIterator<Item = S0>,
S0: AsRef<str>,
I1: IntoIterator<Item = S1>,
S1: AsRef<str>,
I2: IntoIterator<Item = S2>,
S0: AsRef<str>,
S1: AsRef<str>,
S2: AsRef<str>,
{
// make sure that the root column is replaced
Expand All @@ -90,9 +90,9 @@ where
});
polars_ops::pivot::pivot_stable(
df,
values,
index,
columns,
values,
sort_columns,
agg_expr,
separator,
Expand Down
74 changes: 46 additions & 28 deletions crates/polars-ops/src/frame/pivot/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,27 +82,23 @@ fn restore_logical_type(s: &Series, logical_type: &DataType) -> Series {
/// # Note
/// Polars'/arrow memory is not ideal for transposing operations like pivots.
/// If you have a relatively large table, consider using a group_by over a pivot.
pub fn pivot<I0, S0, I1, S1, I2, S2>(
pub fn pivot<I0, I1, I2, S0, S1, S2>(
pivot_df: &DataFrame,
values: I0,
index: I1,
columns: I2,
index: I0,
columns: I1,
values: Option<I2>,
sort_columns: bool,
agg_fn: Option<PivotAgg>,
separator: Option<&str>,
) -> PolarsResult<DataFrame>
where
I0: IntoIterator<Item = S0>,
S0: AsRef<str>,
I1: IntoIterator<Item = S1>,
S1: AsRef<str>,
I2: IntoIterator<Item = S2>,
S0: AsRef<str>,
S1: AsRef<str>,
S2: AsRef<str>,
{
let values = values
.into_iter()
.map(|s| s.as_ref().to_string())
.collect::<Vec<_>>();
let index = index
.into_iter()
.map(|s| s.as_ref().to_string())
Expand All @@ -111,11 +107,12 @@ where
.into_iter()
.map(|s| s.as_ref().to_string())
.collect::<Vec<_>>();
let values = get_values_columns(pivot_df, &index, &columns, values);
pivot_impl(
pivot_df,
&values,
&index,
&columns,
&values,
agg_fn,
sort_columns,
false,
Expand All @@ -128,27 +125,23 @@ where
/// # Note
/// Polars'/arrow memory is not ideal for transposing operations like pivots.
/// If you have a relatively large table, consider using a group_by over a pivot.
pub fn pivot_stable<I0, S0, I1, S1, I2, S2>(
pub fn pivot_stable<I0, I1, I2, S0, S1, S2>(
pivot_df: &DataFrame,
values: I0,
index: I1,
columns: I2,
index: I0,
columns: I1,
values: Option<I2>,
sort_columns: bool,
agg_fn: Option<PivotAgg>,
separator: Option<&str>,
) -> PolarsResult<DataFrame>
where
I0: IntoIterator<Item = S0>,
S0: AsRef<str>,
I1: IntoIterator<Item = S1>,
S1: AsRef<str>,
I2: IntoIterator<Item = S2>,
S0: AsRef<str>,
S1: AsRef<str>,
S2: AsRef<str>,
{
let values = values
.into_iter()
.map(|s| s.as_ref().to_string())
.collect::<Vec<_>>();
let index = index
.into_iter()
.map(|s| s.as_ref().to_string())
Expand All @@ -157,29 +150,54 @@ where
.into_iter()
.map(|s| s.as_ref().to_string())
.collect::<Vec<_>>();

let values = get_values_columns(pivot_df, &index, &columns, values);
pivot_impl(
pivot_df,
&values,
&index,
&columns,
&values,
agg_fn,
sort_columns,
true,
separator,
)
}

/// Determine `values` columns, which is optional in `pivot` calls.
///
/// If not specified (i.e. is `None`), use all remaining columns in the
/// `DataFrame` after `index` and `columns` have been excluded.
fn get_values_columns<I, S>(
df: &DataFrame,
index: &[String],
columns: &[String],
values: Option<I>,
) -> Vec<String>
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
match values {
Some(v) => v.into_iter().map(|s| s.as_ref().to_string()).collect(),
None => df
.get_column_names()
.into_iter()
.map(|c| c.to_string())
.filter(|c| !(index.contains(c) | columns.contains(c)))
.collect(),
}
}

#[allow(clippy::too_many_arguments)]
fn pivot_impl(
pivot_df: &DataFrame,
// these columns will be aggregated in the nested group_by
values: &[String],
// keys of the first group_by operation
index: &[String],
// these columns will be used for a nested group_by
// the rows of this nested group_by will be pivoted as header column values
columns: &[String],
// these columns will be aggregated in the nested group_by
values: &[String],
// aggregation function
agg_fn: Option<PivotAgg>,
sort_columns: bool,
Expand All @@ -206,19 +224,19 @@ fn pivot_impl(
let pivot_df = unsafe { binding.with_column_unchecked(columns_struct) };
pivot_impl_single_column(
pivot_df,
index,
&column,
values,
index,
agg_fn,
sort_columns,
separator,
)
} else {
pivot_impl_single_column(
pivot_df,
index,
unsafe { columns.get_unchecked(0) },
values,
index,
agg_fn,
sort_columns,
separator,
Expand All @@ -228,9 +246,9 @@ fn pivot_impl(

fn pivot_impl_single_column(
pivot_df: &DataFrame,
index: &[String],
column: &str,
values: &[String],
index: &[String],
agg_fn: Option<PivotAgg>,
sort_columns: bool,
separator: Option<&str>,
Expand Down
Loading
Loading