From 4449ee96fe3fd4a0b275da8dd25ce2792699bc98 Mon Sep 17 00:00:00 2001
From: Manish Gill <gill.manish90@gmail.com>
Date: Fri, 14 May 2021 20:42:47 +0200
Subject: [PATCH] Added Decimal support to pretty-print display utility (#230)
 (#273)

* 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>
---
 arrow/src/util/display.rs | 17 ++++++++
 arrow/src/util/pretty.rs  | 89 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs
index e40ababd2338..61f549af7dd2 100644
--- a/arrow/src/util/display.rs
+++ b/arrow/src/util/display.rs
@@ -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
@@ -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)
         }
diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs
index f354899c1dfd..d307406289e2 100644
--- a/arrow/src/util/pretty.rs
+++ b/arrow/src/util/pretty.rs
@@ -115,6 +115,7 @@ mod tests {
     };
 
     use super::*;
+    use crate::array::{DecimalBuilder, Int32Array};
     use std::sync::Arc;
 
     #[test]
@@ -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(())
+    }
 }