This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8c8b2454cb Migrate user_defined tests to insta (#15255)
8c8b2454cb is described below
commit 8c8b2454cbd78204dc6426f9898b79c179486a86
Author: Shruti Sharma <[email protected]>
AuthorDate: Wed Mar 19 00:55:58 2025 +0530
Migrate user_defined tests to insta (#15255)
* migrate user_defined tests to insta
* fix tests
* test
* fmt
* update snaps
* add snaps
* add snaps
* add inline snapshots
* fix
* remove snaps
* Reuse batches_to_stirng
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/core/tests/user_defined/expr_planner.rs | 57 +--
.../tests/user_defined/user_defined_aggregates.rs | 249 +++++++------
.../core/tests/user_defined/user_defined_plan.rs | 100 +++---
.../user_defined/user_defined_scalar_functions.rs | 156 ++++----
.../user_defined/user_defined_table_functions.rs | 59 +--
.../user_defined/user_defined_window_functions.rs | 395 ++++++++++-----------
6 files changed, 514 insertions(+), 502 deletions(-)
diff --git a/datafusion/core/tests/user_defined/expr_planner.rs
b/datafusion/core/tests/user_defined/expr_planner.rs
index 75d890359b..1fc6d14c5b 100644
--- a/datafusion/core/tests/user_defined/expr_planner.rs
+++ b/datafusion/core/tests/user_defined/expr_planner.rs
@@ -16,9 +16,10 @@
// under the License.
use arrow::array::RecordBatch;
+use datafusion::common::test_util::batches_to_string;
use std::sync::Arc;
-use datafusion::common::{assert_batches_eq, DFSchema};
+use datafusion::common::DFSchema;
use datafusion::error::Result;
use datafusion::execution::FunctionRegistry;
use datafusion::logical_expr::Operator;
@@ -76,27 +77,25 @@ async fn plan_and_collect(sql: &str) ->
Result<Vec<RecordBatch>> {
#[tokio::test]
async fn test_custom_operators_arrow() {
let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap();
- let expected = [
- "+----------------------------+",
- "| Utf8(\"foo\") || Utf8(\"bar\") |",
- "+----------------------------+",
- "| foobar |",
- "+----------------------------+",
- ];
- assert_batches_eq!(&expected, &actual);
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +----------------------------+
+ | Utf8("foo") || Utf8("bar") |
+ +----------------------------+
+ | foobar |
+ +----------------------------+
+ "###);
}
#[tokio::test]
async fn test_custom_operators_long_arrow() {
let actual = plan_and_collect("select 1->>2;").await.unwrap();
- let expected = [
- "+---------------------+",
- "| Int64(1) + Int64(2) |",
- "+---------------------+",
- "| 3 |",
- "+---------------------+",
- ];
- assert_batches_eq!(&expected, &actual);
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +---------------------+
+ | Int64(1) + Int64(2) |
+ +---------------------+
+ | 3 |
+ +---------------------+
+ "###);
}
#[tokio::test]
@@ -104,14 +103,13 @@ async fn test_question_select() {
let actual = plan_and_collect("select a ? 2 from (select 1 as a);")
.await
.unwrap();
- let expected = [
- "+--------------+",
- "| a ? Int64(2) |",
- "+--------------+",
- "| true |",
- "+--------------+",
- ];
- assert_batches_eq!(&expected, &actual);
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +--------------+
+ | a ? Int64(2) |
+ +--------------+
+ | true |
+ +--------------+
+ "###);
}
#[tokio::test]
@@ -119,6 +117,11 @@ async fn test_question_filter() {
let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;")
.await
.unwrap();
- let expected = ["+---+", "| a |", "+---+", "| 1 |", "+---+"];
- assert_batches_eq!(&expected, &actual);
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +---+
+ | a |
+ +---+
+ | 1 |
+ +---+
+ "###);
}
diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs
b/datafusion/core/tests/user_defined/user_defined_aggregates.rs
index 7cda6d410f..5cbb05f290 100644
--- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs
+++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs
@@ -30,6 +30,7 @@ use arrow::array::{
};
use arrow::datatypes::{Fields, Schema};
+use datafusion::common::test_util::batches_to_string;
use datafusion::dataframe::DataFrame;
use datafusion::datasource::MemTable;
use datafusion::test_util::plan_and_collect;
@@ -39,7 +40,6 @@ use datafusion::{
datatypes::{DataType, Field, Float64Type, TimeUnit,
TimestampNanosecondType},
record_batch::RecordBatch,
},
- assert_batches_eq,
error::Result,
logical_expr::{
AccumulatorFactoryFunction, AggregateUDF, Signature, TypeSignature,
Volatility,
@@ -48,7 +48,8 @@ use datafusion::{
prelude::SessionContext,
scalar::ScalarValue,
};
-use datafusion_common::{assert_contains, cast::as_primitive_array, exec_err};
+use datafusion_common::assert_contains;
+use datafusion_common::{cast::as_primitive_array, exec_err};
use datafusion_expr::{
col, create_udaf, function::AccumulatorArgs, AggregateUDFImpl,
GroupsAccumulator,
LogicalPlanBuilder, SimpleAggregateUDF,
@@ -60,18 +61,20 @@ use datafusion_functions_aggregate::average::AvgAccumulator;
async fn test_setup() {
let TestContext { ctx, test_state: _ } = TestContext::new();
let sql = "SELECT * from t order by time";
- let expected = [
- "+-------+----------------------------+",
- "| value | time |",
- "+-------+----------------------------+",
- "| 2.0 | 1970-01-01T00:00:00.000002 |",
- "| 3.0 | 1970-01-01T00:00:00.000003 |",
- "| 1.0 | 1970-01-01T00:00:00.000004 |",
- "| 5.0 | 1970-01-01T00:00:00.000005 |",
- "| 5.0 | 1970-01-01T00:00:00.000005 |",
- "+-------+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +-------+----------------------------+
+ | value | time |
+ +-------+----------------------------+
+ | 2.0 | 1970-01-01T00:00:00.000002 |
+ | 3.0 | 1970-01-01T00:00:00.000003 |
+ | 1.0 | 1970-01-01T00:00:00.000004 |
+ | 5.0 | 1970-01-01T00:00:00.000005 |
+ | 5.0 | 1970-01-01T00:00:00.000005 |
+ +-------+----------------------------+
+ "###);
}
/// Basic user defined aggregate
@@ -80,14 +83,17 @@ async fn test_udaf() {
let TestContext { ctx, test_state } = TestContext::new();
assert!(!test_state.update_batch());
let sql = "SELECT time_sum(time) from t";
- let expected = [
- "+----------------------------+",
- "| time_sum(t.time) |",
- "+----------------------------+",
- "| 1970-01-01T00:00:00.000019 |",
- "+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +----------------------------+
+ | time_sum(t.time) |
+ +----------------------------+
+ | 1970-01-01T00:00:00.000019 |
+ +----------------------------+
+ "###);
+
// normal aggregates call update_batch
assert!(test_state.update_batch());
assert!(!test_state.retract_batch());
@@ -98,18 +104,21 @@ async fn test_udaf() {
async fn test_udaf_as_window() {
let TestContext { ctx, test_state } = TestContext::new();
let sql = "SELECT time_sum(time) OVER() as time_sum from t";
- let expected = [
- "+----------------------------+",
- "| time_sum |",
- "+----------------------------+",
- "| 1970-01-01T00:00:00.000019 |",
- "| 1970-01-01T00:00:00.000019 |",
- "| 1970-01-01T00:00:00.000019 |",
- "| 1970-01-01T00:00:00.000019 |",
- "| 1970-01-01T00:00:00.000019 |",
- "+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +----------------------------+
+ | time_sum |
+ +----------------------------+
+ | 1970-01-01T00:00:00.000019 |
+ | 1970-01-01T00:00:00.000019 |
+ | 1970-01-01T00:00:00.000019 |
+ | 1970-01-01T00:00:00.000019 |
+ | 1970-01-01T00:00:00.000019 |
+ +----------------------------+
+ "###);
+
// aggregate over the entire window function call update_batch
assert!(test_state.update_batch());
assert!(!test_state.retract_batch());
@@ -120,18 +129,21 @@ async fn test_udaf_as_window() {
async fn test_udaf_as_window_with_frame() {
let TestContext { ctx, test_state } = TestContext::new();
let sql = "SELECT time_sum(time) OVER(ORDER BY time ROWS BETWEEN 1
PRECEDING AND 1 FOLLOWING) as time_sum from t";
- let expected = [
- "+----------------------------+",
- "| time_sum |",
- "+----------------------------+",
- "| 1970-01-01T00:00:00.000005 |",
- "| 1970-01-01T00:00:00.000009 |",
- "| 1970-01-01T00:00:00.000012 |",
- "| 1970-01-01T00:00:00.000014 |",
- "| 1970-01-01T00:00:00.000010 |",
- "+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +----------------------------+
+ | time_sum |
+ +----------------------------+
+ | 1970-01-01T00:00:00.000005 |
+ | 1970-01-01T00:00:00.000009 |
+ | 1970-01-01T00:00:00.000012 |
+ | 1970-01-01T00:00:00.000014 |
+ | 1970-01-01T00:00:00.000010 |
+ +----------------------------+
+ "###);
+
// user defined aggregates with window frame should be calling retract
batch
assert!(test_state.update_batch());
assert!(test_state.retract_batch());
@@ -155,14 +167,16 @@ async fn
test_udaf_as_window_with_frame_without_retract_batch() {
async fn test_udaf_returning_struct() {
let TestContext { ctx, test_state: _ } = TestContext::new();
let sql = "SELECT first(value, time) from t";
- let expected = [
- "+------------------------------------------------+",
- "| first(t.value,t.time) |",
- "+------------------------------------------------+",
- "| {value: 2.0, time: 1970-01-01T00:00:00.000002} |",
- "+------------------------------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +------------------------------------------------+
+ | first(t.value,t.time) |
+ +------------------------------------------------+
+ | {value: 2.0, time: 1970-01-01T00:00:00.000002} |
+ +------------------------------------------------+
+ "###);
}
/// Demonstrate extracting the fields from a structure using a subquery
@@ -170,14 +184,16 @@ async fn test_udaf_returning_struct() {
async fn test_udaf_returning_struct_subquery() {
let TestContext { ctx, test_state: _ } = TestContext::new();
let sql = "select sq.first['value'], sq.first['time'] from (SELECT
first(value, time) as first from t) as sq";
- let expected = [
- "+-----------------+----------------------------+",
- "| sq.first[value] | sq.first[time] |",
- "+-----------------+----------------------------+",
- "| 2.0 | 1970-01-01T00:00:00.000002 |",
- "+-----------------+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +-----------------+----------------------------+
+ | sq.first[value] | sq.first[time] |
+ +-----------------+----------------------------+
+ | 2.0 | 1970-01-01T00:00:00.000002 |
+ +-----------------+----------------------------+
+ "###);
}
#[tokio::test]
@@ -189,26 +205,29 @@ async fn test_udaf_shadows_builtin_fn() {
let sql = "SELECT sum(arrow_cast(time, 'Int64')) from t";
// compute with builtin `sum` aggregator
- let expected = [
- "+---------------------------------------+",
- "| sum(arrow_cast(t.time,Utf8(\"Int64\"))) |",
- "+---------------------------------------+",
- "| 19000 |",
- "+---------------------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +---------------------------------------+
+ | sum(arrow_cast(t.time,Utf8("Int64"))) |
+ +---------------------------------------+
+ | 19000 |
+ +---------------------------------------+
+ "###);
// Register `TimeSum` with name `sum`. This will shadow the builtin one
- let sql = "SELECT sum(time) from t";
TimeSum::register(&mut ctx, test_state.clone(), "sum");
- let expected = [
- "+----------------------------+",
- "| sum(t.time) |",
- "+----------------------------+",
- "| 1970-01-01T00:00:00.000019 |",
- "+----------------------------+",
- ];
- assert_batches_eq!(expected, &execute(&ctx, sql).await.unwrap());
+ let sql = "SELECT sum(time) from t";
+
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +----------------------------+
+ | sum(t.time) |
+ +----------------------------+
+ | 1970-01-01T00:00:00.000019 |
+ +----------------------------+
+ "###);
}
async fn execute(ctx: &SessionContext, sql: &str) -> Result<Vec<RecordBatch>> {
@@ -248,14 +267,13 @@ async fn simple_udaf() -> Result<()> {
let result = ctx.sql("SELECT MY_AVG(a) FROM t").await?.collect().await?;
- let expected = [
- "+-------------+",
- "| my_avg(t.a) |",
- "+-------------+",
- "| 3.0 |",
- "+-------------+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +-------------+
+ | my_avg(t.a) |
+ +-------------+
+ | 3.0 |
+ +-------------+
+ "###);
Ok(())
}
@@ -315,14 +333,13 @@ async fn
case_sensitive_identifiers_user_defined_aggregates() -> Result<()> {
.collect()
.await?;
- let expected = [
- "+-------------+",
- "| MY_AVG(t.i) |",
- "+-------------+",
- "| 1.0 |",
- "+-------------+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +-------------+
+ | MY_AVG(t.i) |
+ +-------------+
+ | 1.0 |
+ +-------------+
+ "###);
Ok(())
}
@@ -346,19 +363,25 @@ async fn test_user_defined_functions_with_alias() ->
Result<()> {
ctx.register_udaf(my_avg);
- let expected = [
- "+------------+",
- "| dummy(t.i) |",
- "+------------+",
- "| 1.0 |",
- "+------------+",
- ];
-
let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?;
- assert_batches_eq!(expected, &result);
+
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +------------+
+ | dummy(t.i) |
+ +------------+
+ | 1.0 |
+ +------------+
+ "###);
let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM
t").await?;
- assert_batches_eq!(expected, &alias_result);
+
+ insta::assert_snapshot!(batches_to_string(&alias_result), @r###"
+ +------------+
+ | dummy(t.i) |
+ +------------+
+ | 1.0 |
+ +------------+
+ "###);
Ok(())
}
@@ -418,14 +441,14 @@ async fn test_parameterized_aggregate_udf() -> Result<()>
{
);
let actual = DataFrame::new(ctx.state(), plan).collect().await?;
- let expected = [
- "+------+---+---+",
- "| text | a | b |",
- "+------+---+---+",
- "| foo | 1 | 2 |",
- "+------+---+---+",
- ];
- assert_batches_eq!(expected, &actual);
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +------+---+---+
+ | text | a | b |
+ +------+---+---+
+ | foo | 1 | 2 |
+ +------+---+---+
+ "###);
ctx.deregister_table("t")?;
Ok(())
diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs
b/datafusion/core/tests/user_defined/user_defined_plan.rs
index 915d617120..e46940e631 100644
--- a/datafusion/core/tests/user_defined/user_defined_plan.rs
+++ b/datafusion/core/tests/user_defined/user_defined_plan.rs
@@ -155,27 +155,25 @@ const QUERY2: &str = "SELECT 42, arrow_typeof(42)";
// Run the query using the specified execution context and compare it
// to the known result
async fn run_and_compare_query(ctx: SessionContext, description: &str) ->
Result<()> {
- let expected = vec![
- "+-------------+---------+",
- "| customer_id | revenue |",
- "+-------------+---------+",
- "| paul | 300 |",
- "| jorge | 200 |",
- "| andy | 150 |",
- "+-------------+---------+",
- ];
-
let s = exec_sql(&ctx, QUERY).await?;
- let actual = s.lines().collect::<Vec<_>>();
+ let actual = s.lines().collect::<Vec<_>>().join("\n");
+
+ insta::allow_duplicates! {
+ insta::with_settings!({
+ description => description,
+ }, {
+ insta::assert_snapshot!(actual, @r###"
+ +-------------+---------+
+ | customer_id | revenue |
+ +-------------+---------+
+ | paul | 300 |
+ | jorge | 200 |
+ | andy | 150 |
+ +-------------+---------+
+ "###);
+ });
+ }
- assert_eq!(
- expected,
- actual,
- "output mismatch for {}. Expectedn\n{}Actual:\n{}",
- description,
- expected.join("\n"),
- s
- );
Ok(())
}
@@ -185,25 +183,21 @@ async fn run_and_compare_query_with_analyzer_rule(
ctx: SessionContext,
description: &str,
) -> Result<()> {
- let expected = vec![
- "+------------+--------------------------+",
- "| UInt64(42) | arrow_typeof(UInt64(42)) |",
- "+------------+--------------------------+",
- "| 42 | UInt64 |",
- "+------------+--------------------------+",
- ];
-
let s = exec_sql(&ctx, QUERY2).await?;
- let actual = s.lines().collect::<Vec<_>>();
+ let actual = s.lines().collect::<Vec<_>>().join("\n");
+
+ insta::with_settings!({
+ description => description,
+ }, {
+ insta::assert_snapshot!(actual, @r###"
+ +------------+--------------------------+
+ | UInt64(42) | arrow_typeof(UInt64(42)) |
+ +------------+--------------------------+
+ | 42 | UInt64 |
+ +------------+--------------------------+
+ "###);
+ });
- assert_eq!(
- expected,
- actual,
- "output mismatch for {}. Expectedn\n{}Actual:\n{}",
- description,
- expected.join("\n"),
- s
- );
Ok(())
}
@@ -213,27 +207,23 @@ async fn run_and_compare_query_with_auto_schemas(
ctx: SessionContext,
description: &str,
) -> Result<()> {
- let expected = vec![
- "+----------+----------+",
- "| column_1 | column_2 |",
- "+----------+----------+",
- "| andrew | 100 |",
- "| jorge | 200 |",
- "| andy | 150 |",
- "+----------+----------+",
- ];
-
let s = exec_sql(&ctx, QUERY1).await?;
- let actual = s.lines().collect::<Vec<_>>();
+ let actual = s.lines().collect::<Vec<_>>().join("\n");
+
+ insta::with_settings!({
+ description => description,
+ }, {
+ insta::assert_snapshot!(actual, @r###"
+ +----------+----------+
+ | column_1 | column_2 |
+ +----------+----------+
+ | andrew | 100 |
+ | jorge | 200 |
+ | andy | 150 |
+ +----------+----------+
+ "###);
+ });
- assert_eq!(
- expected,
- actual,
- "output mismatch for {}. Expectedn\n{}Actual:\n{}",
- description,
- expected.join("\n"),
- s
- );
Ok(())
}
diff --git
a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
index 68684a81c6..264bd6b66a 100644
--- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs
@@ -26,6 +26,7 @@ use arrow::array::{
};
use arrow::compute::kernels::numeric::add;
use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::test_util::batches_to_string;
use datafusion::execution::context::{FunctionFactory, RegisterFunction,
SessionState};
use datafusion::prelude::*;
use datafusion::{execution::registry::FunctionRegistry, test_util};
@@ -57,14 +58,15 @@ async fn csv_query_custom_udf_with_cast() -> Result<()> {
register_aggregate_csv(&ctx).await?;
let sql = "SELECT avg(custom_sqrt(c11)) FROM aggregate_test_100";
let actual = plan_and_collect(&ctx, sql).await.unwrap();
- let expected = [
- "+------------------------------------------+",
- "| avg(custom_sqrt(aggregate_test_100.c11)) |",
- "+------------------------------------------+",
- "| 0.6584408483418835 |",
- "+------------------------------------------+",
- ];
- assert_batches_eq!(&expected, &actual);
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +------------------------------------------+
+ | avg(custom_sqrt(aggregate_test_100.c11)) |
+ +------------------------------------------+
+ | 0.6584408483418835 |
+ +------------------------------------------+
+ "###);
+
Ok(())
}
@@ -75,14 +77,15 @@ async fn csv_query_avg_sqrt() -> Result<()> {
// Note it is a different column (c12) than above (c11)
let sql = "SELECT avg(custom_sqrt(c12)) FROM aggregate_test_100";
let actual = plan_and_collect(&ctx, sql).await.unwrap();
- let expected = [
- "+------------------------------------------+",
- "| avg(custom_sqrt(aggregate_test_100.c12)) |",
- "+------------------------------------------+",
- "| 0.6706002946036459 |",
- "+------------------------------------------+",
- ];
- assert_batches_eq!(&expected, &actual);
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +------------------------------------------+
+ | avg(custom_sqrt(aggregate_test_100.c12)) |
+ +------------------------------------------+
+ | 0.6706002946036459 |
+ +------------------------------------------+
+ "###);
+
Ok(())
}
@@ -146,17 +149,16 @@ async fn scalar_udf() -> Result<()> {
let result = DataFrame::new(ctx.state(), plan).collect().await?;
- let expected = [
- "+-----+-----+-----------------+",
- "| a | b | my_add(t.a,t.b) |",
- "+-----+-----+-----------------+",
- "| 1 | 2 | 3 |",
- "| 10 | 12 | 22 |",
- "| 10 | 12 | 22 |",
- "| 100 | 120 | 220 |",
- "+-----+-----+-----------------+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +-----+-----+-----------------+
+ | a | b | my_add(t.a,t.b) |
+ +-----+-----+-----------------+
+ | 1 | 2 | 3 |
+ | 10 | 12 | 22 |
+ | 10 | 12 | 22 |
+ | 100 | 120 | 220 |
+ +-----+-----+-----------------+
+ "###);
let batch = &result[0];
let a = as_int32_array(batch.column(0))?;
@@ -272,34 +274,32 @@ async fn scalar_udf_zero_params() -> Result<()> {
ctx.register_udf(ScalarUDF::from(get_100_udf));
let result = plan_and_collect(&ctx, "select get_100() a from t").await?;
- let expected = [
- "+-----+", //
- "| a |", //
- "+-----+", //
- "| 100 |", //
- "| 100 |", //
- "| 100 |", //
- "| 100 |", //
- "+-----+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +-----+
+ | a |
+ +-----+
+ | 100 |
+ | 100 |
+ | 100 |
+ | 100 |
+ +-----+
+ "###);
let result = plan_and_collect(&ctx, "select get_100() a").await?;
- let expected = [
- "+-----+", //
- "| a |", //
- "+-----+", //
- "| 100 |", //
- "+-----+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +-----+
+ | a |
+ +-----+
+ | 100 |
+ +-----+
+ "###);
let result = plan_and_collect(&ctx, "select get_100() from t where
a=999").await?;
- let expected = [
- "++", //
- "++",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ ++
+ ++
+ "###);
+
Ok(())
}
@@ -325,14 +325,14 @@ async fn scalar_udf_override_built_in_scalar_function()
-> Result<()> {
// Make sure that the UDF is used instead of the built-in function
let result = plan_and_collect(&ctx, "select abs(a) a from t").await?;
- let expected = [
- "+---+", //
- "| a |", //
- "+---+", //
- "| 1 |", //
- "+---+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +---+
+ | a |
+ +---+
+ | 1 |
+ +---+
+ "###);
+
Ok(())
}
@@ -427,14 +427,13 @@ async fn
case_sensitive_identifiers_user_defined_functions() -> Result<()> {
// Can call it if you put quotes
let result = plan_and_collect(&ctx, "SELECT \"MY_FUNC\"(i) FROM t").await?;
- let expected = [
- "+--------------+",
- "| MY_FUNC(t.i) |",
- "+--------------+",
- "| 1 |",
- "+--------------+",
- ];
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +--------------+
+ | MY_FUNC(t.i) |
+ +--------------+
+ | 1 |
+ +--------------+
+ "###);
Ok(())
}
@@ -464,18 +463,23 @@ async fn test_user_defined_functions_with_alias() ->
Result<()> {
ctx.register_udf(udf);
- let expected = [
- "+------------+",
- "| dummy(t.i) |",
- "+------------+",
- "| 1 |",
- "+------------+",
- ];
let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?;
- assert_batches_eq!(expected, &result);
+ insta::assert_snapshot!(batches_to_string(&result), @r###"
+ +------------+
+ | dummy(t.i) |
+ +------------+
+ | 1 |
+ +------------+
+ "###);
let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM
t").await?;
- assert_batches_eq!(expected, &alias_result);
+ insta::assert_snapshot!(batches_to_string(&alias_result), @r###"
+ +------------+
+ | dummy(t.i) |
+ +------------+
+ | 1 |
+ +------------+
+ "###);
Ok(())
}
diff --git a/datafusion/core/tests/user_defined/user_defined_table_functions.rs
b/datafusion/core/tests/user_defined/user_defined_table_functions.rs
index 618f0590ab..e4aff0b007 100644
--- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs
@@ -26,6 +26,7 @@ use arrow::csv::ReaderBuilder;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::test_util::batches_to_string;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::TableProvider;
use datafusion::error::Result;
@@ -34,7 +35,7 @@ use datafusion::physical_plan::{collect, ExecutionPlan};
use datafusion::prelude::SessionContext;
use datafusion_catalog::Session;
use datafusion_catalog::TableFunctionImpl;
-use datafusion_common::{assert_batches_eq, DFSchema, ScalarValue};
+use datafusion_common::{DFSchema, ScalarValue};
use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType};
use async_trait::async_trait;
@@ -54,17 +55,17 @@ async fn test_simple_read_csv_udtf() -> Result<()> {
.collect()
.await?;
- let excepted = [
-
"+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+",
- "| n_nationkey | n_name | n_regionkey | n_comment
|",
-
"+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+",
- "| 1 | ARGENTINA | 1 | al foxes promise slyly
according to the regular accounts. bold requests alon
|",
- "| 2 | BRAZIL | 1 | y alongside of the pending
deposits. carefully special packages are about the ironic forges. slyly special
|",
- "| 3 | CANADA | 1 | eas hang ironic, silent
packages. slyly regular packages are furiously over the tithes. fluffily bold
|",
- "| 4 | EGYPT | 4 | y above the carefully
unusual theodolites. final dugouts are quickly across the furiously regular d
|",
- "| 5 | ETHIOPIA | 0 | ven packages wake quickly.
regu
|",
-
"+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+",
];
- assert_batches_eq!(excepted, &rbs);
+ insta::assert_snapshot!(batches_to_string(&rbs), @r###"
+
+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+
+ | n_nationkey | n_name | n_regionkey | n_comment
|
+
+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+
+ | 1 | ARGENTINA | 1 | al foxes promise slyly according
to the regular accounts. bold requests alon |
+ | 2 | BRAZIL | 1 | y alongside of the pending
deposits. carefully special packages are about the ironic forges. slyly special
|
+ | 3 | CANADA | 1 | eas hang ironic, silent
packages. slyly regular packages are furiously over the tithes. fluffily bold
|
+ | 4 | EGYPT | 4 | y above the carefully unusual
theodolites. final dugouts are quickly across the furiously regular d |
+ | 5 | ETHIOPIA | 0 | ven packages wake quickly. regu
|
+
+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+
+ "###);
// just run, return all rows
let rbs = ctx
@@ -72,23 +73,23 @@ async fn test_simple_read_csv_udtf() -> Result<()> {
.await?
.collect()
.await?;
- let excepted = [
-
"+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+",
- "| n_nationkey | n_name | n_regionkey | n_comment
|",
-
"+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+",
- "| 1 | ARGENTINA | 1 | al foxes promise slyly
according to the regular accounts. bold requests alon
|",
- "| 2 | BRAZIL | 1 | y alongside of the pending
deposits. carefully special packages are about the ironic forges. slyly special
|",
- "| 3 | CANADA | 1 | eas hang ironic, silent
packages. slyly regular packages are furiously over the tithes. fluffily bold
|",
- "| 4 | EGYPT | 4 | y above the carefully
unusual theodolites. final dugouts are quickly across the furiously regular d
|",
- "| 5 | ETHIOPIA | 0 | ven packages wake quickly.
regu
|",
- "| 6 | FRANCE | 3 | refully final requests.
regular, ironi
|",
- "| 7 | GERMANY | 3 | l platelets. regular
accounts x-ray: unusual, regular acco
|",
- "| 8 | INDIA | 2 | ss excuses cajole slyly
across the packages. deposits print aroun
|",
- "| 9 | INDONESIA | 2 | slyly express asymptotes.
regular deposits haggle slyly. carefully ironic hockey players sleep blithely.
carefull |",
- "| 10 | IRAN | 4 | efully alongside of the
slyly final dependencies.
|",
-
"+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+"
- ];
- assert_batches_eq!(excepted, &rbs);
+
+ insta::assert_snapshot!(batches_to_string(&rbs), @r###"
+
+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+
+ | n_nationkey | n_name | n_regionkey | n_comment
|
+
+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+
+ | 1 | ARGENTINA | 1 | al foxes promise slyly according
to the regular accounts. bold requests alon
|
+ | 2 | BRAZIL | 1 | y alongside of the pending
deposits. carefully special packages are about the ironic forges. slyly special
|
+ | 3 | CANADA | 1 | eas hang ironic, silent
packages. slyly regular packages are furiously over the tithes. fluffily bold
|
+ | 4 | EGYPT | 4 | y above the carefully unusual
theodolites. final dugouts are quickly across the furiously regular d
|
+ | 5 | ETHIOPIA | 0 | ven packages wake quickly. regu
|
+ | 6 | FRANCE | 3 | refully final requests. regular,
ironi
|
+ | 7 | GERMANY | 3 | l platelets. regular accounts
x-ray: unusual, regular acco
|
+ | 8 | INDIA | 2 | ss excuses cajole slyly across
the packages. deposits print aroun
|
+ | 9 | INDONESIA | 2 | slyly express asymptotes.
regular deposits haggle slyly. carefully ironic hockey players sleep blithely.
carefull |
+ | 10 | IRAN | 4 | efully alongside of the slyly
final dependencies.
|
+
+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+
+ "###);
Ok(())
}
diff --git
a/datafusion/core/tests/user_defined/user_defined_window_functions.rs
b/datafusion/core/tests/user_defined/user_defined_window_functions.rs
index 9acd17493d..28394f0b9d 100644
--- a/datafusion/core/tests/user_defined/user_defined_window_functions.rs
+++ b/datafusion/core/tests/user_defined/user_defined_window_functions.rs
@@ -20,8 +20,9 @@
use arrow::array::{ArrayRef, AsArray, Int64Array, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
-use datafusion::{assert_batches_eq, prelude::SessionContext};
-use datafusion_common::{Result, ScalarValue};
+use datafusion::common::test_util::batches_to_string;
+use datafusion::common::{Result, ScalarValue};
+use datafusion::prelude::SessionContext;
use datafusion_expr::{
PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF,
WindowUDFImpl,
};
@@ -57,30 +58,30 @@ const BOUNDED_WINDOW_QUERY: &str =
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());
+ let actual = execute(&ctx, sql).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+ +---+---+-----+
+ | 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 |
+ +---+---+-----+
+ "###);
}
/// Basic user defined window function
@@ -89,26 +90,25 @@ 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) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |",
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| 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()
- );
+ let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | 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
|
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ "###);
+
// evaluated on two distinct batches
assert_eq!(test_state.evaluate_all_called(), 2);
}
@@ -133,28 +133,26 @@ async fn test_udwf_with_alias() {
let test_state = TestState::new();
let TestContext { ctx, .. } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |",
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| 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_WITH_ALIAS)
- .await
- .unwrap()
- );
+ let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY_WITH_ALIAS)
+ .await
+ .unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | 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
|
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ "###);
}
/// Basic user defined window function with bounded window
@@ -164,26 +162,25 @@ async fn test_udwf_bounded_window_ignores_frame() {
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) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| 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()
- );
+ let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | 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
|
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ "###);
+
// 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);
@@ -195,26 +192,25 @@ 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) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| 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()
- );
+ let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | 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
|
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ "###);
+
// Evaluate is called for each input rows
assert_eq!(test_state.evaluate_called(), 10);
assert_eq!(test_state.evaluate_all_called(), 0);
@@ -228,26 +224,25 @@ async fn test_stateful_udwf() {
.with_uses_window_frame();
let TestContext { ctx, test_state } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |",
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| 1 | a | 0 | 0
|",
- "| 1 | b | 1 | 1
|",
- "| 1 | c | 2 | 1
|",
- "| 2 | d | 3 | 1
|",
- "| 2 | e | 4 | 1
|",
- "| 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()
- );
+ let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | 1 | a | 0 | 0
|
+ | 1 | b | 1 | 1
|
+ | 1 | c | 2 | 1
|
+ | 2 | d | 3 | 1
|
+ | 2 | e | 4 | 1
|
+ | 2 | f | 5 | 2
|
+ | 2 | g | 6 | 2
|
+ | 2 | h | 6 | 2
|
+ | 2 | i | 6 | 2
|
+ | 2 | j | 6 | 2
|
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ "###);
+
assert_eq!(test_state.evaluate_called(), 10);
assert_eq!(test_state.evaluate_all_called(), 0);
}
@@ -260,26 +255,25 @@ async fn test_stateful_udwf_bounded_window() {
.with_uses_window_frame();
let TestContext { ctx, test_state } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| 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()
- );
+ let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | 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
|
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ "###);
+
// Evaluate and update_state is called for each input row
assert_eq!(test_state.evaluate_called(), 10);
assert_eq!(test_state.evaluate_all_called(), 0);
@@ -291,26 +285,25 @@ async fn test_udwf_query_include_rank() {
let test_state = TestState::new().with_include_rank();
let TestContext { ctx, test_state } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |",
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- "| 1 | a | 0 | 3
|",
- "| 1 | b | 1 | 2
|",
- "| 1 | c | 2 | 1
|",
- "| 2 | d | 3 | 7
|",
- "| 2 | e | 4 | 6
|",
- "| 2 | f | 5 | 5
|",
- "| 2 | g | 6 | 4
|",
- "| 2 | h | 6 | 3
|",
- "| 2 | i | 6 | 2
|",
- "| 2 | j | 6 | 1
|",
-
"+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+",
- ];
- assert_batches_eq!(
- expected,
- &execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap()
- );
+ let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW |
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ | 1 | a | 0 | 3
|
+ | 1 | b | 1 | 2
|
+ | 1 | c | 2 | 1
|
+ | 2 | d | 3 | 7
|
+ | 2 | e | 4 | 6
|
+ | 2 | f | 5 | 5
|
+ | 2 | g | 6 | 4
|
+ | 2 | h | 6 | 3
|
+ | 2 | i | 6 | 2
|
+ | 2 | j | 6 | 1
|
+
+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+
+ "###);
+
assert_eq!(test_state.evaluate_called(), 0);
assert_eq!(test_state.evaluate_all_called(), 0);
// evaluated on 2 distinct batches (when x=1 and x=2)
@@ -323,26 +316,25 @@ async fn test_udwf_bounded_query_include_rank() {
let test_state = TestState::new().with_include_rank();
let TestContext { ctx, test_state } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| 1 | a | 0 | 3
|",
- "| 1 | b | 1 | 2
|",
- "| 1 | c | 2 | 1
|",
- "| 2 | d | 3 | 7
|",
- "| 2 | e | 4 | 6
|",
- "| 2 | f | 5 | 5
|",
- "| 2 | g | 6 | 4
|",
- "| 2 | h | 6 | 3
|",
- "| 2 | i | 6 | 2
|",
- "| 2 | j | 6 | 1
|",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- ];
- assert_batches_eq!(
- expected,
- &execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap()
- );
+ let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | 1 | a | 0 | 3
|
+ | 1 | b | 1 | 2
|
+ | 1 | c | 2 | 1
|
+ | 2 | d | 3 | 7
|
+ | 2 | e | 4 | 6
|
+ | 2 | f | 5 | 5
|
+ | 2 | g | 6 | 4
|
+ | 2 | h | 6 | 3
|
+ | 2 | i | 6 | 2
|
+ | 2 | j | 6 | 1
|
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ "###);
+
assert_eq!(test_state.evaluate_called(), 0);
assert_eq!(test_state.evaluate_all_called(), 0);
// evaluated on 2 distinct batches (when x=1 and x=2)
@@ -357,26 +349,25 @@ async fn test_udwf_bounded_window_returns_null() {
.with_null_for_zero();
let TestContext { ctx, test_state } = TestContext::new(test_state);
- let expected = vec![
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y ASC
NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- "| 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 |
|",
- "| 2 | i | 6 |
|",
- "| 2 | j | 6 |
|",
-
"+---+---+-----+--------------------------------------------------------------------------------------------------------------+",
- ];
- assert_batches_eq!(
- expected,
- &execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap()
- );
+ let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap();
+
+ insta::assert_snapshot!(batches_to_string(&actual), @r###"
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | x | y | val | odd_counter(t.val) PARTITION BY [t.x] ORDER BY [t.y
ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING |
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ | 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 |
|
+ | 2 | i | 6 |
|
+ | 2 | j | 6 |
|
+
+---+---+-----+--------------------------------------------------------------------------------------------------------------+
+ "###);
+
// Evaluate is called for each input rows
assert_eq!(test_state.evaluate_called(), 10);
assert_eq!(test_state.evaluate_all_called(), 0);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
