Skip to content

Commit

Permalink
Added Decimal support to pretty-print display utility (#230) (#273)
Browse files Browse the repository at this point in the history
* Added Decimal support to pretty-print display utility (#230)

* Applied cargo fmt to fix linting errors

* Added proper printing for decimals based on scale, moved tests to pretty.rs

* Applied cargo fmt on pretty test

Co-authored-by: Manish Gill <manish.gill@tomtom.com>
  • Loading branch information
mgill25 and manishgill-tomtom authored May 14, 2021
1 parent ce8e67c commit 4449ee9
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 0 deletions.
17 changes: 17 additions & 0 deletions arrow/src/util/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,20 @@ macro_rules! make_string_from_list {
}};
}

macro_rules! make_string_from_decimal {
($array_type: ty, $column: ident, $row: ident, $scale: ident) => {{
let array = $column.as_any().downcast_ref::<$array_type>().unwrap();
let decimal_string = array.value($row).to_string();
let formatted_decimal = if *$scale == 0 {
decimal_string
} else {
let splits = decimal_string.split_at(decimal_string.len() - *$scale);
format!("{}.{}", splits.0, splits.1)
};
Ok(formatted_decimal)
}};
}

/// Get the value at the given row in an array as a String.
///
/// Note this function is quite inefficient and is unlikely to be
Expand All @@ -217,6 +231,9 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
DataType::Float16 => make_string!(array::Float32Array, column, row),
DataType::Float32 => make_string!(array::Float32Array, column, row),
DataType::Float64 => make_string!(array::Float64Array, column, row),
DataType::Decimal(_, scale) => {
make_string_from_decimal!(array::DecimalArray, column, row, scale)
}
DataType::Timestamp(unit, _) if *unit == TimeUnit::Second => {
make_string_datetime!(array::TimestampSecondArray, column, row)
}
Expand Down
89 changes: 89 additions & 0 deletions arrow/src/util/pretty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ mod tests {
};

use super::*;
use crate::array::{DecimalBuilder, Int32Array};
use std::sync::Arc;

#[test]
Expand Down Expand Up @@ -418,4 +419,92 @@ mod tests {
];
check_datetime!(Time64NanosecondArray, 11111111, expected);
}

#[test]
fn test_int_display() -> Result<()> {
let array = Arc::new(Int32Array::from(vec![6, 3])) as ArrayRef;
let actual_one = array_value_to_string(&array, 0).unwrap();
let expected_one = "6";

let actual_two = array_value_to_string(&array, 1).unwrap();
let expected_two = "3";
assert_eq!(actual_one, expected_one);
assert_eq!(actual_two, expected_two);
Ok(())
}

#[test]
fn test_decimal_display() -> Result<()> {
let capacity = 10;
let precision = 10;
let scale = 2;

let mut builder = DecimalBuilder::new(capacity, precision, scale);
builder.append_value(101).unwrap();
builder.append_null().unwrap();
builder.append_value(200).unwrap();
builder.append_value(3040).unwrap();

let dm = Arc::new(builder.finish()) as ArrayRef;

let schema = Arc::new(Schema::new(vec![Field::new(
"f",
dm.data_type().clone(),
true,
)]));

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

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

let expected = vec![
"+-------+",
"| f |",
"+-------+",
"| 1.01 |",
"| |",
"| 2.00 |",
"| 30.40 |",
"+-------+",
];

let actual: Vec<&str> = table.lines().collect();
assert_eq!(expected, actual, "Actual result:\n{}", table);

Ok(())
}

#[test]
fn test_decimal_display_zero_scale() -> Result<()> {
let capacity = 10;
let precision = 5;
let scale = 0;

let mut builder = DecimalBuilder::new(capacity, precision, scale);
builder.append_value(101).unwrap();
builder.append_null().unwrap();
builder.append_value(200).unwrap();
builder.append_value(3040).unwrap();

let dm = Arc::new(builder.finish()) as ArrayRef;

let schema = Arc::new(Schema::new(vec![Field::new(
"f",
dm.data_type().clone(),
true,
)]));

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

let table = pretty_format_batches(&[batch])?;
let expected = vec![
"+------+", "| f |", "+------+", "| 101 |", "| |", "| 200 |",
"| 3040 |", "+------+",
];

let actual: Vec<&str> = table.lines().collect();
assert_eq!(expected, actual, "Actual result:\n{}", table);

Ok(())
}
}

0 comments on commit 4449ee9

Please sign in to comment.