Skip to content

Commit

Permalink
Add variableName and valueName to unpivot
Browse files Browse the repository at this point in the history
  • Loading branch information
sezanzeb committed Jul 15, 2024
1 parent 48bd3b9 commit 7388e3f
Show file tree
Hide file tree
Showing 6 changed files with 103 additions and 16 deletions.
22 changes: 22 additions & 0 deletions __tests__/dataframe.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -664,6 +664,28 @@ describe("dataframe", () => {
});
expect(actual).toFrameEqual(expected);
});
test("unpivot renamed", () => {
const df = pl.DataFrame({
id: [1],
asset_key_1: ["123"],
asset_key_2: ["456"],
asset_key_3: ["abc"],
});
const actual = df.unpivot(
"id",
["asset_key_1", "asset_key_2", "asset_key_3"],
{
variableName: "foo",
valueName: "bar"
}
);
const expected = pl.DataFrame({
id: [1, 1, 1],
foo: ["asset_key_1", "asset_key_2", "asset_key_3"],
bar: ["123", "456", "abc"],
});
expect(actual).toFrameEqual(expected);
});
test("min:axis:0", () => {
const actual = pl
.DataFrame({
Expand Down
24 changes: 24 additions & 0 deletions __tests__/lazyframe.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1347,4 +1347,28 @@ describe("lazyframe", () => {
expect(newDF.sort("foo")).toFrameEqual(actualDf);
fs.rmSync("./test.parquet");
});
test("unpivot renamed", () => {
const ldf = pl
.DataFrame({
id: [1],
asset_key_1: ["123"],
asset_key_2: ["456"],
asset_key_3: ["abc"],
})
.lazy();
const actual = ldf.unpivot(
"id",
["asset_key_1", "asset_key_2", "asset_key_3"],
{
variableName: "foo",
valueName: "bar"
}
);
const expected = pl.DataFrame({
id: [1, 1, 1],
foo: ["asset_key_1", "asset_key_2", "asset_key_3"],
bar: ["123", "456", "abc"],
});
expect(actual.collectSync()).toFrameEqual(expected);
});
});
32 changes: 24 additions & 8 deletions polars/dataframe.ts
Original file line number Diff line number Diff line change
Expand Up @@ -924,10 +924,8 @@ export interface DataFrame
*
* @param idVars - Columns to use as identifier variables.
* @param valueVars - Values to use as value variables.
* @param variableName - Name to give to the `variable` column. Defaults to "variable"
* @param valueName - Name to give to the `value` column. Defaults to "value"
* @param streamable - Allow this node to run in the streaming engine.
If this runs in streaming, the output of the unpivot operation will not have a stable ordering.
* @param options.variableName - Name to give to the `variable` column. Defaults to "variable"
* @param options.valueName - Name to give to the `value` column. Defaults to "value"
* @example
* ```
* > const df1 = pl.DataFrame({
Expand All @@ -951,7 +949,14 @@ export interface DataFrame
* └─────┴─────────────┴───────┘
* ```
*/
unpivot(idVars: ColumnSelection, valueVars: ColumnSelection): DataFrame;
unpivot(
idVars: ColumnSelection,
valueVars: ColumnSelection,
options?: {
variableName?: string | null,
valueName?: string | null,
}
): DataFrame;
/**
* Aggregate the columns of this DataFrame to their minimum value.
* ___
Expand Down Expand Up @@ -1751,7 +1756,7 @@ export interface DataFrame
Or combine them:
- "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
By "calendar day", we mean the corresponding time on the next day (which may not be 24 hours, due to daylight savings).
By "calendar day", we mean the corresponding time on the next day (which may not be 24 hours, due to daylight savings).
Similarly for "calendar week", "calendar month", "calendar quarter", and "calendar year".
Parameters
Expand Down Expand Up @@ -2177,8 +2182,19 @@ export const _DataFrame = (_df: any): DataFrame => {
melt(ids, values) {
return wrap("unpivot", columnOrColumns(ids), columnOrColumns(values));
},
unpivot(ids, values) {
return wrap("unpivot", columnOrColumns(ids), columnOrColumns(values));
unpivot(ids, values, options) {
options = {
variableName: null,
valueName: null,
...options
}
return wrap(
"unpivot",
columnOrColumns(ids),
columnOrColumns(values),
options.variableName,
options.valueName
);
},
min(axis = 0) {
if (axis === 1) {
Expand Down
33 changes: 29 additions & 4 deletions polars/lazy/dataframe.ts
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ export interface LazyDataFrame extends Serialize, GroupByOps<LazyGroupBy> {
.. warning::
Streaming mode is considered **unstable**. It may be changed
at any point without it being considered a breaking change.
*
*
*/
fetch(numRows?: number): Promise<DataFrame>;
fetch(numRows: number, opts: LazyOptions): Promise<DataFrame>;
Expand Down Expand Up @@ -369,7 +369,20 @@ export interface LazyDataFrame extends Serialize, GroupByOps<LazyGroupBy> {
* @deprecated *since 0.13.0* use {@link unpivot}
*/
melt(idVars: ColumnSelection, valueVars: ColumnSelection): LazyDataFrame;
unpivot(idVars: ColumnSelection, valueVars: ColumnSelection): LazyDataFrame;
/**
* @see {@link DataFrame.unpivot}
* @param options.streamable - Allow this node to run in the streaming engine.
* If this runs in streaming, the output of the unpivot operation will not have a stable ordering.
*/
unpivot(
idVars: ColumnSelection,
valueVars: ColumnSelection,
options?: {
variableName?: string | null,
valueName?: string | null,
streamable?: boolean,
}
): LazyDataFrame;
/**
* @see {@link DataFrame.min}
*/
Expand Down Expand Up @@ -982,9 +995,21 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
_ldf.unpivot(columnOrColumnsStrict(ids), columnOrColumnsStrict(values)),
);
},
unpivot(ids, values) {
unpivot(ids, values, options) {
options = {
variableName: null,
valueName: null,
streamable: false,
...options
}
return _LazyDataFrame(
_ldf.unpivot(columnOrColumnsStrict(ids), columnOrColumnsStrict(values)),
_ldf.unpivot(
columnOrColumnsStrict(ids),
columnOrColumnsStrict(values),
options.variableName,
options.valueName,
options.streamable
),
);
},
min() {
Expand Down
4 changes: 2 additions & 2 deletions src/dataframe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -925,15 +925,15 @@ impl JsDataFrame {
&self,
id_vars: Vec<String>,
value_vars: Vec<String>,
value_name: Option<String>,
variable_name: Option<String>,
value_name: Option<String>,
streamable: Option<bool>,
) -> napi::Result<JsDataFrame> {
let args = UnpivotArgs {
index: strings_to_smartstrings(id_vars),
on: strings_to_smartstrings(value_vars),
value_name: value_name.map(|s| s.into()),
variable_name: variable_name.map(|s| s.into()),
value_name: value_name.map(|s| s.into()),
streamable: streamable.unwrap_or(false),
};

Expand Down
4 changes: 2 additions & 2 deletions src/lazy/dataframe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -504,15 +504,15 @@ impl JsLazyFrame {
&self,
id_vars: Vec<&str>,
value_vars: Vec<&str>,
value_name: Option<&str>,
variable_name: Option<&str>,
value_name: Option<&str>,
streamable: Option<bool>,
) -> JsLazyFrame {
let args = UnpivotArgs {
index: strings_to_smartstrings(id_vars),
on: strings_to_smartstrings(value_vars),
value_name: value_name.map(|s| s.into()),
variable_name: variable_name.map(|s| s.into()),
value_name: value_name.map(|s| s.into()),
streamable: streamable.unwrap_or(false),
};
let ldf = self.ldf.clone();
Expand Down

0 comments on commit 7388e3f

Please sign in to comment.