feat: close upstream coverage gaps for DataFusion 55.1.0 #1763
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -55,7 +55,7 @@ use crate::expr::aggregate_expr::PyAggregateFunction; | |
| use crate::expr::binary_expr::PyBinaryExpr; | ||
| use crate::expr::column::PyColumn; | ||
| use crate::expr::literal::PyLiteral; | ||
| use crate::functions::{add_builder_fns_to_window, apply_window_options}; | ||
| use crate::pyarrow_util::scalar_to_pyarrow; | ||
| use crate::sql::logical::PyLogicalPlan; | ||
|
|
||
|
|
@@ -624,44 +624,68 @@ impl PyExpr { | |
|
|
||
| // Expression Function Builder functions | ||
|
|
||
| #[pyo3(signature = (order_by, keep_window_frame=false))] | ||
| pub fn order_by( | ||
| &self, | ||
| order_by: Vec<PySortExpr>, | ||
| keep_window_frame: bool, | ||
| ) -> PyExprFuncBuilder { | ||
| builder_from_expr(&self.expr, keep_window_frame) | ||
| .order_by(to_sort_expressions(order_by)) | ||
| .into() | ||
| } | ||
|
|
||
| #[pyo3(signature = (filter, keep_window_frame=false))] | ||
| pub fn filter(&self, filter: PyExpr, keep_window_frame: bool) -> PyExprFuncBuilder { | ||
| builder_from_expr(&self.expr, keep_window_frame) | ||
| .filter(filter.expr.clone()) | ||
| .into() | ||
| } | ||
|
|
||
| #[pyo3(signature = (keep_window_frame=false))] | ||
| pub fn distinct(&self, keep_window_frame: bool) -> PyExprFuncBuilder { | ||
| builder_from_expr(&self.expr, keep_window_frame) | ||
| .distinct() | ||
| .into() | ||
| } | ||
|
|
||
| #[pyo3(signature = (null_treatment, keep_window_frame=false))] | ||
| pub fn null_treatment( | ||
| &self, | ||
| null_treatment: NullTreatment, | ||
| keep_window_frame: bool, | ||
| ) -> PyExprFuncBuilder { | ||
| builder_from_expr(&self.expr, keep_window_frame) | ||
| .null_treatment(Some(null_treatment.into())) | ||
| .into() | ||
| } | ||
|
|
||
| #[pyo3(signature = (partition_by, keep_window_frame=false))] | ||
| pub fn partition_by( | ||
| &self, | ||
| partition_by: Vec<PyExpr>, | ||
| keep_window_frame: bool, | ||
| ) -> PyExprFuncBuilder { | ||
| let partition_by = partition_by.iter().map(|e| e.expr.clone()).collect(); | ||
| builder_from_expr(&self.expr, keep_window_frame) | ||
| .partition_by(partition_by) | ||
| .into() | ||
| } | ||
|
|
||
| pub fn window_frame(&self, window_frame: PyWindowFrame) -> PyExprFuncBuilder { | ||
| builder_from_expr(&self.expr, false) | ||
| .window_frame(window_frame.into()) | ||
| .into() | ||
| } | ||
|
|
||
| #[pyo3(signature = (partition_by=None, window_frame=None, order_by=None, null_treatment=None, keep_window_frame=false))] | ||
| pub fn over( | ||
| &self, | ||
| partition_by: Option<Vec<PyExpr>>, | ||
| window_frame: Option<PyWindowFrame>, | ||
| order_by: Option<Vec<PySortExpr>>, | ||
| null_treatment: Option<NullTreatment>, | ||
| keep_window_frame: bool, | ||
| ) -> PyDataFusionResult<PyExpr> { | ||
| match &self.expr { | ||
| Expr::AggregateFunction(agg_fn) => { | ||
|
Member
There was a problem hiding this comment. This is #1764, but the new from datafusion import SessionContext, col, functions as f
from datafusion.expr import Window
ctx = SessionContext()
df = ctx.from_pydict({"v": [1.0, 1.0, None, 4.0]}, name="t")
print(df.select(f.mean(col("v"), distinct=True).over(Window()).alias("m")).to_pydict())
# {'m': [2.0, 2.0, 2.0, 2.0]}
print(ctx.sql("SELECT avg(DISTINCT v) OVER () AS m FROM t").to_pydict())
# {'m': [2.5, 2.5, 2.5, 2.5]}Until #1764 lands, a note in the |
||
|
|
@@ -678,8 +702,8 @@ impl PyExpr { | |
| null_treatment, | ||
| ) | ||
| } | ||
| Expr::WindowFunction(_) => apply_window_options( | ||
| builder_from_expr(&self.expr, keep_window_frame), | ||
| partition_by, | ||
| window_frame, | ||
| order_by, | ||
|
|
@@ -743,6 +767,63 @@ impl PyExpr { | |
| } | ||
| } | ||
|
|
||
| /// Start an [`ExprFuncBuilder`] that keeps the options already set on `expr`. | ||
| /// | ||
| /// Upstream's `ExprFunctionExt` methods on an `Expr` start from an empty | ||
| /// builder, so `build()` would reset every option not set again. The Python | ||
| /// function wrappers already apply their keyword options, so chaining another | ||
| /// builder method onto their result must not discard them. | ||
| /// | ||
| /// A built window function always stores a concrete frame, so whether the user | ||
| /// chose it is lost. `keep_window_frame` carries that from the Python side; when | ||
| /// false, a frame equal to the default for the current order-by is treated as | ||
| /// unset. | ||
| fn builder_from_expr(expr: &Expr, keep_window_frame: bool) -> ExprFuncBuilder { | ||
|
Member
There was a problem hiding this comment. Because the seeded builder always carries a function kind, upstream's per-method kind checks in from datafusion import SessionContext, col, functions as f
from datafusion.expr import WindowFrame
ctx = SessionContext()
df = ctx.from_pydict({"g": [1, 1, 1, 2], "v": [3, 1, 2, 4]})
e = f.sum(col("v")).partition_by(col("g")).build()
print(e) # Expr(sum(v))
print(df.aggregate([], [e.alias("s")]).to_pydict()) # {'s': [10]}
print(f.sum(col("v")).window_frame(WindowFrame("rows", 1, 0)).build()) # Expr(sum(v))
print(f.lead(col("v"), order_by="v").distinct().build())
# Expr(lead(DISTINCT v, Int64(1), NULL) ORDER BY [...]) -- runs, DISTINCT ignoredThe new |
||
| match expr { | ||
| Expr::AggregateFunction(agg) => { | ||
| let params = &agg.params; | ||
| let mut builder = expr.clone().null_treatment(params.null_treatment); | ||
| if !params.order_by.is_empty() { | ||
| builder = builder.order_by(params.order_by.clone()); | ||
| } | ||
| if let Some(filter) = ¶ms.filter { | ||
| builder = builder.filter(filter.as_ref().clone()); | ||
| } | ||
| if params.distinct { | ||
| builder = builder.distinct(); | ||
| } | ||
| builder | ||
| } | ||
| Expr::WindowFunction(window) => { | ||
| let params = &window.params; | ||
| let mut builder = expr.clone().null_treatment(params.null_treatment); | ||
| if !params.partition_by.is_empty() { | ||
| builder = builder.partition_by(params.partition_by.clone()); | ||
| } | ||
| let has_order_by = !params.order_by.is_empty(); | ||
| if has_order_by { | ||
| builder = builder.order_by(params.order_by.clone()); | ||
| } | ||
| // A frame equal to the default `build()` derived from the order-by is | ||
| // left unset, so it is derived again from the final order-by. | ||
| if keep_window_frame | ||
| || params.window_frame | ||
| != datafusion::logical_expr::WindowFrame::new(has_order_by.then_some(true)) | ||
|
Member
There was a problem hiding this comment.
from datafusion import SessionContext, col, functions as f
from datafusion.expr import Window
ctx = SessionContext()
df = ctx.from_pydict({"i": [1, 1, 2, 3], "v": [1, 2, 3, 4]})
def run(e):
return df.select(col("i"), e.alias("r")).sort(col("i")).collect_column("r").to_pylist()
print(run(f.sum(col("v")).over(Window(order_by=[])).over(Window(order_by="i"))))
# PR: [3, 3, 6, 10] (RANGE frame kept, ties summed); main: [1, 3, 6, 10]
print(run(f.sum(col("v")).over(Window()).over(Window(order_by="i"))))
# PR and main: [1, 3, 6, 10]An empty sort list is easy to get from a dynamically built one. Also treating |
||
| { | ||
| builder = builder.window_frame(params.window_frame.clone()); | ||
| } | ||
| if let Some(filter) = ¶ms.filter { | ||
| builder = builder.filter(filter.as_ref().clone()); | ||
| } | ||
| if params.distinct { | ||
| builder = builder.distinct(); | ||
| } | ||
| builder | ||
| } | ||
| _ => expr.clone().null_treatment(None), | ||
| } | ||
| } | ||
|
|
||
| #[pyclass( | ||
| from_py_object, | ||
| frozen, | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Upstream
DataFrame::fill_nan(viafill_columns) rebuilds every column withcol(field.name()), which parses the name as an identifier (lowercasing it, splitting on.). Sofill_nanfails on any DataFrame with an uppercase or dotted column name, even when that column isn't insubset:fill_nullhas the same upstream bug, so this is probably worth an upstream issue (ident(field.name())there would fix both). Until then the wrapper could build the projection itself from unparsed column references plusnanvl.