Skip to content

Commit

Permalink
Fix transpose with option columnNames working with arrays (#195)
Browse files Browse the repository at this point in the history
* Fix slice for columnNames in transpose

* Add return type for transpose
  • Loading branch information
gustaferiksson committed Apr 8, 2024
1 parent dda9327 commit a6ab81a
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 3 deletions.
16 changes: 16 additions & 0 deletions __tests__/dataframe.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
6 changes: 3 additions & 3 deletions polars/dataframe.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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) {
Expand All @@ -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) {
Expand Down

0 comments on commit a6ab81a

Please sign in to comment.