forked from pola-rs/polars
-
Notifications
You must be signed in to change notification settings - Fork 0
/
conversion.rs
98 lines (89 loc) · 3.23 KB
/
conversion.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
use crate::prelude::*;
use crate::series::PySeries;
use polars::prelude::AnyValue;
use pyo3::conversion::{FromPyObject, IntoPy};
use pyo3::prelude::*;
use pyo3::types::PySequence;
use pyo3::{PyAny, PyResult};
pub struct Wrap<T>(pub T);
fn get_pyseq(obj: &PyAny) -> PyResult<(&PySequence, usize)> {
let seq = <PySequence as PyTryFrom>::try_from(obj)?;
let len = seq.len()? as usize;
Ok((seq, len))
}
impl<'a, T> FromPyObject<'a> for Wrap<ChunkedArray<T>>
where
T: PyPolarsPrimitiveType,
T::Native: FromPyObject<'a>,
{
fn extract(obj: &'a PyAny) -> PyResult<Self> {
let (seq, len) = get_pyseq(obj)?;
let mut builder = PrimitiveChunkedBuilder::new("", len);
for res in seq.iter()? {
let item = res?;
match item.extract::<T::Native>() {
Ok(val) => builder.append_value(val),
Err(_) => builder.append_null(),
}
}
Ok(Wrap(builder.finish()))
}
}
impl<'a> FromPyObject<'a> for Wrap<BooleanChunked> {
fn extract(obj: &'a PyAny) -> PyResult<Self> {
let (seq, len) = get_pyseq(obj)?;
let mut builder = BooleanChunkedBuilder::new("", len);
for res in seq.iter()? {
let item = res?;
match item.extract::<bool>() {
Ok(val) => builder.append_value(val),
Err(_) => builder.append_null(),
}
}
Ok(Wrap(builder.finish()))
}
}
impl<'a> FromPyObject<'a> for Wrap<Utf8Chunked> {
fn extract(obj: &'a PyAny) -> PyResult<Self> {
let (seq, len) = get_pyseq(obj)?;
let mut builder = Utf8ChunkedBuilder::new("", len, len * 25);
for res in seq.iter()? {
let item = res?;
match item.extract::<&str>() {
Ok(val) => builder.append_value(val),
Err(_) => builder.append_null(),
}
}
Ok(Wrap(builder.finish()))
}
}
impl IntoPy<PyObject> for Wrap<AnyValue<'_>> {
fn into_py(self, py: Python) -> PyObject {
match self.0 {
AnyValue::UInt8(v) => v.into_py(py),
AnyValue::UInt16(v) => v.into_py(py),
AnyValue::UInt32(v) => v.into_py(py),
AnyValue::UInt64(v) => v.into_py(py),
AnyValue::Int8(v) => v.into_py(py),
AnyValue::Int16(v) => v.into_py(py),
AnyValue::Int32(v) => v.into_py(py),
AnyValue::Int64(v) => v.into_py(py),
AnyValue::Float32(v) => v.into_py(py),
AnyValue::Float64(v) => v.into_py(py),
AnyValue::Null => py.None(),
AnyValue::Boolean(v) => v.into_py(py),
AnyValue::Utf8(v) => v.into_py(py),
AnyValue::Date32(v) => v.into_py(py),
AnyValue::Date64(v) => v.into_py(py),
AnyValue::Time64(v, _) => v.into_py(py),
AnyValue::Duration(v, _) => v.into_py(py),
AnyValue::List(v) => {
let pypolars = PyModule::import(py, "polars").expect("polars installed");
let pyseries = PySeries::new(v);
let python_series_wrapper = pypolars.call1("wrap_s", (pyseries,)).unwrap();
python_series_wrapper.into()
}
AnyValue::Object(v) => v.into_py(py),
}
}
}