Skip to content

Add new functions to Series and Expressions, some docs update #319

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.

Sign up for GitHub

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

Merged
merged 4 commits into from
Apr 2, 2025
Merged

Add new functions to Series and Expressions, some docs update #319

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4df666e
add new functions for structs, strings and expression
Mar 20, 2025
8036138
created ListSeries, StringSeries, DatetimeSeries, StructSeries for do…
Mar 20, 2025
8d82131
new functions for structs, fields and expression, update/clarify docu…
Mar 20, 2025
06cb4f7
added tests for new functions in series and expressions
Mar 20, 2025
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
267 changes: 265 additions & 2 deletions __tests__/expr.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,18 @@ describe("expr", () => {
const actual = df.select(col("nrs").explode());
expect(actual).toFrameEqual(expected);
});
test("implode", () => {
const df = pl.DataFrame({
nrs: [1, 2, 1, 3],
strs: ["a", "b", null, "d"],
});
const expected = pl.DataFrame({
nrs: [[1, 2, 1, 3]],
strs: [["a", "b", null, "d"]],
});
const actual = df.select(col("nrs").implode(), col("strs").implode());
expect(actual).toFrameEqual(expected);
});
test("extend", () => {
const df = pl.DataFrame({
a: [1, 2, 3, 4, 5],
Expand Down Expand Up @@ -1075,6 +1087,56 @@ describe("expr.str", () => {
);
expect(actual).toFrameEqual(expected);
});
test("endsWith", () => {
const df = pl.DataFrame({
fruits: ["apple", "mango", null],
});
const expected = df.withColumn(
pl.Series("has_suffix", [false, true, null], pl.Bool),
);
const actual = df.withColumn(
col("fruits").str.endsWith("go").as("has_suffix"),
);
expect(actual).toFrameEqual(expected);
});
test("endsWith:expr", () => {
const df = pl.DataFrame({
fruits: ["apple", "mango", "banana"],
suffix: ["le", "go", "nu"],
});
const expected = df.withColumn(
pl.Series("has_suffix", [true, true, false], pl.Bool),
);
const actual = df.withColumn(
col("fruits").str.endsWith(pl.col("suffix")).as("has_suffix"),
);
expect(actual).toFrameEqual(expected);
});
test("startsWith", () => {
const df = pl.DataFrame({
fruits: ["apple", "mango", null],
});
const expected = df.withColumn(
pl.Series("has_prefix", [true, false, null], pl.Bool),
);
const actual = df.withColumn(
col("fruits").str.startsWith("app").as("has_prefix"),
);
expect(actual).toFrameEqual(expected);
});
test("startsWith:expr", () => {
const df = pl.DataFrame({
fruits: ["apple", "mango", "banana"],
prefix: ["app", "na", "ba"],
});
const expected = df.withColumn(
pl.Series("has_prefix", [true, false, true], pl.Bool),
);
const actual = df.withColumn(
col("fruits").str.startsWith(pl.col("prefix")).as("has_prefix"),
);
expect(actual).toFrameEqual(expected);
});
test("split", () => {
const df = pl.DataFrame({ a: ["ab,cd", "e,fg", "h"] });
const expected = pl.DataFrame({
Expand Down Expand Up @@ -1196,10 +1258,10 @@ describe("expr.str", () => {
});
test("str.replace", () => {
const df = pl.DataFrame({
os: ["kali-linux", "debian-linux", "ubuntu-linux", "mac-sierra"],
os: ["kali-linux", "debian-linux", null, "mac-sierra"],
});
const expected = pl.DataFrame({
os: ["kali:linux", "debian:linux", "ubuntu:linux", "mac:sierra"],
os: ["kali:linux", "debian:linux", null, "mac:sierra"],
});
const seriesActual = df
.getColumn("os")
Expand All @@ -1210,6 +1272,48 @@ describe("expr.str", () => {
expect(actual).toFrameEqual(expected);
expect(seriesActual).toFrameEqual(expected);
});
test("str.replace:Expr1", () => {
const df = pl.DataFrame({
os: ["kali-linux", "debian-linux", null, "mac-sierra"],
val: ["windows", "acorn", "atari", null],
});
const expected = pl.DataFrame({
os: ["kali-windows", "debian-acorn", null, null],
});
const actual = df.select(
col("os").str.replace("linux", col("val")).as("os"),
);
expect(actual).toFrameEqual(expected);
});
test("str.replace:Expr2", () => {
const df = pl.DataFrame({
cost: ["#12.34", "#56.78"],
text: ["123abc", "abc456"],
});
const expected = pl.DataFrame({
expr: ["123#12.34", "#56.78456"],
});
const actual = df.select(
col("text").str.replace("abc", pl.col("cost")).alias("expr"),
);
expect(actual).toFrameEqual(expected);
});
// TODO: Remove skip when polars-plan will support for "dynamic pattern length in 'str.replace' expressions"
test.skip("str.replace:Expr3", () => {
const df = pl.DataFrame({
os: ["kali-linux", "debian-linux", "ubuntu-linux", "mac-sierra"],
pat: ["linux", "linux", "linux", "mac"],
val: ["windows", "acorn", "atari", "arm"],
});
const expected = pl.DataFrame({
os: ["kali-windows", "debian-acorn", "ubuntu-atari", "arm-sierra"],
});
const actual = df.select(
col("os").str.replace(col("pat"), col("val")).as("os"),
);
expect(actual).toFrameEqual(expected);
});

test("str.replaceAll", () => {
const df = pl.DataFrame({
os: [
Expand All @@ -1236,6 +1340,108 @@ describe("expr.str", () => {
expect(actual).toFrameEqual(expected);
expect(seriesActual).toFrameEqual(expected);
});
test("str.replaceAll:Expr", () => {
const df = pl.DataFrame({
os: [
"kali-linux-2021.3a",
null,
"ubuntu-linux-16.04",
"mac-sierra-10.12.1",
],
val: [":", ":", null, "_"],
});
const expected = pl.DataFrame({
os: ["kali:linux:2021.3a", null, null, "mac_sierra_10.12.1"],
});
const actual = df.select(
col("os").str.replaceAll("-", col("val")).as("os"),
);
expect(actual).toFrameEqual(expected);
});
// TODO: Remove skip when polars-plan will support for "dynamic pattern length in 'str.replace' expressions"
test.skip("str.replaceAll:Expr2", () => {
const df = pl.DataFrame({
os: [
"kali-linux-2021.3a",
null,
"ubuntu-linux-16.04",
"mac-sierra-10.12.1",
],
pat: ["-", "-", "-", "."],
val: [":", ":", null, "_"],
});
const expected = pl.DataFrame({
os: ["kali:linux:2021.3a", null, null, "mac-sierra-10_12_1"],
});
const actual = df.select(
col("os").str.replaceAll(col("pat"), col("val")).as("os"),
);
expect(actual).toFrameEqual(expected);
});
test("struct:field", () => {
const df = pl.DataFrame({
objs: [
{ a: 1, b: 2.0, c: "abc" },
{ a: 10, b: 20.0, c: "def" },
],
});
const expected = pl.DataFrame({
b: [2.0, 20.0],
last: ["abc", "def"],
});
const actual = df.select(
col("objs").struct.field("b"),
col("objs").struct.field("c").as("last"),
);
expect(actual).toFrameStrictEqual(expected);
});
test("struct:nth", () => {
const df = pl.DataFrame({
objs: [
{ a: 1, b: 2.0, c: "abc" },
{ a: 10, b: 20.0, c: "def" },
],
});
const expected = pl.DataFrame({
b: [2.0, 20.0],
last: ["abc", "def"],
});
const actual = df.select(
col("objs").struct.nth(1),
col("objs").struct.nth(2).as("last"),
);
expect(actual).toFrameStrictEqual(expected);
});
test("struct:withFields", () => {
const df = pl.DataFrame({
objs: [
{ a: 1, b: 2.0, c: "abc" },
{ a: 10, b: 20.0, c: "def" },
],
more: ["text1", "text2"],
final: [100, null],
});
const expected = pl.DataFrame({
objs: [
{ a: 1, b: 2.0, c: "abc", d: null, e: "text" },
{ a: 10, b: 20.0, c: "def", d: null, e: "text" },
],
new: [
{ a: 1, b: 2.0, c: "abc", more: "text1", final: 100 },
{ a: 10, b: 20.0, c: "def", more: "text2", final: null },
],
});
const actual = df.select(
col("objs").struct.withFields([
pl.lit(null).alias("d"),
pl.lit("text").alias("e"),
]),
col("objs")
.struct.withFields([col("more"), col("final")])
.alias("new"),
);
expect(actual).toFrameStrictEqual(expected);
});
test("expr.replace", () => {
const df = pl.DataFrame({ a: [1, 2, 2, 3], b: ["a", "b", "c", "d"] });
{
Expand Down Expand Up @@ -1443,6 +1649,63 @@ describe("expr.str", () => {
expect(actual).toFrameEqual(expected);
expect(seriesActual).toFrameEqual(expected);
});

test("stripChars:Expr", () => {
const df = pl.DataFrame({
os: [
"#Kali-Linux###",
"$$$Debian-Linux$",
null,
"Ubuntu-Linux ",
" Mac-Sierra",
],
chars: ["#", "$", " ", " ", null],
});
const expected = pl.DataFrame({
os: ["Kali-Linux", "Debian-Linux", null, "Ubuntu-Linux", "Mac-Sierra"],
});
const actual = df.select(col("os").str.stripChars(col("chars")).as("os"));
expect(actual).toFrameEqual(expected);
});
test("stripCharsStart:Expr", () => {
const df = pl.DataFrame({
os: [
"#Kali-Linux###",
"$$$Debian-Linux$",
null,
" Ubuntu-Linux ",
"Mac-Sierra",
],
chars: ["#", "$", " ", null, "Mac-"],
});
const expected = pl.DataFrame({
os: ["Kali-Linux###", "Debian-Linux$", null, "Ubuntu-Linux ", "Sierra"],
});
const actual = df.select(
col("os").str.stripCharsStart(col("chars")).as("os"),
);
expect(actual).toFrameEqual(expected);
});
test("stripCharsEnd:Expr", () => {
const df = pl.DataFrame({
os: [
"#Kali-Linux###",
"$$$Debian-Linux$",
null,
"Ubuntu-Linux ",
" Mac-Sierra",
],
chars: ["#", "$", " ", null, "-Sierra"],
});
const expected = pl.DataFrame({
os: ["#Kali-Linux", "$$$Debian-Linux", null, "Ubuntu-Linux", " Mac"],
});
const actual = df.select(
col("os").str.stripCharsEnd(col("chars")).as("os"),
);
expect(actual).toFrameEqual(expected);
});

test("padStart", () => {
const df = pl.DataFrame({
foo: ["a", "b", "cow", "longer"],
Expand Down
16 changes: 16 additions & 0 deletions __tests__/series.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -871,6 +871,12 @@ describe("StringFunctions", () => {
const serString = s.toString();
expect(actualInspect).toStrictEqual(serString);
});
test("str contains", () => {
const s = pl.Series(["linux-kali", "linux-debian", "windows-vista"]);
const expected = pl.Series([true, true, false]);
const encoded = s.str.contains("linux");
expect(encoded).toSeriesEqual(expected);
});
});
describe("series struct", () => {
test("struct:fields", () => {
Expand Down Expand Up @@ -904,6 +910,16 @@ describe("series struct", () => {
.toArray();
expect(actual).toEqual(expected);
});
test("struct:nth", () => {
const arr = [
{ foo: 1, bar: 2, ham: "c" },
{ foo: null, bar: 10, ham: null },
{ foo: 2, bar: 0, ham: "z" },
];
const expected = [1, null, 2];
const actual = pl.Series(arr).struct.nth(0).toArray();
expect(actual).toEqual(expected);
});
});
describe("generics", () => {
const series = pl.Series([1, 2, 3]);
Expand Down
4 changes: 2 additions & 2 deletions polars/lazy/expr/datetime.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@ import type { DateFunctions } from "../../shared_traits";
import { type Expr, _Expr } from "../expr";

/**
* DateTime functions
* DateTime functions for Lazy dataframes
*/
export type ExprDateTime = DateFunctions<Expr>;
export interface ExprDateTime extends DateFunctions<Expr> {}

export const ExprDateTimeFunctions = (_expr: any): ExprDateTime => {
const wrap = (method, ...args: any[]): Expr => {
Expand Down
8 changes: 6 additions & 2 deletions polars/lazy/expr/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import * as dt from "./datetime";
import * as lst from "./list";
import * as str from "./string";
import * as struct from "./struct";
export type { StringNamespace } from "./string";
export type { ExprString as StringNamespace } from "./string";
export type { ExprList as ListNamespace } from "./list";
export type { ExprDateTime as DatetimeNamespace } from "./datetime";
export type { ExprStruct as StructNamespace } from "./struct";
Expand Down Expand Up @@ -53,7 +53,7 @@ export interface Expr
/**
* String namespace
*/
get str(): str.StringNamespace;
get str(): str.ExprString;
/**
* List namespace
*/
Expand Down Expand Up @@ -471,6 +471,7 @@ export interface Expr
/** Take the first n values. */
head(length?: number): Expr;
head({ length }: { length: number }): Expr;
implode(): Expr;
inner(): any;
/** Interpolate intermediate values. The interpolation method is linear. */
interpolate(): Expr;
Expand Down Expand Up @@ -1528,6 +1529,9 @@ export const _Expr = (_expr: any): Expr => {

return wrap("head", length.length);
},
implode() {
return _Expr(_expr.implode());
},
interpolate(method: InterpolationMethod = "linear") {
return _Expr(_expr.interpolate(method));
},
Expand Down
Loading