From cb087f9a571bcc44040af2a5519277294c5b599c Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Sat, 14 Aug 2021 06:41:10 +0000 Subject: [PATCH] Added display of tz and interval. --- src/array/display.rs | 107 ++++++++++++++++++++++++++++++------------- src/io/print.rs | 55 ++++++++++++++-------- 2 files changed, 110 insertions(+), 52 deletions(-) diff --git a/src/array/display.rs b/src/array/display.rs index 4d7483db40a..b90fe527953 100644 --- a/src/array/display.rs +++ b/src/array/display.rs @@ -1,7 +1,6 @@ use crate::{ array::*, datatypes::{DataType, IntervalUnit, TimeUnit}, - error::{ArrowError, Result}, temporal_conversions, }; @@ -25,18 +24,16 @@ macro_rules! dyn_dict { .downcast_ref::>() .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 String + 'a>> { +pub fn get_value_display<'a>(array: &'a dyn Array) -> Box 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::().unwrap(); @@ -61,32 +58,76 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result { 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)), @@ -123,7 +164,7 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result { let f = |x: Box| { - 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::>(); format!("[{}]", string_values.join(", ")) }; @@ -131,7 +172,7 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result { let f = |x: Box| { - 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::>(); format!("[{}]", string_values.join(", ")) }; @@ -139,13 +180,13 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result { let f = |x: Box| { - 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::>(); format!("[{}]", string_values.join(", ")) }; dyn_display!(array, ListArray, 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), @@ -162,7 +203,7 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result>>()?; + .collect::>(); Box::new(move |row: usize| { let mut string = displays .iter() @@ -181,19 +222,19 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Result 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 String + 'a>> { - let value_display = get_value_display(array)?; - Ok(Box::new(move |row| { +pub fn get_display<'a>(array: &'a dyn Array) -> Box String + 'a> { + let value_display = get_value_display(array); + Box::new(move |row| { if array.is_null(row) { "".to_string() } else { value_display(row) } - })) + }) } diff --git a/src/io/print.rs b/src/io/print.rs index 102ab355956..6a7b2c533de 100644 --- a/src/io/print.rs +++ b/src/io/print.rs @@ -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 { - 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 { +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(); @@ -52,7 +51,7 @@ fn create_table(results: &[RecordBatch]) -> Result
{ .columns() .iter() .map(|array| get_display(array.as_ref())) - .collect::>>()?; + .collect::>(); for row in 0..batch.num_rows() { let mut cells = Vec::new(); @@ -63,13 +62,12 @@ fn create_table(results: &[RecordBatch]) -> Result
{ 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; @@ -96,7 +94,7 @@ mod tests { ], )?; - let table = write(&[batch])?; + let table = write(&[batch]); let expected = vec![ "+---+-----+", @@ -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), @@ -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![ "+---+---+---+", @@ -150,6 +148,7 @@ mod tests { let actual: Vec<&str> = table.lines().collect(); assert_eq!(expected, actual, "Actual result:\n{:#?}", table); + Ok(()) } #[test] @@ -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![ "+-------+", @@ -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(); @@ -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![ @@ -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![ "+--------------+",