fix: canonicalize NaN in binary serializers

Author: zjw1111

src/paimon/common/data/variant/generic_variant_test.cpp

@@ -19,6 +19,8 @@
 
 #include "paimon/common/data/variant/generic_variant.h"
 
+#include <cstdint>
+#include <cstring>
 #include <functional>
 #include <limits>
 #include <string>
@@ -31,6 +33,17 @@
 #include "paimon/testing/utils/testharness.h"
 
 namespace paimon::test {
+namespace {
+
+template <typename FloatingPoint, typename UInt>
+FloatingPoint FloatingPointFromBits(UInt bits) {
+    static_assert(sizeof(FloatingPoint) == sizeof(UInt));
+    FloatingPoint value;
+    std::memcpy(&value, &bits, sizeof(value));
+    return value;
+}
+
+}  // namespace
 
 class GenericVariantTest : public ::testing::Test {
  public:
@@ -350,6 +363,24 @@ TEST_F(GenericVariantTest, NonFiniteDoubleToJson) {
     ASSERT_EQ(json, "\"Infinity\"");
 }
 
+TEST_F(GenericVariantTest, CanonicalizesFloatingPointNaN) {
+    {
+        VariantBuilder builder(false);
+        ASSERT_OK(builder.AppendFloat(FloatingPointFromBits<float>(uint32_t{0xffc12345})));
+        ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> variant, builder.Build(pool_));
+        ASSERT_OK_AND_ASSIGN(std::string_view value, variant->Value());
+        ASSERT_EQ(ToHex(value), "380000c07f");
+    }
+    {
+        VariantBuilder builder(false);
+        ASSERT_OK(
+            builder.AppendDouble(FloatingPointFromBits<double>(uint64_t{0xfff8123456789abc})));
+        ASSERT_OK_AND_ASSIGN(std::shared_ptr<GenericVariant> variant, builder.Build(pool_));
+        ASSERT_OK_AND_ASSIGN(std::string_view value, variant->Value());
+        ASSERT_EQ(ToHex(value), "1c000000000000f87f");
+    }
+}
+
 TEST_F(GenericVariantTest, GetTypeInfoReturnsHeaderBits) {
     // GetTypeInfo exposes the primitive header's type-info bits; 42 is encoded as an int1.
     auto v = FromJson("42");

src/paimon/common/data/variant/variant_builder.cpp

@@ -30,6 +30,7 @@
 
 #include "fmt/format.h"
 #include "paimon/common/data/variant/variant_defs.h"
+#include "paimon/common/utils/math.h"
 #include "rapidjson/error/en.h"
 #include "rapidjson/memorystream.h"
 #include "rapidjson/reader.h"
@@ -339,8 +340,7 @@ Status VariantBuilder::AppendLong(int64_t l) {
 Status VariantBuilder::AppendDouble(double d) {
     PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 8));
     write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kDouble);
-    int64_t bits;
-    memcpy(&bits, &d, sizeof(bits));
+    const int64_t bits = CanonicalizeDoubleToLongBits(d);
     VariantBinaryUtil::WriteLong(bits, 8, write_buffer_.data(), write_pos_);
     write_pos_ += 8;
     return Status::OK();
@@ -409,8 +409,7 @@ Status VariantBuilder::AppendTimestampNtz(int64_t micros_since_epoch) {
 Status VariantBuilder::AppendFloat(float f) {
     PAIMON_RETURN_NOT_OK(CheckCapacity(1 + 4));
     write_buffer_[write_pos_++] = VariantBinaryUtil::PrimitiveHeader(VariantDefs::kFloat);
-    int32_t bits;
-    memcpy(&bits, &f, sizeof(bits));
+    const int32_t bits = CanonicalizeFloatToIntBits(f);
     VariantBinaryUtil::WriteLong(bits, 4, write_buffer_.data(), write_pos_);
     write_pos_ += 4;
     return Status::OK();

src/paimon/common/global_index/btree/key_serializer.cpp

@@ -27,6 +27,7 @@
 #include "paimon/common/utils/date_time_utils.h"
 #include "paimon/common/utils/field_type_utils.h"
 #include "paimon/common/utils/fields_comparator.h"
+#include "paimon/common/utils/math.h"
 #include "paimon/common/utils/preconditions.h"
 #include "paimon/common/utils/var_length_int_utils.h"
 #include "paimon/data/decimal.h"
@@ -164,19 +165,13 @@ Result<std::shared_ptr<Bytes>> KeySerializer::SerializeKey(
         case FieldType::FLOAT: {
             MemorySliceOutput output(4, pool);
             output.Reset();
-            auto fvalue = literal.GetValue<float>();
-            int32_t ivalue;
-            memcpy(&ivalue, &fvalue, sizeof(float));
-            output.WriteValue<int32_t>(ivalue);
+            output.WriteValue<int32_t>(CanonicalizeFloatToIntBits(literal.GetValue<float>()));
             return output.ToSlice().CopyBytes(pool);
         }
         case FieldType::DOUBLE: {
             MemorySliceOutput output(8, pool);
             output.Reset();
-            auto dvalue = literal.GetValue<double>();
-            int64_t ivalue;
-            memcpy(&ivalue, &dvalue, sizeof(double));
-            output.WriteValue<int64_t>(ivalue);
+            output.WriteValue<int64_t>(CanonicalizeDoubleToLongBits(literal.GetValue<double>()));
             return output.ToSlice().CopyBytes(pool);
         }
         case FieldType::STRING: {

src/paimon/common/global_index/btree/key_serializer_test.cpp

@@ -19,12 +19,27 @@
 
 #include "paimon/common/global_index/btree/key_serializer.h"
 
+#include <cstdint>
+#include <cstring>
+#include <string>
+
 #include "gtest/gtest.h"
 #include "paimon/data/decimal.h"
 #include "paimon/data/timestamp.h"
 #include "paimon/testing/utils/testharness.h"
 
 namespace paimon::test {
+namespace {
+
+template <typename FloatingPoint, typename UInt>
+FloatingPoint FloatingPointFromBits(UInt bits) {
+    static_assert(sizeof(FloatingPoint) == sizeof(UInt));
+    FloatingPoint value;
+    std::memcpy(&value, &bits, sizeof(value));
+    return value;
+}
+
+}  // namespace
 
 class KeySerializerTest : public ::testing::Test {
  protected:
@@ -208,6 +223,30 @@ TEST_F(KeySerializerTest, SerializeAndDeserializeAllTypes) {
     }
 }
 
+TEST_F(KeySerializerTest, CanonicalizesFloatingPointNaN) {
+    const float float_nan = FloatingPointFromBits<float>(uint32_t{0xffc12345});
+    const float canonical_float_nan = FloatingPointFromBits<float>(uint32_t{0x7fc00000});
+    ASSERT_OK_AND_ASSIGN(
+        std::shared_ptr<Bytes> float_bytes,
+        KeySerializer::SerializeKey(Literal(float_nan), arrow::float32(), pool_.get()));
+    ASSERT_OK_AND_ASSIGN(
+        std::shared_ptr<Bytes> canonical_float_bytes,
+        KeySerializer::SerializeKey(Literal(canonical_float_nan), arrow::float32(), pool_.get()));
+    ASSERT_EQ(std::string(float_bytes->data(), float_bytes->size()),
+              std::string(canonical_float_bytes->data(), canonical_float_bytes->size()));
+
+    const double double_nan = FloatingPointFromBits<double>(uint64_t{0xfff8123456789abc});
+    const double canonical_double_nan = FloatingPointFromBits<double>(uint64_t{0x7ff8000000000000});
+    ASSERT_OK_AND_ASSIGN(
+        std::shared_ptr<Bytes> double_bytes,
+        KeySerializer::SerializeKey(Literal(double_nan), arrow::float64(), pool_.get()));
+    ASSERT_OK_AND_ASSIGN(
+        std::shared_ptr<Bytes> canonical_double_bytes,
+        KeySerializer::SerializeKey(Literal(canonical_double_nan), arrow::float64(), pool_.get()));
+    ASSERT_EQ(std::string(double_bytes->data(), double_bytes->size()),
+              std::string(canonical_double_bytes->data(), canonical_double_bytes->size()));
+}
+
 TEST_F(KeySerializerTest, RejectsMalformedSerializedKeys) {
     auto wrap = [this](const std::string& value) {
         return MemorySlice::Wrap(std::make_shared<Bytes>(value, pool_.get()));

src/paimon/common/global_index/global_index_result.cpp

@@ -22,6 +22,7 @@
 #include "fmt/format.h"
 #include "paimon/common/io/memory_segment_output_stream.h"
 #include "paimon/common/memory/memory_segment_utils.h"
+#include "paimon/common/utils/math.h"
 #include "paimon/global_index/bitmap_global_index_result.h"
 #include "paimon/global_index/bitmap_scored_global_index_result.h"
 #include "paimon/io/byte_array_input_stream.h"
@@ -37,8 +38,8 @@ void WriteBitmapAndScores(const RoaringBitmap64* bitmap, const std::vector<float
     out->WriteBytes(bitmap_bytes);
 
     out->WriteValue<int32_t>(scores.size());
-    for (auto score : scores) {
-        out->WriteValue<float>(score);
+    for (float score : scores) {
+        out->WriteValue<float>(CanonicalizeFloatingPoint(score));
     }
 }
 

src/paimon/common/global_index/global_index_result_test.cpp

@@ -19,6 +19,9 @@
 
 #include "paimon/global_index/global_index_result.h"
 
+#include <cmath>
+#include <cstdint>
+#include <cstring>
 #include <utility>
 
 #include "gtest/gtest.h"
@@ -144,6 +147,28 @@ TEST_F(GlobalIndexResultTest, TestSerializeAndDeserializeWithScore) {
                                                 serialize_bytes->data() + serialize_bytes->size()));
 }
 
+TEST_F(GlobalIndexResultTest, TestSerializeCanonicalizesNaNScore) {
+    auto pool = GetDefaultPool();
+    uint32_t payload_bits = 0xffc12345;
+    float payload_nan;
+    std::memcpy(&payload_nan, &payload_bits, sizeof(payload_nan));
+    auto index_result = std::make_shared<BitmapScoredGlobalIndexResult>(
+        RoaringBitmap64::From({1}), std::vector<float>{payload_nan});
+
+    ASSERT_OK_AND_ASSIGN(PAIMON_UNIQUE_PTR<Bytes> serialized,
+                         GlobalIndexResult::Serialize(index_result, pool));
+    ASSERT_GE(serialized->size(), sizeof(float));
+    ASSERT_EQ(std::string(serialized->data() + serialized->size() - sizeof(float), sizeof(float)),
+              std::string("\x7f\xc0\x00\x00", sizeof(float)));
+
+    ASSERT_OK_AND_ASSIGN(
+        std::shared_ptr<GlobalIndexResult> deserialized,
+        GlobalIndexResult::Deserialize(serialized->data(), serialized->size(), pool));
+    auto scored_result = std::dynamic_pointer_cast<BitmapScoredGlobalIndexResult>(deserialized);
+    ASSERT_TRUE(scored_result);
+    ASSERT_TRUE(std::isnan(scored_result->GetScores()[0]));
+}
+
 TEST_F(GlobalIndexResultTest, TestInvalidSerialize) {
     auto pool = GetDefaultPool();
     auto result = std::make_shared<FakeGlobalIndexResult>(std::vector<int64_t>({1, 3, 5, 100}));

src/paimon/core/global_index/indexed_split_test.cpp

@@ -17,6 +17,8 @@
  * under the License.
  */
 
+#include <cstdint>
+#include <cstring>
 #include <memory>
 #include <string>
 #include <utility>
@@ -155,6 +157,27 @@ TEST(IndexedSplitTest, TestIndexedSplitWithScore) {
         << roundtrip_indexed_split->ToString();
 }
 
+TEST(IndexedSplitTest, TestSerializeCanonicalizesNaNScore) {
+    auto pool = GetDefaultPool();
+    DataSplitImpl::Builder builder(
+        /*partition=*/BinaryRow::EmptyRow(),
+        /*bucket=*/0, /*bucket_path=*/"bucket-0",
+        /*data_files=*/{});
+    ASSERT_OK_AND_ASSIGN(std::shared_ptr<Split> data_split, builder.Build());
+
+    uint32_t payload_bits = 0xffc12345;
+    float payload_nan;
+    std::memcpy(&payload_nan, &payload_bits, sizeof(payload_nan));
+    auto indexed_split = std::make_shared<IndexedSplitImpl>(
+        std::dynamic_pointer_cast<DataSplitImpl>(data_split), std::vector<Range>{Range(0, 0)},
+        std::vector<float>{payload_nan});
+
+    ASSERT_OK_AND_ASSIGN(std::string serialized, Split::Serialize(indexed_split, pool));
+    ASSERT_GE(serialized.size(), sizeof(float));
+    ASSERT_EQ(serialized.substr(serialized.size() - sizeof(float)),
+              std::string("\x7f\xc0\x00\x00", sizeof(float)));
+}
+
 TEST(IndexedSplitTest, TestValidate) {
     auto meta = std::make_shared<DataFileMeta>(
         "file.orc", 1l, 200l, BinaryRow::EmptyRow(), BinaryRow::EmptyRow(),

src/paimon/core/table/source/split.cpp

@@ -23,6 +23,7 @@
 #include "paimon/common/data/binary_row.h"
 #include "paimon/common/io/memory_segment_output_stream.h"
 #include "paimon/common/memory/memory_segment_utils.h"
+#include "paimon/common/utils/math.h"
 #include "paimon/common/utils/serialization_utils.h"
 #include "paimon/core/global_index/indexed_split_impl.h"
 #include "paimon/core/io/data_file_meta_serializer.h"
@@ -159,8 +160,8 @@ Result<std::string> Split::Serialize(const std::shared_ptr<Split>& split,
         if (!scores.empty()) {
             out.WriteValue<bool>(true);
             out.WriteValue<int32_t>(scores.size());
-            for (const auto& score : scores) {
-                out.WriteValue<float>(score);
+            for (float score : scores) {
+                out.WriteValue<float>(CanonicalizeFloatingPoint(score));
             }
         } else {
             out.WriteValue<bool>(false);