You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/08/02 21:14:19 UTC
[GitHub] [arrow-datafusion] jgoday opened a new pull request #819: Draft: python bindings for window functions
jgoday opened a new pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819
# Which issue does this PR close?
Closes #572.
# Rationale for this change
Allows to use window functions from python bindings (inside a select expression).
```python
df = df.select(
f.col("a"),
f.alias(f.window("lead", [f.col("b")], order_by=[f.order_by(f.col("b"))]), "a_next"),
)
```
# What changes are included in this PR?
# Are there any user-facing changes?
Changes to python crate, new exported functions (sort_by, alias and window).
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] houqp commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
houqp commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-891533471
cc @Jimexist
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday edited a comment on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday edited a comment on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-891341609
There is some pending issues:
1. Currently, only a generic way of calling window functions by name is exported. Should we export all window functions individually ?
```rust
fn window(
name: &str,
args: Vec<expression::Expression>,
partition_by: Option<Vec<expression::Expression>>,
order_by: Option<Vec<expression::Expression>>,
) -> PyResult<expression::Expression> {
```
2. There is some functionality copied from internal datafusion::sql::utils crate,
to find window expressions and wrap a logical plan into a LogicalPlan::Window.
I have changed python crate datafusion dependency (to local project).
```toml
datafusion = { path = "../datafusion" }
```
**Obviously this change is completely wrong and temporal**,
just to allow using internal datafusion::sql::utils from python crate (find_window_exprs and group_window_expr_by_sort_keys).
- (Edit) I have reverted it in the last update (to allowing running the test without changing the python crate dependencies)
- Would it be better to duplicate the code so as not to change the visibility of that internal crate and functions currently public(crate) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] Jimexist merged pull request #819: Python bindings for window functions
Posted by GitBox <gi...@apache.org>.
Jimexist merged pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday closed pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday closed pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#discussion_r753901912
##########
File path: python/src/functions.rs
##########
@@ -85,6 +86,63 @@ fn concat_ws(sep: String, args: Vec<PyExpr>) -> PyResult<PyExpr> {
Ok(logical_plan::concat_ws(sep, &args).into())
}
+/// Creates a new Sort expression
+#[pyfunction]
+fn order_by(
+ expr: PyExpr,
+ asc: Option<bool>,
+ nulls_first: Option<bool>,
+) -> PyResult<PyExpr> {
+ Ok(PyExpr {
+ expr: datafusion::logical_plan::Expr::Sort {
+ expr: Box::new(expr.expr),
+ asc: asc.unwrap_or(true),
+ nulls_first: nulls_first.unwrap_or(true),
+ },
+ })
+}
+
+/// Creates a new Alias expression
+#[pyfunction]
+fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
+ Ok(PyExpr {
+ expr: datafusion::logical_plan::Expr::Alias(
+ Box::new(expr.expr),
+ String::from(name),
+ ),
+ })
+}
Review comment:
do you think alias should be a method on the expression or column object?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday edited a comment on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday edited a comment on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-891341609
There is some pending issues:
1. Currently, only a generic way of calling window functions by name is exported. Should we export all window functions individually ?
```rust
fn window(
name: &str,
args: Vec<expression::Expression>,
partition_by: Option<Vec<expression::Expression>>,
order_by: Option<Vec<expression::Expression>>,
) -> PyResult<expression::Expression> {
```
2. There is some functionality copied from internal datafusion::sql::utils crate,
to find window expressions and wrap a logical plan into a LogicalPlan::Window.
I have changed python crate datafusion dependency (to local project).
```toml
datafusion = { path = "../datafusion" }
```
**Obviously this change is completely wrong and temporal**,
just to allow using internal datafusion::sql::utils from python crate (find_window_exprs and group_window_expr_by_sort_keys).
- I have reverted it in the last update (to allowing running the test without changing the python crate dependencies)
- Would it be better to duplicate the code so as not to change the visibility of that internal crate and functions currently public(crate) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-891341609
There is some pending issues:
1. Currently, only a generic way of calling window functions by name is exported. Should we export all window functions individually ?
```rust
fn window(
name: &str,
args: Vec<expression::Expression>,
partition_by: Option<Vec<expression::Expression>>,
order_by: Option<Vec<expression::Expression>>,
) -> PyResult<expression::Expression> {
```
2. There is some functionality copied from internal datafusion::sql::utils crate,
to find window expressions and wrap a logical plan into a LogicalPlan::Window.
I have changed python crate datafusion dependency (to local project).
```toml
datafusion = { path = "../datafusion" }
```
Obviously this change is completely wrong and temporal,
just to allow using internal datafusion::sql::utils from python crate (find_window_exprs and group_window_expr_by_sort_keys).
- Would it be better to duplicate the code so as not to change the visibility of that internal crate and functions currently public(crate) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday edited a comment on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday edited a comment on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-891341609
There is some pending issues:
1. Currently, only a generic way of calling window functions by name is exported. Should we export all window functions individually ?
```rust
fn window(
name: &str,
args: Vec<expression::Expression>,
partition_by: Option<Vec<expression::Expression>>,
order_by: Option<Vec<expression::Expression>>,
) -> PyResult<expression::Expression> {
```
2. There is some functionality copied from internal datafusion::sql::utils crate,
to find window expressions and wrap a logical plan into a LogicalPlan::Window.
I have changed python crate datafusion dependency (to local project).
```toml
datafusion = { path = "../datafusion" }
```
**Obviously this change is completely wrong and temporal**,
just to allow using internal datafusion::sql::utils from python crate (find_window_exprs and group_window_expr_by_sort_keys).
- Would it be better to duplicate the code so as not to change the visibility of that internal crate and functions currently public(crate) ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#discussion_r683531698
##########
File path: python/src/dataframe.rs
##########
@@ -45,16 +46,111 @@ impl DataFrame {
}
}
+/// TODO: this functions belongs to datafusion::sql::planner and datafusion::sql::utils but right now there are private to the datafusion crate
+pub(crate) fn find_window_exprs(exprs: &[Expr]) -> Vec<Expr> {
+ exprs
+ .into_iter()
+ .map(|e| match e {
+ Expr::WindowFunction { .. } => vec![e.clone()],
+ Expr::Alias(e, _) => find_window_exprs(&vec![*e.clone()]),
+ _ => vec![],
+ })
+ .flatten()
+ .collect::<Vec<_>>()
+}
+
+type WindowSortKey = Vec<Expr>;
+
+/// Generate a sort key for a given window expr's partition_by and order_bu expr
+fn generate_sort_key(partition_by: &[Expr], order_by: &[Expr]) -> WindowSortKey {
+ let mut sort_key = vec![];
+ partition_by.iter().for_each(|e| {
+ let e = e.clone().sort(true, true);
+ if !sort_key.contains(&e) {
+ sort_key.push(e);
+ }
+ });
+ order_by.iter().for_each(|e| {
+ if !sort_key.contains(e) {
+ sort_key.push(e.clone());
+ }
+ });
+ sort_key
+}
+
+/// group a slice of window expression expr by their order by expressions
+pub fn group_window_expr_by_sort_keys(
+ window_expr: &[Expr],
+) -> Result<Vec<(WindowSortKey, Vec<&Expr>)>, DataFusionError> {
+ let mut result = vec![];
+ window_expr.iter().try_for_each(|expr| match expr {
+ Expr::WindowFunction { partition_by, order_by, .. } => {
+ let sort_key = generate_sort_key(partition_by, order_by);
+ if let Some((_, values)) = result.iter_mut().find(
+ |group: &&mut (WindowSortKey, Vec<&Expr>)| matches!(group, (key, _) if *key == sort_key),
+ ) {
+ values.push(expr);
+ } else {
+ result.push((sort_key, vec![expr]))
+ }
+ Ok(())
+ }
+ other => Err(DataFusionError::Common(format!(
+ "Impossibly got non-window expr {:?}",
+ other,
+ ))),
+ })?;
+ Ok(result)
+}
+
+fn wrap_window(
+ input: LogicalPlan,
+ window_exprs: Vec<Expr>,
+) -> Result<LogicalPlan, DataFusionError> {
+ let mut plan = input;
+ let mut groups = group_window_expr_by_sort_keys(&window_exprs[..])?;
+ // sort by sort_key len descending, so that more deeply sorted plans gets nested further
+ // down as children; to further mimic the behavior of PostgreSQL, we want stable sort
+ // and a reverse so that tieing sort keys are reversed in order; note that by this rule
+ // if there's an empty over, it'll be at the top level
Review comment:
i wonder if there can be a way to reuse this part of the code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
houqp commented on a change in pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#discussion_r753916429
##########
File path: python/src/functions.rs
##########
@@ -85,6 +86,63 @@ fn concat_ws(sep: String, args: Vec<PyExpr>) -> PyResult<PyExpr> {
Ok(logical_plan::concat_ws(sep, &args).into())
}
+/// Creates a new Sort expression
+#[pyfunction]
+fn order_by(
+ expr: PyExpr,
+ asc: Option<bool>,
+ nulls_first: Option<bool>,
+) -> PyResult<PyExpr> {
+ Ok(PyExpr {
+ expr: datafusion::logical_plan::Expr::Sort {
+ expr: Box::new(expr.expr),
+ asc: asc.unwrap_or(true),
+ nulls_first: nulls_first.unwrap_or(true),
+ },
+ })
+}
+
+/// Creates a new Alias expression
+#[pyfunction]
+fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
+ Ok(PyExpr {
+ expr: datafusion::logical_plan::Expr::Alias(
+ Box::new(expr.expr),
+ String::from(name),
+ ),
+ })
+}
Review comment:
expression already has an alias method, but perhaps we can support both APIs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-953333876
> FYI #1167 added some support in DataFrame API
Thank you @Jimexist, I have updated the code to avoid code duplication and therefore use the DataFrame API as you pointed me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] jgoday commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
jgoday commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-972099356
Hi @Jimexist, rebasing done.
Now, thanks to the last changes from master, there's no need to adapt the PyDataFrame select expressions (there is no changes in python/src/dataframe.rs now), much clearer now !
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] Jimexist commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
Jimexist commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-951521671
FYI #1167 added some support in DataFrame API
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] Jimexist commented on pull request #819: Draft: python bindings for window functions
Posted by GitBox <gi...@apache.org>.
Jimexist commented on pull request #819:
URL: https://github.com/apache/arrow-datafusion/pull/819#issuecomment-971252078
hi @jgoday do you mind rebasing first?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org