diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index 2cdf94bfab95..98cc7d8254f9 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -2558,6 +2558,32 @@ TEST(TestPrimitiveAdHoc, FloatingUlpDistanceEquality) { CheckFloatApproxEqualsWithUlpDistance(); } +TEST(TestRunEncodedBuilder, MergeConsecutiveNaNs) { + auto run_end_encoded_type = run_end_encoded(int32(), float32()); + std::shared_ptr result; + auto run_end_builder = std::make_shared(); + auto value_builder = std::make_shared(); + RunEndEncodedBuilder builder(default_memory_pool(), run_end_builder, value_builder, + run_end_encoded_type); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), 1), 15)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), 1), 10)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), 2), 20)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), std::nanf("")), 25)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), std::nanf("")), 5)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), 3), 15)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), std::nanf("")), 20)); + ASSERT_OK(builder.AppendScalar(**MakeScalar(float32(), 4), 25)); + ASSERT_OK(builder.Finish(&result)); + + std::shared_ptr expected_run_end_array; + std::shared_ptr expected_value_array; + ArrayFromVector({25, 45, 75, 90, 110, 135}, &expected_run_end_array); + ArrayFromVector({1.0, 2.0, std::nanf(""), 3.0, std::nanf(""), 4.0}, + &expected_value_array); + AssertArraysEqual(*expected_run_end_array, *result->run_ends()); + AssertArraysEqual(*expected_value_array, *result->values()); +} + // ---------------------------------------------------------------------- // FixedSizeBinary tests diff --git a/cpp/src/arrow/array/builder_run_end.cc b/cpp/src/arrow/array/builder_run_end.cc index 2edeaff504d2..8b38783eb898 100644 --- a/cpp/src/arrow/array/builder_run_end.cc +++ b/cpp/src/arrow/array/builder_run_end.cc @@ -23,6 +23,8 @@ #include #include +#include "arrow/array/array_run_end.h" +#include "arrow/array/util.h" #include "arrow/scalar.h" #include "arrow/util/checked_cast.h" #include "arrow/util/int_util_overflow.h" @@ -102,7 +104,7 @@ Status RunCompressorBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeat current_value_ = scalar.is_valid ? scalar.shared_from_this() : NULLPTR; current_run_length_ = n_repeats; } else if ((current_value_ == NULLPTR && !scalar.is_valid) || - (current_value_ != NULLPTR && current_value_->Equals(scalar))) { + (current_value_ != NULLPTR && current_value_->Equals(scalar, options_))) { // Extend the currently open run current_run_length_ += n_repeats; } else { diff --git a/cpp/src/arrow/array/builder_run_end.h b/cpp/src/arrow/array/builder_run_end.h index ac92efbd0dbe..43e59df1c6ef 100644 --- a/cpp/src/arrow/array/builder_run_end.h +++ b/cpp/src/arrow/array/builder_run_end.h @@ -23,8 +23,10 @@ #include #include -#include "arrow/array.h" +#include "arrow/array/array_base.h" +#include "arrow/array/array_run_end.h" #include "arrow/array/builder_base.h" +#include "arrow/compare.h" namespace arrow { @@ -161,6 +163,7 @@ class RunCompressorBuilder : public ArrayBuilder { std::shared_ptr inner_builder_; std::shared_ptr current_value_ = NULLPTR; int64_t current_run_length_ = 0; + const EqualOptions options_ = EqualOptions::Defaults().nans_equal(true); }; } // namespace internal