Skip to content

Commit

Permalink
Adding Pretty Print Support For Fixed Size List (#958)
Browse files Browse the repository at this point in the history
* Inferring 2. as Float64 for issue #929

* Adding pretty print support for fixed size list array

* fixing linting errors

* adding null row to test
  • Loading branch information
brianrackle authored and alamb committed Nov 22, 2021
1 parent c96e8de commit 41d59eb
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 1 deletion.
17 changes: 17 additions & 0 deletions arrow/src/util/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,22 @@ macro_rules! make_string_from_list {
}};
}

macro_rules! make_string_from_fixed_size_list {
($column: ident, $row: ident) => {{
let list = $column
.as_any()
.downcast_ref::<array::FixedSizeListArray>()
.ok_or(ArrowError::InvalidArgumentError(format!(
"Repl error: could not convert list column to list array."
)))?
.value($row);
let string_values = (0..list.len())
.map(|i| array_value_to_string(&list.clone(), i))
.collect::<Result<Vec<String>>>()?;
Ok(format!("[{}]", string_values.join(", ")))
}};
}

#[inline(always)]
pub fn make_string_from_decimal(column: &Arc<dyn Array>, row: usize) -> Result<String> {
let array = column
Expand Down Expand Up @@ -308,6 +324,7 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
column.data_type()
))),
},
DataType::FixedSizeList(_, _) => make_string_from_fixed_size_list!(column, row),
DataType::Struct(_) => {
let st = column
.as_any()
Expand Down
42 changes: 41 additions & 1 deletion arrow/src/util/pretty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ mod tests {
};

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

#[test]
Expand Down Expand Up @@ -263,6 +263,46 @@ mod tests {
Ok(())
}

#[test]
fn test_pretty_format_fixed_size_list() -> Result<()> {
// define a schema.
let field_type = DataType::FixedSizeList(
Box::new(Field::new("item", DataType::Int32, true)),
3,
);
let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));

let keys_builder = Int32Array::builder(3);
let mut builder = FixedSizeListBuilder::new(keys_builder, 3);

builder.values().append_slice(&[1, 2, 3]).unwrap();
builder.append(true).unwrap();
builder.values().append_slice(&[4, 5, 6]).unwrap();
builder.append(false).unwrap();
builder.values().append_slice(&[7, 8, 9]).unwrap();
builder.append(true).unwrap();

let array = Arc::new(builder.finish());

let batch = RecordBatch::try_new(schema, vec![array])?;
let table = pretty_format_batches(&[batch])?;
let expected = vec![
"+-----------+",
"| d1 |",
"+-----------+",
"| [1, 2, 3] |",
"| |",
"| [7, 8, 9] |",
"+-----------+",
];

let actual: Vec<&str> = table.lines().collect();

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

Ok(())
}

/// Generate an array with type $ARRAYTYPE with a numeric value of
/// $VALUE, and compare $EXPECTED_RESULT to the output of
/// formatting that array with `pretty_format_batches`
Expand Down

0 comments on commit 41d59eb

Please sign in to comment.