Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Second batch of updating sql tests to use assert_batches_eq #1180

Merged
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 119 additions & 42 deletions datafusion/tests/sql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ extern crate datafusion;

use arrow::{
array::*, datatypes::*, record_batch::RecordBatch,
util::display::array_value_to_string,
util::display::array_value_to_string, util::pretty::pretty_format_batches,
};

use datafusion::assert_batches_eq;
Expand Down Expand Up @@ -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();
Expand All @@ -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]
Expand All @@ -1886,29 +1889,47 @@ 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]
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]
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]
Expand All @@ -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(())
}

Expand All @@ -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(())
}

Expand All @@ -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(())
}

Expand All @@ -1961,12 +2009,23 @@ 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(())
}

// --- End Test Porting ---

fn create_case_context() -> Result<ExecutionContext> {
let mut ctx = ExecutionContext::new();
let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)]));
Expand All @@ -1992,24 +2051,36 @@ 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()?;
let equivalent_sql = [
"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(())
}
Expand All @@ -2024,13 +2095,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(())
}
Expand All @@ -2041,6 +2116,8 @@ async fn equijoin_and_other_condition() -> Result<()> {
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 a = execute_to_batches(&mut ctx, sql).await;
// println!("{}", pretty_format_batches(&a).unwrap());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you either remove this or put a todo on what to do next?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i just caught up the tracking comment that is used for test porting:

// --- End Test Porting ---

Let me know when you think its enough updates on this batch / PR to merge and ill clean up those helpers so it can be merged.

let expected = vec![vec!["11", "a", "z"], vec!["22", "b", "y"]];
assert_eq!(expected, actual);
Ok(())
Expand Down