Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Added display of tz and interval.
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorgecarleitao
jorgecarleitao committed Aug 14, 2021
1 parent 8a5df12 commit cb087f9
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 52 deletions.
107 changes: 74 additions & 33 deletions src/array/display.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use crate::{
array::*,
datatypes::{DataType, IntervalUnit, TimeUnit},
error::{ArrowError, Result},
temporal_conversions,
};

Expand All @@ -25,18 +24,16 @@ macro_rules! dyn_dict {
.downcast_ref::<DictionaryArray<$ty>>()
.unwrap();
let keys = a.keys();
let display = get_value_display(a.values().as_ref())?;
let display = get_value_display(a.values().as_ref());
Box::new(move |row: usize| display(keys.value(row) as usize))
}};
}

/// Returns a function of index returning the string representation of the _value_ of `array`.
/// This does not take nulls into account.
/// # Errors
/// This function errors iff the datatype is not yet supported for printing.
pub fn get_value_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -> String + 'a>> {
pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> String + 'a> {
use DataType::*;
Ok(match array.data_type() {
match array.data_type() {
Null => Box::new(|_: usize| "".to_string()),
Boolean => {
let a = array.as_any().downcast_ref::<BooleanArray>().unwrap();
Expand All @@ -61,32 +58,76 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -
Time32(TimeUnit::Millisecond) => {
dyn_primitive!(array, i32, temporal_conversions::time32ms_to_time)
}
Time32(_) => unreachable!(), // remaining are not valid
Time64(TimeUnit::Microsecond) => {
dyn_primitive!(array, i64, temporal_conversions::time64us_to_time)
}
Time64(TimeUnit::Nanosecond) => {
dyn_primitive!(array, i64, temporal_conversions::time64ns_to_time)
}
Timestamp(TimeUnit::Second, None) => {
dyn_primitive!(array, i64, temporal_conversions::timestamp_s_to_datetime)
}
Timestamp(TimeUnit::Millisecond, None) => {
dyn_primitive!(array, i64, temporal_conversions::timestamp_ms_to_datetime)
}
Timestamp(TimeUnit::Microsecond, None) => {
dyn_primitive!(array, i64, temporal_conversions::timestamp_us_to_datetime)
}
Timestamp(TimeUnit::Nanosecond, None) => {
dyn_primitive!(array, i64, temporal_conversions::timestamp_ns_to_datetime)
}
Timestamp(_, _) => {
return Err(ArrowError::NotYetImplemented(
"Priting of timestamps with timezones is not yet implemented.".to_string(),
))
Time64(_) => unreachable!(), // remaining are not valid
Timestamp(TimeUnit::Second, tz) => {
if let Some(tz) = tz {
dyn_primitive!(array, i64, |x| {
format!(
"{} {}",
temporal_conversions::timestamp_s_to_datetime(x),
tz
)
})
} else {
dyn_primitive!(array, i64, temporal_conversions::timestamp_s_to_datetime)
}
}
Timestamp(TimeUnit::Millisecond, tz) => {
if let Some(tz) = tz {
dyn_primitive!(array, i64, |x| {
format!(
"{} {}",
temporal_conversions::timestamp_ms_to_datetime(x),
tz
)
})
} else {
dyn_primitive!(array, i64, temporal_conversions::timestamp_ms_to_datetime)
}
}
Timestamp(TimeUnit::Microsecond, tz) => {
if let Some(tz) = tz {
dyn_primitive!(array, i64, |x| {
format!(
"{} {}",
temporal_conversions::timestamp_us_to_datetime(x),
tz
)
})
} else {
dyn_primitive!(array, i64, temporal_conversions::timestamp_us_to_datetime)
}
}
Timestamp(TimeUnit::Nanosecond, tz) => {
if let Some(tz) = tz {
dyn_primitive!(array, i64, |x| {
format!(
"{} {}",
temporal_conversions::timestamp_ns_to_datetime(x),
tz
)
})
} else {
dyn_primitive!(array, i64, temporal_conversions::timestamp_ns_to_datetime)
}
}
Interval(IntervalUnit::YearMonth) => {
dyn_primitive!(array, i32, |x| format!("{}m", x))
}
Interval(IntervalUnit::DayTime) => {
dyn_primitive!(array, days_ms, |x: days_ms| format!(
"{}d{}ms",
x.days(),
x.milliseconds()
))
}
Duration(TimeUnit::Second) => dyn_primitive!(array, i64, |x| format!("{}s", x)),
Duration(TimeUnit::Millisecond) => dyn_primitive!(array, i64, |x| format!("{}ms", x)),
Duration(TimeUnit::Microsecond) => dyn_primitive!(array, i64, |x| format!("{}us", x)),
Expand Down Expand Up @@ -123,29 +164,29 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -
}
List(_) => {
let f = |x: Box<dyn Array>| {
let display = get_value_display(x.as_ref()).unwrap();
let display = get_value_display(x.as_ref());
let string_values = (0..x.len()).map(|i| display(i)).collect::<Vec<String>>();
format!("[{}]", string_values.join(", "))
};
dyn_display!(array, ListArray<i32>, f)
}
FixedSizeList(_, _) => {
let f = |x: Box<dyn Array>| {
let display = get_value_display(x.as_ref()).unwrap();
let display = get_value_display(x.as_ref());
let string_values = (0..x.len()).map(|i| display(i)).collect::<Vec<String>>();
format!("[{}]", string_values.join(", "))
};
dyn_display!(array, FixedSizeListArray, f)
}
LargeList(_) => {
let f = |x: Box<dyn Array>| {
let display = get_value_display(x.as_ref()).unwrap();
let display = get_value_display(x.as_ref());
let string_values = (0..x.len()).map(|i| display(i)).collect::<Vec<String>>();
format!("[{}]", string_values.join(", "))
};
dyn_display!(array, ListArray<i64>, f)
}
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => dyn_dict!(array, i8),
DataType::Int16 => dyn_dict!(array, i16),
DataType::Int32 => dyn_dict!(array, i32),
Expand All @@ -162,7 +203,7 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -
.values()
.iter()
.map(|x| get_value_display(x.as_ref()))
.collect::<Result<Vec<_>>>()?;
.collect::<Vec<_>>();
Box::new(move |row: usize| {
let mut string = displays
.iter()
Expand All @@ -181,19 +222,19 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -
string
})
}
_ => unreachable!(),
})
Union(_) => todo!(),
}
}

/// Returns a function of index returning the string representation of the item of `array`.
/// This outputs an empty string on nulls.
pub fn get_display<'a>(array: &'a dyn Array) -> Result<Box<dyn Fn(usize) -> String + 'a>> {
let value_display = get_value_display(array)?;
Ok(Box::new(move |row| {
pub fn get_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> String + 'a> {
let value_display = get_value_display(array);
Box::new(move |row| {
if array.is_null(row) {
"".to_string()
} else {
value_display(row)
}
}))
})
}
55 changes: 36 additions & 19 deletions src/io/print.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,28 +15,27 @@
// specific language governing permissions and limitations
// under the License.

use crate::{array::*, error::Result, record_batch::RecordBatch};
use crate::{array::*, record_batch::RecordBatch};

use comfy_table::{Cell, Table};

/// Returns a visual representation of multiple [`RecordBatch`]es.
pub fn write(batches: &[RecordBatch]) -> Result<String> {
Ok(create_table(batches)?.to_string())
pub fn write(batches: &[RecordBatch]) -> String {
create_table(batches).to_string()
}

/// Prints a visual representation of record batches to stdout
pub fn print(results: &[RecordBatch]) -> Result<()> {
println!("{}", create_table(results)?);
Ok(())
pub fn print(results: &[RecordBatch]) {
println!("{}", create_table(results))
}

/// Convert a series of record batches into a table
fn create_table(results: &[RecordBatch]) -> Result<Table> {
fn create_table(results: &[RecordBatch]) -> Table {
let mut table = Table::new();
table.load_preset("||--+-++| ++++++");

if results.is_empty() {
return Ok(table);
return table;
}

let schema = results[0].schema();
Expand All @@ -52,7 +51,7 @@ fn create_table(results: &[RecordBatch]) -> Result<Table> {
.columns()
.iter()
.map(|array| get_display(array.as_ref()))
.collect::<Result<Vec<_>>>()?;
.collect::<Vec<_>>();

for row in 0..batch.num_rows() {
let mut cells = Vec::new();
Expand All @@ -63,13 +62,12 @@ fn create_table(results: &[RecordBatch]) -> Result<Table> {
table.add_row(cells);
}
}

Ok(table)
table
}

#[cfg(test)]
mod tests {
use crate::{array::*, bitmap::Bitmap, datatypes::*};
use crate::{array::*, bitmap::Bitmap, datatypes::*, error::Result};

use super::*;
use std::sync::Arc;
Expand All @@ -96,7 +94,7 @@ mod tests {
],
)?;

let table = write(&[batch])?;
let table = write(&[batch]);

let expected = vec![
"+---+-----+",
Expand All @@ -117,7 +115,7 @@ mod tests {
}

#[test]
fn test_write_null() {
fn test_write_null() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::Utf8, true),
Field::new("b", DataType::Int32, true),
Expand All @@ -132,9 +130,9 @@ mod tests {
.collect();

// define data (null)
let batch = RecordBatch::try_new(schema, arrays).unwrap();
let batch = RecordBatch::try_new(schema, arrays)?;

let table = write(&[batch]).unwrap();
let table = write(&[batch]);

let expected = vec![
"+---+---+---+",
Expand All @@ -150,6 +148,7 @@ mod tests {
let actual: Vec<&str> = table.lines().collect();

assert_eq!(expected, actual, "Actual result:\n{:#?}", table);
Ok(())
}

#[test]
Expand All @@ -165,7 +164,7 @@ mod tests {

let batch = RecordBatch::try_new(schema, vec![array])?;

let table = write(&[batch])?;
let table = write(&[batch]);

let expected = vec![
"+-------+",
Expand Down Expand Up @@ -198,7 +197,7 @@ mod tests {
)]));
let batch = RecordBatch::try_new(schema, vec![array]).unwrap();

let table = write(&[batch]).expect("formatting batches");
let table = write(&[batch]);

let expected = $EXPECTED_RESULT;
let actual: Vec<&str> = table.lines().collect();
Expand All @@ -225,6 +224,24 @@ mod tests {
);
}

#[test]
fn test_write_timestamp_second_with_tz() {
let expected = vec![
"+-------------------------+",
"| f |",
"+-------------------------+",
"| 1970-05-09 14:25:11 UTC |",
"| |",
"+-------------------------+",
];
check_datetime!(
i64,
DataType::Timestamp(TimeUnit::Second, Some("UTC".to_string())),
11111111,
expected
);
}

#[test]
fn test_write_timestamp_millisecond() {
let expected = vec![
Expand Down Expand Up @@ -391,7 +408,7 @@ mod tests {

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)])?;

let table = write(&[batch])?;
let table = write(&[batch]);

let expected = vec![
"+--------------+",
Expand Down

0 comments on commit cb087f9

Please sign in to comment.