Skip to content

Commit 8a17c18

Browse files
alambalamb
authored
Move assert_batches_eq! macros to test_utils.rs (#746)
* Move assert_batches_eq! macros to test_utils.rs * port test
1 parent 3fb600d commit 8a17c18

File tree

3 files changed

+92
-82
lines changed

3 files changed

+92
-82
lines changed

‎datafusion/src/test/mod.rs‎

Lines changed: 0 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -280,72 +280,3 @@ pub fn make_timestamps() -> RecordBatch {
280280
pub mod exec;
281281
pub mod user_defined;
282282
pub mod variable;
283-
284-
/// Compares formatted output of a record batch with an expected
285-
/// vector of strings, with the result of pretty formatting record
286-
/// batches. This is a macro so errors appear on the correct line
287-
///
288-
/// Designed so that failure output can be directly copy/pasted
289-
/// into the test code as expected results.
290-
///
291-
/// Expects to be called about like this:
292-
///
293-
/// `assert_batch_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
294-
#[macro_export]
295-
macro_rules! assert_batches_eq {
296-
($EXPECTED_LINES: expr, $CHUNKS: expr) => {
297-
let expected_lines: Vec<String> =
298-
$EXPECTED_LINES.iter().map(|&s| s.into()).collect();
299-
300-
let formatted = arrow::util::pretty::pretty_format_batches($CHUNKS).unwrap();
301-
302-
let actual_lines: Vec<&str> = formatted.trim().lines().collect();
303-
304-
assert_eq!(
305-
expected_lines, actual_lines,
306-
"\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
307-
expected_lines, actual_lines
308-
);
309-
};
310-
}
311-
312-
/// Compares formatted output of a record batch with an expected
313-
/// vector of strings in a way that order does not matter.
314-
/// This is a macro so errors appear on the correct line
315-
///
316-
/// Designed so that failure output can be directly copy/pasted
317-
/// into the test code as expected results.
318-
///
319-
/// Expects to be called about like this:
320-
///
321-
/// `assert_batch_sorted_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
322-
#[macro_export]
323-
macro_rules! assert_batches_sorted_eq {
324-
($EXPECTED_LINES: expr, $CHUNKS: expr) => {
325-
let mut expected_lines: Vec<String> =
326-
$EXPECTED_LINES.iter().map(|&s| s.into()).collect();
327-
328-
// sort except for header + footer
329-
let num_lines = expected_lines.len();
330-
if num_lines > 3 {
331-
expected_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
332-
}
333-
334-
let formatted = arrow::util::pretty::pretty_format_batches($CHUNKS).unwrap();
335-
// fix for windows: \r\n -->
336-
337-
let mut actual_lines: Vec<&str> = formatted.trim().lines().collect();
338-
339-
// sort except for header + footer
340-
let num_lines = actual_lines.len();
341-
if num_lines > 3 {
342-
actual_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
343-
}
344-
345-
assert_eq!(
346-
expected_lines, actual_lines,
347-
"\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
348-
expected_lines, actual_lines
349-
);
350-
};
351-
}

‎datafusion/src/test_util.rs‎

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,79 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
//! Utils to make testing easier
18+
//! Utility functions to make testing DataFusion based crates easier
1919
2020
use std::{env, error::Error, path::PathBuf};
2121

22+
/// Compares formatted output of a record batch with an expected
23+
/// vector of strings, with the result of pretty formatting record
24+
/// batches. This is a macro so errors appear on the correct line
25+
///
26+
/// Designed so that failure output can be directly copy/pasted
27+
/// into the test code as expected results.
28+
///
29+
/// Expects to be called about like this:
30+
///
31+
/// `assert_batch_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
32+
#[macro_export]
33+
macro_rules! assert_batches_eq {
34+
($EXPECTED_LINES: expr, $CHUNKS: expr) => {
35+
let expected_lines: Vec<String> =
36+
$EXPECTED_LINES.iter().map(|&s| s.into()).collect();
37+
38+
let formatted = arrow::util::pretty::pretty_format_batches($CHUNKS).unwrap();
39+
40+
let actual_lines: Vec<&str> = formatted.trim().lines().collect();
41+
42+
assert_eq!(
43+
expected_lines, actual_lines,
44+
"\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
45+
expected_lines, actual_lines
46+
);
47+
};
48+
}
49+
50+
/// Compares formatted output of a record batch with an expected
51+
/// vector of strings in a way that order does not matter.
52+
/// This is a macro so errors appear on the correct line
53+
///
54+
/// Designed so that failure output can be directly copy/pasted
55+
/// into the test code as expected results.
56+
///
57+
/// Expects to be called about like this:
58+
///
59+
/// `assert_batch_sorted_eq!(expected_lines: &[&str], batches: &[RecordBatch])`
60+
#[macro_export]
61+
macro_rules! assert_batches_sorted_eq {
62+
($EXPECTED_LINES: expr, $CHUNKS: expr) => {
63+
let mut expected_lines: Vec<String> =
64+
$EXPECTED_LINES.iter().map(|&s| s.into()).collect();
65+
66+
// sort except for header + footer
67+
let num_lines = expected_lines.len();
68+
if num_lines > 3 {
69+
expected_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
70+
}
71+
72+
let formatted = arrow::util::pretty::pretty_format_batches($CHUNKS).unwrap();
73+
// fix for windows: \r\n -->
74+
75+
let mut actual_lines: Vec<&str> = formatted.trim().lines().collect();
76+
77+
// sort except for header + footer
78+
let num_lines = actual_lines.len();
79+
if num_lines > 3 {
80+
actual_lines.as_mut_slice()[2..num_lines - 1].sort_unstable()
81+
}
82+
83+
assert_eq!(
84+
expected_lines, actual_lines,
85+
"\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
86+
expected_lines, actual_lines
87+
);
88+
};
89+
}
90+
2291
/// Returns the arrow test data directory, which is by default stored
2392
/// in a git submodule rooted at `testing/data`.
2493
///

‎datafusion/tests/sql.rs‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ use arrow::{
3535
util::display::array_value_to_string,
3636
};
3737

38+
use datafusion::assert_batches_eq;
3839
use datafusion::logical_plan::LogicalPlan;
3940
use datafusion::prelude::*;
4041
use datafusion::{
@@ -112,19 +113,23 @@ async fn parquet_query() {
112113
// NOTE that string_col is actually a binary column and does not have the UTF8 logical type
113114
// so we need an explicit cast
114115
let sql = "SELECT id, CAST(string_col AS varchar) FROM alltypes_plain";
115-
let actual = execute(&mut ctx, sql).await;
116+
let actual = execute_to_batches(&mut ctx, sql).await;
116117
let expected = vec![
117-
vec!["4", "0"],
118-
vec!["5", "1"],
119-
vec!["6", "0"],
120-
vec!["7", "1"],
121-
vec!["2", "0"],
122-
vec!["3", "1"],
123-
vec!["0", "0"],
124-
vec!["1", "1"],
118+
"+----+--------------------------+",
119+
"| id | CAST(string_col AS Utf8) |",
120+
"+----+--------------------------+",
121+
"| 4 | 0 |",
122+
"| 5 | 1 |",
123+
"| 6 | 0 |",
124+
"| 7 | 1 |",
125+
"| 2 | 0 |",
126+
"| 3 | 1 |",
127+
"| 0 | 0 |",
128+
"| 1 | 1 |",
129+
"+----+--------------------------+",
125130
];
126131

127-
assert_eq!(expected, actual);
132+
assert_batches_eq!(expected, &actual);
128133
}
129134

130135
#[tokio::test]
@@ -2487,7 +2492,7 @@ fn register_alltypes_parquet(ctx: &mut ExecutionContext) {
24872492

24882493
/// Execute query and return result set as 2-d table of Vecs
24892494
/// `result[row][column]`
2490-
async fn execute(ctx: &mut ExecutionContext, sql: &str) -> Vec<Vec<String>> {
2495+
async fn execute_to_batches(ctx: &mut ExecutionContext, sql: &str) -> Vec<RecordBatch> {
24912496
let msg = format!("Creating logical plan for '{}'", sql);
24922497
let plan = ctx.create_logical_plan(sql).expect(&msg);
24932498
let logical_schema = plan.schema();
@@ -2503,8 +2508,13 @@ async fn execute(ctx: &mut ExecutionContext, sql: &str) -> Vec<Vec<String>> {
25032508
let results = collect(plan).await.expect(&msg);
25042509

25052510
assert_eq!(logical_schema.as_ref(), optimized_logical_schema.as_ref());
2511+
results
2512+
}
25062513

2507-
result_vec(&results)
2514+
/// Execute query and return result set as 2-d table of Vecs
2515+
/// `result[row][column]`
2516+
async fn execute(ctx: &mut ExecutionContext, sql: &str) -> Vec<Vec<String>> {
2517+
result_vec(&execute_to_batches(ctx, sql).await)
25082518
}
25092519

25102520
/// Specialised String representation

0 commit comments

Comments
 (0)