Skip to content

Commit

Permalink
Add support for arrow2::types::f16 (#104)
Browse files Browse the repository at this point in the history
* Add support for arrow2::types::f16

* Add a few more exotic values for half round-trip

* cargo fmt
  • Loading branch information
jleibs authored Mar 18, 2023
1 parent 6450acb commit d517f07
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 0 deletions.
1 change: 1 addition & 0 deletions arrow2_convert/src/deserialize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ impl_arrow_deserialize_primitive!(i8);
impl_arrow_deserialize_primitive!(i16);
impl_arrow_deserialize_primitive!(i32);
impl_arrow_deserialize_primitive!(i64);
impl_arrow_deserialize_primitive!(arrow2::types::f16);
impl_arrow_deserialize_primitive!(f32);
impl_arrow_deserialize_primitive!(f64);

Expand Down
1 change: 1 addition & 0 deletions arrow2_convert/src/field.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ impl_numeric_type_full!(i8, Int8);
impl_numeric_type_full!(i16, Int16);
impl_numeric_type_full!(i32, Int32);
impl_numeric_type_full!(i64, Int64);
impl_numeric_type_full!(arrow2::types::f16, Float16);
impl_numeric_type_full!(f32, Float32);
impl_numeric_type_full!(f64, Float64);

Expand Down
1 change: 1 addition & 0 deletions arrow2_convert/src/serialize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ impl_numeric_type!(i8);
impl_numeric_type!(i16);
impl_numeric_type!(i32);
impl_numeric_type!(i64);
impl_numeric_type!(arrow2::types::f16);
impl_numeric_type!(f32);
impl_numeric_type!(f64);

Expand Down
29 changes: 29 additions & 0 deletions arrow2_convert/tests/test_round_trip.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use arrow2_convert::{
ArrowDeserialize, ArrowField, ArrowSerialize,
};
use std::borrow::Borrow;
use std::f32::INFINITY;
use std::sync::Arc;

#[test]
Expand Down Expand Up @@ -194,6 +195,34 @@ fn test_primitive_type_vec() {
test_float_type!(f32);
test_float_type!(f64);

// `arrow2::types::f16` isn't a native type so we can't just use `as`
{
let original_array: Vec<arrow2::types::f16> =
vec![1.0, 2.5, 47800.0, 0.000012, -0.0, 0.0, INFINITY]
.iter()
.map(|f| arrow2::types::f16::from_f32(*f))
.collect();
let b: Box<dyn Array> = original_array.try_into_arrow().unwrap();
let round_trip: Vec<arrow2::types::f16> = b.try_into_collection().unwrap();
assert_eq!(original_array, round_trip);

let original_array: Vec<Option<arrow2::types::f16>> = vec![Some(1.), None, Some(3.)]
.iter()
.map(|f| f.map(arrow2::types::f16::from_f32))
.collect();
let b: Box<dyn Array> = original_array.try_into_arrow().unwrap();
let round_trip: Vec<Option<arrow2::types::f16>> = b.try_into_collection().unwrap();
assert_eq!(original_array, round_trip);

let original_array: Vec<Option<arrow2::types::f16>> = vec![Some(1.), None, Some(3.)]
.iter()
.map(|f| f.map(arrow2::types::f16::from_f32))
.collect();
let b: Arc<dyn Array> = original_array.try_into_arrow().unwrap();
let round_trip: Vec<Option<arrow2::types::f16>> = b.try_into_collection().unwrap();
assert_eq!(original_array, round_trip);
};

// i128
// i128 is special since we need to require precision and scale so the TryIntoArrow trait
// is not implemented for Vec<i128>.
Expand Down

0 comments on commit d517f07

Please sign in to comment.