You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2023/06/16 20:03:09 UTC
[GitHub] [arrow-datafusion] alamb opened a new pull request, #6703: Support User Defined Window Functions
alamb opened a new pull request, #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703
Draft:
- [ ] Builds on https://github.com/apache/arrow-datafusion/pull/6690
- [ ] Complete tests
# Which issue does this PR close?
Closes https://github.com/apache/arrow-datafusion/issues/5781
Closes https://github.com/apache/arrow-datafusion/pull/6617
# Rationale for this change
See https://github.com/apache/arrow-datafusion/issues/5781
This is the final PR that integrates all the work we have done this week with @mustafarepo and others to support User Defined Window Functions
It basically takes the parts of https://github.com/apache/arrow-datafusion/pull/6617 that have not been merged already, and adds some more tests.
# What changes are included in this PR?
1. Introduce `WindowUDF` and connect it through datafusion_expr and contexts
2. Add `simple_udwf.rs` example
3. Add tests in user_defined_functions.rs integration test
# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
4. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
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] alamb merged pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238418202
##########
datafusion/core/tests/user_defined/user_defined_window_functions.rs:
##########
@@ -0,0 +1,548 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This module contains end to end tests of creating
+//! user defined window functions
+
+use std::{
+ ops::Range,
+ sync::{
+ atomic::{AtomicUsize, Ordering},
+ Arc,
+ },
+};
+
+use arrow::array::AsArray;
+use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow_schema::DataType;
+use datafusion::{assert_batches_eq, prelude::SessionContext};
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::{
+ function::PartitionEvaluatorFactory, window_state::WindowAggState,
+ PartitionEvaluator, ReturnTypeFunction, Signature, Volatility, WindowUDF,
+};
+
+/// A query with a window function evaluated over the entire partition
+const UNBOUNDED_WINDOW_QUERY: &str = "SELECT x, y, val, \
+ odd_counter(val) OVER (PARTITION BY x ORDER BY y) \
+ from t ORDER BY x, y";
+
+/// A query with a window function evaluated over a moving window
+const BOUNDED_WINDOW_QUERY: &str =
+ "SELECT x, y, val, \
+ odd_counter(val) OVER (PARTITION BY x ORDER BY y ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) \
+ from t ORDER BY x, y";
+
+/// Test to show the contents of the setup
+#[tokio::test]
+async fn test_setup() {
+ let test_state = TestState::new();
+ let TestContext { ctx, test_state: _ } = TestContext::new(test_state);
+
+ let sql = "SELECT * from t order by x, y";
+ let expected = vec![
+ "+---+---+-----+",
+ "| x | y | val |",
+ "+---+---+-----+",
+ "| 1 | a | 0 |",
+ "| 1 | b | 1 |",
+ "| 1 | c | 2 |",
+ "| 2 | d | 3 |",
+ "| 2 | e | 4 |",
+ "| 2 | f | 5 |",
+ "| 2 | g | 6 |",
+ "| 2 | h | 6 |",
+ "| 2 | i | 6 |",
+ "| 2 | j | 6 |",
+ "+---+---+-----+",
+ ];
+ assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+}
+
+/// Basic user defined window function
+#[tokio::test]
+async fn test_udwf() {
+ let test_state = TestState::new();
+ let TestContext { ctx, test_state } = TestContext::new(test_state);
+
+ let expected = vec![
+ "+---+---+-----+--------------------+",
+ "| x | y | val | odd_counter(t.val) |",
+ "+---+---+-----+--------------------+",
+ "| 1 | a | 0 | 1 |",
+ "| 1 | b | 1 | 1 |",
+ "| 1 | c | 2 | 1 |",
+ "| 2 | d | 3 | 2 |",
+ "| 2 | e | 4 | 2 |",
+ "| 2 | f | 5 | 2 |",
+ "| 2 | g | 6 | 2 |",
+ "| 2 | h | 6 | 2 |",
+ "| 2 | i | 6 | 2 |",
+ "| 2 | j | 6 | 2 |",
+ "+---+---+-----+--------------------+",
+ ];
+ assert_batches_eq!(
+ expected,
+ &execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap()
+ );
+ // evaluated on three distinct batches
Review Comment:
```suggestion
// evaluated on two distinct batches
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238394899
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
+ // creating a new `PartitionEvaluator`)
+ //
+ // `ORDER BY time`: within each partition ('green' or 'red') the
+ // rows will be be orderd by the value in the `time` column
+ //
+ // `evaluate_inside_range` is invoked with a window defined by the
+ // SQL. In this case:
+ //
+ // The first invocation will be passed row 0, the first row in the
+ // partition.
+ //
+ // The second invocation will be passed rows 0 and 1, the first
+ // two rows in the partition.
+ //
+ // etc.
+ let df = ctx
+ .sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ )
+ .await?;
+ // print the results
+ df.show().await?;
+
+ // this time, call the new widow function with an explicit
+ // window so evaluate will be invoked with each window.
+ //
+ // `ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING`: each invocation
Review Comment:
```suggestion
// `ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING`: each invocation
```
--
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] alamb commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238711464
##########
datafusion/expr/src/expr_fn.rs:
##########
@@ -791,6 +793,25 @@ pub fn create_udaf(
)
}
+/// Creates a new UDWF with a specific signature, state type and return type.
+///
+/// The signature and state type must match the ]`PartitionEvaluator`]'s implementation`.
Review Comment:
thank you -- this is a great catch
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
+ // creating a new `PartitionEvaluator`)
+ //
+ // `ORDER BY time`: within each partition ('green' or 'red') the
+ // rows will be be orderd by the value in the `time` column
+ //
+ // `evaluate_inside_range` is invoked with a window defined by the
+ // SQL. In this case:
+ //
+ // The first invocation will be passed row 0, the first row in the
+ // partition.
+ //
+ // The second invocation will be passed rows 0 and 1, the first
+ // two rows in the partition.
+ //
+ // etc.
+ let df = ctx
+ .sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ )
+ .await?;
+ // print the results
+ df.show().await?;
+
+ // this time, call the new widow function with an explicit
+ // window so evaluate will be invoked with each window.
+ //
+ // `ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING`: each invocation
+ // sees at most 3 rows: the row before, the current row, and the 1
+ // row afterward.
+ let df = ctx.sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ ).await?;
+ // print the results
+ df.show().await?;
+
+ // Now, run the function using the DataFrame API:
+ let window_expr = smooth_it().call(
+ vec![col("speed")], // smooth_it(speed)
+ vec![col("car")], // PARTITION BY car
+ vec![col("time").sort(true, true)], // ORDER BY time ASC
+ WindowFrame::new(false),
+ );
+ let df = ctx.table("cars").await?.window(vec![window_expr])?;
+
+ // print the results
+ df.show().await?;
+
+ Ok(())
+}
+
+fn smooth_it() -> WindowUDF {
+ WindowUDF {
+ name: String::from("smooth_it"),
+ // it will take 1 arguments -- the column to smooth
+ signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),
Review Comment:
I thought it was the default type that the CSV parser picked, but apparently not:
```sql
❯ describe './datafusion/core/tests/data/cars.csv';
describe './datafusion/core/tests/data/cars.csv';
+-------------+-----------------------------+-------------+
| column_name | data_type | is_nullable |
+-------------+-----------------------------+-------------+
| car | Utf8 | YES |
| speed | Float64 | YES |
| time | Timestamp(Nanosecond, None) | YES |
+-------------+-----------------------------+-------------+
3 rows in set. Query took 0.031 seconds.
```
I'lll update this to take a Float instead.
##########
datafusion/core/src/dataframe.rs:
##########
@@ -218,6 +218,14 @@ impl DataFrame {
Ok(DataFrame::new(self.session_state, plan))
}
+ /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+ pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {
Review Comment:
Good call -- I tried to make one, and it turns out to be non trivial. I will do it in a follow on PR
##########
datafusion/execution/src/task.rs:
##########
@@ -92,13 +100,15 @@ impl TaskContext {
config.set(&k, &v)?;
}
let session_config = SessionConfig::from(config);
+ let window_functions = HashMap::new();
Review Comment:
Yes, exactly. I also tried to lead people to `SessionContext::task_ctx` with another doc comment above.
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238390315
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
Review Comment:
```suggestion
// declare a new context. In spark API, this corresponds to a new spark SQL session
```
--
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] alamb commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1239927124
##########
datafusion/core/src/dataframe.rs:
##########
@@ -218,6 +218,14 @@ impl DataFrame {
Ok(DataFrame::new(self.session_state, plan))
}
+ /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+ pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {
Review Comment:
@stuartcarnie I wonder if you have any thoughts about the API here for making window functions: https://github.com/apache/arrow-datafusion/pull/6746 (it is somewhat complex at the moment)
--
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] alamb commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1237006176
##########
datafusion/core/tests/user_defined_integration.rs:
##########
@@ -0,0 +1,26 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+/// Run all tests that are found in the `user_defined` directory
Review Comment:
All the user defined tests now run as part of the same test binaries `user_defined_integration`
##########
datafusion/core/src/physical_plan/windows/mod.rs:
##########
@@ -206,6 +212,71 @@ fn create_built_in_window_expr(
})
}
+/// Creates a `BuiltInWindowFunctionExpr` suitable for a user defined window function
Review Comment:
This mirrors the code and structure for `AggregateUDF`
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
This is a user facing example of how to use User Defined Window Functions
--
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] alamb commented on pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#issuecomment-1599464261
Ok, I think this PR is basically ready for review (once https://github.com/apache/arrow-datafusion/pull/6690 is merged).
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238397589
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
+ // creating a new `PartitionEvaluator`)
+ //
+ // `ORDER BY time`: within each partition ('green' or 'red') the
+ // rows will be be orderd by the value in the `time` column
+ //
+ // `evaluate_inside_range` is invoked with a window defined by the
+ // SQL. In this case:
+ //
+ // The first invocation will be passed row 0, the first row in the
+ // partition.
+ //
+ // The second invocation will be passed rows 0 and 1, the first
+ // two rows in the partition.
+ //
+ // etc.
+ let df = ctx
+ .sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ )
+ .await?;
+ // print the results
+ df.show().await?;
+
+ // this time, call the new widow function with an explicit
+ // window so evaluate will be invoked with each window.
+ //
+ // `ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING`: each invocation
+ // sees at most 3 rows: the row before, the current row, and the 1
+ // row afterward.
+ let df = ctx.sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ ).await?;
+ // print the results
+ df.show().await?;
+
+ // Now, run the function using the DataFrame API:
+ let window_expr = smooth_it().call(
+ vec![col("speed")], // smooth_it(speed)
+ vec![col("car")], // PARTITION BY car
+ vec![col("time").sort(true, true)], // ORDER BY time ASC
+ WindowFrame::new(false),
+ );
+ let df = ctx.table("cars").await?.window(vec![window_expr])?;
+
+ // print the results
+ df.show().await?;
+
+ Ok(())
+}
+
+fn smooth_it() -> WindowUDF {
+ WindowUDF {
+ name: String::from("smooth_it"),
+ // it will take 1 arguments -- the column to smooth
+ signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),
+ return_type: Arc::new(return_type),
+ partition_evaluator_factory: Arc::new(make_partition_evaluator),
+ }
+}
+
+/// Compute the return type of the smooth_it window function given
+/// arguments of `arg_types`.
+fn return_type(arg_types: &[DataType]) -> Result<Arc<DataType>> {
+ if arg_types.len() != 1 {
+ return Err(DataFusionError::Plan(format!(
+ "my_udwf expects 1 argument, got {}: {:?}",
+ arg_types.len(),
+ arg_types
+ )));
+ }
+ Ok(Arc::new(arg_types[0].clone()))
+}
+
+/// Create a `PartitionEvalutor` to evaluate this function on a new
+/// partition.
+fn make_partition_evaluator() -> Result<Box<dyn PartitionEvaluator>> {
+ Ok(Box::new(MyPartitionEvaluator::new()))
+}
+
+/// This implements the lowest level evaluation for a window function
+///
+/// It handles calculating the value of the window function for each
+/// distinct values of `PARTITION BY` (each car type in our example)
+#[derive(Clone, Debug)]
+struct MyPartitionEvaluator {}
+
+impl MyPartitionEvaluator {
+ fn new() -> Self {
+ Self {}
+ }
+}
+
+/// Different evaluation methods are called depending on the various
+/// settings of WindowUDF. This example uses the simplest and most
+/// general, `evaluate`. See `PartitionEvaluator` for the other more
+/// advanced uses.
+impl PartitionEvaluator for MyPartitionEvaluator {
+ /// Tell DataFusion the window function varies based on the value
+ /// of the window frame.
+ fn uses_window_frame(&self) -> bool {
+ true
+ }
+
+ /// This function is called once per input row.
+ ///
+ /// `range`specifies which indexes of `values` should be
+ /// considered for the calculation.
+ ///
+ /// Note this is the SLOWEST, but simplest, way to evaluate a
+ /// window function. It is much faster to implement
+ /// evaluate_all or evaluate_all_with_rank, if possible
+ fn evaluate(
+ &mut self,
+ values: &[ArrayRef],
+ range: &std::ops::Range<usize>,
+ ) -> Result<ScalarValue> {
+ //println!("evaluate_inside_range(). range: {range:#?}, values: {values:#?}");
Review Comment:
I think this is leftover
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238390315
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
Review Comment:
```suggestion
// declare a new context. In spark API, this corresponds to a new spark SQL session
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238393725
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
Review Comment:
```suggestion
// should be treated as a separate partition (and will result in
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238447930
##########
datafusion/expr/src/expr_fn.rs:
##########
@@ -791,6 +793,25 @@ pub fn create_udaf(
)
}
+/// Creates a new UDWF with a specific signature, state type and return type.
+///
+/// The signature and state type must match the ]`PartitionEvaluator`]'s implementation`.
Review Comment:
```suggestion
/// The signature and state type must match the [`PartitionEvaluator`]'s implementation`.
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238404625
##########
datafusion/core/src/execution/context.rs:
##########
@@ -786,6 +786,20 @@ impl SessionContext {
.insert(f.name.clone(), Arc::new(f));
}
+ /// Registers an window UDF within this context.
Review Comment:
```suggestion
/// Registers a window UDF within this context.
```
--
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] stuartcarnie commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "stuartcarnie (via GitHub)" <gi...@apache.org>.
stuartcarnie commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238004128
##########
datafusion-examples/examples/simple_udwf.rs:
##########
Review Comment:
The documentation in this example is excellent 👏🏻
##########
datafusion-examples/README.md:
##########
@@ -57,6 +57,7 @@ cargo run --example csv_sql
- [`rewrite_expr.rs`](examples/rewrite_expr.rs): Define and invoke a custom Query Optimizer pass
- [`simple_udaf.rs`](examples/simple_udaf.rs): Define and invoke a User Defined Aggregate Function (UDAF)
- [`simple_udf.rs`](examples/simple_udf.rs): Define and invoke a User Defined (scalar) Function (UDF)
+- [`simple_udfw.rs`](examples/simple_udwf.rs): Define and invoke a User Defined Window Function (UDWF)
Review Comment:
These examples are 💯
##########
datafusion/core/src/dataframe.rs:
##########
@@ -218,6 +218,14 @@ impl DataFrame {
Ok(DataFrame::new(self.session_state, plan))
}
+ /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+ pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {
Review Comment:
I was just looking for a similar API today – great to know it is coming
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238445524
##########
datafusion/execution/src/task.rs:
##########
@@ -92,13 +100,15 @@ impl TaskContext {
config.set(&k, &v)?;
}
let session_config = SessionConfig::from(config);
+ let window_functions = HashMap::new();
Review Comment:
I guess, since this API is deprecated anyway, you didn't update its arguments.
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238453481
##########
datafusion/expr/src/udwf.rs:
##########
@@ -0,0 +1,118 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Support for user-defined window (UDWF) window functions
+
+use std::{
+ fmt::{self, Debug, Display, Formatter},
+ sync::Arc,
+};
+
+use crate::{
+ Expr, PartitionEvaluatorFactory, ReturnTypeFunction, Signature, WindowFrame,
+};
+
+/// Logical representation of a user-defined window function (UDWF)
+/// A UDAF is different from a UDF in that it is stateful across batches.
Review Comment:
```suggestion
/// A UDWF is different from a UDF in that it is stateful across batches.
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238394073
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
+ // creating a new `PartitionEvaluator`)
+ //
+ // `ORDER BY time`: within each partition ('green' or 'red') the
+ // rows will be be orderd by the value in the `time` column
Review Comment:
```suggestion
// rows will be be ordered by the value in the `time` column
```
--
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] alamb commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1235805257
##########
datafusion/core/tests/user_defined/user_defined_window_functions.rs:
##########
@@ -0,0 +1,548 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This module contains end to end tests of creating
+//! user defined window functions
+
+use std::{
+ ops::Range,
+ sync::{
+ atomic::{AtomicUsize, Ordering},
+ Arc,
+ },
+};
+
+use arrow::array::AsArray;
+use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow_schema::DataType;
+use datafusion::{assert_batches_eq, prelude::SessionContext};
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::{
+ function::PartitionEvaluatorFactory, window_state::WindowAggState,
+ PartitionEvaluator, ReturnTypeFunction, Signature, Volatility, WindowUDF,
+};
+
+/// A query with a window function evaluated over the entire partition
+const UNBOUNDED_WINDOW_QUERY: &str = "SELECT x, y, val, \
+ odd_counter(val) OVER (PARTITION BY x ORDER BY y) \
+ from t ORDER BY x, y";
+
+/// A query with a window function evaluated over a moving window
+const BOUNDED_WINDOW_QUERY: &str =
+ "SELECT x, y, val, \
+ odd_counter(val) OVER (PARTITION BY x ORDER BY y ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) \
+ from t ORDER BY x, y";
+
+/// Test to show the contents of the setup
+#[tokio::test]
+async fn test_setup() {
+ let test_state = TestState::new();
+ let TestContext { ctx, test_state: _ } = TestContext::new(test_state);
+
+ let sql = "SELECT * from t order by x, y";
+ let expected = vec![
+ "+---+---+-----+",
+ "| x | y | val |",
+ "+---+---+-----+",
+ "| 1 | a | 0 |",
+ "| 1 | b | 1 |",
+ "| 1 | c | 2 |",
+ "| 2 | d | 3 |",
+ "| 2 | e | 4 |",
+ "| 2 | f | 5 |",
+ "| 2 | g | 6 |",
+ "| 2 | h | 6 |",
+ "| 2 | i | 6 |",
+ "| 2 | j | 6 |",
+ "+---+---+-----+",
+ ];
+ assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+}
+
+/// Basic user defined window function
+#[tokio::test]
+async fn test_udwf() {
+ let test_state = TestState::new();
+ let TestContext { ctx, test_state } = TestContext::new(test_state);
+
+ let expected = vec![
+ "+---+---+-----+--------------------+",
+ "| x | y | val | odd_counter(t.val) |",
+ "+---+---+-----+--------------------+",
+ "| 1 | a | 0 | 1 |",
+ "| 1 | b | 1 | 1 |",
+ "| 1 | c | 2 | 1 |",
+ "| 2 | d | 3 | 2 |",
+ "| 2 | e | 4 | 2 |",
+ "| 2 | f | 5 | 2 |",
+ "| 2 | g | 6 | 2 |",
+ "| 2 | h | 6 | 2 |",
+ "| 2 | i | 6 | 2 |",
+ "| 2 | j | 6 | 2 |",
+ "+---+---+-----+--------------------+",
+ ];
+ assert_batches_eq!(
+ expected,
+ &execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap()
+ );
+ // evaluated on three distinct batches
+ assert_eq!(test_state.evaluate_all_called(), 2);
+}
+
+/// Basic user defined window function with bounded window
+#[tokio::test]
+async fn test_udwf_bounded_window_ignores_frame() {
+ let test_state = TestState::new();
+ let TestContext { ctx, test_state } = TestContext::new(test_state);
+
+ // Since the UDWF doesn't say it needs the window frame, the frame is ignored
+ let expected = vec![
+ "+---+---+-----+--------------------+",
+ "| x | y | val | odd_counter(t.val) |",
+ "+---+---+-----+--------------------+",
+ "| 1 | a | 0 | 1 |",
+ "| 1 | b | 1 | 1 |",
+ "| 1 | c | 2 | 1 |",
+ "| 2 | d | 3 | 2 |",
+ "| 2 | e | 4 | 2 |",
+ "| 2 | f | 5 | 2 |",
+ "| 2 | g | 6 | 2 |",
+ "| 2 | h | 6 | 2 |",
+ "| 2 | i | 6 | 2 |",
+ "| 2 | j | 6 | 2 |",
+ "+---+---+-----+--------------------+",
+ ];
+ assert_batches_eq!(
+ expected,
+ &execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap()
+ );
+ // evaluated on 2 distinct batches (when x=1 and x=2)
+ assert_eq!(test_state.evaluate_called(), 0);
+ assert_eq!(test_state.evaluate_all_called(), 2);
+}
+
+/// Basic user defined window function with bounded window
+#[tokio::test]
+async fn test_udwf_bounded_window() {
+ let test_state = TestState::new().with_uses_window_frame();
+ let TestContext { ctx, test_state } = TestContext::new(test_state);
+
+ let expected = vec![
+ "+---+---+-----+--------------------+",
+ "| x | y | val | odd_counter(t.val) |",
+ "+---+---+-----+--------------------+",
+ "| 1 | a | 0 | 1 |",
+ "| 1 | b | 1 | 1 |",
+ "| 1 | c | 2 | 1 |",
+ "| 2 | d | 3 | 1 |",
+ "| 2 | e | 4 | 2 |",
+ "| 2 | f | 5 | 1 |",
+ "| 2 | g | 6 | 1 |",
+ "| 2 | h | 6 | 0 |",
+ "| 2 | i | 6 | 0 |",
+ "| 2 | j | 6 | 0 |",
+ "+---+---+-----+--------------------+",
+ ];
+ assert_batches_eq!(
+ expected,
+ &execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap()
+ );
+ // Evaluate is called for each input rows
+ assert_eq!(test_state.evaluate_called(), 10);
+ assert_eq!(test_state.evaluate_all_called(), 0);
+}
+
+/// Basic stateful user defined window function
+#[tokio::test]
+async fn test_stateful_udwf() {
Review Comment:
I couldn't think of anything else to test with a stateful function other than `update_state` is called. I would be interested in anyone else's opinions in this matter
##########
datafusion/core/tests/user_defined/user_defined_window_functions.rs:
##########
@@ -0,0 +1,548 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This module contains end to end tests of creating
+//! user defined window functions
+
+use std::{
+ ops::Range,
+ sync::{
+ atomic::{AtomicUsize, Ordering},
+ Arc,
+ },
+};
+
+use arrow::array::AsArray;
+use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use arrow_schema::DataType;
+use datafusion::{assert_batches_eq, prelude::SessionContext};
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::{
+ function::PartitionEvaluatorFactory, window_state::WindowAggState,
+ PartitionEvaluator, ReturnTypeFunction, Signature, Volatility, WindowUDF,
+};
+
+/// A query with a window function evaluated over the entire partition
Review Comment:
Here are the new tests
--
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] alamb commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238821133
##########
datafusion/core/src/dataframe.rs:
##########
@@ -218,6 +218,14 @@ impl DataFrame {
Ok(DataFrame::new(self.session_state, plan))
}
+ /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+ pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {
Review Comment:
Here is what I was thinking: https://github.com/apache/arrow-datafusion/pull/6746 if you like it I will polish it up
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238401273
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
+ ctx.register_udwf(smooth_it());
+
+ // Use SQL to run the new window function
+ let df = ctx.sql("SELECT * from cars").await?;
+ // print the results
+ df.show().await?;
+
+ // Use SQL to run the new window function:
+ //
+ // `PARTITION BY car`:each distinct value of car (red, and green)
+ // should be treated as a seprate partition (and will result in
+ // creating a new `PartitionEvaluator`)
+ //
+ // `ORDER BY time`: within each partition ('green' or 'red') the
+ // rows will be be orderd by the value in the `time` column
+ //
+ // `evaluate_inside_range` is invoked with a window defined by the
+ // SQL. In this case:
+ //
+ // The first invocation will be passed row 0, the first row in the
+ // partition.
+ //
+ // The second invocation will be passed rows 0 and 1, the first
+ // two rows in the partition.
+ //
+ // etc.
+ let df = ctx
+ .sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ )
+ .await?;
+ // print the results
+ df.show().await?;
+
+ // this time, call the new widow function with an explicit
+ // window so evaluate will be invoked with each window.
+ //
+ // `ROWS BETWEEN 2 PRECEDING AND 2 FOLLOWING`: each invocation
+ // sees at most 3 rows: the row before, the current row, and the 1
+ // row afterward.
+ let df = ctx.sql(
+ "SELECT \
+ car, \
+ speed, \
+ smooth_it(speed) OVER (PARTITION BY car ORDER BY time ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING),\
+ time \
+ from cars \
+ ORDER BY \
+ car",
+ ).await?;
+ // print the results
+ df.show().await?;
+
+ // Now, run the function using the DataFrame API:
+ let window_expr = smooth_it().call(
+ vec![col("speed")], // smooth_it(speed)
+ vec![col("car")], // PARTITION BY car
+ vec![col("time").sort(true, true)], // ORDER BY time ASC
+ WindowFrame::new(false),
+ );
+ let df = ctx.table("cars").await?.window(vec![window_expr])?;
+
+ // print the results
+ df.show().await?;
+
+ Ok(())
+}
+
+fn smooth_it() -> WindowUDF {
+ WindowUDF {
+ name: String::from("smooth_it"),
+ // it will take 1 arguments -- the column to smooth
+ signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),
Review Comment:
It doesn't seem to be effecting result, but I don't understand why signature takes `DataType::Int32`?
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238391268
##########
datafusion-examples/examples/simple_udwf.rs:
##########
@@ -0,0 +1,213 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{ArrayRef, AsArray, Float64Array},
+ datatypes::Float64Type,
+};
+use arrow_schema::DataType;
+use datafusion::datasource::file_format::options::CsvReadOptions;
+
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowUDF,
+};
+
+// create local execution context with `cars.csv` registered as a table named `cars`
+async fn create_context() -> Result<SessionContext> {
+ // declare a new context. In spark API, this corresponds to a new spark SQLsession
+ let ctx = SessionContext::new();
+
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ println!("pwd: {}", std::env::current_dir().unwrap().display());
+ let csv_path = "datafusion/core/tests/data/cars.csv".to_string();
+ let read_options = CsvReadOptions::default().has_header(true);
+
+ ctx.register_csv("cars", &csv_path, read_options).await?;
+ Ok(ctx)
+}
+
+/// In this example we will declare a user defined window function that computes a moving average and then run it using SQL
+#[tokio::main]
+async fn main() -> Result<()> {
+ let ctx = create_context().await?;
+
+ // register the window function with DataFusion so wecan call it
Review Comment:
```suggestion
// register the window function with DataFusion so we can call it
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238405451
##########
datafusion/core/src/execution/context.rs:
##########
@@ -786,6 +786,20 @@ impl SessionContext {
.insert(f.name.clone(), Arc::new(f));
}
+ /// Registers an window UDF within this context.
+ ///
+ /// Note in SQL queries, window function names are looked up using
+ /// lowercase unless the query uses quotes. For example,
+ ///
+ /// - `SELECT MY_UDWF(x)...` will look for a window function named `"my_uwaf"`
Review Comment:
```suggestion
/// - `SELECT MY_UDWF(x)...` will look for a window function named `"my_udwf"`
```
--
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] mustafasrepo commented on a diff in pull request #6703: Support User Defined Window Functions
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on code in PR #6703:
URL: https://github.com/apache/arrow-datafusion/pull/6703#discussion_r1238403292
##########
datafusion/core/src/dataframe.rs:
##########
@@ -218,6 +218,14 @@ impl DataFrame {
Ok(DataFrame::new(self.session_state, plan))
}
+ /// Apply one or more window functions ([`Expr::WindowFunction`]) to extend the schema
+ pub fn window(self, window_exprs: Vec<Expr>) -> Result<DataFrame> {
Review Comment:
Maybe we can add an example usage to the docstring. such as in the `pub fn aggregate`. We can do so in following PRs also.
--
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