diff --git a/__tests__/series.test.ts b/__tests__/series.test.ts index e5a73329c..5ade5301e 100644 --- a/__tests__/series.test.ts +++ b/__tests__/series.test.ts @@ -2,6 +2,41 @@ import pl, { DataType } from "@polars"; import Chance from "chance"; +describe("mapElements", () => { + test("mapElements string", () => { + const mapping: Record = { + A: "AA", + B: "BB", + C: "CC", + D: "OtherD", + }; + const funcMap = (k: string): string => mapping[k] ?? ""; + const actual = pl + .Series("foo", ["A", "B", "C", "D", "F", null], pl.String) + .mapElements(funcMap); + const expected = pl.Series( + "foo", + ["AA", "BB", "CC", "OtherD", "", ""], + pl.String, + ); + expect(actual).toSeriesEqual(expected); + }); + test("mapElements int", () => { + const mapping: Record = { 1: 11, 2: 22, 3: 33, 4: 44 }; + const funcMap = (k: number): number => mapping[k] ?? ""; + let actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(funcMap); + let expected = pl.Series("foo", [11, 22, 33, null], pl.Int32); + expect(actual).toSeriesEqual(expected); + const multiFunc = (k: number): number => k * 2; + actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(multiFunc); + expected = pl.Series("foo", [2, 4, 6, 10], pl.Int32); + expect(actual).toSeriesEqual(expected); + const funcStr = (k: number): string => `${k}x`; + actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(funcStr); + expected = pl.Series("foo", ["1x", "2x", "3x", "5x"], pl.String); + expect(actual).toSeriesEqual(expected); + }); +}); describe("from lists", () => { test("bool", () => { const expected = [[true, false], [true], [null], []]; diff --git a/polars/series/index.ts b/polars/series/index.ts index 9234966a1..5ff942117 100644 --- a/polars/series/index.ts +++ b/polars/series/index.ts @@ -628,6 +628,28 @@ export interface Series * ``` */ limit(n?: number): Series; + /** + * Map a custom/user-defined function (UDF) over elements in this Series. + * @param fn - Custom function to call + * + * @example + * ``` + * > const mapping: Record = { A: 'AA', B: 'BB', C: 'CC', D: 'OtherD', }; + * > const funcMap = (k: string): string => mapping[k] ?? ''; + * > pl.Series("foo", ["A", "B", "C", "D", "F", null], pl.String).mapElements(funcMap); + shape: (6,) + Series: 'foo' [str] + [ + "AA" + "BB" + "CC" + "OtherD" + "" + "" + ] + * ``` + */ + mapElements(fn: (v: any) => any): Series; /** * Get the maximum value in this Series. * @example @@ -1612,6 +1634,9 @@ export function _Series(_s: any): Series { limit(n = 10) { return wrap("limit", n); }, + mapElements(fn: (v: any) => any) { + return Series(this.name, [...this.values()].map(fn)); + }, max() { return _s.max() as any; },