From a4f4de82981edd1cab9bacd59f3203173b892e83 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Tue, 2 Nov 2021 16:02:28 -0400 Subject: [PATCH] Second batch of updating sql tests to use assert_batches_eq (#1180) * First test updated * More tests * More tests * More tests * Move test porting comment * More tests * More tests * Move tracker --- datafusion/tests/sql.rs | 293 +++++++++++++++++++++++++++------------- 1 file changed, 200 insertions(+), 93 deletions(-) diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs index cf099193085e..dd9198c926a3 100644 --- a/datafusion/tests/sql.rs +++ b/datafusion/tests/sql.rs @@ -1809,8 +1809,6 @@ async fn csv_query_limit_zero() -> Result<()> { Ok(()) } -// --- End Test Porting --- - #[tokio::test] async fn csv_query_create_external_table() { let mut ctx = ExecutionContext::new(); @@ -1832,11 +1830,16 @@ async fn csv_query_external_table_count() { let mut ctx = ExecutionContext::new(); register_aggregate_csv_by_sql(&mut ctx).await; let sql = "SELECT COUNT(c12) FROM aggregate_test_100"; - let actual = execute(&mut ctx, sql).await; - // let a = execute_to_batches(&mut ctx, sql).await; - // println!("{}", pretty_format_batches(&a).unwrap()); - let expected = vec![vec!["100"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-------------------------------+", + "| COUNT(aggregate_test_100.c12) |", + "+-------------------------------+", + "| 100 |", + "+-------------------------------+", + ]; + + assert_batches_eq!(expected, &actual); } #[tokio::test] @@ -1846,9 +1849,15 @@ async fn csv_query_external_table_sum() { register_aggregate_csv_by_sql(&mut ctx).await; let sql = "SELECT SUM(CAST(c7 AS BIGINT)), SUM(CAST(c8 AS BIGINT)) FROM aggregate_test_100"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["13060", "3017641"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-------------------------------------------+-------------------------------------------+", + "| SUM(CAST(aggregate_test_100.c7 AS Int64)) | SUM(CAST(aggregate_test_100.c8 AS Int64)) |", + "+-------------------------------------------+-------------------------------------------+", + "| 13060 | 3017641 |", + "+-------------------------------------------+-------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); } #[tokio::test] @@ -1856,9 +1865,15 @@ async fn csv_query_count_star() { let mut ctx = ExecutionContext::new(); register_aggregate_csv_by_sql(&mut ctx).await; let sql = "SELECT COUNT(*) FROM aggregate_test_100"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["100"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-----------------+", + "| COUNT(UInt8(1)) |", + "+-----------------+", + "| 100 |", + "+-----------------+", + ]; + assert_batches_eq!(expected, &actual); } #[tokio::test] @@ -1866,9 +1881,15 @@ async fn csv_query_count_one() { let mut ctx = ExecutionContext::new(); register_aggregate_csv_by_sql(&mut ctx).await; let sql = "SELECT COUNT(1) FROM aggregate_test_100"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["100"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-----------------+", + "| COUNT(UInt8(1)) |", + "+-----------------+", + "| 100 |", + "+-----------------+", + ]; + assert_batches_eq!(expected, &actual); } #[tokio::test] @@ -1879,9 +1900,18 @@ async fn case_when() -> Result<()> { WHEN c1 = 'b' THEN 2 \ END \ FROM t1"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["1"], vec!["2"], vec!["NULL"], vec!["NULL"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+--------------------------------------------------------------------------------------+", + "| CASE WHEN #t1.c1 = Utf8(\"a\") THEN Int64(1) WHEN #t1.c1 = Utf8(\"b\") THEN Int64(2) END |", + "+--------------------------------------------------------------------------------------+", + "| 1 |", + "| 2 |", + "| |", + "| |", + "+--------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); Ok(()) } @@ -1893,9 +1923,18 @@ async fn case_when_else() -> Result<()> { WHEN c1 = 'b' THEN 2 \ ELSE 999 END \ FROM t1"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["1"], vec!["2"], vec!["999"], vec!["999"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+------------------------------------------------------------------------------------------------------+", + "| CASE WHEN #t1.c1 = Utf8(\"a\") THEN Int64(1) WHEN #t1.c1 = Utf8(\"b\") THEN Int64(2) ELSE Int64(999) END |", + "+------------------------------------------------------------------------------------------------------+", + "| 1 |", + "| 2 |", + "| 999 |", + "| 999 |", + "+------------------------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); Ok(()) } @@ -1907,9 +1946,18 @@ async fn case_when_with_base_expr() -> Result<()> { WHEN 'b' THEN 2 \ END \ FROM t1"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["1"], vec!["2"], vec!["NULL"], vec!["NULL"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+---------------------------------------------------------------------------+", + "| CASE #t1.c1 WHEN Utf8(\"a\") THEN Int64(1) WHEN Utf8(\"b\") THEN Int64(2) END |", + "+---------------------------------------------------------------------------+", + "| 1 |", + "| 2 |", + "| |", + "| |", + "+---------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); Ok(()) } @@ -1921,9 +1969,18 @@ async fn case_when_else_with_base_expr() -> Result<()> { WHEN 'b' THEN 2 \ ELSE 999 END \ FROM t1"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["1"], vec!["2"], vec!["999"], vec!["999"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-------------------------------------------------------------------------------------------+", + "| CASE #t1.c1 WHEN Utf8(\"a\") THEN Int64(1) WHEN Utf8(\"b\") THEN Int64(2) ELSE Int64(999) END |", + "+-------------------------------------------------------------------------------------------+", + "| 1 |", + "| 2 |", + "| 999 |", + "| 999 |", + "+-------------------------------------------------------------------------------------------+", + ]; + assert_batches_eq!(expected, &actual); Ok(()) } @@ -1952,13 +2009,17 @@ async fn equijoin() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 JOIN t2 ON t2_id = t1_id ORDER BY t1_id", ]; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["44", "d", "x"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 44 | d | x |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } let mut ctx = create_join_context_qualified()?; @@ -1966,10 +2027,18 @@ async fn equijoin() -> Result<()> { "SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t1.a = t2.a ORDER BY t1.a", "SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t2.a = t1.a ORDER BY t1.a", ]; - let expected = vec![vec!["1", "100"], vec!["2", "200"], vec!["4", "400"]]; + let expected = vec![ + "+---+-----+", + "| a | b |", + "+---+-----+", + "| 1 | 100 |", + "| 2 | 200 |", + "| 4 | 400 |", + "+---+-----+", + ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) } @@ -1984,13 +2053,17 @@ async fn equijoin_multiple_condition_ordering() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 JOIN t2 ON t2_id = t1_id AND t2_name <> t1_name ORDER BY t1_id", ]; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["44", "d", "x"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 44 | d | x |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) } @@ -2000,9 +2073,16 @@ async fn equijoin_and_other_condition() -> Result<()> { let mut ctx = create_join_context("t1_id", "t2_id")?; let sql = "SELECT t1_id, t1_name, t2_name FROM t1 JOIN t2 ON t1_id = t2_id AND t2_name >= 'y' ORDER BY t1_id"; - let actual = execute(&mut ctx, sql).await; - let expected = vec![vec!["11", "a", "z"], vec!["22", "b", "y"]]; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + let expected = vec![ + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "+-------+---------+---------+", + ]; + assert_batches_eq!(expected, &actual); Ok(()) } @@ -2013,15 +2093,18 @@ async fn equijoin_left_and_condition_from_right() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id AND t2_name >= 'y' ORDER BY t1_id"; let res = ctx.create_logical_plan(sql); assert!(res.is_ok()); - let actual = execute(&mut ctx, sql).await; - + let actual = execute_to_batches(&mut ctx, sql).await; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["33", "c", "NULL"], - vec!["44", "d", "NULL"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 33 | c | |", + "| 44 | d | |", + "+-------+---------+---------+", ]; - assert_eq!(expected, actual); + assert_batches_eq!(expected, &actual); Ok(()) } @@ -2033,16 +2116,18 @@ async fn equijoin_right_and_condition_from_left() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 RIGHT JOIN t2 ON t1_id = t2_id AND t1_id >= 22 ORDER BY t2_name"; let res = ctx.create_logical_plan(sql); assert!(res.is_ok()); - let actual = execute(&mut ctx, sql).await; - + let actual = execute_to_batches(&mut ctx, sql).await; let expected = vec![ - vec!["NULL", "NULL", "w"], - vec!["44", "d", "x"], - vec!["22", "b", "y"], - vec!["NULL", "NULL", "z"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| | | w |", + "| 44 | d | x |", + "| 22 | b | y |", + "| | | z |", + "+-------+---------+---------+", ]; - assert_eq!(expected, actual); - + assert_batches_eq!(expected, &actual); Ok(()) } @@ -2067,14 +2152,18 @@ async fn left_join() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t2_id = t1_id ORDER BY t1_id", ]; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["33", "c", "NULL"], - vec!["44", "d", "x"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 33 | c | |", + "| 44 | d | x |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) } @@ -2088,15 +2177,19 @@ async fn left_join_unbalanced() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t2_id = t1_id ORDER BY t1_id", ]; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["33", "c", "NULL"], - vec!["44", "d", "x"], - vec!["77", "e", "NULL"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 33 | c | |", + "| 44 | d | x |", + "| 77 | e | |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) } @@ -2109,14 +2202,18 @@ async fn right_join() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 RIGHT JOIN t2 ON t2_id = t1_id ORDER BY t1_id" ]; let expected = vec![ - vec!["NULL", "NULL", "w"], - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["44", "d", "x"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| | | w |", + "| 11 | a | z |", + "| 22 | b | y |", + "| 44 | d | x |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) } @@ -2129,15 +2226,19 @@ async fn full_join() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 FULL JOIN t2 ON t2_id = t1_id ORDER BY t1_id", ]; let expected = vec![ - vec!["NULL", "NULL", "w"], - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["33", "c", "NULL"], - vec!["44", "d", "x"], + "+-------+---------+---------+", + "| t1_id | t1_name | t2_name |", + "+-------+---------+---------+", + "| | | w |", + "| 11 | a | z |", + "| 22 | b | y |", + "| 33 | c | |", + "| 44 | d | x |", + "+-------+---------+---------+", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } let equivalent_sql = [ @@ -2145,8 +2246,8 @@ async fn full_join() -> Result<()> { "SELECT t1_id, t1_name, t2_name FROM t1 FULL OUTER JOIN t2 ON t2_id = t1_id ORDER BY t1_id", ]; for sql in equivalent_sql.iter() { - let actual = execute(&mut ctx, sql).await; - assert_eq!(expected, actual); + let actual = execute_to_batches(&mut ctx, sql).await; + assert_batches_eq!(expected, &actual); } Ok(()) @@ -2156,17 +2257,23 @@ async fn full_join() -> Result<()> { async fn left_join_using() -> Result<()> { let mut ctx = create_join_context("id", "id")?; let sql = "SELECT id, t1_name, t2_name FROM t1 LEFT JOIN t2 USING (id) ORDER BY id"; - let actual = execute(&mut ctx, sql).await; + let actual = execute_to_batches(&mut ctx, sql).await; let expected = vec![ - vec!["11", "a", "z"], - vec!["22", "b", "y"], - vec!["33", "c", "NULL"], - vec!["44", "d", "x"], + "+----+---------+---------+", + "| id | t1_name | t2_name |", + "+----+---------+---------+", + "| 11 | a | z |", + "| 22 | b | y |", + "| 33 | c | |", + "| 44 | d | x |", + "+----+---------+---------+", ]; - assert_eq!(expected, actual); + assert_batches_eq!(expected, &actual); Ok(()) } +// --- End Test Porting --- + #[tokio::test] async fn equijoin_implicit_syntax() -> Result<()> { let mut ctx = create_join_context("t1_id", "t2_id")?;