Fix pivot() to correctly accept options object as first parameter (#318)

Fix pivot() to correctly accept options object as first parameter to
close #317

Author: elegantcoder

__tests__/dataframe.test.ts

@@ -1412,7 +1412,7 @@ Series: 'attributes' [struct[5]]
       rows.map((e) => e["attributes"]),
     );
   });
-  test("pivot", () => {
+  test("pivot:values-with-options", () => {
     {
       const df = pl.DataFrame({
         a: pl.Series([1, 2, 3]).cast(pl.Int32),
@@ -1507,6 +1507,115 @@ Series: 'attributes' [struct[5]]
         separator: "/",
       });
 
+      const expected = pl.DataFrame({
+        ix: [1, 2],
+        "foo/a": [1, 4],
+        "foo/b": [7, 1],
+        "bar/a": [2, 0],
+        "bar/b": [9, 4],
+      });
+      expect(actual).toFrameEqual(expected, true);
+    }
+  });
+  test("pivot:options-only", () => {
+    {
+      const df = pl.DataFrame({
+        a: pl.Series([1, 2, 3]).cast(pl.Int32),
+        b: pl
+          .Series([
+            [1, 1],
+            [2, 2],
+            [3, 3],
+          ])
+          .cast(pl.List(pl.Int32)),
+      });
+
+      const expected = pl
+        .DataFrame({
+          a: pl.Series([1, 2, 3]).cast(pl.Int32),
+          "1": pl.Series([[1, 1], null, null]).cast(pl.List(pl.Int32)),
+          "2": pl.Series([null, [2, 2], null]).cast(pl.List(pl.Int32)),
+          "3": pl.Series([null, null, [3, 3]]).cast(pl.List(pl.Int32)),
+        })
+        .select("a", "1", "2", "3");
+
+      const actual = df.pivot({
+        values: "b",
+        index: "a",
+        on: "a",
+        aggregateFunc: "first",
+        sortColumns: true,
+      });
+
+      expect(actual).toFrameEqual(expected, true);
+    }
+
+    {
+      const df = pl.DataFrame({
+        a: ["beep", "bop"],
+        b: ["a", "b"],
+        c: ["s", "f"],
+        d: [7, 8],
+        e: ["x", "y"],
+      });
+      const actual = df.pivot({
+        values: ["a", "e"],
+        index: "b",
+        on: ["b"],
+        aggregateFunc: "first",
+        separator: "|",
+        maintainOrder: true,
+      });
+
+      const expected = pl.DataFrame({
+        b: ["a", "b"],
+        "a|a": ["beep", null],
+        "a|b": [null, "bop"],
+        "e|a": ["x", null],
+        "e|b": [null, "y"],
+      });
+      expect(actual).toFrameEqual(expected, true);
+    }
+    {
+      const df = pl.DataFrame({
+        foo: ["A", "A", "B", "B", "C"],
+        N: [1, 2, 2, 4, 2],
+        bar: ["k", "l", "m", "n", "o"],
+      });
+      const actual = df.pivot({
+        values: ["N"],
+        index: "foo",
+        on: "bar",
+        aggregateFunc: "first",
+      });
+
+      const expected = pl.DataFrame({
+        foo: ["A", "B", "C"],
+        k: [1, null, null],
+        l: [2, null, null],
+        m: [null, 2, null],
+        n: [null, 4, null],
+        o: [null, null, 2],
+      });
+
+      expect(actual).toFrameEqual(expected, true);
+    }
+    {
+      const df = pl.DataFrame({
+        ix: [1, 1, 2, 2, 1, 2],
+        col: ["a", "a", "a", "a", "b", "b"],
+        foo: [0, 1, 2, 2, 7, 1],
+        bar: [0, 2, 0, 0, 9, 4],
+      });
+
+      const actual = df.pivot({
+        values: ["foo", "bar"],
+        index: "ix",
+        on: "col",
+        aggregateFunc: "sum",
+        separator: "/",
+      });
+
       const expected = pl.DataFrame({
         ix: [1, 2],
         "foo/a": [1, 4],

polars/dataframe.ts

@@ -2348,6 +2348,9 @@ export const _DataFrame = (_df: any): DataFrame => {
         .map((d) => mapFn(_DataFrame(d)));
     },
     pivot(values, options?) {
+      if (values && !options) {
+        options = values;
+      }
       let {
         values: values_,
         index,