Fix transpose with option columnNames working with arrays (#195)

* Fix slice for columnNames in transpose

* Add return type for transpose

Author: gustaferiksson

__tests__/dataframe.test.ts

@@ -1104,6 +1104,22 @@ describe("dataframe", () => {
     const actual = df.transpose({ includeHeader: false, columnNames: "abc" });
     expect(actual).toFrameEqual(expected);
   });
+  test("transpose:columnNames:array", () => {
+    const expected = pl.DataFrame({
+      a: [1, 1],
+      b: [2, 2],
+      c: [3, 3],
+    });
+    const df = pl.DataFrame({
+      a: [1, 2, 3],
+      b: [1, 2, 3],
+    });
+    const actual = df.transpose({
+      includeHeader: false,
+      columnNames: ["a", "b", "c"],
+    });
+    expect(actual).toFrameEqual(expected);
+  });
   test("transpose:columnNames:generator", () => {
     const expected = pl.DataFrame({
       col_0: [1, 1],

polars/dataframe.ts

@@ -1594,7 +1594,7 @@ export interface DataFrame
     includeHeader?: boolean;
     headerName?: string;
     columnNames?: Iterable<string>;
-  });
+  }): DataFrame;
   /**
    * Drop duplicate rows from this DataFrame.
    * Note that this fails if there is a column of type `List` in the DataFrame.
@@ -2521,7 +2521,7 @@ export const _DataFrame = (_df: any): DataFrame => {
       const headeName = options?.headerName ?? "column";
       const keep_names_as = includeHeader ? headeName : undefined;
       if (options?.columnNames) {
-        function takeNItems(iterable: Iterable<string>, n) {
+        function takeNItems(iterable: Iterable<string>, n: number) {
           const result: Array<string> = [];
           let i = 0;
           for (const item of iterable) {
@@ -2534,7 +2534,7 @@ export const _DataFrame = (_df: any): DataFrame => {
           return result;
         }
         options.columnNames = Array.isArray(options.columnNames)
-          ? options.columnNames.slice(this.height)
+          ? options.columnNames.slice(0, this.height)
           : takeNItems(options.columnNames, this.height);
       }
       if (!options?.columnNames) {