Skip to content

Commit

Permalink
Adapt to latest IR updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@wence-
wence- committed Apr 29, 2024
1 parent 1c32963 commit 0afbe5a
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 41 deletions.
12 changes: 12 additions & 0 deletions py-polars/src/conversion/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ pub(crate) mod any_value;
pub(crate) mod chunked_array;
use std::fmt::{Display, Formatter};
use std::hash::{Hash, Hasher};
use std::sync::Mutex;

#[cfg(feature = "object")]
use polars::chunked_array::object::PolarsObjectSafe;
Expand Down Expand Up @@ -463,6 +464,17 @@ impl IntoPy<PyObject> for Wrap<&Schema> {
}
}

impl IntoPy<PyObject> for Wrap<&Mutex<Option<SchemaRef>>> {
fn into_py(self, py: Python<'_>) -> PyObject {
let guard = self.0.lock().unwrap();
if let Some(schema) = &*guard {
Wrap(schema.as_ref()).into_py(py)
} else {
py.None()
}
}
}

#[derive(Clone, Debug)]
#[repr(transparent)]
pub struct ObjectValue {
Expand Down
40 changes: 32 additions & 8 deletions py-polars/src/lazyframe/visitor/expr_nodes.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
use polars_core::series::IsSorted;
use polars_plan::dsl::function_expr::rolling::RollingFunction;
use polars_plan::dsl::function_expr::trigonometry::TrigonometricFunction;
use polars_plan::dsl::BooleanFunction;
use polars_plan::prelude::{
AAggExpr, AExpr, FunctionExpr, GroupbyOptions, LiteralValue, Operator, PowFunction,
WindowMapping, WindowType,
};
use polars_rs::series::IsSorted;
use polars_time::prelude::RollingGroupOptions;
use pyo3::exceptions::PyNotImplementedError;
use pyo3::prelude::*;
Expand Down Expand Up @@ -168,7 +168,8 @@ pub struct SortBy {
#[pyo3(get)]
by: Vec<usize>,
#[pyo3(get)]
descending: Vec<bool>,
/// descending, nulls_last, maintain_order
sort_options: (Vec<bool>, bool, bool),
}

#[pyclass]
Expand Down Expand Up @@ -334,6 +335,10 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
use LiteralValue::*;
let dtype: PyObject = Wrap(lit.get_datatype()).to_object(py);
match lit {
Float(v) => Literal {
value: v.to_object(py),
dtype,
},
Float32(v) => Literal {
value: v.to_object(py),
dtype,
Expand All @@ -342,6 +347,10 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
value: v.to_object(py),
dtype,
},
Int(v) => Literal {
value: v.to_object(py),
dtype,
},
Int8(v) => Literal {
value: v.to_object(py),
dtype,
Expand Down Expand Up @@ -378,6 +387,10 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
value: v.to_object(py),
dtype,
},
StrCat(v) => Literal {
value: v.to_object(py),
dtype,
},
String(v) => Literal {
value: v.to_object(py),
dtype,
Expand Down Expand Up @@ -442,11 +455,15 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
AExpr::SortBy {
expr,
by,
descending,
sort_options,
} => SortBy {
expr: expr.0,
by: by.iter().map(|n| n.0).collect(),
descending: descending.clone(),
sort_options: (
sort_options.descending.clone(),
sort_options.nulls_last,
sort_options.maintain_order,
),
}
.into_py(py),
AExpr::Agg(aggexpr) => match aggexpr {
Expand Down Expand Up @@ -606,9 +623,7 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
.to_object(py),
FunctionExpr::Atan2 => ("atan2",).to_object(py),
FunctionExpr::Sign => ("sign",).to_object(py),
FunctionExpr::FillNull { super_type: _ } => {
return Err(PyNotImplementedError::new_err("fill null"))
},
FunctionExpr::FillNull => return Err(PyNotImplementedError::new_err("fill null")),
FunctionExpr::RollingExpr(rolling) => match rolling {
RollingFunction::Min(_) => {
return Err(PyNotImplementedError::new_err("rolling min"))
Expand Down Expand Up @@ -679,7 +694,9 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
has_max: _,
} => return Err(PyNotImplementedError::new_err("clip")),
FunctionExpr::AsStruct => return Err(PyNotImplementedError::new_err("as struct")),
FunctionExpr::TopK(_) => return Err(PyNotImplementedError::new_err("top k")),
FunctionExpr::TopK { sort_options: _ } => {
return Err(PyNotImplementedError::new_err("top k"))
},
FunctionExpr::CumCount { reverse } => ("cumcount", reverse).to_object(py),
FunctionExpr::CumSum { reverse } => ("cumsum", reverse).to_object(py),
FunctionExpr::CumProd { reverse } => ("cumprod", reverse).to_object(py),
Expand Down Expand Up @@ -806,6 +823,13 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<PyObject> {
FunctionExpr::Business(_) => {
return Err(PyNotImplementedError::new_err("business"))
},
FunctionExpr::TopKBy { sort_options: _ } => {
return Err(PyNotImplementedError::new_err("top_k_by"))
},
FunctionExpr::EwmMeanBy {
half_life: _,
check_sorted: _,
} => return Err(PyNotImplementedError::new_err("ewm_mean_by")),
},
options: py.None(),
}
Expand Down
53 changes: 20 additions & 33 deletions py-polars/src/lazyframe/visitor/nodes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,9 @@ pub struct Sort {
#[pyo3(get)]
by_column: Vec<PyExprIR>,
#[pyo3(get)]
args: PyObject,
sort_options: (Vec<bool>, bool, bool),
#[pyo3(get)]
slice: Option<(i64, usize)>,
}

#[pyclass]
Expand Down Expand Up @@ -239,7 +241,7 @@ pub struct Union {
#[pyo3(get)]
inputs: Vec<usize>,
#[pyo3(get)]
options: PyObject,
options: Option<(i64, usize)>,
}
#[pyclass]
/// Horizontal concatenation of multiple plans
Expand Down Expand Up @@ -384,17 +386,17 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
IR::Sort {
input,
by_column,
args,
slice,
sort_options,
} => Sort {
input: input.0,
by_column: by_column.iter().map(|e| e.into()).collect(),
args: (
args.maintain_order,
args.nulls_last,
&args.descending,
args.slice,
)
.to_object(py),
sort_options: (
sort_options.descending.clone(),
sort_options.nulls_last,
sort_options.maintain_order,
),
slice: *slice,
}
.into_py(py),
IR::Cache {
Expand Down Expand Up @@ -448,28 +450,16 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
match options.args.how {
JoinType::Left => "left",
JoinType::Inner => "inner",
JoinType::Outer { coalesce } => {
if coalesce {
"outer_coalesce"
} else {
"outer"
}
},
JoinType::Outer => "outer",
JoinType::AsOf(_) => return Err(PyNotImplementedError::new_err("asof join")),
JoinType::Cross => "cross",
JoinType::Semi => "leftsemi",
JoinType::Anti => "leftanti",
},
options.args.join_nulls,
options
.args
.slice
.map_or_else(|| py.None(), |f| f.to_object(py)),
options
.args
.suffix
.as_ref()
.map_or_else(|| py.None(), |f| f.to_object(py)),
options.args.slice,
options.args.suffix.clone(),
options.args.coalesce.coalesce(&options.args.how),
)
.to_object(py),
}
Expand Down Expand Up @@ -502,7 +492,7 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
.as_ref()
.map_or_else(|| py.None(), |f| f.to_object(py)),
options.maintain_order,
options.slice.map_or_else(|| py.None(), |f| f.to_object(py)),
options.slice,
)
.to_object(py),
}
Expand Down Expand Up @@ -536,11 +526,6 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
columns.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
)
.to_object(py),
FunctionNode::DropNulls { subset } => (
"drop_nulls",
subset.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
)
.to_object(py),
FunctionNode::Rechunk => ("rechunk",).to_object(py),
FunctionNode::MergeSorted { column } => {
("merge_sorted", column.to_string()).to_object(py)
Expand All @@ -549,11 +534,13 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
existing,
new,
swapping,
schema,
} => (
"rename",
existing.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
new.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
*swapping,
Wrap(schema.as_ref()).into_py(py),
)
.to_object(py),
FunctionNode::Explode { columns, schema } => (
Expand Down Expand Up @@ -600,7 +587,7 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
IR::Union { inputs, options } => Union {
inputs: inputs.iter().map(|n| n.0).collect(),
// TODO: rest of options
options: options.slice.map_or_else(|| py.None(), |f| f.to_object(py)),
options: options.slice,
}
.into_py(py),
IR::HConcat {
Expand Down

0 comments on commit 0afbe5a

Please sign in to comment.