Add variableName and valueName to unpivot (#340)

Add variableName and valueName to unpivot to close #246

---------

Co-authored-by: sezanzeb <[email protected]>
Co-authored-by: Bidek56 <[email protected]>

Author: Bidek56

__tests__/dataframe.test.ts

@@ -673,6 +673,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({

__tests__/lazyframe.test.ts

@@ -1351,4 +1351,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);
+  });
 });

polars/dataframe.ts

@@ -1000,6 +1000,8 @@ export interface DataFrame<T extends Record<string, Series> = any>
    *
    * @param idVars - Columns to use as identifier variables.
    * @param valueVars - Values to use as value variables.
+   * @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({
@@ -1023,7 +1025,14 @@ export interface DataFrame<T extends Record<string, Series> = any>
    * └─────┴─────────────┴───────┘
    * ```
    */
-  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.
    * ___
@@ -2309,9 +2318,23 @@ export const _DataFrame = (_df: any): DataFrame => {
     median() {
       return this.lazy().median().collectSync();
     },
-    unpivot(ids, values) {
+    melt(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) {
         return _Series(_df.hmin() as any) as any;

polars/lazy/dataframe.ts

@@ -107,7 +107,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, opts: LazyOptions): Promise<DataFrame>;
   fetch(numRows?: number): Promise<DataFrame>;
@@ -351,7 +351,18 @@ export interface LazyDataFrame extends Serialize, GroupByOps<LazyGroupBy> {
   /**
    * @see {@link DataFrame.unpivot}
    */
-  unpivot(idVars: ColumnSelection, valueVars: ColumnSelection): LazyDataFrame;
+  melt(idVars: ColumnSelection, valueVars: ColumnSelection): LazyDataFrame;
+  /**
+   * @see {@link DataFrame.unpivot}
+   */
+  unpivot(
+    idVars: ColumnSelection,
+    valueVars: ColumnSelection,
+    options?: {
+      variableName?: string | null;
+      valueName?: string | null;
+    },
+  ): LazyDataFrame;
   /**
    * @see {@link DataFrame.min}
    */
@@ -947,11 +958,26 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
     median() {
       return _LazyDataFrame(_ldf.median());
     },
-    unpivot(ids, values) {
+    melt(ids, values) {
       return _LazyDataFrame(
         _ldf.unpivot(columnOrColumnsStrict(ids), columnOrColumnsStrict(values)),
       );
     },
+    unpivot(ids, values, options) {
+      options = {
+        variableName: null,
+        valueName: null,
+        ...options,
+      };
+      return _LazyDataFrame(
+        _ldf.unpivot(
+          columnOrColumnsStrict(ids),
+          columnOrColumnsStrict(values),
+          options.variableName,
+          options.valueName,
+        ),
+      );
+    },
     min() {
       return _LazyDataFrame(_ldf.min());
     },