Skip to content

Commit 181cf47

Browse files
authored
Adding trigonometry functions (#198)
* adding trigonometry fn * Adding more trigonometry functions * Fixing linting errors
1 parent ff60d5a commit 181cf47

File tree

3 files changed

+380
-24
lines changed

3 files changed

+380
-24
lines changed

__tests__/expr.test.ts

+47-7
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,18 @@ describe("expr", () => {
9696
const fn = () => df.select(col("a").cast(pl.Int16, true));
9797
expect(fn).toThrow();
9898
});
99+
test("cos", () => {
100+
const df = pl.DataFrame({ a: [1, 2, 3] });
101+
const expected = pl.DataFrame({ cos: [0.540302, -0.416147, -0.989992] });
102+
const actual = df.select(col("a").cos().round(6).as("cos"));
103+
expect(actual).toFrameEqual(expected);
104+
});
105+
test("cot", () => {
106+
const df = pl.DataFrame({ a: [1, 2, 3] });
107+
const expected = pl.DataFrame({ cot: [0.642093, -0.457658, -7.015253] });
108+
const actual = df.select(col("a").cot().round(6).as("cot"));
109+
expect(actual).toFrameEqual(expected);
110+
});
99111
test("count", () => {
100112
const df = pl.DataFrame({ a: [1, 0, 3, 4, 6, 0] });
101113
const expected = pl.DataFrame({ a: [6] });
@@ -208,6 +220,12 @@ describe("expr", () => {
208220
const actual = df.select(col("*").exclude("b", "c"));
209221
expect(actual).toFrameEqual(expected);
210222
});
223+
test("exp", () => {
224+
const df = pl.DataFrame({ a: [1.0] });
225+
const actual = df.select(pl.col("a").exp());
226+
const expected = pl.DataFrame({ a: [Math.E] });
227+
expect(actual).toFrameEqual(expected);
228+
});
211229
test("explode", () => {
212230
const df = pl.DataFrame({
213231
letters: ["c", "a"],
@@ -329,6 +347,15 @@ describe("expr", () => {
329347
const actual = df.select(col("a").gtEq(0));
330348
expect(actual).toFrameEqual(expected);
331349
});
350+
test("gatherEvery", () => {
351+
const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] });
352+
let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] });
353+
let actual = df.select(col("a").gatherEvery(2).as("everyother"));
354+
expect(actual).toFrameEqual(expected);
355+
expected = pl.DataFrame({ everyother: [2, 3, 8, 1] });
356+
actual = df.select(col("a").gatherEvery(2, 2).as("everyother"));
357+
expect(actual).toFrameEqual(expected);
358+
});
332359
test.each`
333360
args | hashValue
334361
${[0]} | ${7355865757046787768n}
@@ -514,6 +541,16 @@ describe("expr", () => {
514541
const actual = df.select(col("a").ltEq(2).as("lt"));
515542
expect(actual).toFrameEqual(expected);
516543
});
544+
test("log", () => {
545+
let df = pl.DataFrame({ a: [1, 2, 3] });
546+
let actual = df.select(col("a").log(2).round(6).as("log"));
547+
let expected = pl.DataFrame({ log: [0.0, 1.0, 1.584963] });
548+
expect(actual).toFrameEqual(expected);
549+
df = pl.DataFrame({ a: [2] });
550+
actual = df.select(col("a").log().as("log"));
551+
expected = pl.DataFrame({ log: [Math.LN2] });
552+
expect(actual).toFrameEqual(expected);
553+
});
517554
test("max", () => {
518555
const df = pl.DataFrame({ a: [1, 5, 3] });
519556
const expected = pl.DataFrame({ max: [5] });
@@ -709,6 +746,12 @@ describe("expr", () => {
709746
const actual = df.select(col("a"), ...shifts);
710747
expect(actual).toFrameStrictEqual(expected);
711748
});
749+
test("sin", () => {
750+
const df = pl.DataFrame({ a: [1, 2, 3] });
751+
const expected = pl.DataFrame({ sin: [0.841471, 0.909297, 0.14112] });
752+
const actual = df.select(col("a").sin().round(6).as("sin"));
753+
expect(actual).toFrameEqual(expected);
754+
});
712755
test("skew", () => {
713756
const df = pl.DataFrame({ a: [1, 2, 3, 3] });
714757
const expected = pl.DataFrame({
@@ -870,13 +913,10 @@ describe("expr", () => {
870913
);
871914
expect(actual).toFrameEqual(expected);
872915
});
873-
test("gatherEvery", () => {
874-
const df = pl.DataFrame({ a: [1, 1, 2, 2, 3, 3, 8, null, 1] });
875-
let expected = pl.DataFrame({ everyother: [1, 2, 3, 8, 1] });
876-
let actual = df.select(col("a").gatherEvery(2).as("everyother"));
877-
expect(actual).toFrameEqual(expected);
878-
expected = pl.DataFrame({ everyother: [2, 3, 8, 1] });
879-
actual = df.select(col("a").gatherEvery(2, 2).as("everyother"));
916+
test("tan", () => {
917+
const df = pl.DataFrame({ a: [1, 2, 3] });
918+
const expected = pl.DataFrame({ tan: [1.557408, -2.18504, -0.142547] });
919+
const actual = df.select(col("a").tan().round(6).as("tan"));
880920
expect(actual).toFrameEqual(expected);
881921
});
882922
test("unique", () => {

0 commit comments

Comments
 (0)