forked from pola-rs/polars
-
Notifications
You must be signed in to change notification settings - Fork 0
/
lib.rs
109 lines (96 loc) · 2.55 KB
/
lib.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#[macro_use]
extern crate polars;
use crate::lazy::dsl::PyExpr;
use crate::{
dataframe::PyDataFrame,
lazy::{
dataframe::{PyLazyFrame, PyLazyGroupBy},
dsl,
},
series::PySeries,
};
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
pub mod arrow_interop;
pub mod conversion;
pub mod dataframe;
pub mod datatypes;
pub mod dispatch;
pub mod error;
pub mod file;
pub mod lazy;
pub mod npy;
pub mod prelude;
pub mod series;
pub mod utils;
#[pyfunction]
fn col(name: &str) -> dsl::PyExpr {
dsl::col(name)
}
#[pyfunction]
fn except_(name: &str) -> dsl::PyExpr {
dsl::except(name)
}
#[pyfunction]
fn lit(value: &PyAny) -> dsl::PyExpr {
dsl::lit(value)
}
#[pyfunction]
fn range(low: i64, high: i64, dtype: &PyAny) -> dsl::PyExpr {
dsl::range(low, high, dtype)
}
#[pyfunction]
fn binary_expr(l: dsl::PyExpr, op: u8, r: dsl::PyExpr) -> dsl::PyExpr {
dsl::binary_expr(l, op, r)
}
#[pyfunction]
fn binary_function(
a: dsl::PyExpr,
b: dsl::PyExpr,
lambda: PyObject,
output_type: &PyAny,
) -> dsl::PyExpr {
dsl::binary_function(a, b, lambda, output_type)
}
#[pyfunction]
fn pearson_corr(a: dsl::PyExpr, b: dsl::PyExpr) -> dsl::PyExpr {
polars::lazy::functions::pearson_corr(a.inner, b.inner).into()
}
#[pyfunction]
fn cov(a: dsl::PyExpr, b: dsl::PyExpr) -> dsl::PyExpr {
polars::lazy::functions::cov(a.inner, b.inner).into()
}
#[pyfunction]
fn when(predicate: PyExpr) -> dsl::When {
dsl::when(predicate)
}
const VERSION: &str = env!("CARGO_PKG_VERSION");
#[pyfunction]
fn version() -> &'static str {
VERSION
}
#[pyfunction]
fn toggle_string_cache(toggle: bool) {
polars::toggle_string_cache(toggle)
}
#[pymodule]
fn polars(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<PySeries>().unwrap();
m.add_class::<PyDataFrame>().unwrap();
m.add_class::<PyLazyFrame>().unwrap();
m.add_class::<PyLazyGroupBy>().unwrap();
m.add_class::<dsl::PyExpr>().unwrap();
m.add_wrapped(wrap_pyfunction!(col)).unwrap();
m.add_wrapped(wrap_pyfunction!(lit)).unwrap();
m.add_wrapped(wrap_pyfunction!(binary_expr)).unwrap();
m.add_wrapped(wrap_pyfunction!(binary_function)).unwrap();
m.add_wrapped(wrap_pyfunction!(pearson_corr)).unwrap();
m.add_wrapped(wrap_pyfunction!(cov)).unwrap();
m.add_wrapped(wrap_pyfunction!(when)).unwrap();
m.add_wrapped(wrap_pyfunction!(version)).unwrap();
m.add_wrapped(wrap_pyfunction!(toggle_string_cache))
.unwrap();
m.add_wrapped(wrap_pyfunction!(except_)).unwrap();
m.add_wrapped(wrap_pyfunction!(range)).unwrap();
Ok(())
}