use crate::record_batch::RecordBatch;
use prettytable::format;
use prettytable::{Cell, Row, Table};
use crate::error::Result;
use super::display::array_value_to_string;
pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<String> {
Ok(create_table(results)?.to_string())
}
pub fn print_batches(results: &[RecordBatch]) -> Result<()> {
create_table(results)?.printstd();
Ok(())
}
fn create_table(results: &[RecordBatch]) -> Result<Table> {
let mut table = Table::new();
table.set_format(*format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
if results.is_empty() {
return Ok(table);
}
let schema = results[0].schema();
let mut header = Vec::new();
for field in schema.fields() {
header.push(Cell::new(&field.name()));
}
table.set_titles(Row::new(header));
for batch in results {
for row in 0..batch.num_rows() {
let mut cells = Vec::new();
for col in 0..batch.num_columns() {
let column = batch.column(col);
cells.push(Cell::new(&array_value_to_string(&column, row)?));
}
table.add_row(Row::new(cells));
}
}
Ok(table)
}
#[cfg(test)]
mod tests {
use crate::array::{self, PrimitiveBuilder, StringBuilder, StringDictionaryBuilder};
use super::*;
use crate::datatypes::{DataType, Field, Int32Type, Schema};
use std::sync::Arc;
#[test]
fn test_pretty_format_batches() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Utf8, true),
Field::new("b", DataType::Int32, true),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(array::StringArray::from(vec![
Some("a"),
Some("b"),
None,
Some("d"),
])),
Arc::new(array::Int32Array::from(vec![
Some(1),
None,
Some(10),
Some(100),
])),
],
)?;
let table = pretty_format_batches(&[batch])?;
let expected = vec![
"+---+-----+",
"| a | b |",
"+---+-----+",
"| a | 1 |",
"| b | |",
"| | 10 |",
"| d | 100 |",
"+---+-----+",
];
let actual: Vec<&str> = table.lines().collect();
assert_eq!(expected, actual, "Actual result:\n{}", table);
Ok(())
}
#[test]
fn test_pretty_format_dictionary() -> Result<()> {
let field_type =
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
let keys_builder = PrimitiveBuilder::<Int32Type>::new(10);
let values_builder = StringBuilder::new(10);
let mut builder = StringDictionaryBuilder::new(keys_builder, values_builder);
builder.append("one")?;
builder.append_null()?;
builder.append("three")?;
let array = Arc::new(builder.finish());
let batch = RecordBatch::try_new(schema, vec![array])?;
let table = pretty_format_batches(&[batch])?;
let expected = vec![
"+-------+",
"| d1 |",
"+-------+",
"| one |",
"| |",
"| three |",
"+-------+",
];
let actual: Vec<&str> = table.lines().collect();
assert_eq!(expected, actual, "Actual result:\n{}", table);
Ok(())
}
}