diff --git a/__tests__/dataframe.test.ts b/__tests__/dataframe.test.ts index 4dfc46254..b66a1c277 100644 --- a/__tests__/dataframe.test.ts +++ b/__tests__/dataframe.test.ts @@ -1038,6 +1038,28 @@ describe("dataframe", () => { }); expect(actual).toFrameEqual(expected); }); + test("sort:nulls_last:false", () => { + const actual = pl + .DataFrame({ + foo: [1, null, 2, 3], + }) + .sort({ by: "foo", nulls_last: false }); + const expected = pl.DataFrame({ + foo: [null, 1, 2, 3], + }); + expect(actual).toFrameEqual(expected); + }); + test("sort:nulls_last:true", () => { + const actual = pl + .DataFrame({ + foo: [1, null, 2, 3], + }) + .sort({ by: "foo", nulls_last: true }); + const expected = pl.DataFrame({ + foo: [1, 2, 3, null], + }); + expect(actual).toFrameEqual(expected); + }); test("std", () => { const actual = pl .DataFrame({ diff --git a/__tests__/lazyframe.test.ts b/__tests__/lazyframe.test.ts index 0fdeb442b..fead12808 100644 --- a/__tests__/lazyframe.test.ts +++ b/__tests__/lazyframe.test.ts @@ -997,6 +997,32 @@ describe("lazyframe", () => { }); expect(actual).toFrameEqual(expected); }); + test("sort:nulls_last:false", () => { + const actual = pl + .DataFrame({ + foo: [1, null, 2, 3], + }) + .lazy() + .sort({ by: "foo", nulls_last: false }) + .collectSync(); + const expected = pl.DataFrame({ + foo: [null, 1, 2, 3], + }); + expect(actual).toFrameEqual(expected); + }); + test("sort:nulls_last:true", () => { + const actual = pl + .DataFrame({ + foo: [1, null, 2, 3], + }) + .lazy() + .sort({ by: "foo", nulls_last: true }) + .collectSync(); + const expected = pl.DataFrame({ + foo: [1, 2, 3, null], + }); + expect(actual).toFrameEqual(expected); + }); test("sum", () => { const actual = pl .DataFrame({ diff --git a/polars/dataframe.ts b/polars/dataframe.ts index 3ac473368..78f22478e 100644 --- a/polars/dataframe.ts +++ b/polars/dataframe.ts @@ -1336,6 +1336,7 @@ export interface DataFrame sort( by: ColumnsOrExpr, descending?: boolean, + nulls_last?: boolean, maintain_order?: boolean, ): DataFrame; sort({ @@ -1345,6 +1346,7 @@ export interface DataFrame }: { by: ColumnsOrExpr; descending?: boolean; + nulls_last?: boolean; maintain_order?: boolean; }): DataFrame; /** @@ -2332,17 +2334,22 @@ export const _DataFrame = (_df: any): DataFrame => { } return wrap("slice", opts.offset, opts.length); }, - sort(arg, descending = false, maintain_order = false) { + sort(arg, descending = false, nulls_last = false, maintain_order = false) { if (arg?.by !== undefined) { - return this.sort(arg.by, arg.descending); + return this.sort( + arg.by, + arg.descending, + arg.nulls_last, + arg.maintain_order, + ); } if (Array.isArray(arg) || Expr.isExpr(arg)) { return _DataFrame(_df) .lazy() - .sort(arg, descending, maintain_order) + .sort(arg, descending, nulls_last, maintain_order) .collectSync({ noOptimization: true, stringCache: false }); } - return wrap("sort", arg, descending, true, maintain_order); + return wrap("sort", arg, descending, nulls_last, maintain_order); }, std() { return this.lazy().std().collectSync(); diff --git a/polars/lazy/dataframe.ts b/polars/lazy/dataframe.ts index 7c5e697f9..b7d7d4a48 100644 --- a/polars/lazy/dataframe.ts +++ b/polars/lazy/dataframe.ts @@ -401,11 +401,13 @@ export interface LazyDataFrame extends Serialize, GroupByOps { sort( by: ColumnsOrExpr, descending?: ValueOrArray, + nulls_last?: boolean, maintain_order?: boolean, ): LazyDataFrame; sort(opts: { by: ColumnsOrExpr; descending?: ValueOrArray; + nulls_last?: boolean; maintain_order?: boolean; }): LazyDataFrame; /** @@ -965,15 +967,20 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => { } return _LazyDataFrame(_ldf.slice(opt, len)); }, - sort(arg, descending = false, maintain_order = false) { + sort(arg, descending = false, nulls_last = false, maintain_order = false) { if (arg?.by !== undefined) { - return this.sort(arg.by, arg.descending, arg.maintain_order); + return this.sort( + arg.by, + arg.descending, + arg.nulls_last, + arg.maintain_order, + ); } if (typeof arg === "string") { - return wrap("sort", arg, descending, true, maintain_order); + return wrap("sort", arg, descending, nulls_last, maintain_order); } const by = selectionToExprList(arg, false); - return wrap("sortByExprs", by, descending, true, maintain_order); + return wrap("sortByExprs", by, descending, nulls_last, maintain_order); }, std() { return _LazyDataFrame(_ldf.std());