diff --git a/__tests__/expr.test.ts b/__tests__/expr.test.ts index 35649672..cd108264 100644 --- a/__tests__/expr.test.ts +++ b/__tests__/expr.test.ts @@ -96,6 +96,18 @@ describe("expr", () => { const fn = () => df.select(col("a").cast(pl.Int16, true)); expect(fn).toThrow(); }); + test("cos", () => { + const df = pl.DataFrame({ a: [1, 2, 3] }); + const expected = pl.DataFrame({ cos: [0.540302, -0.416147, -0.989992] }); + const actual = df.select(col("a").cos().round(6).as("cos")); + expect(actual).toFrameEqual(expected); + }); + test("cot", () => { + const df = pl.DataFrame({ a: [1, 2, 3] }); + const expected = pl.DataFrame({ cot: [0.642093, -0.457658, -7.015253] }); + const actual = df.select(col("a").cot().round(6).as("cot")); + expect(actual).toFrameEqual(expected); + }); test("count", () => { const df = pl.DataFrame({ a: [1, 0, 3, 4, 6, 0] }); const expected = pl.DataFrame({ a: [6] }); @@ -208,6 +220,12 @@ describe("expr", () => { const actual = df.select(col("*").exclude("b", "c")); expect(actual).toFrameEqual(expected); }); + test("exp", () => { + const df = pl.DataFrame({ a: [1.0] }); + const actual = df.select(pl.col("a").exp()); + const expected = pl.DataFrame({ a: [Math.E] }); + expect(actual).toFrameEqual(expected); + }); test("explode", () => { const df = pl.DataFrame({ letters: ["c", "a"], @@ -329,6 +347,15 @@ describe("expr", () => { const actual = df.select(col("a").gtEq(0)); expect(actual).toFrameEqual(expected); }); + test("gatherEvery", () => { + const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] }); + let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] }); + let actual = df.select(col("a").gatherEvery(2).as("everyother")); + expect(actual).toFrameEqual(expected); + expected = pl.DataFrame({ everyother: [2, 3, 8, 1] }); + actual = df.select(col("a").gatherEvery(2, 2).as("everyother")); + expect(actual).toFrameEqual(expected); + }); test.each` args | hashValue ${[0]} | ${7355865757046787768n} @@ -514,6 +541,16 @@ describe("expr", () => { const actual = df.select(col("a").ltEq(2).as("lt")); expect(actual).toFrameEqual(expected); }); + test("log", () => { + let df = pl.DataFrame({ a: [1, 2, 3] }); + let actual = df.select(col("a").log(2).round(6).as("log")); + let expected = pl.DataFrame({ log: [0.0, 1.0, 1.584963] }); + expect(actual).toFrameEqual(expected); + df = pl.DataFrame({ a: [2] }); + actual = df.select(col("a").log().as("log")); + expected = pl.DataFrame({ log: [Math.LN2] }); + expect(actual).toFrameEqual(expected); + }); test("max", () => { const df = pl.DataFrame({ a: [1, 5, 3] }); const expected = pl.DataFrame({ max: [5] }); @@ -709,6 +746,12 @@ describe("expr", () => { const actual = df.select(col("a"), ...shifts); expect(actual).toFrameStrictEqual(expected); }); + test("sin", () => { + const df = pl.DataFrame({ a: [1, 2, 3] }); + const expected = pl.DataFrame({ sin: [0.841471, 0.909297, 0.14112] }); + const actual = df.select(col("a").sin().round(6).as("sin")); + expect(actual).toFrameEqual(expected); + }); test("skew", () => { const df = pl.DataFrame({ a: [1, 2, 3, 3] }); const expected = pl.DataFrame({ @@ -870,13 +913,10 @@ describe("expr", () => { ); expect(actual).toFrameEqual(expected); }); - test("gatherEvery", () => { - const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] }); - let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] }); - let actual = df.select(col("a").gatherEvery(2).as("everyother")); - expect(actual).toFrameEqual(expected); - expected = pl.DataFrame({ everyother: [2, 3, 8, 1] }); - actual = df.select(col("a").gatherEvery(2, 2).as("everyother")); + test("tan", () => { + const df = pl.DataFrame({ a: [1, 2, 3] }); + const expected = pl.DataFrame({ tan: [1.557408, -2.18504, -0.142547] }); + const actual = df.select(col("a").tan().round(6).as("tan")); expect(actual).toFrameEqual(expected); }); test("unique", () => { diff --git a/polars/lazy/expr/index.ts b/polars/lazy/expr/index.ts index 6eb565de..14b55f79 100644 --- a/polars/lazy/expr/index.ts +++ b/polars/lazy/expr/index.ts @@ -69,6 +69,36 @@ export interface Expr toJSON(): string; /** Take absolute values */ abs(): Expr; + /** + * Get the group indexes of the group by operation. + * Should be used in aggregation context only. + * @example + * ``` + >>> const df = pl.DataFrame( + ... { + ... "group": [ + ... "one", + ... "one", + ... "one", + ... "two", + ... "two", + ... "two", + ... ], + ... "value": [94, 95, 96, 97, 97, 99], + ... } + ... ) + >>> df.group_by("group", maintain_order=True).agg(pl.col("value").aggGroups()) + shape: (2, 2) + ┌───────┬───────────┐ + │ group ┆ value │ + │ --- ┆ --- │ + │ str ┆ list[u32] │ + ╞═══════╪═══════════╡ + │ one ┆ [0, 1, 2] │ + │ two ┆ [3, 4, 5] │ + └───────┴───────────┘ + *``` + */ aggGroups(): Expr; /** * Rename the output of an expression. @@ -137,6 +167,60 @@ export interface Expr backwardFill(): Expr; /** Cast between data types. */ cast(dtype: DataType, strict?: boolean): Expr; + /** + * Compute the element-wise value for the cosine. + * @returns Expression of data type :class:`Float64`. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [0.0]}) + >>> df.select(pl.col("a").cos()) + shape: (1, 1) + ┌─────┐ + │ a │ + │ --- │ + │ f64 │ + ╞═════╡ + │ 1.0 │ + └─────┘ + * ``` + */ + cos(): Expr; + /** + * Compute the element-wise value for the hyperbolic cosine. + * @returns Expression of data type :class:`Float64`. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [1.0]}) + >>> df.select(pl.col("a").cosh()) + shape: (1, 1) + ┌──────────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 1.543081 │ + └──────────┘ + * ``` + */ + cosh(): Expr; + /** + * Compute the element-wise value for the cotangent. + * @returns Expression of data type :class:`Float64`. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [1.0]}) + >>> df.select(pl.col("a").cot().round(2)) + shape: (1, 1) + ┌──────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════╡ + │ 0.64 │ + └──────┘ + * ``` + */ + cot(): Expr; /** Count the number of values in this expression */ count(): Expr; /** Calculate the n-th discrete difference. @@ -151,7 +235,6 @@ export interface Expr * @param other Expression to compute dot product with */ dot(other: any): Expr; - /** * Exclude certain columns from a wildcard/regex selection. * @@ -196,6 +279,25 @@ export interface Expr * ``` */ exclude(column: string, ...columns: string[]): Expr; + /** + * Compute the exponential, element-wise. + * @example + * ``` + >>> const df = pl.DataFrame({"values": [1.0, 2.0, 4.0]}) + >>> df.select(pl.col("values").exp()) + shape: (3, 1) + ┌──────────┐ + │ values │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 2.718282 │ + │ 7.389056 │ + │ 54.59815 │ + └──────────┘ + * ``` + */ + exp(): Expr; /** * Explode a list or utf8 Series. * @@ -236,6 +338,14 @@ export interface Expr flatten(): Expr; /** Fill missing values with the latest seen values */ forwardFill(): Expr; + /** + * Take values by index. + * @param index An expression that leads to a UInt32 dtyped Series. + */ + gather(index: Expr | number[] | Series): Expr; + gather({ index }: { index: Expr | number[] | Series }): Expr; + /** Take every nth value in the Series and return as a new Series. */ + gatherEvery(n: number, offset?: number): Expr; /** Hash the Series. */ hash(k0?: number, k1?: number, k2?: number, k3?: number): Expr; hash({ @@ -356,6 +466,46 @@ export interface Expr last(): Expr; /** Aggregate to list. */ list(): Expr; + /*** + * Compute the natural logarithm of each element plus one. + * This computes `log(1 + x)` but is more numerically stable for `x` close to zero. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [1, 2, 3]}) + >>> df.select(pl.col("a").log1p()) + shape: (3, 1) + ┌──────────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 0.693147 │ + │ 1.098612 │ + │ 1.386294 │ + └──────────┘ + * ``` + */ + log1p(): Expr; + /** + * Compute the logarithm to a given base. + * @param base - Given base, defaults to `e` + * @example + * ``` + >>> const df = pl.DataFrame({"a": [1, 2, 3]}) + >>> df.select(pl.col("a").log(base=2)) + shape: (3, 1) + ┌──────────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 0.0 │ + │ 1.0 │ + │ 1.584963 │ + └──────────┘ + * ``` + */ + log(base?: number): Expr; /** Returns a unit Series with the lowest value possible for the dtype of this expression. */ lowerBound(): Expr; peakMax(): Expr; @@ -503,6 +653,42 @@ export interface Expr periods, fillValue, }: { periods: number; fillValue: number }): Expr; + /** + * Compute the element-wise value for the sine. + * @returns Expression of data type :class:`Float64`. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [0.0]}) + >>> df.select(pl.col("a").sin()) + shape: (1, 1) + ┌─────┐ + │ a │ + │ --- │ + │ f64 │ + ╞═════╡ + │ 0.0 │ + └─────┘ + *``` + */ + sin(): Expr; + /** + * Compute the element-wise value for the hyperbolic sine. + * @returns Expression of data type :class:`Float64`. + * @example + * ``` + >>> const df = pl.DataFrame({"a": [1.0]}) + >>> df.select(pl.col("a").sinh()) + shape: (1, 1) + ┌──────────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 1.175201 │ + └──────────┘ + *``` + */ + sinh(): Expr; /** * Compute the sample skewness of a data set. * For normally distributed data, the skewness should be about zero. For @@ -565,13 +751,41 @@ export interface Expr tail(length?: number): Expr; tail({ length }: { length: number }): Expr; /** - * Take values by index. - * @param index An expression that leads to a UInt32 dtyped Series. + * Compute the element-wise value for the tangent. + * @returns Expression of data type :class:`Float64`. + * @example + *``` + >>> const df = pl.DataFrame({"a": [1.0]}) + >>> df.select(pl.col("a").tan().round(2)) + shape: (1, 1) + ┌──────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════╡ + │ 1.56 │ + └──────┘ + *``` */ - gather(index: Expr | number[] | Series): Expr; - gather({ index }: { index: Expr | number[] | Series }): Expr; - /** Take every nth value in the Series and return as a new Series. */ - gatherEvery(n: number, offset?: number): Expr; + tan(): Expr; + /** + * Compute the element-wise value for the hyperbolic tangent. + * @returns Expression of data type :class:`Float64`. + * @example + *``` + >>> const df = pl.DataFrame({"a": [1.0]}) + >>> df.select(pl.col("a").tanh()) + shape: (1, 1) + ┌──────────┐ + │ a │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 0.761594 │ + └──────────┘ + *``` + */ + tanh(): Expr; /** * Get the unique values of this expression; * @param maintainOrder Maintain order of data. This requires more work. @@ -709,6 +923,15 @@ export const _Expr = (_expr: any): Expr => { ), ); }, + cos() { + return _Expr(_expr.cos()); + }, + cosh() { + return _Expr(_expr.cosh()); + }, + cot() { + return _Expr(_expr.cot()); + }, count() { return _Expr(_expr.count()); }, @@ -859,6 +1082,9 @@ export const _Expr = (_expr: any): Expr => { explode() { return _Expr(_expr.explode()); }, + exp() { + return _Expr(_expr.exp()); + }, extend(o, n?) { if (n !== null && typeof n === "number") { return _Expr(_expr.extendConstant(o, n)); @@ -908,6 +1134,17 @@ export const _Expr = (_expr: any): Expr => { forwardFill() { return _Expr(_expr.forwardFill()); }, + gather(indices) { + if (Array.isArray(indices)) { + indices = pli.lit(Series(indices).inner()); + } else { + indices = indices.inner(); + } + return wrap("gather", indices); + }, + gatherEvery(n, offset = 0) { + return _Expr(_expr.gatherEvery(n, offset)); + }, hash(obj: any = 0, k1 = 1, k2 = 2, k3 = 3) { if (typeof obj === "number" || typeof obj === "bigint") { return wrap("hash", BigInt(obj), BigInt(k1), BigInt(k2), BigInt(k3)); @@ -974,7 +1211,6 @@ export const _Expr = (_expr: any): Expr => { kurtosis(obj?, bias = true) { const fisher = obj?.["fisher"] ?? (typeof obj === "boolean" ? obj : true); bias = obj?.["bias"] ?? bias; - return _Expr(_expr.kurtosis(fisher, bias)); }, last() { @@ -983,6 +1219,13 @@ export const _Expr = (_expr: any): Expr => { list() { return _Expr(_expr.list()); }, + log1p() { + console.log(_expr.log1p); + return _Expr(_expr.log1p()); + }, + log(base?: number) { + return _Expr(_expr.log(base ?? Math.E)); + }, lowerBound() { return _Expr(_expr.lowerBound()); }, @@ -1139,6 +1382,12 @@ export const _Expr = (_expr: any): Expr => { skew(bias) { return wrap("skew", bias?.bias ?? bias ?? true); }, + sin() { + return _Expr(_expr.sin()); + }, + sinh() { + return _Expr(_expr.sinh()); + }, slice(arg, len?) { if (typeof arg === "number") { return wrap("slice", pli.lit(arg), pli.lit(len)); @@ -1181,16 +1430,11 @@ export const _Expr = (_expr: any): Expr => { tail(length) { return _Expr(_expr.tail(length)); }, - gather(indices) { - if (Array.isArray(indices)) { - indices = pli.lit(Series(indices).inner()); - } else { - indices = indices.inner(); - } - return wrap("gather", indices); + tan() { + return _Expr(_expr.tan()); }, - gatherEvery(n, offset = 0) { - return _Expr(_expr.gatherEvery(n, offset)); + tanh() { + return _Expr(_expr.tanh()); }, unique(opt?) { if (opt || opt?.maintainOrder) { diff --git a/src/lazy/dsl.rs b/src/lazy/dsl.rs index 1b314b7a..a455152a 100644 --- a/src/lazy/dsl.rs +++ b/src/lazy/dsl.rs @@ -491,6 +491,70 @@ impl JsExpr { self.clone().inner.abs().into() } #[napi(catch_unwind)] + pub fn sin(&self) -> Self { + self.inner.clone().sin().into() + } + #[napi(catch_unwind)] + pub fn cos(&self) -> Self { + self.inner.clone().cos().into() + } + #[napi(catch_unwind)] + pub fn tan(&self) -> Self { + self.inner.clone().tan().into() + } + #[napi(catch_unwind)] + pub fn cot(&self) -> Self { + self.inner.clone().cot().into() + } + #[napi(catch_unwind)] + pub fn arcsin(&self) -> Self { + self.inner.clone().arcsin().into() + } + #[napi(catch_unwind)] + pub fn arccos(&self) -> Self { + self.inner.clone().arccos().into() + } + #[napi(catch_unwind)] + pub fn arctan(&self) -> Self { + self.inner.clone().arctan().into() + } + #[napi(catch_unwind)] + pub fn arctan2(&self, y: &JsExpr) -> Self { + self.inner.clone().arctan2(y.inner.clone()).into() + } + #[napi(catch_unwind)] + pub fn sinh(&self) -> Self { + self.inner.clone().sinh().into() + } + #[napi(catch_unwind)] + pub fn cosh(&self) -> Self { + self.inner.clone().cosh().into() + } + #[napi(catch_unwind)] + pub fn tanh(&self) -> Self { + self.inner.clone().tanh().into() + } + #[napi(catch_unwind)] + pub fn arcsinh(&self) -> Self { + self.inner.clone().arcsinh().into() + } + #[napi(catch_unwind)] + pub fn arccosh(&self) -> Self { + self.inner.clone().arccosh().into() + } + #[napi(catch_unwind)] + pub fn arctanh(&self) -> Self { + self.inner.clone().arctanh().into() + } + #[napi(catch_unwind)] + pub fn degrees(&self) -> Self { + self.inner.clone().degrees().into() + } + #[napi(catch_unwind)] + pub fn radians(&self) -> Self { + self.inner.clone().radians().into() + } + #[napi(catch_unwind)] pub fn is_duplicated(&self) -> JsExpr { self.clone().inner.is_duplicated().into() } @@ -1436,6 +1500,14 @@ impl JsExpr { self.inner.clone().log(base).into() } #[napi(catch_unwind)] + pub fn log1p(&self) -> JsExpr { + self.inner.clone().log1p().into() + } + #[napi(catch_unwind)] + pub fn exp(&self) -> JsExpr { + self.inner.clone().exp().into() + } + #[napi(catch_unwind)] pub fn entropy(&self, base: f64, normalize: bool) -> JsExpr { self.inner.clone().entropy(base, normalize).into() }