From 93f9f85b55c3a51ea1bab7ee534d9fa7ae40d70d Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Tue, 28 Feb 2023 19:01:05 +0100 Subject: [PATCH 1/3] Add support for arrow2::types::f16 --- arrow2_convert/src/deserialize.rs | 1 + arrow2_convert/src/field.rs | 1 + arrow2_convert/src/serialize.rs | 1 + arrow2_convert/tests/test_round_trip.rs | 27 +++++++++++++++++++++++++ 4 files changed, 30 insertions(+) diff --git a/arrow2_convert/src/deserialize.rs b/arrow2_convert/src/deserialize.rs index 00fc5b8..13f5c52 100644 --- a/arrow2_convert/src/deserialize.rs +++ b/arrow2_convert/src/deserialize.rs @@ -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); diff --git a/arrow2_convert/src/field.rs b/arrow2_convert/src/field.rs index f581ddc..5a0e8fc 100644 --- a/arrow2_convert/src/field.rs +++ b/arrow2_convert/src/field.rs @@ -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); diff --git a/arrow2_convert/src/serialize.rs b/arrow2_convert/src/serialize.rs index bf0e2ea..3460445 100644 --- a/arrow2_convert/src/serialize.rs +++ b/arrow2_convert/src/serialize.rs @@ -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); diff --git a/arrow2_convert/tests/test_round_trip.rs b/arrow2_convert/tests/test_round_trip.rs index d6bd0f3..b147e47 100644 --- a/arrow2_convert/tests/test_round_trip.rs +++ b/arrow2_convert/tests/test_round_trip.rs @@ -194,6 +194,33 @@ 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 = vec![1., 2., 3.] + .iter() + .map(|f| arrow2::types::f16::from_f32(*f)) + .collect(); + let b: Box = original_array.try_into_arrow().unwrap(); + let round_trip: Vec = b.try_into_collection().unwrap(); + assert_eq!(original_array, round_trip); + + let original_array: Vec> = vec![Some(1.), None, Some(3.)] + .iter() + .map(|f| f.map(arrow2::types::f16::from_f32)) + .collect(); + let b: Box = original_array.try_into_arrow().unwrap(); + let round_trip: Vec> = b.try_into_collection().unwrap(); + assert_eq!(original_array, round_trip); + + let original_array: Vec> = vec![Some(1.), None, Some(3.)] + .iter() + .map(|f| f.map(arrow2::types::f16::from_f32)) + .collect(); + let b: Arc = original_array.try_into_arrow().unwrap(); + let round_trip: Vec> = 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. From 0cd962699ff1754d89dc9db1e8a1b438d42a6b63 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 1 Mar 2023 00:40:07 +0100 Subject: [PATCH 2/3] Add a few more exotic values for half round-trip --- arrow2_convert/tests/test_round_trip.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/arrow2_convert/tests/test_round_trip.rs b/arrow2_convert/tests/test_round_trip.rs index b147e47..184385b 100644 --- a/arrow2_convert/tests/test_round_trip.rs +++ b/arrow2_convert/tests/test_round_trip.rs @@ -9,6 +9,7 @@ use arrow2_convert::{ ArrowDeserialize, ArrowField, ArrowSerialize, }; use std::borrow::Borrow; +use std::f32::INFINITY; use std::sync::Arc; #[test] @@ -196,7 +197,7 @@ fn test_primitive_type_vec() { // `arrow2::types::f16` isn't a native type so we can't just use `as` { - let original_array: Vec = vec![1., 2., 3.] + let original_array: Vec = vec![1.0, 2.5, 47800.0, 0.000012, -0.0, 0.0, INFINITY] .iter() .map(|f| arrow2::types::f16::from_f32(*f)) .collect(); From 90e4fabf2ce267483816ea1fd4e60a59999f9e67 Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Wed, 1 Mar 2023 19:53:07 +0100 Subject: [PATCH 3/3] cargo fmt --- arrow2_convert/tests/test_round_trip.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/arrow2_convert/tests/test_round_trip.rs b/arrow2_convert/tests/test_round_trip.rs index 184385b..d9a6ed3 100644 --- a/arrow2_convert/tests/test_round_trip.rs +++ b/arrow2_convert/tests/test_round_trip.rs @@ -197,10 +197,11 @@ fn test_primitive_type_vec() { // `arrow2::types::f16` isn't a native type so we can't just use `as` { - let original_array: Vec = 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 original_array: Vec = + 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 = original_array.try_into_arrow().unwrap(); let round_trip: Vec = b.try_into_collection().unwrap(); assert_eq!(original_array, round_trip);