bugfix in display of float16 array (#1194)

Due to a typo the float16 array was being cast to a float32 array, causing a crash when pretty printing a record batch containing float16.
apache · Jan 18, 2022 · 9d637a4 · 9d637a4
1 parent 0cc0c05
commit 9d637a4
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs
@@ -313,7 +313,7 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
         DataType::UInt16 => make_string!(array::UInt16Array, column, row),
         DataType::UInt32 => make_string!(array::UInt32Array, column, row),
         DataType::UInt64 => make_string!(array::UInt64Array, column, row),
-        DataType::Float16 => make_string!(array::Float32Array, column, row),
+        DataType::Float16 => make_string!(array::Float16Array, column, row),
         DataType::Float32 => make_string!(array::Float32Array, column, row),
         DataType::Float64 => make_string!(array::Float64Array, column, row),
         DataType::Decimal(..) => make_string_from_decimal(column, row),

diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs
@@ -109,8 +109,8 @@ mod tests {
     use crate::{
         array::{
             self, new_null_array, Array, Date32Array, Date64Array,
-            FixedSizeBinaryBuilder, PrimitiveBuilder, StringArray, StringBuilder,
-            StringDictionaryBuilder, StructArray, Time32MillisecondArray,
+            FixedSizeBinaryBuilder, Float16Array, PrimitiveBuilder, StringArray,
+            StringBuilder, StringDictionaryBuilder, StructArray, Time32MillisecondArray,
             Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
             TimestampMicrosecondArray, TimestampMillisecondArray,
             TimestampNanosecondArray, TimestampSecondArray,
@@ -123,6 +123,8 @@ mod tests {
     use std::fmt::Write;
     use std::sync::Arc;
 
+    use half::f16;
+
     #[test]
     fn test_pretty_format_batches() -> Result<()> {
         // define a schema.
@@ -692,4 +694,34 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_float16_display() -> Result<()> {
+        let values = vec![
+            Some(f16::from_f32(f32::NAN)),
+            Some(f16::from_f32(4.0)),
+            Some(f16::from_f32(f32::NEG_INFINITY)),
+        ];
+        let array = Arc::new(values.into_iter().collect::<Float16Array>()) as ArrayRef;
+
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "f16",
+            array.data_type().clone(),
+            true,
+        )]));
+
+        let batch = RecordBatch::try_new(schema, vec![array])?;
+
+        let table = pretty_format_batches(&[batch])?.to_string();
+
+        let expected = vec![
+            "+------+", "| f16  |", "+------+", "| NaN  |", "| 4    |", "| -inf |",
+            "+------+",
+        ];
+
+        let actual: Vec<&str> = table.lines().collect();
+        assert_eq!(expected, actual, "Actual result:\n{}", table);
+
+        Ok(())
+    }
 }