Adding pivot separator option (#181)

* Adding pivot separator option * Fixing documentation
pola-rs · Mar 14, 2024 · 137a2d9 · 137a2d9
1 parent 307c034
commit 137a2d9
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/__tests__/dataframe.test.ts b/__tests__/dataframe.test.ts
@@ -1243,7 +1243,7 @@ describe("dataframe", () => {
     expect(actual).toFrameEqual(expected);
   });
   test("pivot", () => {
-    const df = pl.DataFrame({
+    let df = pl.DataFrame({
       a: pl.Series([1, 2, 3]).cast(pl.Int32),
       b: pl
         .Series([
@@ -1254,7 +1254,7 @@ describe("dataframe", () => {
         .cast(pl.List(pl.Int32)),
     });
 
-    const expected = pl
+    let expected = pl
       .DataFrame({
         a: pl.Series([1, 2, 3]).cast(pl.Int32),
         "1": pl.Series([[1, 1], null, null]).cast(pl.List(pl.Int32)),
@@ -1263,14 +1263,38 @@ describe("dataframe", () => {
       })
       .select("a", "1", "2", "3");
 
-    const actual = df.pivot("b", {
+    let actual = df.pivot("b", {
       index: "a",
       columns: "a",
       aggregateFunc: "first",
       sortColumns: true,
     });
 
     expect(actual).toFrameEqual(expected, true);
+
+    df = pl.DataFrame({
+      a: ["beep", "bop"],
+      b: ["a", "b"],
+      c: ["s", "f"],
+      d: [7, 8],
+      e: ["x", "y"],
+    });
+    actual = df.pivot(["a", "e"], {
+      index: "b",
+      columns: ["c"],
+      aggregateFunc: "first",
+      separator: "|",
+      maintainOrder: true,
+    });
+
+    expected = pl.DataFrame({
+      b: ["a", "b"],
+      "a|c|s": ["beep", null],
+      "a|c|f": [null, "bop"],
+      "e|c|s": ["x", null],
+      "e|c|f": [null, "y"],
+    });
+    expect(actual).toFrameEqual(expected, true);
   });
 });
 describe("join", () => {

diff --git a/bun.lockb b/bun.lockb
diff --git a/polars/dataframe.ts b/polars/dataframe.ts
@@ -1008,6 +1008,7 @@ export interface DataFrame
    *     Defaults to "first"
    *   @param options.maintainOrder Sort the grouped keys so that the output order is predictable.
    *   @param options.sortColumns Sort the transposed columns by name. Default is by order of discovery.
+   *   @param options.separator Used as separator/delimiter in generated column names.
    *   @example
    * ```
    *   > const df = pl.DataFrame(
@@ -1017,12 +1018,12 @@ export interface DataFrame
    *   ...         "baz": [1, 2, 3, 4, 5, 6],
    *   ...     }
    *   ... );
-   *   > df.pivot({values:"baz", index:"foo", columns:"bar"});
+   *   > df.pivot(values:"baz", {index:"foo", columns:"bar"});
    *   shape: (2, 4)
    *   ┌─────┬─────┬─────┬─────┐
    *   │ foo ┆ A   ┆ B   ┆ C   │
    *   │ --- ┆ --- ┆ --- ┆ --- │
-   *   │ str ┆ i64 ┆ i64 ┆ i64 │
+   *   │ str ┆ f64 ┆ f64 ┆ f64 │
    *   ╞═════╪═════╪═════╪═════╡
    *   │ one ┆ 1   ┆ 2   ┆ 3   │
    *   ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
@@ -1047,6 +1048,7 @@ export interface DataFrame
         | Expr;
       maintainOrder?: boolean;
       sortColumns?: boolean;
+      separator?: string;
     },
   ): DataFrame;
   pivot(options: {
@@ -1065,6 +1067,7 @@ export interface DataFrame
       | Expr;
     maintainOrder?: boolean;
     sortColumns?: boolean;
+    separator?: string;
   }): DataFrame;
   // TODO!
   // /**
@@ -2188,6 +2191,7 @@ export const _DataFrame = (_df: any): DataFrame => {
         maintainOrder = true,
         sortColumns = false,
         aggregateFunc = "first",
+        separator,
       } = options;
       values = values_ ?? values;
       values = typeof values === "string" ? [values] : values;
@@ -2216,7 +2220,15 @@ export const _DataFrame = (_df: any): DataFrame => {
       }
 
       return _DataFrame(
-        _df.pivotExpr(values, index, columns, fn, maintainOrder, sortColumns),
+        _df.pivotExpr(
+          values,
+          index,
+          columns,
+          fn,
+          maintainOrder,
+          sortColumns,
+          separator,
+        ),
       );
     },
     quantile(quantile) {