Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding string to parameters for replace function #230

Merged
merged 1 commit into from
Jun 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 38 additions & 5 deletions __tests__/expr.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1178,17 +1178,43 @@ describe("expr.str", () => {
expect(seriesActual).toFrameEqual(expected);
});
test("expr.replace", () => {
const df = pl.DataFrame({ a: [1, 2, 2, 3] });
const df = pl.DataFrame({ a: [1, 2, 2, 3], b: ["a", "b", "c", "d"] });
let actual = df.withColumns(pl.col("a").replace(2, 100).alias("replaced"));
let expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: [1, 100, 100, 3],
});
expect(actual).toFrameEqual(expected);
actual = df.withColumns(
pl.col("a").replace([2, 3], [100, 200], -1, pl.Float64).alias("replaced"),
);
expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [-1, 100, 100, 200] });
expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: [-1, 100, 100, 200],
});
expect(actual).toFrameEqual(expected);
actual = df.withColumns(
pl.col("b").replace("a", "c", "e", pl.Utf8).alias("replaced"),
);
expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: ["c", "e", "e", "e"],
});
expect(actual).toFrameEqual(expected);
actual = df.withColumns(
pl
.col("b")
.replace(["a", "b"], ["c", "d"], "e", pl.Utf8)
.alias("replaced"),
);
Comment on lines +1208 to +1212
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it'd be a bit more readable if we took in a object/map here instead of the array. (this could totally be done in a follow up pr though.

pl.col("a").replace({ 1: 11, 2: 12 }, 10, pl.Int64)
pl.col("b").replace({ a: "c", b: "d" }, "e", pl.Utf8)

also, i think using an object for the additional parameters would make things more readable

pl.col("a").replace({ 1: 11, 2: 12 }, {
  defaultValue: 10,
  returnType: pl.Int64,
})

Copy link
Collaborator Author

@Bidek56 Bidek56 Jun 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be more readable but I want to test list as param.
Python version supports both list and dict, hence we need to support both and test for it.
We can add another test using object if you like.

expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: ["c", "d", "e", "e"],
});
expect(actual).toFrameEqual(expected);
const mapping = { 2: 100, 3: 200 };
actual = df.withColumns(
Expand All @@ -1197,13 +1223,20 @@ describe("expr.str", () => {
.replace({ old: mapping, default_: -1, returnDtype: pl.Int64 })
.alias("replaced"),
);
expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [-1, 100, 100, 200] });
expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: [-1, 100, 100, 200],
});
expect(actual).toFrameEqual(expected);

actual = df.withColumns(
pl.col("a").replace({ old: mapping }).alias("replaced"),
);
expected = pl.DataFrame({ a: [1, 2, 2, 3], replaced: [1, 100, 100, 200] });
expected = pl.DataFrame({
a: [1, 2, 2, 3],
b: ["a", "b", "c", "d"],
replaced: [1, 100, 100, 200],
});
expect(actual).toFrameEqual(expected);
});
test("slice", () => {
Expand Down
12 changes: 6 additions & 6 deletions polars/lazy/expr/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -871,9 +871,9 @@ export interface Expr
* ```
*/
replace(
old: Expr | number | number[],
new_: Expr | number | number[],
default_?: Expr | number | number[],
old: Expr | string | number | (number | string)[],
new_: Expr | string | number | (number | string)[],
default_?: Expr | string | number | (number | string)[],
returnDtype?: DataType,
): Expr;
replace({
Expand All @@ -882,9 +882,9 @@ export interface Expr
default_,
returnDtype,
}: {
old: unknown | Expr | number | number[];
new_?: Expr | number | number[];
default_?: Expr | number | number[];
old: unknown | Expr | string | number | (number | string)[];
new_?: Expr | string | number | (number | string)[];
default_?: Expr | string | number | (number | string)[];
returnDtype?: DataType;
}): Expr;
/** Reverse the arrays in the list */
Expand Down