diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index 75d890359ba8..1fc6d14c5b22 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> { #[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 7cda6d410f4e..5cbb05f290a7 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> { @@ -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 915d61712074..e46940e63154 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::>(); + let actual = s.lines().collect::>().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::>(); + let actual = s.lines().collect::>().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::>(); + let actual = s.lines().collect::>().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 68684a81c656..264bd6b66a60 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 618f0590ab3d..e4aff0b00705 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 9acd17493da4..28394f0b9dfa 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);