ARROW-10162: [Rust] Add pretty print support for DictionaryArray

rust/arrow/src/util/pretty.rs

@@ -19,9 +19,13 @@
 
 use crate::array;
 use crate::array::{Array, PrimitiveArrayOps};
-use crate::datatypes::{DataType, TimeUnit};
+use crate::datatypes::{
+ ArrowNativeType, ArrowPrimitiveType, DataType, Int16Type, Int32Type, Int64Type,
+ Int8Type, TimeUnit, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+};
 use crate::record_batch::RecordBatch;
 
+use array::DictionaryArray;
 use prettytable::format;
 use prettytable::{Cell, Row, Table};
 
@@ -60,7 +64,7 @@ fn create_table(results: &[RecordBatch]) -> Result<Table> {
  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.clone(), row)?));
+ cells.push(Cell::new(&array_value_to_string(&column, row)?));
  }
  table.add_row(Row::new(cells));
  }
@@ -83,8 +87,8 @@ macro_rules! make_string {
  }};
 }
 
-/// Get the value at the given row in an array as a string
-fn array_value_to_string(column: array::ArrayRef, row: usize) -> Result<String> {
+/// Get the value at the given row in an array as a String
+pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<String> {
  match column.data_type() {
  DataType::Utf8 => make_string!(array::StringArray, column, row),
  DataType::Boolean => make_string!(array::BooleanArray, column, row),
@@ -124,15 +128,55 @@ fn array_value_to_string(column: array::ArrayRef, row: usize) -> Result<String>
  DataType::Time64(unit) if *unit == TimeUnit::Nanosecond => {
  make_string!(array::Time64NanosecondArray, column, row)
  }
+ DataType::Dictionary(index_type, _value_type) => match **index_type {
+ DataType::Int8 => dict_array_value_to_string::<Int8Type>(column, row),
+ DataType::Int16 => dict_array_value_to_string::<Int16Type>(column, row),
+ DataType::Int32 => dict_array_value_to_string::<Int32Type>(column, row),
+ DataType::Int64 => dict_array_value_to_string::<Int64Type>(column, row),
+ DataType::UInt8 => dict_array_value_to_string::<UInt8Type>(column, row),
+ DataType::UInt16 => dict_array_value_to_string::<UInt16Type>(column, row),
+ DataType::UInt32 => dict_array_value_to_string::<UInt32Type>(column, row),
+ DataType::UInt64 => dict_array_value_to_string::<UInt64Type>(column, row),
+ _ => Err(ArrowError::InvalidArgumentError(format!(
+ "Pretty printing not supported for {:?} due to index type",
+ column.data_type()
+ ))),
+ },
  _ => Err(ArrowError::InvalidArgumentError(format!(
- "Unsupported {:?} type for repl.",
+ "Pretty printing not implemented for {:?} type",
  column.data_type()
  ))),
  }
 }
 
+/// Converts the value of the dictionary array at `row` to a String
+fn dict_array_value_to_string<K: ArrowPrimitiveType>(
+ colum: &array::ArrayRef,
+ row: usize,
+) -> Result<String> {
+ let dict_array = colum.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
+
+ let keys_array = dict_array.keys_array();
+
+ if keys_array.is_null(row) {
+ return Ok(String::from(""));
+ }
+
+ let dict_index = keys_array.value(row).to_usize().ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Can not convert value {:?} at index {:?} to usize for repl.",
+ keys_array.value(row),
+ row
+ ))
+ })?;
+
+ array_value_to_string(&dict_array.values(), dict_index)
+}
+
 #[cfg(test)]
 mod tests {
+ use array::{PrimitiveBuilder, StringBuilder, StringDictionaryBuilder};
+
  use super::*;
  use crate::datatypes::{Field, Schema};
  use std::sync::Arc;
@@ -183,4 +227,41 @@ mod tests {
 
  Ok(())
  }
+
+ #[test]
+ fn test_pretty_format_dictionary() -> Result<()> {
+ // define a schema.
+ 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.clone(), 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(())
+ }
 }