From 1d7de54849bdc563f6b55813c2cb470fe99de65c Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 02:49:46 +0530 Subject: [PATCH 01/11] migrate user_defined tests to insta --- Cargo.lock | 1 + datafusion/core/Cargo.toml | 1 + .../core/tests/user_defined/expr_planner.rs | 64 +-- .../tests/user_defined/insert_operation.rs | 39 +- .../user_defined/user_defined_aggregates.rs | 256 ++++++----- .../tests/user_defined/user_defined_plan.rs | 125 +++--- .../user_defined_scalar_functions.rs | 163 +++---- .../user_defined_table_functions.rs | 66 +-- .../user_defined_window_functions.rs | 400 +++++++++--------- 9 files changed, 582 insertions(+), 533 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e772b2466c341..c7183557c9b7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1779,6 +1779,7 @@ dependencies = [ "env_logger", "flate2", "futures", + "insta", "itertools 0.14.0", "log", "nix", diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index fd1fd4164da0d..438c00b6b8faf 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -152,6 +152,7 @@ datafusion-functions-window-common = { workspace = true } datafusion-physical-optimizer = { workspace = true } doc-comment = { workspace = true } env_logger = { workspace = true } +insta = "1.42.2" paste = "^1.0" rand = { workspace = true, features = ["small_rng"] } rand_distr = "0.4.3" diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index 75d890359ba82..8d5e363213d35 100644 --- a/datafusion/core/tests/user_defined/expr_planner.rs +++ b/datafusion/core/tests/user_defined/expr_planner.rs @@ -18,7 +18,7 @@ use arrow::array::RecordBatch; 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; @@ -73,30 +73,36 @@ async fn plan_and_collect(sql: &str) -> Result> { ctx.sql(sql).await?.collect().await } +fn fmt_batches(batches: &[RecordBatch]) -> String { + use arrow::util::pretty::pretty_format_batches; + match pretty_format_batches(batches) { + Ok(formatted) => formatted.to_string(), + Err(e) => format!("Error formatting record batches: {}", e), + } +} + #[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!(fmt_batches(&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!(fmt_batches(&actual), @r###" + +---------------------+ + | Int64(1) + Int64(2) | + +---------------------+ + | 3 | + +---------------------+ + "###); } #[tokio::test] @@ -104,14 +110,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!(fmt_batches(&actual), @r###" + +--------------+ + | a ? Int64(2) | + +--------------+ + | true | + +--------------+ + "###); } #[tokio::test] @@ -119,6 +124,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!(fmt_batches(&actual), @r###" + +---+ + | a | + +---+ + | 1 | + +---+ + "###); } diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 12f700ce572ba..839c78c990559 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -30,6 +30,7 @@ use datafusion_physical_plan::{ execution_plan::{Boundedness, EmissionType}, DisplayAs, ExecutionPlan, PlanProperties, }; +use insta::assert_debug_snapshot; #[tokio::test] async fn insert_operation_is_passed_correctly_to_table_provider() { @@ -39,24 +40,26 @@ async fn insert_operation_is_passed_correctly_to_table_provider() { ctx.register_table("testing", table_provider.clone()) .unwrap(); - let sql = "INSERT INTO testing (column) VALUES (1)"; - assert_insert_op(&ctx, sql, InsertOp::Append).await; - - let sql = "INSERT OVERWRITE testing (column) VALUES (1)"; - assert_insert_op(&ctx, sql, InsertOp::Overwrite).await; - - let sql = "REPLACE INTO testing (column) VALUES (1)"; - assert_insert_op(&ctx, sql, InsertOp::Replace).await; - - let sql = "INSERT OR REPLACE INTO testing (column) VALUES (1)"; - assert_insert_op(&ctx, sql, InsertOp::Replace).await; -} - -async fn assert_insert_op(ctx: &SessionContext, sql: &str, insert_op: InsertOp) { - let df = ctx.sql(sql).await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - let exec = plan.as_any().downcast_ref::().unwrap(); - assert_eq!(exec.op, insert_op); + let test_cases = vec![ + ("INSERT INTO testing (column) VALUES (1)", "insert_append"), + ( + "INSERT OVERWRITE testing (column) VALUES (1)", + "insert_overwrite", + ), + ("REPLACE INTO testing (column) VALUES (1)", "insert_replace"), + ( + "INSERT OR REPLACE INTO testing (column) VALUES (1)", + "insert_or_replace", + ), + ]; + + for (sql, snapshot_name) in test_cases { + let df = ctx.sql(sql).await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let exec = plan.as_any().downcast_ref::().unwrap(); + + assert_debug_snapshot!(snapshot_name, exec.op); + } } fn session_ctx_with_dialect(dialect: impl Into) -> SessionContext { diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index 7cda6d410f4e4..d45df33de9300 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -39,7 +39,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,30 +47,41 @@ 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, }; use datafusion_functions_aggregate::average::AvgAccumulator; +fn fmt_batches(batches: &[RecordBatch]) -> String { + use arrow::util::pretty::pretty_format_batches; + match pretty_format_batches(batches) { + Ok(formatted) => formatted.to_string(), + Err(e) => format!("Error formatting record batches: {}", e), + } +} + /// Test to show the contents of the setup #[tokio::test] 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!(fmt_batches(&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 +90,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!(fmt_batches(&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 +111,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!(fmt_batches(&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 +136,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!(fmt_batches(&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 +174,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!(fmt_batches(&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 +191,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!(fmt_batches(&actual), @r###" + +-----------------+----------------------------+ + | sq.first[value] | sq.first[time] | + +-----------------+----------------------------+ + | 2.0 | 1970-01-01T00:00:00.000002 | + +-----------------+----------------------------+ + "###); } #[tokio::test] @@ -189,26 +212,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!(fmt_batches(&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!(fmt_batches(&actual), @r###" + +----------------------------+ + | sum(t.time) | + +----------------------------+ + | 1970-01-01T00:00:00.000019 | + +----------------------------+ + "###); } async fn execute(ctx: &SessionContext, sql: &str) -> Result> { @@ -248,14 +274,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!(fmt_batches(&result), @r###" + +-------------+ + | my_avg(t.a) | + +-------------+ + | 3.0 | + +-------------+ + "###); Ok(()) } @@ -315,14 +340,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!(fmt_batches(&result), @r###" + +-------------+ + | MY_AVG(t.i) | + +-------------+ + | 1.0 | + +-------------+ + "###); Ok(()) } @@ -346,19 +370,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!(fmt_batches(&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!(fmt_batches(&alias_result), @r###"" + +------------+ + | dummy(t.i) | + +------------+ + | 1.0 | + +------------+ + "###); Ok(()) } @@ -418,14 +448,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!(fmt_batches(&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 915d617120744..1d99e6dcd91cf 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -120,12 +120,13 @@ async fn setup_table(ctx: SessionContext) -> Result { OPTIONS('format.has_header' 'false') "; - let expected = vec!["++", "++"]; + let result = exec_sql(&ctx, sql).await?; - let s = exec_sql(&ctx, sql).await?; - let actual = s.lines().collect::>(); + insta::assert_snapshot!(result, @r###" + ++ + ++ + "###); - assert_eq!(expected, actual, "Creating table"); Ok(ctx) } @@ -136,12 +137,13 @@ async fn setup_table_without_schemas(ctx: SessionContext) -> Result>(); + insta::assert_snapshot!(result, @r###" + ++ + ++ + "###); - assert_eq!(expected, actual, "Creating table"); Ok(ctx) } @@ -155,27 +157,22 @@ 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::>(); - - assert_eq!( - expected, - actual, - "output mismatch for {}. Expectedn\n{}Actual:\n{}", - description, - expected.join("\n"), - s - ); + let result = exec_sql(&ctx, QUERY).await?; + + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(result, @r###" + +-------------+---------+ + | customer_id | revenue | + +-------------+---------+ + | paul | 300 | + | jorge | 200 | + | andy | 150 | + +-------------+---------+ + "###); + }); + Ok(()) } @@ -185,25 +182,20 @@ 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::>(); - - assert_eq!( - expected, - actual, - "output mismatch for {}. Expectedn\n{}Actual:\n{}", - description, - expected.join("\n"), - s - ); + let result = exec_sql(&ctx, QUERY2).await?; + + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(result, @r###" + +------------+--------------------------+ + | UInt64(42) | arrow_typeof(UInt64(42)) | + +------------+--------------------------+ + | 42 | UInt64 | + +------------+--------------------------+ + "###); + }); + Ok(()) } @@ -213,27 +205,22 @@ 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::>(); - - assert_eq!( - expected, - actual, - "output mismatch for {}. Expectedn\n{}Actual:\n{}", - description, - expected.join("\n"), - s - ); + let result = exec_sql(&ctx, QUERY1).await?; + + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(result, @r###" + +----------+----------+ + | column_1 | column_2 | + +----------+----------+ + | andrew | 100 | + | jorge | 200 | + | andy | 150 | + +----------+----------+ + "###); + }); + 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 87694d83bb99d..0aa356a77a4ba 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -48,6 +48,14 @@ use regex::Regex; use sqlparser::ast::Ident; use sqlparser::tokenizer::Span; +fn fmt_batches(batches: &[RecordBatch]) -> String { + use arrow::util::pretty::pretty_format_batches; + match pretty_format_batches(batches) { + Ok(formatted) => formatted.to_string(), + Err(e) => format!("Error formatting record batches: {}", e), + } +} + /// test that casting happens on udfs. /// c11 is f32, but `custom_sqrt` requires f64. Casting happens but the logical plan and /// physical plan have the same schema. @@ -57,14 +65,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!(fmt_batches(&actual), @r###" + +------------------------------------------+ + | avg(custom_sqrt(aggregate_test_100.c11)) | + +------------------------------------------+ + | 0.6584408483418835 | + +------------------------------------------+ + "###); + Ok(()) } @@ -75,14 +84,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!(fmt_batches(&actual), @r###" + +------------------------------------------+ + | avg(custom_sqrt(aggregate_test_100.c12)) | + +------------------------------------------+ + | 0.6706002946036459 | + +------------------------------------------+ + "###); + Ok(()) } @@ -146,17 +156,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!(fmt_batches(&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 +281,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!(fmt_batches(&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!(fmt_batches(&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!(fmt_batches(&result), @r###" + ++ + ++ + "###); + Ok(()) } @@ -325,14 +332,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!(fmt_batches(&result), @r###" + +---+ + | a | + +---+ + | 1 | + +---+ + "###); + Ok(()) } @@ -427,14 +434,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!(fmt_batches(&result), @r###" + +--------------+ + | MY_FUNC(t.i) | + +--------------+ + | 1 | + +--------------+ + "###); Ok(()) } @@ -464,18 +470,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!(fmt_batches(&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!(fmt_batches(&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 618f0590ab3d6..08aeee8a3275a 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -34,11 +34,19 @@ 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; +fn fmt_batches(batches: &[RecordBatch]) -> String { + use arrow::util::pretty::pretty_format_batches; + match pretty_format_batches(batches) { + Ok(formatted) => formatted.to_string(), + Err(e) => format!("Error formatting record batches: {}", e), + } +} + /// test simple udtf with define read_csv with parameters #[tokio::test] async fn test_simple_read_csv_udtf() -> Result<()> { @@ -54,17 +62,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!(fmt_batches(&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 +80,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!(fmt_batches(&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 9acd17493da43..bd5bbd302f720 100644 --- a/datafusion/core/tests/user_defined/user_defined_window_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_window_functions.rs @@ -20,7 +20,7 @@ use arrow::array::{ArrayRef, AsArray, Int64Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::{assert_batches_eq, prelude::SessionContext}; +use datafusion::prelude::SessionContext; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF, WindowUDFImpl, @@ -57,30 +57,38 @@ 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 +fn fmt_batches(batches: &[RecordBatch]) -> String { + use arrow::util::pretty::pretty_format_batches; + match pretty_format_batches(batches) { + Ok(formatted) => formatted.to_string(), + Err(e) => format!("Error formatting record batches: {}", e), + } +} + #[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!(fmt_batches(&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 +97,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!(fmt_batches(&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 +140,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!(fmt_batches(&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 +169,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!(fmt_batches(&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 +199,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!(fmt_batches(&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 +231,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!(fmt_batches(&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 +262,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!(fmt_batches(&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 +292,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!(fmt_batches(&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 +323,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!(fmt_batches(&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 +356,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!(fmt_batches(&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); From b0d18df36ba2150dab998a20c6ba8751c9096c2e Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 04:18:07 +0530 Subject: [PATCH 02/11] fix tests --- .../tests/user_defined/insert_operation.rs | 39 +++++++++---------- .../user_defined/user_defined_aggregates.rs | 2 +- .../tests/user_defined/user_defined_plan.rs | 33 ++++++++-------- 3 files changed, 36 insertions(+), 38 deletions(-) diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 839c78c990559..12f700ce572ba 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -30,7 +30,6 @@ use datafusion_physical_plan::{ execution_plan::{Boundedness, EmissionType}, DisplayAs, ExecutionPlan, PlanProperties, }; -use insta::assert_debug_snapshot; #[tokio::test] async fn insert_operation_is_passed_correctly_to_table_provider() { @@ -40,26 +39,24 @@ async fn insert_operation_is_passed_correctly_to_table_provider() { ctx.register_table("testing", table_provider.clone()) .unwrap(); - let test_cases = vec![ - ("INSERT INTO testing (column) VALUES (1)", "insert_append"), - ( - "INSERT OVERWRITE testing (column) VALUES (1)", - "insert_overwrite", - ), - ("REPLACE INTO testing (column) VALUES (1)", "insert_replace"), - ( - "INSERT OR REPLACE INTO testing (column) VALUES (1)", - "insert_or_replace", - ), - ]; - - for (sql, snapshot_name) in test_cases { - let df = ctx.sql(sql).await.unwrap(); - let plan = df.create_physical_plan().await.unwrap(); - let exec = plan.as_any().downcast_ref::().unwrap(); - - assert_debug_snapshot!(snapshot_name, exec.op); - } + let sql = "INSERT INTO testing (column) VALUES (1)"; + assert_insert_op(&ctx, sql, InsertOp::Append).await; + + let sql = "INSERT OVERWRITE testing (column) VALUES (1)"; + assert_insert_op(&ctx, sql, InsertOp::Overwrite).await; + + let sql = "REPLACE INTO testing (column) VALUES (1)"; + assert_insert_op(&ctx, sql, InsertOp::Replace).await; + + let sql = "INSERT OR REPLACE INTO testing (column) VALUES (1)"; + assert_insert_op(&ctx, sql, InsertOp::Replace).await; +} + +async fn assert_insert_op(ctx: &SessionContext, sql: &str, insert_op: InsertOp) { + let df = ctx.sql(sql).await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let exec = plan.as_any().downcast_ref::().unwrap(); + assert_eq!(exec.op, insert_op); } fn session_ctx_with_dialect(dialect: impl Into) -> SessionContext { diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index d45df33de9300..7392f3d4479d8 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -382,7 +382,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result), @r###"" + insta::assert_snapshot!(fmt_batches(&alias_result), @r###" +------------+ | dummy(t.i) | +------------+ diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 1d99e6dcd91cf..200ca2b7006c2 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -120,13 +120,12 @@ async fn setup_table(ctx: SessionContext) -> Result { OPTIONS('format.has_header' 'false') "; - let result = exec_sql(&ctx, sql).await?; + let expected = vec!["++", "++"]; - insta::assert_snapshot!(result, @r###" - ++ - ++ - "###); + let s = exec_sql(&ctx, sql).await?; + let actual = s.lines().collect::>(); + assert_eq!(expected, actual, "Creating table"); Ok(ctx) } @@ -137,13 +136,12 @@ async fn setup_table_without_schemas(ctx: SessionContext) -> Result>(); + assert_eq!(expected, actual, "Creating table"); Ok(ctx) } @@ -157,12 +155,13 @@ 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 result = exec_sql(&ctx, QUERY).await?; + let s = exec_sql(&ctx, QUERY).await?; + let actual = s.lines().collect::>().join("\n"); insta::with_settings!({ description => description, }, { - insta::assert_snapshot!(result, @r###" + insta::assert_snapshot!(actual, @r###" +-------------+---------+ | customer_id | revenue | +-------------+---------+ @@ -182,12 +181,13 @@ async fn run_and_compare_query_with_analyzer_rule( ctx: SessionContext, description: &str, ) -> Result<()> { - let result = exec_sql(&ctx, QUERY2).await?; + let s = exec_sql(&ctx, QUERY2).await?; + let actual = s.lines().collect::>().join("\n"); insta::with_settings!({ description => description, }, { - insta::assert_snapshot!(result, @r###" + insta::assert_snapshot!(actual, @r###" +------------+--------------------------+ | UInt64(42) | arrow_typeof(UInt64(42)) | +------------+--------------------------+ @@ -205,12 +205,13 @@ async fn run_and_compare_query_with_auto_schemas( ctx: SessionContext, description: &str, ) -> Result<()> { - let result = exec_sql(&ctx, QUERY1).await?; + let s = exec_sql(&ctx, QUERY1).await?; + let actual = s.lines().collect::>().join("\n"); insta::with_settings!({ description => description, }, { - insta::assert_snapshot!(result, @r###" + insta::assert_snapshot!(actual, @r###" +----------+----------+ | column_1 | column_2 | +----------+----------+ From 46bb92f946fa89024ee4ebc6a4bcda52ad4402c0 Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 05:57:58 +0530 Subject: [PATCH 03/11] test --- .../tests/user_defined/user_defined_plan.rs | 98 ++++++++++--------- 1 file changed, 52 insertions(+), 46 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 200ca2b7006c2..96315a848babf 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -152,30 +152,32 @@ const QUERY1: &str = "SELECT * FROM sales limit 3"; const QUERY2: &str = "SELECT 42, arrow_typeof(42)"; -// Run the query using the specified execution context and compare it +// 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 s = exec_sql(&ctx, QUERY).await?; let actual = s.lines().collect::>().join("\n"); - - insta::with_settings!({ - description => description, - }, { - insta::assert_snapshot!(actual, @r###" - +-------------+---------+ - | customer_id | revenue | - +-------------+---------+ - | paul | 300 | - | jorge | 200 | - | andy | 150 | - +-------------+---------+ - "###); - }); - + + insta::allow_duplicates!(|| { + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(actual, @r###" + +-------------+---------+ + | customer_id | revenue | + +-------------+---------+ + | paul | 300 | + | jorge | 200 | + | andy | 150 | + +-------------+---------+ + "###); + }); + })(); + Ok(()) } -// Run the query using the specified execution context and compare it +// Run the query using the specified execution context and compare it // to the known result async fn run_and_compare_query_with_analyzer_rule( ctx: SessionContext, @@ -183,23 +185,25 @@ async fn run_and_compare_query_with_analyzer_rule( ) -> Result<()> { let s = exec_sql(&ctx, QUERY2).await?; 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 | - +------------+--------------------------+ - "###); - }); - + + insta::allow_duplicates!(|| { + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(actual, @r###" + +------------+--------------------------+ + | UInt64(42) | arrow_typeof(UInt64(42)) | + +------------+--------------------------+ + | 42 | UInt64 | + +------------+--------------------------+ + "###); + }); + })(); + Ok(()) } -// Run the query using the specified execution context and compare it +// Run the query using the specified execution context and compare it // to the known result async fn run_and_compare_query_with_auto_schemas( ctx: SessionContext, @@ -207,21 +211,23 @@ async fn run_and_compare_query_with_auto_schemas( ) -> Result<()> { let s = exec_sql(&ctx, QUERY1).await?; 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 | - +----------+----------+ - "###); - }); - + + insta::allow_duplicates!(|| { + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(actual, @r###" + +----------+----------+ + | column_1 | column_2 | + +----------+----------+ + | andrew | 100 | + | jorge | 200 | + | andy | 150 | + +----------+----------+ + "###); + }); + })(); + Ok(()) } From 40b9a12e28aef455508a63ff3ffa88d2f677237b Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 05:58:58 +0530 Subject: [PATCH 04/11] fmt --- .../tests/user_defined/user_defined_plan.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 96315a848babf..9a604e2bf881b 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -152,12 +152,12 @@ const QUERY1: &str = "SELECT * FROM sales limit 3"; const QUERY2: &str = "SELECT 42, arrow_typeof(42)"; -// Run the query using the specified execution context and compare it +// 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 s = exec_sql(&ctx, QUERY).await?; let actual = s.lines().collect::>().join("\n"); - + insta::allow_duplicates!(|| { insta::with_settings!({ description => description, @@ -173,11 +173,11 @@ async fn run_and_compare_query(ctx: SessionContext, description: &str) -> Result "###); }); })(); - + Ok(()) } -// Run the query using the specified execution context and compare it +// Run the query using the specified execution context and compare it // to the known result async fn run_and_compare_query_with_analyzer_rule( ctx: SessionContext, @@ -185,7 +185,7 @@ async fn run_and_compare_query_with_analyzer_rule( ) -> Result<()> { let s = exec_sql(&ctx, QUERY2).await?; let actual = s.lines().collect::>().join("\n"); - + insta::allow_duplicates!(|| { insta::with_settings!({ description => description, @@ -199,11 +199,11 @@ async fn run_and_compare_query_with_analyzer_rule( "###); }); })(); - + Ok(()) } -// Run the query using the specified execution context and compare it +// Run the query using the specified execution context and compare it // to the known result async fn run_and_compare_query_with_auto_schemas( ctx: SessionContext, @@ -211,7 +211,7 @@ async fn run_and_compare_query_with_auto_schemas( ) -> Result<()> { let s = exec_sql(&ctx, QUERY1).await?; let actual = s.lines().collect::>().join("\n"); - + insta::allow_duplicates!(|| { insta::with_settings!({ description => description, @@ -227,7 +227,7 @@ async fn run_and_compare_query_with_auto_schemas( "###); }); })(); - + Ok(()) } From 1e4d20fe44ff2153818f271f52600c40d10caca3 Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 07:59:09 +0530 Subject: [PATCH 05/11] update snaps --- .../core/tests/user_defined/expr_planner.rs | 32 +--- ..._expr_planner__custom_operators_arrow.snap | 9 ++ ..._planner__custom_operators_long_arrow.snap | 9 ++ ...efined__expr_planner__question_filter.snap | 9 ++ ...efined__expr_planner__question_select.snap | 9 ++ ...e_identifiers_user_defined_aggregates.snap | 9 ++ ...gregates__parameterized_aggregate_udf.snap | 9 ++ ...fined__user_defined_aggregates__setup.snap | 13 ++ ..._user_defined_aggregates__simple_udaf.snap | 9 ++ ...efined__user_defined_aggregates__udaf.snap | 9 ++ ...er_defined_aggregates__udaf_as_window.snap | 13 ++ ...aggregates__udaf_as_window_with_frame.snap | 13 ++ ...ned_aggregates__udaf_returning_struct.snap | 9 ++ ...gates__udaf_returning_struct_subquery.snap | 9 ++ ...d_aggregates__udaf_shadows_builtin_fn.snap | 9 ++ ...es__user_defined_functions_with_alias.snap | 9 ++ ...r_defined_plan__query_Default_context.snap | 12 ++ ...query_Default_context_without_schemas.snap | 12 ++ ...er_defined_plan__query_MyAnalyzerRule.snap | 10 ++ ...user_defined_plan__query_Topk_context.snap | 12 ++ ...ve_identifiers_user_defined_functions.snap | 9 ++ ..._defined_scalar_functions__scalar_udf.snap | 12 ++ ...udf_override_built_in_scalar_function.snap | 9 ++ ...lar_functions__scalar_udf_zero_params.snap | 12 ++ ...ns__user_defined_functions_with_alias.snap | 9 ++ ...ble_functions__simple_read_csv_udtf-2.snap | 18 +++ ...table_functions__simple_read_csv_udtf.snap | 13 ++ ...fined_window_functions__stateful_udwf.snap | 18 +++ ...nctions__stateful_udwf_bounded_window.snap | 18 +++ ...__user_defined_window_functions__udwf.snap | 18 +++ ...ions__udwf_bounded_query_include_rank.snap | 18 +++ ...window_functions__udwf_bounded_window.snap | 18 +++ ...ns__udwf_bounded_window_ignores_frame.snap | 18 +++ ...ons__udwf_bounded_window_returns_null.snap | 18 +++ ...ow_functions__udwf_query_include_rank.snap | 18 +++ ...ned_window_functions__udwf_with_alias.snap | 18 +++ .../user_defined/user_defined_aggregates.rs | 116 ++----------- .../tests/user_defined/user_defined_plan.rs | 60 ++----- .../user_defined_scalar_functions.rs | 83 ++-------- .../user_defined_table_functions.rs | 29 +--- .../user_defined_window_functions.rs | 153 ++---------------- 41 files changed, 491 insertions(+), 419 deletions(-) create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index 8d5e363213d35..f99665a49f3df 100644 --- a/datafusion/core/tests/user_defined/expr_planner.rs +++ b/datafusion/core/tests/user_defined/expr_planner.rs @@ -84,25 +84,13 @@ fn fmt_batches(batches: &[RecordBatch]) -> String { #[tokio::test] async fn test_custom_operators_arrow() { let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +----------------------------+ - | Utf8("foo") || Utf8("bar") | - +----------------------------+ - | foobar | - +----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } #[tokio::test] async fn test_custom_operators_long_arrow() { let actual = plan_and_collect("select 1->>2;").await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +---------------------+ - | Int64(1) + Int64(2) | - +---------------------+ - | 3 | - +---------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } #[tokio::test] @@ -110,13 +98,7 @@ async fn test_question_select() { let actual = plan_and_collect("select a ? 2 from (select 1 as a);") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +--------------+ - | a ? Int64(2) | - +--------------+ - | true | - +--------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } #[tokio::test] @@ -124,11 +106,5 @@ async fn test_question_filter() { let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +---+ - | a | - +---+ - | 1 | - +---+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap new file mode 100644 index 0000000000000..5ebfd6881d44d --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/expr_planner.rs +expression: fmt_batches(&actual) +--- ++----------------------------+ +| Utf8("foo") || Utf8("bar") | ++----------------------------+ +| foobar | ++----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap new file mode 100644 index 0000000000000..a036d3485dc87 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/expr_planner.rs +expression: fmt_batches(&actual) +--- ++---------------------+ +| Int64(1) + Int64(2) | ++---------------------+ +| 3 | ++---------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap new file mode 100644 index 0000000000000..694e189fe5868 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/expr_planner.rs +expression: fmt_batches(&actual) +--- ++---+ +| a | ++---+ +| 1 | ++---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap new file mode 100644 index 0000000000000..2f4be644dbb28 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/expr_planner.rs +expression: fmt_batches(&actual) +--- ++--------------+ +| a ? Int64(2) | ++--------------+ +| true | ++--------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap new file mode 100644 index 0000000000000..62e93175e7f88 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&result) +--- ++-------------+ +| MY_AVG(t.i) | ++-------------+ +| 1.0 | ++-------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap new file mode 100644 index 0000000000000..9dbfe8779eda5 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++------+---+---+ +| text | a | b | ++------+---+---+ +| foo | 1 | 2 | ++------+---+---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap new file mode 100644 index 0000000000000..c39125a536f30 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap @@ -0,0 +1,13 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++-------+----------------------------+ +| 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 | ++-------+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap new file mode 100644 index 0000000000000..da3045a9607d9 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&result) +--- ++-------------+ +| my_avg(t.a) | ++-------------+ +| 3.0 | ++-------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap new file mode 100644 index 0000000000000..93da02ea5b50a --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++----------------------------+ +| time_sum(t.time) | ++----------------------------+ +| 1970-01-01T00:00:00.000019 | ++----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap new file mode 100644 index 0000000000000..a82c5fab76323 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap @@ -0,0 +1,13 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++----------------------------+ +| 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 | ++----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap new file mode 100644 index 0000000000000..b7f577e3f96cb --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap @@ -0,0 +1,13 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++----------------------------+ +| 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 | ++----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap new file mode 100644 index 0000000000000..a622e52545309 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++------------------------------------------------+ +| first(t.value,t.time) | ++------------------------------------------------+ +| {value: 2.0, time: 1970-01-01T00:00:00.000002} | ++------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap new file mode 100644 index 0000000000000..57590e867f66e --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++-----------------+----------------------------+ +| sq.first[value] | sq.first[time] | ++-----------------+----------------------------+ +| 2.0 | 1970-01-01T00:00:00.000002 | ++-----------------+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap new file mode 100644 index 0000000000000..51a2bb10e0457 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++---------------------------------------+ +| sum(arrow_cast(t.time,Utf8("Int64"))) | ++---------------------------------------+ +| 19000 | ++---------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap new file mode 100644 index 0000000000000..3ebc95bd52c7f --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&result) +--- ++------------+ +| dummy(t.i) | ++------------+ +| 1.0 | ++------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap new file mode 100644 index 0000000000000..274798d80ef73 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap @@ -0,0 +1,12 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_plan.rs +description: Default context +expression: actual +--- ++-------------+---------+ +| customer_id | revenue | ++-------------+---------+ +| paul | 300 | +| jorge | 200 | +| andy | 150 | ++-------------+---------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap new file mode 100644 index 0000000000000..d441732da4f5d --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap @@ -0,0 +1,12 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_plan.rs +description: Default context without schemas +expression: actual +--- ++----------+----------+ +| column_1 | column_2 | ++----------+----------+ +| andrew | 100 | +| jorge | 200 | +| andy | 150 | ++----------+----------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap new file mode 100644 index 0000000000000..7c4b3e4d75847 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap @@ -0,0 +1,10 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_plan.rs +description: MyAnalyzerRule +expression: actual +--- ++------------+--------------------------+ +| UInt64(42) | arrow_typeof(UInt64(42)) | ++------------+--------------------------+ +| 42 | UInt64 | ++------------+--------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap new file mode 100644 index 0000000000000..d7a8dd0622362 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap @@ -0,0 +1,12 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_plan.rs +description: Topk context +expression: actual +--- ++-------------+---------+ +| customer_id | revenue | ++-------------+---------+ +| paul | 300 | +| jorge | 200 | +| andy | 150 | ++-------------+---------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap new file mode 100644 index 0000000000000..c0768c5770b35 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++--------------+ +| MY_FUNC(t.i) | ++--------------+ +| 1 | ++--------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap new file mode 100644 index 0000000000000..d649f1f733787 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap @@ -0,0 +1,12 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++-----+-----+-----------------+ +| a | b | my_add(t.a,t.b) | ++-----+-----+-----------------+ +| 1 | 2 | 3 | +| 10 | 12 | 22 | +| 10 | 12 | 22 | +| 100 | 120 | 220 | ++-----+-----+-----------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap new file mode 100644 index 0000000000000..f28352148bdd6 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++---+ +| a | ++---+ +| 1 | ++---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap new file mode 100644 index 0000000000000..c06b649a3cb16 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap @@ -0,0 +1,12 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++-----+ +| a | ++-----+ +| 100 | +| 100 | +| 100 | +| 100 | ++-----+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap new file mode 100644 index 0000000000000..8c4683042aee2 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++------------+ +| dummy(t.i) | ++------------+ +| 1 | ++------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap new file mode 100644 index 0000000000000..66c9bcee59ace --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_table_functions.rs +expression: fmt_batches(&rbs) +--- ++-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ +| 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. | ++-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap new file mode 100644 index 0000000000000..7391dfa1943af --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap @@ -0,0 +1,13 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_table_functions.rs +expression: fmt_batches(&rbs) +--- ++-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ +| 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 | ++-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap new file mode 100644 index 0000000000000..0553f35043d6c --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap new file mode 100644 index 0000000000000..cd97fc388c678 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap new file mode 100644 index 0000000000000..55b53cccef701 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap new file mode 100644 index 0000000000000..d5f3cb1c1d9c1 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap new file mode 100644 index 0000000000000..cd97fc388c678 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap new file mode 100644 index 0000000000000..1834edba2e385 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap new file mode 100644 index 0000000000000..b34427e622ff5 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ +| 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 | | ++---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap new file mode 100644 index 0000000000000..2cb430d2553bf --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap new file mode 100644 index 0000000000000..55b53cccef701 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap @@ -0,0 +1,18 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_window_functions.rs +expression: fmt_batches(&actual) +--- ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ +| 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 | ++---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index 7392f3d4479d8..69ec3f431fc1a 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -71,17 +71,7 @@ async fn test_setup() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +-------+----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } /// Basic user defined aggregate @@ -93,13 +83,7 @@ async fn test_udaf() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +----------------------------+ - | time_sum(t.time) | - +----------------------------+ - | 1970-01-01T00:00:00.000019 | - +----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // normal aggregates call update_batch assert!(test_state.update_batch()); @@ -114,17 +98,7 @@ async fn test_udaf_as_window() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // aggregate over the entire window function call update_batch assert!(test_state.update_batch()); @@ -139,17 +113,7 @@ async fn test_udaf_as_window_with_frame() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // user defined aggregates with window frame should be calling retract batch assert!(test_state.update_batch()); @@ -177,13 +141,7 @@ async fn test_udaf_returning_struct() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +------------------------------------------------+ - | first(t.value,t.time) | - +------------------------------------------------+ - | {value: 2.0, time: 1970-01-01T00:00:00.000002} | - +------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } /// Demonstrate extracting the fields from a structure using a subquery @@ -194,13 +152,7 @@ async fn test_udaf_returning_struct_subquery() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +-----------------+----------------------------+ - | sq.first[value] | sq.first[time] | - +-----------------+----------------------------+ - | 2.0 | 1970-01-01T00:00:00.000002 | - +-----------------+----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } #[tokio::test] @@ -214,13 +166,7 @@ async fn test_udaf_shadows_builtin_fn() { // compute with builtin `sum` aggregator let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +---------------------------------------+ - | sum(arrow_cast(t.time,Utf8("Int64"))) | - +---------------------------------------+ - | 19000 | - +---------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // Register `TimeSum` with name `sum`. This will shadow the builtin one TimeSum::register(&mut ctx, test_state.clone(), "sum"); @@ -228,13 +174,7 @@ async fn test_udaf_shadows_builtin_fn() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +----------------------------+ - | sum(t.time) | - +----------------------------+ - | 1970-01-01T00:00:00.000019 | - +----------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } async fn execute(ctx: &SessionContext, sql: &str) -> Result> { @@ -274,13 +214,7 @@ async fn simple_udaf() -> Result<()> { let result = ctx.sql("SELECT MY_AVG(a) FROM t").await?.collect().await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +-------------+ - | my_avg(t.a) | - +-------------+ - | 3.0 | - +-------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); Ok(()) } @@ -340,13 +274,7 @@ async fn case_sensitive_identifiers_user_defined_aggregates() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +-------------+ - | MY_AVG(t.i) | - +-------------+ - | 1.0 | - +-------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); Ok(()) } @@ -372,23 +300,11 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +------------+ - | dummy(t.i) | - +------------+ - | 1.0 | - +------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result), @r###" - +------------+ - | dummy(t.i) | - +------------+ - | 1.0 | - +------------+ - "###); + insta::assert_snapshot!(fmt_batches(&alias_result)); Ok(()) } @@ -449,13 +365,7 @@ async fn test_parameterized_aggregate_udf() -> Result<()> { let actual = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +------+---+---+ - | text | a | b | - +------+---+---+ - | foo | 1 | 2 | - +------+---+---+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); 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 9a604e2bf881b..9cc1106f6c25c 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -158,21 +158,11 @@ async fn run_and_compare_query(ctx: SessionContext, description: &str) -> Result let s = exec_sql(&ctx, QUERY).await?; 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 | - +-------------+---------+ - "###); - }); - })(); + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); + }); Ok(()) } @@ -186,19 +176,11 @@ async fn run_and_compare_query_with_analyzer_rule( let s = exec_sql(&ctx, QUERY2).await?; let actual = s.lines().collect::>().join("\n"); - insta::allow_duplicates!(|| { - insta::with_settings!({ - description => description, - }, { - insta::assert_snapshot!(actual, @r###" - +------------+--------------------------+ - | UInt64(42) | arrow_typeof(UInt64(42)) | - +------------+--------------------------+ - | 42 | UInt64 | - +------------+--------------------------+ - "###); - }); - })(); + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); + }); Ok(()) } @@ -212,21 +194,11 @@ async fn run_and_compare_query_with_auto_schemas( let s = exec_sql(&ctx, QUERY1).await?; let actual = s.lines().collect::>().join("\n"); - insta::allow_duplicates!(|| { - insta::with_settings!({ - description => description, - }, { - insta::assert_snapshot!(actual, @r###" - +----------+----------+ - | column_1 | column_2 | - +----------+----------+ - | andrew | 100 | - | jorge | 200 | - | andy | 150 | - +----------+----------+ - "###); - }); - })(); + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); + }); Ok(()) } @@ -235,7 +207,7 @@ async fn run_and_compare_query_with_auto_schemas( // Run the query using default planners and optimizer async fn normal_query_without_schemas() -> Result<()> { let ctx = setup_table_without_schemas(SessionContext::new()).await?; - run_and_compare_query_with_auto_schemas(ctx, "Default context").await + run_and_compare_query_with_auto_schemas(ctx, "Default context without schemas").await } #[tokio::test] 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 0aa356a77a4ba..3ada72d71f77e 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -66,13 +66,7 @@ async fn csv_query_custom_udf_with_cast() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c11)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +------------------------------------------+ - | avg(custom_sqrt(aggregate_test_100.c11)) | - +------------------------------------------+ - | 0.6584408483418835 | - +------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); Ok(()) } @@ -85,13 +79,7 @@ async fn csv_query_avg_sqrt() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c12)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" - +------------------------------------------+ - | avg(custom_sqrt(aggregate_test_100.c12)) | - +------------------------------------------+ - | 0.6706002946036459 | - +------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); Ok(()) } @@ -156,16 +144,7 @@ async fn scalar_udf() -> Result<()> { let result = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +-----+-----+-----------------+ - | a | b | my_add(t.a,t.b) | - +-----+-----+-----------------+ - | 1 | 2 | 3 | - | 10 | 12 | 22 | - | 10 | 12 | 22 | - | 100 | 120 | 220 | - +-----+-----+-----------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); let batch = &result[0]; let a = as_int32_array(batch.column(0))?; @@ -281,31 +260,13 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +-----+ - | a | - +-----+ - | 100 | - | 100 | - | 100 | - | 100 | - +-----+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); let result = plan_and_collect(&ctx, "select get_100() a").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +-----+ - | a | - +-----+ - | 100 | - +-----+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); let result = plan_and_collect(&ctx, "select get_100() from t where a=999").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - ++ - ++ - "###); + insta::assert_snapshot!(fmt_batches(&result)); Ok(()) } @@ -332,13 +293,7 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +---+ - | a | - +---+ - | 1 | - +---+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); Ok(()) } @@ -434,13 +389,7 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +--------------+ - | MY_FUNC(t.i) | - +--------------+ - | 1 | - +--------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); Ok(()) } @@ -471,22 +420,10 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { ctx.register_udf(udf); let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" - +------------+ - | dummy(t.i) | - +------------+ - | 1 | - +------------+ - "###); + insta::assert_snapshot!(fmt_batches(&result)); let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result), @r###" - +------------+ - | dummy(t.i) | - +------------+ - | 1 | - +------------+ - "###); + insta::assert_snapshot!(fmt_batches(&alias_result)); 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 08aeee8a3275a..efc23dcce2ced 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -62,17 +62,7 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&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 | - +-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&rbs)); // just run, return all rows let rbs = ctx @@ -81,22 +71,7 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&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. | - +-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&rbs)); 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 bd5bbd302f720..e7b82467f95be 100644 --- a/datafusion/core/tests/user_defined/user_defined_window_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_window_functions.rs @@ -99,22 +99,7 @@ async fn test_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // evaluated on two distinct batches assert_eq!(test_state.evaluate_all_called(), 2); @@ -144,22 +129,7 @@ async fn test_udwf_with_alias() { .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); } /// Basic user defined window function with bounded window @@ -171,22 +141,7 @@ async fn test_udwf_bounded_window_ignores_frame() { // Since the UDWF doesn't say it needs the window frame, the frame is ignored let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+--------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // evaluated on 2 distinct batches (when x=1 and x=2) assert_eq!(test_state.evaluate_called(), 0); @@ -201,22 +156,7 @@ async fn test_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+--------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // Evaluate is called for each input rows assert_eq!(test_state.evaluate_called(), 10); @@ -233,22 +173,7 @@ async fn test_stateful_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); assert_eq!(test_state.evaluate_called(), 10); assert_eq!(test_state.evaluate_all_called(), 0); @@ -264,22 +189,7 @@ async fn test_stateful_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+--------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // Evaluate and update_state is called for each input row assert_eq!(test_state.evaluate_called(), 10); @@ -294,22 +204,7 @@ async fn test_udwf_query_include_rank() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); assert_eq!(test_state.evaluate_called(), 0); assert_eq!(test_state.evaluate_all_called(), 0); @@ -325,22 +220,7 @@ async fn test_udwf_bounded_query_include_rank() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | - +---+---+-----+--------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); assert_eq!(test_state.evaluate_called(), 0); assert_eq!(test_state.evaluate_all_called(), 0); @@ -358,22 +238,7 @@ async fn test_udwf_bounded_window_returns_null() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&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 | | - +---+---+-----+--------------------------------------------------------------------------------------------------------------+ - "###); + insta::assert_snapshot!(fmt_batches(&actual)); // Evaluate is called for each input rows assert_eq!(test_state.evaluate_called(), 10); From 08ebef51cadb91c81c1989aa12e09b0006e00ff9 Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 08:38:31 +0530 Subject: [PATCH 06/11] add snaps --- ...er_defined_aggregates__udaf_shadows_builtin_fn-2.snap | 9 +++++++++ ..._aggregates__user_defined_functions_with_alias-2.snap | 9 +++++++++ ...fined_scalar_functions__scalar_udf_zero_params-2.snap | 9 +++++++++ ...fined_scalar_functions__scalar_udf_zero_params-3.snap | 6 ++++++ ...r_functions__user_defined_functions_with_alias-2.snap | 9 +++++++++ 5 files changed, 42 insertions(+) create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap new file mode 100644 index 0000000000000..b73ec68d76d85 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&actual) +--- ++----------------------------+ +| sum(t.time) | ++----------------------------+ +| 1970-01-01T00:00:00.000019 | ++----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap new file mode 100644 index 0000000000000..d6baca2de76c4 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_aggregates.rs +expression: fmt_batches(&alias_result) +--- ++------------+ +| dummy(t.i) | ++------------+ +| 1.0 | ++------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap new file mode 100644 index 0000000000000..8058fc0c09fdd --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- ++-----+ +| a | ++-----+ +| 100 | ++-----+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap new file mode 100644 index 0000000000000..df3f626376514 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap @@ -0,0 +1,6 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&result) +--- +++ +++ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap new file mode 100644 index 0000000000000..cee733014f634 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&alias_result) +--- ++------------+ +| dummy(t.i) | ++------------+ +| 1 | ++------------+ From d695e5381f2f9155ca1c0abf591429615f303ecf Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Sun, 16 Mar 2025 09:23:43 +0530 Subject: [PATCH 07/11] add snaps --- ...ser_defined_scalar_functions__csv_query_avg_sqrt.snap | 9 +++++++++ ...scalar_functions__csv_query_custom_udf_with_cast.snap | 9 +++++++++ 2 files changed, 18 insertions(+) create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap create mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap new file mode 100644 index 0000000000000..bd3a996844b27 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&actual) +--- ++------------------------------------------+ +| avg(custom_sqrt(aggregate_test_100.c12)) | ++------------------------------------------+ +| 0.6706002946036459 | ++------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap new file mode 100644 index 0000000000000..f486c40e17609 --- /dev/null +++ b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap @@ -0,0 +1,9 @@ +--- +source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +expression: fmt_batches(&actual) +--- ++------------------------------------------+ +| avg(custom_sqrt(aggregate_test_100.c11)) | ++------------------------------------------+ +| 0.6584408483418835 | ++------------------------------------------+ From 259aedd3df38da27a035ae5a902d1a6981286140 Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Tue, 18 Mar 2025 00:21:02 +0530 Subject: [PATCH 08/11] add inline snapshots --- .../core/tests/user_defined/expr_planner.rs | 32 +++- .../user_defined/user_defined_aggregates.rs | 116 +++++++++++-- .../tests/user_defined/user_defined_plan.rs | 44 +++-- .../user_defined_scalar_functions.rs | 83 ++++++++-- .../user_defined_table_functions.rs | 29 +++- .../user_defined_window_functions.rs | 153 ++++++++++++++++-- 6 files changed, 409 insertions(+), 48 deletions(-) diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index f99665a49f3df..8d5e363213d35 100644 --- a/datafusion/core/tests/user_defined/expr_planner.rs +++ b/datafusion/core/tests/user_defined/expr_planner.rs @@ -84,13 +84,25 @@ fn fmt_batches(batches: &[RecordBatch]) -> String { #[tokio::test] async fn test_custom_operators_arrow() { let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +---------------------+ + | Int64(1) + Int64(2) | + +---------------------+ + | 3 | + +---------------------+ + "###); } #[tokio::test] @@ -98,7 +110,13 @@ async fn test_question_select() { let actual = plan_and_collect("select a ? 2 from (select 1 as a);") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +--------------+ + | a ? Int64(2) | + +--------------+ + | true | + +--------------+ + "###); } #[tokio::test] @@ -106,5 +124,11 @@ async fn test_question_filter() { let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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 69ec3f431fc1a..7392f3d4479d8 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -71,7 +71,17 @@ async fn test_setup() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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 @@ -83,7 +93,13 @@ async fn test_udaf() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +----------------------------+ + | time_sum(t.time) | + +----------------------------+ + | 1970-01-01T00:00:00.000019 | + +----------------------------+ + "###); // normal aggregates call update_batch assert!(test_state.update_batch()); @@ -98,7 +114,17 @@ async fn test_udaf_as_window() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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()); @@ -113,7 +139,17 @@ async fn test_udaf_as_window_with_frame() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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()); @@ -141,7 +177,13 @@ async fn test_udaf_returning_struct() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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 @@ -152,7 +194,13 @@ async fn test_udaf_returning_struct_subquery() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +-----------------+----------------------------+ + | sq.first[value] | sq.first[time] | + +-----------------+----------------------------+ + | 2.0 | 1970-01-01T00:00:00.000002 | + +-----------------+----------------------------+ + "###); } #[tokio::test] @@ -166,7 +214,13 @@ async fn test_udaf_shadows_builtin_fn() { // compute with builtin `sum` aggregator let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +---------------------------------------+ + | sum(arrow_cast(t.time,Utf8("Int64"))) | + +---------------------------------------+ + | 19000 | + +---------------------------------------+ + "###); // Register `TimeSum` with name `sum`. This will shadow the builtin one TimeSum::register(&mut ctx, test_state.clone(), "sum"); @@ -174,7 +228,13 @@ async fn test_udaf_shadows_builtin_fn() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +----------------------------+ + | sum(t.time) | + +----------------------------+ + | 1970-01-01T00:00:00.000019 | + +----------------------------+ + "###); } async fn execute(ctx: &SessionContext, sql: &str) -> Result> { @@ -214,7 +274,13 @@ async fn simple_udaf() -> Result<()> { let result = ctx.sql("SELECT MY_AVG(a) FROM t").await?.collect().await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +-------------+ + | my_avg(t.a) | + +-------------+ + | 3.0 | + +-------------+ + "###); Ok(()) } @@ -274,7 +340,13 @@ async fn case_sensitive_identifiers_user_defined_aggregates() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +-------------+ + | MY_AVG(t.i) | + +-------------+ + | 1.0 | + +-------------+ + "###); Ok(()) } @@ -300,11 +372,23 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +------------+ + | dummy(t.i) | + +------------+ + | 1.0 | + +------------+ + "###); let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result)); + insta::assert_snapshot!(fmt_batches(&alias_result), @r###" + +------------+ + | dummy(t.i) | + +------------+ + | 1.0 | + +------------+ + "###); Ok(()) } @@ -365,7 +449,13 @@ async fn test_parameterized_aggregate_udf() -> Result<()> { let actual = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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 9cc1106f6c25c..3da51b8206991 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -158,11 +158,21 @@ async fn run_and_compare_query(ctx: SessionContext, description: &str) -> Result let s = exec_sql(&ctx, QUERY).await?; let actual = s.lines().collect::>().join("\n"); - insta::with_settings!({ - description => description, - }, { - insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); - }); + insta::allow_duplicates! { + insta::with_settings!({ + description => description, + }, { + insta::assert_snapshot!(actual, @r###" + +-------------+---------+ + | customer_id | revenue | + +-------------+---------+ + | paul | 300 | + | jorge | 200 | + | andy | 150 | + +-------------+---------+ + "###); + }); + } Ok(()) } @@ -179,7 +189,13 @@ async fn run_and_compare_query_with_analyzer_rule( insta::with_settings!({ description => description, }, { - insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); + insta::assert_snapshot!(actual, @r###" + +------------+--------------------------+ + | UInt64(42) | arrow_typeof(UInt64(42)) | + +------------+--------------------------+ + | 42 | UInt64 | + +------------+--------------------------+ + "###); }); Ok(()) @@ -195,9 +211,17 @@ async fn run_and_compare_query_with_auto_schemas( let actual = s.lines().collect::>().join("\n"); insta::with_settings!({ - description => description, - }, { - insta::assert_snapshot!(format!("query_{}", description.replace(" ", "_")), actual); + description => description, + }, { + insta::assert_snapshot!(actual, @r###" + +-------------+---------+ + | customer_id | revenue | + +-------------+---------+ + | paul | 300 | + | jorge | 200 | + | andy | 150 | + +-------------+---------+ + "###); }); Ok(()) @@ -207,7 +231,7 @@ async fn run_and_compare_query_with_auto_schemas( // Run the query using default planners and optimizer async fn normal_query_without_schemas() -> Result<()> { let ctx = setup_table_without_schemas(SessionContext::new()).await?; - run_and_compare_query_with_auto_schemas(ctx, "Default context without schemas").await + run_and_compare_query_with_auto_schemas(ctx, "Default context").await } #[tokio::test] 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 3ada72d71f77e..0aa356a77a4ba 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -66,7 +66,13 @@ async fn csv_query_custom_udf_with_cast() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c11)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +------------------------------------------+ + | avg(custom_sqrt(aggregate_test_100.c11)) | + +------------------------------------------+ + | 0.6584408483418835 | + +------------------------------------------+ + "###); Ok(()) } @@ -79,7 +85,13 @@ async fn csv_query_avg_sqrt() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c12)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&actual), @r###" + +------------------------------------------+ + | avg(custom_sqrt(aggregate_test_100.c12)) | + +------------------------------------------+ + | 0.6706002946036459 | + +------------------------------------------+ + "###); Ok(()) } @@ -144,7 +156,16 @@ async fn scalar_udf() -> Result<()> { let result = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&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))?; @@ -260,13 +281,31 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +-----+ + | a | + +-----+ + | 100 | + | 100 | + | 100 | + | 100 | + +-----+ + "###); let result = plan_and_collect(&ctx, "select get_100() a").await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +-----+ + | a | + +-----+ + | 100 | + +-----+ + "###); let result = plan_and_collect(&ctx, "select get_100() from t where a=999").await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + ++ + ++ + "###); Ok(()) } @@ -293,7 +332,13 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +---+ + | a | + +---+ + | 1 | + +---+ + "###); Ok(()) } @@ -389,7 +434,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?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +--------------+ + | MY_FUNC(t.i) | + +--------------+ + | 1 | + +--------------+ + "###); Ok(()) } @@ -420,10 +471,22 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { ctx.register_udf(udf); let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result)); + insta::assert_snapshot!(fmt_batches(&result), @r###" + +------------+ + | dummy(t.i) | + +------------+ + | 1 | + +------------+ + "###); let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result)); + insta::assert_snapshot!(fmt_batches(&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 efc23dcce2ced..08aeee8a3275a 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -62,7 +62,17 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&rbs)); + insta::assert_snapshot!(fmt_batches(&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 @@ -71,7 +81,22 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&rbs)); + insta::assert_snapshot!(fmt_batches(&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 e7b82467f95be..bd5bbd302f720 100644 --- a/datafusion/core/tests/user_defined/user_defined_window_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_window_functions.rs @@ -99,7 +99,22 @@ async fn test_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -129,7 +144,22 @@ async fn test_udwf_with_alias() { .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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 @@ -141,7 +171,22 @@ async fn test_udwf_bounded_window_ignores_frame() { // Since the UDWF doesn't say it needs the window frame, the frame is ignored let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -156,7 +201,22 @@ async fn test_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -173,7 +233,22 @@ async fn test_stateful_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -189,7 +264,22 @@ async fn test_stateful_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -204,7 +294,22 @@ async fn test_udwf_query_include_rank() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -220,7 +325,22 @@ async fn test_udwf_bounded_query_include_rank() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); @@ -238,7 +358,22 @@ async fn test_udwf_bounded_window_returns_null() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual)); + insta::assert_snapshot!(fmt_batches(&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); From fe1447c9d07f80ab190963ce00cecc4729c12ce2 Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Tue, 18 Mar 2025 00:37:29 +0530 Subject: [PATCH 09/11] fix --- .../core/tests/user_defined/user_defined_plan.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 3da51b8206991..e46940e631542 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -214,13 +214,13 @@ async fn run_and_compare_query_with_auto_schemas( description => description, }, { insta::assert_snapshot!(actual, @r###" - +-------------+---------+ - | customer_id | revenue | - +-------------+---------+ - | paul | 300 | - | jorge | 200 | - | andy | 150 | - +-------------+---------+ + +----------+----------+ + | column_1 | column_2 | + +----------+----------+ + | andrew | 100 | + | jorge | 200 | + | andy | 150 | + +----------+----------+ "###); }); From 5d4ada3b9ac3cfbc4950aa6e6afddeb3a77b975b Mon Sep 17 00:00:00 2001 From: shruti2522 Date: Tue, 18 Mar 2025 00:43:23 +0530 Subject: [PATCH 10/11] remove snaps --- ...__expr_planner__custom_operators_arrow.snap | 9 --------- ...r_planner__custom_operators_long_arrow.snap | 9 --------- ...defined__expr_planner__question_filter.snap | 9 --------- ...defined__expr_planner__question_select.snap | 9 --------- ...ve_identifiers_user_defined_aggregates.snap | 9 --------- ...ggregates__parameterized_aggregate_udf.snap | 9 --------- ...efined__user_defined_aggregates__setup.snap | 13 ------------- ...__user_defined_aggregates__simple_udaf.snap | 9 --------- ...defined__user_defined_aggregates__udaf.snap | 9 --------- ...ser_defined_aggregates__udaf_as_window.snap | 13 ------------- ..._aggregates__udaf_as_window_with_frame.snap | 13 ------------- ...ined_aggregates__udaf_returning_struct.snap | 9 --------- ...egates__udaf_returning_struct_subquery.snap | 9 --------- ..._aggregates__udaf_shadows_builtin_fn-2.snap | 9 --------- ...ed_aggregates__udaf_shadows_builtin_fn.snap | 9 --------- ...s__user_defined_functions_with_alias-2.snap | 9 --------- ...tes__user_defined_functions_with_alias.snap | 9 --------- ...er_defined_plan__query_Default_context.snap | 12 ------------ ..._query_Default_context_without_schemas.snap | 12 ------------ ...ser_defined_plan__query_MyAnalyzerRule.snap | 10 ---------- ..._user_defined_plan__query_Topk_context.snap | 12 ------------ ...ive_identifiers_user_defined_functions.snap | 9 --------- ...d_scalar_functions__csv_query_avg_sqrt.snap | 9 --------- ...ctions__csv_query_custom_udf_with_cast.snap | 9 --------- ...r_defined_scalar_functions__scalar_udf.snap | 12 ------------ ..._udf_override_built_in_scalar_function.snap | 9 --------- ...ar_functions__scalar_udf_zero_params-2.snap | 9 --------- ...ar_functions__scalar_udf_zero_params-3.snap | 6 ------ ...alar_functions__scalar_udf_zero_params.snap | 12 ------------ ...s__user_defined_functions_with_alias-2.snap | 9 --------- ...ons__user_defined_functions_with_alias.snap | 9 --------- ...able_functions__simple_read_csv_udtf-2.snap | 18 ------------------ ..._table_functions__simple_read_csv_udtf.snap | 13 ------------- ...efined_window_functions__stateful_udwf.snap | 18 ------------------ ...unctions__stateful_udwf_bounded_window.snap | 18 ------------------ ...d__user_defined_window_functions__udwf.snap | 18 ------------------ ...tions__udwf_bounded_query_include_rank.snap | 18 ------------------ ..._window_functions__udwf_bounded_window.snap | 18 ------------------ ...ons__udwf_bounded_window_ignores_frame.snap | 18 ------------------ ...ions__udwf_bounded_window_returns_null.snap | 18 ------------------ ...dow_functions__udwf_query_include_rank.snap | 18 ------------------ ...ined_window_functions__udwf_with_alias.snap | 18 ------------------ 42 files changed, 497 deletions(-) delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap delete mode 100644 datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap deleted file mode 100644 index 5ebfd6881d44d..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_arrow.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/expr_planner.rs -expression: fmt_batches(&actual) ---- -+----------------------------+ -| Utf8("foo") || Utf8("bar") | -+----------------------------+ -| foobar | -+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap deleted file mode 100644 index a036d3485dc87..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__custom_operators_long_arrow.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/expr_planner.rs -expression: fmt_batches(&actual) ---- -+---------------------+ -| Int64(1) + Int64(2) | -+---------------------+ -| 3 | -+---------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap deleted file mode 100644 index 694e189fe5868..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_filter.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/expr_planner.rs -expression: fmt_batches(&actual) ---- -+---+ -| a | -+---+ -| 1 | -+---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap deleted file mode 100644 index 2f4be644dbb28..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__expr_planner__question_select.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/expr_planner.rs -expression: fmt_batches(&actual) ---- -+--------------+ -| a ? Int64(2) | -+--------------+ -| true | -+--------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap deleted file mode 100644 index 62e93175e7f88..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__case_sensitive_identifiers_user_defined_aggregates.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&result) ---- -+-------------+ -| MY_AVG(t.i) | -+-------------+ -| 1.0 | -+-------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap deleted file mode 100644 index 9dbfe8779eda5..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__parameterized_aggregate_udf.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+------+---+---+ -| text | a | b | -+------+---+---+ -| foo | 1 | 2 | -+------+---+---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap deleted file mode 100644 index c39125a536f30..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__setup.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+-------+----------------------------+ -| 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 | -+-------+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap deleted file mode 100644 index da3045a9607d9..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__simple_udaf.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&result) ---- -+-------------+ -| my_avg(t.a) | -+-------------+ -| 3.0 | -+-------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap deleted file mode 100644 index 93da02ea5b50a..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+----------------------------+ -| time_sum(t.time) | -+----------------------------+ -| 1970-01-01T00:00:00.000019 | -+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap deleted file mode 100644 index a82c5fab76323..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+----------------------------+ -| 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 | -+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap deleted file mode 100644 index b7f577e3f96cb..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_as_window_with_frame.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+----------------------------+ -| 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 | -+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap deleted file mode 100644 index a622e52545309..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+------------------------------------------------+ -| first(t.value,t.time) | -+------------------------------------------------+ -| {value: 2.0, time: 1970-01-01T00:00:00.000002} | -+------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap deleted file mode 100644 index 57590e867f66e..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_returning_struct_subquery.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+-----------------+----------------------------+ -| sq.first[value] | sq.first[time] | -+-----------------+----------------------------+ -| 2.0 | 1970-01-01T00:00:00.000002 | -+-----------------+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap deleted file mode 100644 index b73ec68d76d85..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn-2.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+----------------------------+ -| sum(t.time) | -+----------------------------+ -| 1970-01-01T00:00:00.000019 | -+----------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap deleted file mode 100644 index 51a2bb10e0457..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__udaf_shadows_builtin_fn.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&actual) ---- -+---------------------------------------+ -| sum(arrow_cast(t.time,Utf8("Int64"))) | -+---------------------------------------+ -| 19000 | -+---------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap deleted file mode 100644 index d6baca2de76c4..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias-2.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&alias_result) ---- -+------------+ -| dummy(t.i) | -+------------+ -| 1.0 | -+------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap deleted file mode 100644 index 3ebc95bd52c7f..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_aggregates__user_defined_functions_with_alias.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_aggregates.rs -expression: fmt_batches(&result) ---- -+------------+ -| dummy(t.i) | -+------------+ -| 1.0 | -+------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap deleted file mode 100644 index 274798d80ef73..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_plan.rs -description: Default context -expression: actual ---- -+-------------+---------+ -| customer_id | revenue | -+-------------+---------+ -| paul | 300 | -| jorge | 200 | -| andy | 150 | -+-------------+---------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap deleted file mode 100644 index d441732da4f5d..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Default_context_without_schemas.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_plan.rs -description: Default context without schemas -expression: actual ---- -+----------+----------+ -| column_1 | column_2 | -+----------+----------+ -| andrew | 100 | -| jorge | 200 | -| andy | 150 | -+----------+----------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap deleted file mode 100644 index 7c4b3e4d75847..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_MyAnalyzerRule.snap +++ /dev/null @@ -1,10 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_plan.rs -description: MyAnalyzerRule -expression: actual ---- -+------------+--------------------------+ -| UInt64(42) | arrow_typeof(UInt64(42)) | -+------------+--------------------------+ -| 42 | UInt64 | -+------------+--------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap deleted file mode 100644 index d7a8dd0622362..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_plan__query_Topk_context.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_plan.rs -description: Topk context -expression: actual ---- -+-------------+---------+ -| customer_id | revenue | -+-------------+---------+ -| paul | 300 | -| jorge | 200 | -| andy | 150 | -+-------------+---------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap deleted file mode 100644 index c0768c5770b35..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__case_sensitive_identifiers_user_defined_functions.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+--------------+ -| MY_FUNC(t.i) | -+--------------+ -| 1 | -+--------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap deleted file mode 100644 index bd3a996844b27..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_avg_sqrt.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&actual) ---- -+------------------------------------------+ -| avg(custom_sqrt(aggregate_test_100.c12)) | -+------------------------------------------+ -| 0.6706002946036459 | -+------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap deleted file mode 100644 index f486c40e17609..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__csv_query_custom_udf_with_cast.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&actual) ---- -+------------------------------------------+ -| avg(custom_sqrt(aggregate_test_100.c11)) | -+------------------------------------------+ -| 0.6584408483418835 | -+------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap deleted file mode 100644 index d649f1f733787..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+-----+-----+-----------------+ -| a | b | my_add(t.a,t.b) | -+-----+-----+-----------------+ -| 1 | 2 | 3 | -| 10 | 12 | 22 | -| 10 | 12 | 22 | -| 100 | 120 | 220 | -+-----+-----+-----------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap deleted file mode 100644 index f28352148bdd6..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_override_built_in_scalar_function.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+---+ -| a | -+---+ -| 1 | -+---+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap deleted file mode 100644 index 8058fc0c09fdd..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-2.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+-----+ -| a | -+-----+ -| 100 | -+-----+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap deleted file mode 100644 index df3f626376514..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params-3.snap +++ /dev/null @@ -1,6 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -++ -++ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap deleted file mode 100644 index c06b649a3cb16..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__scalar_udf_zero_params.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+-----+ -| a | -+-----+ -| 100 | -| 100 | -| 100 | -| 100 | -+-----+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap deleted file mode 100644 index cee733014f634..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias-2.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&alias_result) ---- -+------------+ -| dummy(t.i) | -+------------+ -| 1 | -+------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap deleted file mode 100644 index 8c4683042aee2..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_scalar_functions__user_defined_functions_with_alias.snap +++ /dev/null @@ -1,9 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_scalar_functions.rs -expression: fmt_batches(&result) ---- -+------------+ -| dummy(t.i) | -+------------+ -| 1 | -+------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap deleted file mode 100644 index 66c9bcee59ace..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf-2.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_table_functions.rs -expression: fmt_batches(&rbs) ---- -+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ -| 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. | -+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap deleted file mode 100644 index 7391dfa1943af..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_table_functions__simple_read_csv_udtf.snap +++ /dev/null @@ -1,13 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_table_functions.rs -expression: fmt_batches(&rbs) ---- -+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ -| 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 | -+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap deleted file mode 100644 index 0553f35043d6c..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap deleted file mode 100644 index cd97fc388c678..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__stateful_udwf_bounded_window.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap deleted file mode 100644 index 55b53cccef701..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap deleted file mode 100644 index d5f3cb1c1d9c1..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_query_include_rank.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap deleted file mode 100644 index cd97fc388c678..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap deleted file mode 100644 index 1834edba2e385..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_ignores_frame.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap deleted file mode 100644 index b34427e622ff5..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_bounded_window_returns_null.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ -| 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 | | -+---+---+-----+--------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap deleted file mode 100644 index 2cb430d2553bf..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_query_include_rank.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ diff --git a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap b/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap deleted file mode 100644 index 55b53cccef701..0000000000000 --- a/datafusion/core/tests/user_defined/snapshots/user_defined_integration__user_defined__user_defined_window_functions__udwf_with_alias.snap +++ /dev/null @@ -1,18 +0,0 @@ ---- -source: datafusion/core/tests/user_defined/user_defined_window_functions.rs -expression: fmt_batches(&actual) ---- -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ -| 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 | -+---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ From c6cdabdf9b329cb4a1acdf111cc1724c331c091f Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Tue, 18 Mar 2025 13:56:40 -0400 Subject: [PATCH 11/11] Reuse batches_to_stirng --- .../core/tests/user_defined/expr_planner.rs | 17 +++------ .../user_defined/user_defined_aggregates.rs | 35 ++++++++----------- .../user_defined_scalar_functions.rs | 29 ++++++--------- .../user_defined_table_functions.rs | 13 ++----- .../user_defined_window_functions.rs | 31 +++++++--------- 5 files changed, 45 insertions(+), 80 deletions(-) diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index 8d5e363213d35..1fc6d14c5b229 100644 --- a/datafusion/core/tests/user_defined/expr_planner.rs +++ b/datafusion/core/tests/user_defined/expr_planner.rs @@ -16,6 +16,7 @@ // under the License. use arrow::array::RecordBatch; +use datafusion::common::test_util::batches_to_string; use std::sync::Arc; use datafusion::common::DFSchema; @@ -73,18 +74,10 @@ async fn plan_and_collect(sql: &str) -> Result> { ctx.sql(sql).await?.collect().await } -fn fmt_batches(batches: &[RecordBatch]) -> String { - use arrow::util::pretty::pretty_format_batches; - match pretty_format_batches(batches) { - Ok(formatted) => formatted.to_string(), - Err(e) => format!("Error formatting record batches: {}", e), - } -} - #[tokio::test] async fn test_custom_operators_arrow() { let actual = plan_and_collect("select 'foo'->'bar';").await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +----------------------------+ | Utf8("foo") || Utf8("bar") | +----------------------------+ @@ -96,7 +89,7 @@ async fn test_custom_operators_arrow() { #[tokio::test] async fn test_custom_operators_long_arrow() { let actual = plan_and_collect("select 1->>2;").await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +---------------------+ | Int64(1) + Int64(2) | +---------------------+ @@ -110,7 +103,7 @@ async fn test_question_select() { let actual = plan_and_collect("select a ? 2 from (select 1 as a);") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +--------------+ | a ? Int64(2) | +--------------+ @@ -124,7 +117,7 @@ async fn test_question_filter() { let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;") .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +---+ | a | +---+ diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index 7392f3d4479d8..5cbb05f290a70 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; @@ -55,14 +56,6 @@ use datafusion_expr::{ }; use datafusion_functions_aggregate::average::AvgAccumulator; -fn fmt_batches(batches: &[RecordBatch]) -> String { - use arrow::util::pretty::pretty_format_batches; - match pretty_format_batches(batches) { - Ok(formatted) => formatted.to_string(), - Err(e) => format!("Error formatting record batches: {}", e), - } -} - /// Test to show the contents of the setup #[tokio::test] async fn test_setup() { @@ -71,7 +64,7 @@ async fn test_setup() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +-------+----------------------------+ | value | time | +-------+----------------------------+ @@ -93,7 +86,7 @@ async fn test_udaf() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +----------------------------+ | time_sum(t.time) | +----------------------------+ @@ -114,7 +107,7 @@ async fn test_udaf_as_window() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +----------------------------+ | time_sum | +----------------------------+ @@ -139,7 +132,7 @@ async fn test_udaf_as_window_with_frame() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +----------------------------+ | time_sum | +----------------------------+ @@ -177,7 +170,7 @@ async fn test_udaf_returning_struct() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +------------------------------------------------+ | first(t.value,t.time) | +------------------------------------------------+ @@ -194,7 +187,7 @@ async fn test_udaf_returning_struct_subquery() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +-----------------+----------------------------+ | sq.first[value] | sq.first[time] | +-----------------+----------------------------+ @@ -214,7 +207,7 @@ async fn test_udaf_shadows_builtin_fn() { // compute with builtin `sum` aggregator let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +---------------------------------------+ | sum(arrow_cast(t.time,Utf8("Int64"))) | +---------------------------------------+ @@ -228,7 +221,7 @@ async fn test_udaf_shadows_builtin_fn() { let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +----------------------------+ | sum(t.time) | +----------------------------+ @@ -274,7 +267,7 @@ async fn simple_udaf() -> Result<()> { let result = ctx.sql("SELECT MY_AVG(a) FROM t").await?.collect().await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +-------------+ | my_avg(t.a) | +-------------+ @@ -340,7 +333,7 @@ async fn case_sensitive_identifiers_user_defined_aggregates() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +-------------+ | MY_AVG(t.i) | +-------------+ @@ -372,7 +365,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +------------+ | dummy(t.i) | +------------+ @@ -382,7 +375,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result), @r###" + insta::assert_snapshot!(batches_to_string(&alias_result), @r###" +------------+ | dummy(t.i) | +------------+ @@ -449,7 +442,7 @@ async fn test_parameterized_aggregate_udf() -> Result<()> { let actual = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +------+---+---+ | text | a | b | +------+---+---+ 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 0aa356a77a4ba..41f91d6b8336f 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}; @@ -48,14 +49,6 @@ use regex::Regex; use sqlparser::ast::Ident; use sqlparser::tokenizer::Span; -fn fmt_batches(batches: &[RecordBatch]) -> String { - use arrow::util::pretty::pretty_format_batches; - match pretty_format_batches(batches) { - Ok(formatted) => formatted.to_string(), - Err(e) => format!("Error formatting record batches: {}", e), - } -} - /// test that casting happens on udfs. /// c11 is f32, but `custom_sqrt` requires f64. Casting happens but the logical plan and /// physical plan have the same schema. @@ -66,7 +59,7 @@ async fn csv_query_custom_udf_with_cast() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c11)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +------------------------------------------+ | avg(custom_sqrt(aggregate_test_100.c11)) | +------------------------------------------+ @@ -85,7 +78,7 @@ async fn csv_query_avg_sqrt() -> Result<()> { let sql = "SELECT avg(custom_sqrt(c12)) FROM aggregate_test_100"; let actual = plan_and_collect(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +------------------------------------------+ | avg(custom_sqrt(aggregate_test_100.c12)) | +------------------------------------------+ @@ -156,7 +149,7 @@ async fn scalar_udf() -> Result<()> { let result = DataFrame::new(ctx.state(), plan).collect().await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +-----+-----+-----------------+ | a | b | my_add(t.a,t.b) | +-----+-----+-----------------+ @@ -281,7 +274,7 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +-----+ | a | +-----+ @@ -293,7 +286,7 @@ async fn scalar_udf_zero_params() -> Result<()> { "###); let result = plan_and_collect(&ctx, "select get_100() a").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +-----+ | a | +-----+ @@ -302,7 +295,7 @@ async fn scalar_udf_zero_params() -> Result<()> { "###); let result = plan_and_collect(&ctx, "select get_100() from t where a=999").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" ++ ++ "###); @@ -332,7 +325,7 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +---+ | a | +---+ @@ -434,7 +427,7 @@ 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?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +--------------+ | MY_FUNC(t.i) | +--------------+ @@ -471,7 +464,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { ctx.register_udf(udf); let result = plan_and_collect(&ctx, "SELECT dummy(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&result), @r###" + insta::assert_snapshot!(batches_to_string(&result), @r###" +------------+ | dummy(t.i) | +------------+ @@ -480,7 +473,7 @@ async fn test_user_defined_functions_with_alias() -> Result<()> { "###); let alias_result = plan_and_collect(&ctx, "SELECT dummy_alias(i) FROM t").await?; - insta::assert_snapshot!(fmt_batches(&alias_result), @r###" + insta::assert_snapshot!(batches_to_string(&alias_result), @r###" +------------+ | dummy(t.i) | +------------+ 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 08aeee8a3275a..e4aff0b00705d 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; @@ -39,14 +40,6 @@ use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; use async_trait::async_trait; -fn fmt_batches(batches: &[RecordBatch]) -> String { - use arrow::util::pretty::pretty_format_batches; - match pretty_format_batches(batches) { - Ok(formatted) => formatted.to_string(), - Err(e) => format!("Error formatting record batches: {}", e), - } -} - /// test simple udtf with define read_csv with parameters #[tokio::test] async fn test_simple_read_csv_udtf() -> Result<()> { @@ -62,7 +55,7 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&rbs), @r###" + insta::assert_snapshot!(batches_to_string(&rbs), @r###" +-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ | n_nationkey | n_name | n_regionkey | n_comment | +-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+ @@ -81,7 +74,7 @@ async fn test_simple_read_csv_udtf() -> Result<()> { .collect() .await?; - insta::assert_snapshot!(fmt_batches(&rbs), @r###" + insta::assert_snapshot!(batches_to_string(&rbs), @r###" +-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ | n_nationkey | n_name | n_regionkey | n_comment | +-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+ 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 bd5bbd302f720..28394f0b9dfaf 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::common::test_util::batches_to_string; +use datafusion::common::{Result, ScalarValue}; use datafusion::prelude::SessionContext; -use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ PartitionEvaluator, Signature, TypeSignature, Volatility, WindowUDF, WindowUDFImpl, }; @@ -57,14 +58,6 @@ 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"; -fn fmt_batches(batches: &[RecordBatch]) -> String { - use arrow::util::pretty::pretty_format_batches; - match pretty_format_batches(batches) { - Ok(formatted) => formatted.to_string(), - Err(e) => format!("Error formatting record batches: {}", e), - } -} - #[tokio::test] async fn test_setup() { let test_state = TestState::new(); @@ -73,7 +66,7 @@ async fn test_setup() { let sql = "SELECT * from t order by x, y"; let actual = execute(&ctx, sql).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + insta::assert_snapshot!(batches_to_string(&actual), @r###" +---+---+-----+ | x | y | val | +---+---+-----+ @@ -99,7 +92,7 @@ async fn test_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ @@ -144,7 +137,7 @@ async fn test_udwf_with_alias() { .await .unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ @@ -171,7 +164,7 @@ async fn test_udwf_bounded_window_ignores_frame() { // Since the UDWF doesn't say it needs the window frame, the frame is ignored let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+--------------------------------------------------------------------------------------------------------------+ @@ -201,7 +194,7 @@ async fn test_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+--------------------------------------------------------------------------------------------------------------+ @@ -233,7 +226,7 @@ async fn test_stateful_udwf() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ @@ -264,7 +257,7 @@ async fn test_stateful_udwf_bounded_window() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+--------------------------------------------------------------------------------------------------------------+ @@ -294,7 +287,7 @@ async fn test_udwf_query_include_rank() { let actual = execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+-----------------------------------------------------------------------------------------------------------------------+ @@ -325,7 +318,7 @@ async fn test_udwf_bounded_query_include_rank() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+--------------------------------------------------------------------------------------------------------------+ @@ -358,7 +351,7 @@ async fn test_udwf_bounded_window_returns_null() { let actual = execute(&ctx, BOUNDED_WINDOW_QUERY).await.unwrap(); - insta::assert_snapshot!(fmt_batches(&actual), @r###" + 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 | +---+---+-----+--------------------------------------------------------------------------------------------------------------+