diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs index e4841529ab83..9998fc1483aa 100644 --- a/datafusion/tests/sql.rs +++ b/datafusion/tests/sql.rs @@ -1849,8 +1849,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(); @@ -1872,11 +1870,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] @@ -1886,9 +1889,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] @@ -1896,9 +1905,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] @@ -1906,9 +1921,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] @@ -1919,9 +1940,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(()) } @@ -1933,9 +1963,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(()) } @@ -1947,9 +1986,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(()) } @@ -1961,9 +2009,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(()) } @@ -1992,13 +2049,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()?; @@ -2006,10 +2067,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(()) } @@ -2024,13 +2093,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(()) } @@ -2040,9 +2113,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(()) } @@ -2053,15 +2133,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(()) } @@ -2073,16 +2156,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(()) } @@ -2107,14 +2192,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(()) } @@ -2128,15 +2217,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(()) } @@ -2149,14 +2242,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(()) } @@ -2169,15 +2266,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 = [ @@ -2185,8 +2286,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(()) @@ -2196,17 +2297,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")?;