Skip to content

Commit

Permalink
Add nth fn binding (#212)
Browse files Browse the repository at this point in the history
The [nth fn](https://docs.rs/polars/latest/polars/prelude/fn.nth.html)
in rust polar is different from what is available in the [python
repo](https://docs.pola.rs/py-polars/html/reference/expressions/api/polars.nth.html#polars-nth).
Not sure what's the convention for this nodejs library, but just
exposing what the rust library does is easier and satisfies our needs.
  • Loading branch information
froxCZ authored Jun 5, 2024
1 parent 6e9caaf commit c7af942
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 0 deletions.
22 changes: 22 additions & 0 deletions __tests__/dataframe.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,28 @@ describe("dataframe", () => {
});
expect(actual.nChunks()).toEqual(1);
});
test("nth", () => {
const actual = pl
.DataFrame({
foo: [1, 2, null],
bar: [6, 2, 8],
apple: [6, 2, 8],
pizza: [null, null, 8],
})
.select(pl.nth(2));
expect(actual.columns).toEqual(["apple"]);
});
test("nth:multiple", () => {
const actual = pl
.DataFrame({
foo: [1, 2, null],
bar: [6, 2, 8],
apple: [6, 2, 8],
pizza: [null, null, 8],
})
.select(pl.nth(2), pl.nth(3));
expect(actual.columns).toEqual(["apple", "pizza"]);
});
test("nullCount", () => {
const actual = pl
.DataFrame({
Expand Down
1 change: 1 addition & 0 deletions polars/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ export namespace pl {

// lazy
export import col = lazy.col;
export import nth = lazy.nth;
export import cols = lazy.cols;
export import lit = lazy.lit;
export import intRange = lazy.intRange;
Expand Down
28 changes: 28 additions & 0 deletions polars/lazy/functions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,34 @@ export function cols(...cols): Expr {
return col(cols.flat());
}

/**
* Select nth column index in a DataFrame.
* @param n - Column index to select, starting at 0.
* @example
* ```
* > df = pl.DataFrame({
* > "ham": [1, 2, 3],
* > "hamburger": [11, 22, 33],
* > "foo": [3, 2, 1]})
* > df.select(nth(2))
* shape: (3, 1)
* ╭─────╮
* │ foo │
* │ --- │
* │ i64 │
* ╞═════╡
* │ 3 │
* ├╌╌╌╌╌┤
* │ 2 │
* ├╌╌╌╌╌┤
* │ 1 │
* ╰─────╯
* ```
*/
export function nth(n: number): Expr {
return _Expr(pli.nth(n));
}

export function lit(value: any): Expr {
if (Array.isArray(value)) {
value = Series(value);
Expand Down
5 changes: 5 additions & 0 deletions src/lazy/dsl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1625,6 +1625,11 @@ pub fn last() -> JsExpr {
dsl::last().into()
}

#[napi(catch_unwind)]
pub fn nth(n: i64) -> JsExpr {
Expr::Nth(n).into()
}

#[napi(catch_unwind)]
pub fn cols(names: Vec<String>) -> JsExpr {
dsl::cols(names).into()
Expand Down

0 comments on commit c7af942

Please sign in to comment.