diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 5277a50b85ca..320fa21c8665 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -467,11 +467,12 @@ mod tests { use arrow::array::*; use arrow::compute::SortOptions; use arrow::datatypes::*; + use datafusion_common::test_util::batches_to_string; use futures::FutureExt; + use insta::allow_duplicates; + use insta::assert_snapshot; use itertools::Itertools; - use datafusion_common::assert_batches_eq; - use crate::collect; use crate::expressions::col; use crate::expressions::PhysicalSortExpr; @@ -522,20 +523,21 @@ mod tests { let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; - let expected_after_sort = [ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| 0 | 1 | 0 |", - "| 0 | 1 | 1 |", - "| 0 | 2 | 5 |", - "| 1 | 2 | 4 |", - "| 1 | 3 | 2 |", - "| 1 | 3 | 3 |", - "+---+---+---+", - ]; assert_eq!(2, result.len()); - assert_batches_eq!(expected_after_sort, &result); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r#" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 0 | + | 0 | 1 | 1 | + | 0 | 2 | 5 | + | 1 | 2 | 4 | + | 1 | 3 | 2 | + | 1 | 3 | 3 | + +---+---+---+ + "#); + } assert_eq!( task_ctx.runtime_env().memory_pool.reserved(), 0, @@ -588,18 +590,19 @@ mod tests { let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; - let expected_after_sort = [ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| 0 | 1 | 4 |", - "| 0 | 2 | 3 |", - "| 1 | 2 | 2 |", - "| 1 | 3 | 0 |", - "+---+---+---+", - ]; assert_eq!(2, result.len()); - assert_batches_eq!(expected_after_sort, &result); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r#" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 4 | + | 0 | 2 | 3 | + | 1 | 2 | 2 | + | 1 | 3 | 0 | + +---+---+---+ + "#); + } assert_eq!( task_ctx.runtime_env().memory_pool.reserved(), 0, @@ -663,21 +666,22 @@ mod tests { 0, "The sort should have returned all memory used back to the memory manager" ); - let expected = [ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| 0 | 1 | 6 |", - "| 0 | 1 | 7 |", - "| 0 | 3 | 4 |", - "| 0 | 3 | 5 |", - "| 1 | 2 | 0 |", - "| 1 | 2 | 1 |", - "| 1 | 4 | 2 |", - "| 1 | 4 | 3 |", - "+---+---+---+", - ]; - assert_batches_eq!(expected, &result); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r#" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 6 | + | 0 | 1 | 7 | + | 0 | 3 | 4 | + | 0 | 3 | 5 | + | 1 | 2 | 0 | + | 1 | 2 | 1 | + | 1 | 4 | 2 | + | 1 | 4 | 3 | + +---+---+---+ + "#); + } } Ok(()) } @@ -1000,21 +1004,6 @@ mod tests { 2, )); - let expected = [ - "+-----+------+-------+", - "| a | b | c |", - "+-----+------+-------+", - "| 1.0 | 20.0 | 20.0 |", - "| 1.0 | 20.0 | 10.0 |", - "| 1.0 | 40.0 | 10.0 |", - "| 2.0 | 40.0 | 100.0 |", - "| 2.0 | NaN | NaN |", - "| 3.0 | | |", - "| 3.0 | | 100.0 |", - "| 3.0 | NaN | NaN |", - "+-----+------+-------+", - ]; - assert_eq!( DataType::Float32, *partial_sort_exec.schema().field(0).data_type() @@ -1033,7 +1022,20 @@ mod tests { task_ctx, ) .await?; - assert_batches_eq!(expected, &result); + assert_snapshot!(batches_to_string(&result), @r#" + +-----+------+-------+ + | a | b | c | + +-----+------+-------+ + | 1.0 | 20.0 | 20.0 | + | 1.0 | 20.0 | 10.0 | + | 1.0 | 40.0 | 10.0 | + | 2.0 | 40.0 | 100.0 | + | 2.0 | NaN | NaN | + | 3.0 | | | + | 3.0 | | 100.0 | + | 3.0 | NaN | NaN | + +-----+------+-------+ + "#); assert_eq!(result.len(), 2); let metrics = partial_sort_exec.metrics().unwrap(); assert!(metrics.elapsed_compute().unwrap() > 0); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index e2d665e1d814..731e24b53c60 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1244,7 +1244,8 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::*; use datafusion_common::cast::as_primitive_array; - use datafusion_common::{assert_batches_eq, Result, ScalarValue}; + use datafusion_common::test_util::batches_to_string; + use datafusion_common::{Result, ScalarValue}; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_execution::RecordBatchStream; @@ -1252,6 +1253,7 @@ mod tests { use datafusion_physical_expr::EquivalenceProperties; use futures::{FutureExt, Stream}; + use insta::assert_snapshot; #[derive(Debug, Clone)] pub struct SortedUnboundedExec { @@ -1913,22 +1915,21 @@ mod tests { plan = plan.with_fetch(Some(9)); let batches = collect(Arc::new(plan), task_ctx).await?; - #[rustfmt::skip] - let expected = [ - "+----+", - "| c1 |", - "+----+", - "| 0 |", - "| 1 |", - "| 2 |", - "| 3 |", - "| 4 |", - "| 5 |", - "| 6 |", - "| 7 |", - "| 8 |", - "+----+",]; - assert_batches_eq!(expected, &batches); + assert_snapshot!(batches_to_string(&batches), @r#" + +----+ + | c1 | + +----+ + | 0 | + | 1 | + | 2 | + | 3 | + | 4 | + | 5 | + | 6 | + | 7 | + | 8 | + +----+ + "#); Ok(()) } } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index ca06a029e8db..b987dff36441 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -412,6 +412,7 @@ mod tests { }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::test_util::batches_to_string; use datafusion_common::{assert_batches_eq, assert_contains, DataFusionError}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::config::SessionConfig; @@ -423,6 +424,7 @@ mod tests { use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::{FutureExt, Stream, StreamExt}; + use insta::assert_snapshot; use tokio::time::timeout; // The number in the function is highly related to the memory limit we are testing @@ -992,25 +994,22 @@ mod tests { let collected = collect(merge, task_ctx).await.unwrap(); assert_eq!(collected.len(), 1); - assert_batches_eq!( - &[ - "+---+---+-------------------------------+", - "| a | b | c |", - "+---+---+-------------------------------+", - "| 1 | | 1970-01-01T00:00:00.000000008 |", - "| 1 | | 1970-01-01T00:00:00.000000008 |", - "| 2 | a | |", - "| 7 | b | 1970-01-01T00:00:00.000000006 |", - "| 2 | b | |", - "| 9 | d | |", - "| 3 | e | 1970-01-01T00:00:00.000000004 |", - "| 3 | g | 1970-01-01T00:00:00.000000005 |", - "| 4 | h | |", - "| 5 | i | 1970-01-01T00:00:00.000000004 |", - "+---+---+-------------------------------+", - ], - collected.as_slice() - ); + assert_snapshot!(batches_to_string(collected.as_slice()), @r#" + +---+---+-------------------------------+ + | a | b | c | + +---+---+-------------------------------+ + | 1 | | 1970-01-01T00:00:00.000000008 | + | 1 | | 1970-01-01T00:00:00.000000008 | + | 2 | a | | + | 7 | b | 1970-01-01T00:00:00.000000006 | + | 2 | b | | + | 9 | d | | + | 3 | e | 1970-01-01T00:00:00.000000004 | + | 3 | g | 1970-01-01T00:00:00.000000005 | + | 4 | h | | + | 5 | i | 1970-01-01T00:00:00.000000004 | + +---+---+-------------------------------+ + "#); } #[tokio::test] @@ -1035,17 +1034,14 @@ mod tests { let collected = collect(merge, task_ctx).await.unwrap(); assert_eq!(collected.len(), 1); - assert_batches_eq!( - &[ - "+---+---+", - "| a | b |", - "+---+---+", - "| 1 | a |", - "| 2 | b |", - "+---+---+", - ], - collected.as_slice() - ); + assert_snapshot!(batches_to_string(collected.as_slice()), @r#" + +---+---+ + | a | b | + +---+---+ + | 1 | a | + | 2 | b | + +---+---+ + "#); } #[tokio::test] @@ -1069,20 +1065,17 @@ mod tests { let collected = collect(merge, task_ctx).await.unwrap(); assert_eq!(collected.len(), 1); - assert_batches_eq!( - &[ - "+---+---+", - "| a | b |", - "+---+---+", - "| 1 | a |", - "| 2 | b |", - "| 7 | c |", - "| 9 | d |", - "| 3 | e |", - "+---+---+", - ], - collected.as_slice() - ); + assert_snapshot!(batches_to_string(collected.as_slice()), @r#" + +---+---+ + | a | b | + +---+---+ + | 1 | a | + | 2 | b | + | 7 | c | + | 9 | d | + | 3 | e | + +---+---+ + "#); } #[tokio::test] @@ -1179,17 +1172,16 @@ mod tests { let collected = collect(Arc::clone(&merge) as Arc, task_ctx) .await .unwrap(); - let expected = [ - "+----+---+", - "| a | b |", - "+----+---+", - "| 1 | a |", - "| 10 | b |", - "| 2 | c |", - "| 20 | d |", - "+----+---+", - ]; - assert_batches_eq!(expected, collected.as_slice()); + assert_snapshot!(batches_to_string(collected.as_slice()), @r#" + +----+---+ + | a | b | + +----+---+ + | 1 | a | + | 10 | b | + | 2 | c | + | 20 | d | + +----+---+ + "#); // Now, validate metrics let metrics = merge.metrics().unwrap(); @@ -1293,35 +1285,32 @@ mod tests { // Expect the data to be sorted first by "batch_number" (because // that was the order it was fed in, even though only "value" // is in the sort key) - assert_batches_eq!( - &[ - "+--------------+-------+", - "| batch_number | value |", - "+--------------+-------+", - "| 0 | A |", - "| 1 | A |", - "| 2 | A |", - "| 3 | A |", - "| 4 | A |", - "| 5 | A |", - "| 6 | A |", - "| 7 | A |", - "| 8 | A |", - "| 9 | A |", - "| 0 | B |", - "| 1 | B |", - "| 2 | B |", - "| 3 | B |", - "| 4 | B |", - "| 5 | B |", - "| 6 | B |", - "| 7 | B |", - "| 8 | B |", - "| 9 | B |", - "+--------------+-------+", - ], - collected.as_slice() - ); + assert_snapshot!(batches_to_string(collected.as_slice()), @r#" + +--------------+-------+ + | batch_number | value | + +--------------+-------+ + | 0 | A | + | 1 | A | + | 2 | A | + | 3 | A | + | 4 | A | + | 5 | A | + | 6 | A | + | 7 | A | + | 8 | A | + | 9 | A | + | 0 | B | + | 1 | B | + | 2 | B | + | 3 | B | + | 4 | B | + | 5 | B | + | 6 | B | + | 7 | B | + | 8 | B | + | 9 | B | + +--------------+-------+ + "#); } /// It returns pending for the 2nd partition until the 3rd partition is polled. The 1st diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index f9f4b78686db..92138bf6a7a1 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -1222,9 +1222,8 @@ mod tests { }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use datafusion_common::{ - assert_batches_eq, exec_datafusion_err, Result, ScalarValue, - }; + use datafusion_common::test_util::batches_to_string; + use datafusion_common::{exec_datafusion_err, Result, ScalarValue}; use datafusion_execution::config::SessionConfig; use datafusion_execution::{ RecordBatchStream, SendableRecordBatchStream, TaskContext, @@ -1241,6 +1240,7 @@ mod tests { use futures::future::Shared; use futures::{pin_mut, ready, FutureExt, Stream, StreamExt}; + use insta::assert_snapshot; use itertools::Itertools; use tokio::time::timeout; @@ -1664,22 +1664,21 @@ mod tests { "\n**Optimized Plan Mismatch\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" ); - let expected = [ - "+---+------+---------------+---------------+", - "| a | last | nth_value(-1) | nth_value(-2) |", - "+---+------+---------------+---------------+", - "| 1 | 1 | 1 | |", - "| 2 | 2 | 2 | 1 |", - "| 3 | 3 | 3 | 2 |", - "| 1 | 1 | 1 | 3 |", - "| 2 | 2 | 2 | 1 |", - "| 3 | 3 | 3 | 2 |", - "| 1 | 1 | 1 | 3 |", - "| 2 | 2 | 2 | 1 |", - "| 3 | 3 | 3 | 2 |", - "+---+------+---------------+---------------+", - ]; - assert_batches_eq!(expected, &batches); + assert_snapshot!(batches_to_string(&batches), @r#" + +---+------+---------------+---------------+ + | a | last | nth_value(-1) | nth_value(-2) | + +---+------+---------------+---------------+ + | 1 | 1 | 1 | | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + | 1 | 1 | 1 | 3 | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + | 1 | 1 | 1 | 3 | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + +---+------+---------------+---------------+ + "#); Ok(()) } @@ -1792,21 +1791,20 @@ mod tests { let task_ctx = task_context(); let batches = collect_with_timeout(plan, task_ctx, timeout_duration).await?; - let expected = [ - "+----+------+-------+", - "| sn | hash | col_2 |", - "+----+------+-------+", - "| 0 | 2 | 2 |", - "| 1 | 2 | 2 |", - "| 2 | 2 | 2 |", - "| 3 | 2 | 1 |", - "| 4 | 1 | 2 |", - "| 5 | 1 | 2 |", - "| 6 | 1 | 2 |", - "| 7 | 1 | 1 |", - "+----+------+-------+", - ]; - assert_batches_eq!(expected, &batches); + assert_snapshot!(batches_to_string(&batches), @r#" + +----+------+-------+ + | sn | hash | col_2 | + +----+------+-------+ + | 0 | 2 | 2 | + | 1 | 2 | 2 | + | 2 | 2 | 2 | + | 3 | 2 | 1 | + | 4 | 1 | 2 | + | 5 | 1 | 2 | + | 6 | 1 | 2 | + | 7 | 1 | 1 | + +----+------+-------+ + "#); Ok(()) }