diff --git a/cpp/src/arrow/compare.cc b/cpp/src/arrow/compare.cc index 26f56d9b588..7a0fd1ec2ac 100644 --- a/cpp/src/arrow/compare.cc +++ b/cpp/src/arrow/compare.cc @@ -1249,8 +1249,9 @@ bool StridedFloatTensorContentEquals(const int dim_index, int64_t left_offset, int64_t right_offset, const Tensor& left, const Tensor& right, const EqualOptions& opts) { using c_type = typename DataType::c_type; - static_assert(std::is_floating_point::value, - "DataType must be a floating point type"); + static_assert(std::is_floating_point::value || + std::is_same::value, + "DataType must be a floating point type or half-float"); const auto n = left.shape()[dim_index]; const auto left_stride = left.strides()[dim_index]; @@ -1309,8 +1310,9 @@ bool TensorEquals(const Tensor& left, const Tensor& right, const EqualOptions& o } switch (left.type_id()) { - // TODO: Support half-float tensors - // case Type::HALF_FLOAT: + case Type::HALF_FLOAT: + return FloatTensorEquals(left, right, opts); + case Type::FLOAT: return FloatTensorEquals(left, right, opts); @@ -1347,8 +1349,9 @@ bool FloatSparseTensorDataEquals(const typename DataType::c_type* left_data, const typename DataType::c_type* right_data, const int64_t length, const EqualOptions& opts) { using c_type = typename DataType::c_type; - static_assert(std::is_floating_point::value, - "DataType must be a floating point type"); + static_assert(std::is_floating_point::value || + std::is_same::value, + "DataType must be a floating point type or half-float"); if (opts.nans_equal()) { if (left_data == right_data) { return true; @@ -1395,8 +1398,11 @@ struct SparseTensorEqualsImpl { const uint8_t* left_data = left.data()->data(); const uint8_t* right_data = right.data()->data(); switch (left.type()->id()) { - // TODO: Support half-float tensors - // case Type::HALF_FLOAT: + case Type::HALF_FLOAT: + return FloatSparseTensorDataEquals( + reinterpret_cast(left_data), + reinterpret_cast(right_data), length, opts); + case Type::FLOAT: return FloatSparseTensorDataEquals( reinterpret_cast(left_data), diff --git a/cpp/src/arrow/tensor_test.cc b/cpp/src/arrow/tensor_test.cc index 05dcf38e1e6..a6511948162 100644 --- a/cpp/src/arrow/tensor_test.cc +++ b/cpp/src/arrow/tensor_test.cc @@ -31,6 +31,7 @@ #include "arrow/tensor.h" #include "arrow/testing/gtest_util.h" #include "arrow/type.h" +#include "arrow/util/float16.h" namespace arrow { @@ -505,8 +506,18 @@ TYPED_TEST_P(TestFloatTensor, Equals) { std::vector shape = {4, 4}; - std::vector c_values = {1, 2, 3, 4, 5, 6, 7, 8, + constexpr int c_vals_as_int_arr[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}; + std::vector c_values; + if constexpr (std::is_same_v) { + for (int i : c_vals_as_int_arr) { + c_values.push_back(util::Float16::FromFloat(static_cast(i)).bits()); + } + } else { + for (int i : c_vals_as_int_arr) { + c_values.push_back(static_cast(i)); + } + } std::vector c_strides = {unit_size * shape[1], unit_size}; Tensor tc1(TypeTraits::type_singleton(), Buffer::Wrap(c_values), shape, c_strides); @@ -515,8 +526,18 @@ TYPED_TEST_P(TestFloatTensor, Equals) { Tensor tc2(TypeTraits::type_singleton(), Buffer::Wrap(c_values_2), shape, c_strides); - std::vector f_values = {1, 5, 9, 13, 2, 6, 10, 14, + constexpr int f_vals_as_int_arr[] = {1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15, 4, 8, 12, 16}; + std::vector f_values; + if constexpr (std::is_same_v) { + for (int i : f_vals_as_int_arr) { + f_values.push_back(util::Float16::FromFloat(static_cast(i)).bits()); + } + } else { + for (int i : f_vals_as_int_arr) { + f_values.push_back(static_cast(i)); + } + } Tensor tc3(TypeTraits::type_singleton(), Buffer::Wrap(f_values), shape, c_strides); @@ -531,9 +552,19 @@ TYPED_TEST_P(TestFloatTensor, Equals) { Tensor tf3(TypeTraits::type_singleton(), Buffer::Wrap(c_values), shape, f_strides); - std::vector nc_values = {1, 0, 5, 0, 9, 0, 13, 0, 2, 0, 6, + constexpr int nc_vals_as_int_arr[] = {1, 0, 5, 0, 9, 0, 13, 0, 2, 0, 6, 0, 10, 0, 14, 0, 3, 0, 7, 0, 11, 0, 15, 0, 4, 0, 8, 0, 12, 0, 16, 0}; + std::vector nc_values; + if constexpr (std::is_same_v) { + for (int i : nc_vals_as_int_arr) { + nc_values.push_back(util::Float16::FromFloat(static_cast(i)).bits()); + } + } else { + for (int i : nc_vals_as_int_arr) { + nc_values.push_back(static_cast(i)); + } + } std::vector nc_strides = {unit_size * 2, unit_size * 2 * shape[0]}; Tensor tnc(TypeTraits::type_singleton(), Buffer::Wrap(nc_values), shape, nc_strides); @@ -569,29 +600,54 @@ TYPED_TEST_P(TestFloatTensor, Equals) { EXPECT_FALSE(tf3.Equals(tnc)); // signed zeros - c_values[0] = -0.0; - c_values_2[0] = 0.0; + if constexpr (std::is_same_v) { + c_values[0] = 0x8000; // -0.0 + c_values_2[0] = 0x0000; // +0.0 + } else { + c_values[0] = -0.0; + c_values_2[0] = 0.0; + } EXPECT_TRUE(tc1.Equals(tc2)); EXPECT_FALSE(tc1.Equals(tc2, EqualOptions().signed_zeros_equal(false))); // tensors with NaNs - const c_data_type nan_value = static_cast(NAN); + c_data_type nan_value; + if constexpr (std::is_same_v) { + nan_value = 0x7e00; // NaN + } else { + nan_value = static_cast(NAN); + } + c_values[0] = nan_value; - EXPECT_TRUE(std::isnan(tc1.Value({0, 0}))); + if constexpr (std::is_same_v) { + EXPECT_TRUE(util::Float16::FromBits(tc1.Value({0, 0})).is_nan()); + } else { + EXPECT_TRUE(std::isnan(tc1.Value({0, 0}))); + } EXPECT_FALSE(tc1.Equals(tc1)); // same object EXPECT_TRUE(tc1.Equals(tc1, EqualOptions().nans_equal(true))); // same object - EXPECT_FALSE(std::isnan(tc2.Value({0, 0}))); + + if constexpr (std::is_same_v) { + EXPECT_FALSE(util::Float16::FromBits(tc2.Value({0, 0})).is_nan()); + } else { + EXPECT_FALSE(std::isnan(tc2.Value({0, 0}))); + } EXPECT_FALSE(tc1.Equals(tc2)); // different memory EXPECT_FALSE(tc1.Equals(tc2, EqualOptions().nans_equal(true))); // different memory c_values_2[0] = nan_value; - EXPECT_TRUE(std::isnan(tc2.Value({0, 0}))); + if constexpr (std::is_same_v) { + EXPECT_TRUE(util::Float16::FromBits(tc2.Value({0, 0})).is_nan()); + } else { + EXPECT_TRUE(std::isnan(tc2.Value({0, 0}))); + } EXPECT_FALSE(tc1.Equals(tc2)); // different memory EXPECT_TRUE(tc1.Equals(tc2, EqualOptions().nans_equal(true))); // different memory } REGISTER_TYPED_TEST_SUITE_P(TestFloatTensor, Equals); +INSTANTIATE_TYPED_TEST_SUITE_P(Float16, TestFloatTensor, HalfFloatType); INSTANTIATE_TYPED_TEST_SUITE_P(Float32, TestFloatTensor, FloatType); INSTANTIATE_TYPED_TEST_SUITE_P(Float64, TestFloatTensor, DoubleType);