alamb commented on code in PR #8896:
URL: https://github.com/apache/arrow-datafusion/pull/8896#discussion_r1455802246
##########
datafusion-cli/src/print_format.rs:
##########
@@ -161,152 +161,263 @@ impl PrintFormat {
maxrows: MaxRows,
with_header: bool,
) -> Result<()> {
- if batches.is_empty() || batches[0].num_rows() == 0 {
+ // filter out any empty batches
+ let batches: Vec<_> = batches
+ .iter()
+ .filter(|b| b.num_rows() > 0)
+ .cloned()
+ .collect();
+ if batches.is_empty() {
return Ok(());
}
match self {
Self::Csv | Self::Automatic => {
- print_batches_with_sep(writer, batches, b',', with_header)
+ print_batches_with_sep(writer, &batches, b',', with_header)
}
- Self::Tsv => print_batches_with_sep(writer, batches, b'\t',
with_header),
+ Self::Tsv => print_batches_with_sep(writer, &batches, b'\t',
with_header),
Self::Table => {
if maxrows == MaxRows::Limited(0) {
return Ok(());
}
- format_batches_with_maxrows(writer, batches, maxrows)
+ format_batches_with_maxrows(writer, &batches, maxrows)
}
- Self::Json => batches_to_json!(ArrayWriter, writer, batches),
- Self::NdJson => batches_to_json!(LineDelimitedWriter, writer,
batches),
+ Self::Json => batches_to_json!(ArrayWriter, writer, &batches),
+ Self::NdJson => batches_to_json!(LineDelimitedWriter, writer,
&batches),
}
}
}
#[cfg(test)]
mod tests {
- use std::io::{Cursor, Read, Write};
- use std::sync::Arc;
-
use super::*;
+ use std::sync::Arc;
- use arrow::array::Int32Array;
+ use arrow::array::{ArrayRef, Int32Array};
use arrow::datatypes::{DataType, Field, Schema};
- use datafusion::error::Result;
-
- fn run_test<F>(batches: &[RecordBatch], test_fn: F) -> Result<String>
- where
- F: Fn(&mut Cursor<Vec<u8>>, &[RecordBatch]) -> Result<()>,
- {
- let mut buffer = Cursor::new(Vec::new());
- test_fn(&mut buffer, batches)?;
- buffer.set_position(0);
- let mut contents = String::new();
- buffer.read_to_string(&mut contents)?;
- Ok(contents)
- }
#[test]
- fn test_print_batches_with_sep() -> Result<()> {
- let contents = run_test(&[], |buffer, batches| {
- print_batches_with_sep(buffer, batches, b',', true)
- })?;
- assert_eq!(contents, "");
-
- let schema = Arc::new(Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, false),
- Field::new("c", DataType::Int32, false),
- ]));
- let batch = RecordBatch::try_new(
- schema,
- vec![
- Arc::new(Int32Array::from(vec![1, 2, 3])),
- Arc::new(Int32Array::from(vec![4, 5, 6])),
- Arc::new(Int32Array::from(vec![7, 8, 9])),
- ],
- )?;
-
- let contents = run_test(&[batch], |buffer, batches| {
- print_batches_with_sep(buffer, batches, b',', true)
- })?;
- assert_eq!(contents, "a,b,c\n1,4,7\n2,5,8\n3,6,9\n");
-
- Ok(())
+ fn print_empty() {
+ for format in [
+ PrintFormat::Csv,
+ PrintFormat::Tsv,
+ PrintFormat::Table,
+ PrintFormat::Json,
+ PrintFormat::NdJson,
+ PrintFormat::Automatic,
+ ] {
+ // no output for empty batches, even with header set
+ PrintBatchesTest::new()
Review Comment:
I think this style of test is clearer what is expected and what is covered
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
