Skip to content

Commit f3873fc

Browse files
authored
feat: Adding series mapElements (#417)
Adding series mapElements to close #46 This feature will allow for DF column map thru series, for example: ``` const df = pl.DataFrame({ a: [1, 5, 3] }); const mapping: Record<number, number> = { 1: 11, 2: 22, 3: 33, 4: 44 }; const funcMap = (k: number): number => mapping[k] ?? ''; const mappedSeries: pl.Series = df.select(pl.col("a")).toSeries().mapElements(funcMap); const df2 = df.withColumn(mappedSeries.alias("mappedSeries")); ``` Co-authored-by: Bidek56 <[email protected]>
1 parent 8b01bbf commit f3873fc

File tree

2 files changed

+60
-0
lines changed

2 files changed

+60
-0
lines changed

__tests__/series.test.ts

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,41 @@
22
import pl, { DataType } from "@polars";
33
import Chance from "chance";
44

5+
describe("mapElements", () => {
6+
test("mapElements string", () => {
7+
const mapping: Record<string, string> = {
8+
A: "AA",
9+
B: "BB",
10+
C: "CC",
11+
D: "OtherD",
12+
};
13+
const funcMap = (k: string): string => mapping[k] ?? "";
14+
const actual = pl
15+
.Series("foo", ["A", "B", "C", "D", "F", null], pl.String)
16+
.mapElements(funcMap);
17+
const expected = pl.Series(
18+
"foo",
19+
["AA", "BB", "CC", "OtherD", "", ""],
20+
pl.String,
21+
);
22+
expect(actual).toSeriesEqual(expected);
23+
});
24+
test("mapElements int", () => {
25+
const mapping: Record<number, number> = { 1: 11, 2: 22, 3: 33, 4: 44 };
26+
const funcMap = (k: number): number => mapping[k] ?? "";
27+
let actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(funcMap);
28+
let expected = pl.Series("foo", [11, 22, 33, null], pl.Int32);
29+
expect(actual).toSeriesEqual(expected);
30+
const multiFunc = (k: number): number => k * 2;
31+
actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(multiFunc);
32+
expected = pl.Series("foo", [2, 4, 6, 10], pl.Int32);
33+
expect(actual).toSeriesEqual(expected);
34+
const funcStr = (k: number): string => `${k}x`;
35+
actual = pl.Series("foo", [1, 2, 3, 5], pl.Int32).mapElements(funcStr);
36+
expected = pl.Series("foo", ["1x", "2x", "3x", "5x"], pl.String);
37+
expect(actual).toSeriesEqual(expected);
38+
});
39+
});
540
describe("from lists", () => {
641
test("bool", () => {
742
const expected = [[true, false], [true], [null], []];

polars/series/index.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,6 +628,28 @@ export interface Series<T extends DataType = any, Name extends string = string>
628628
* ```
629629
*/
630630
limit(n?: number): Series;
631+
/**
632+
* Map a custom/user-defined function (UDF) over elements in this Series.
633+
* @param fn - Custom function to call
634+
*
635+
* @example
636+
* ```
637+
* > const mapping: Record<string, string> = { A: 'AA', B: 'BB', C: 'CC', D: 'OtherD', };
638+
* > const funcMap = (k: string): string => mapping[k] ?? '';
639+
* > pl.Series("foo", ["A", "B", "C", "D", "F", null], pl.String).mapElements(funcMap);
640+
shape: (6,)
641+
Series: 'foo' [str]
642+
[
643+
"AA"
644+
"BB"
645+
"CC"
646+
"OtherD"
647+
""
648+
""
649+
]
650+
* ```
651+
*/
652+
mapElements(fn: (v: any) => any): Series;
631653
/**
632654
* Get the maximum value in this Series.
633655
* @example
@@ -1612,6 +1634,9 @@ export function _Series(_s: any): Series {
16121634
limit(n = 10) {
16131635
return wrap("limit", n);
16141636
},
1637+
mapElements(fn: (v: any) => any) {
1638+
return Series(this.name, [...this.values()].map(fn));
1639+
},
16151640
max() {
16161641
return _s.max() as any;
16171642
},

0 commit comments

Comments
 (0)