From 7effad7113bc1b01bf3b8abbcdda7ca21498afb5 Mon Sep 17 00:00:00 2001
From: gx <xuan.guo@atomdatatech.com>
Date: Fri, 7 Aug 2026 17:27:51 +0800
Subject: [PATCH] fix(cpp): handle TS2DIFF float prefixes in batch decode

---
 cpp/src/encoding/ts2diff_decoder.h      | 26 +++++----------
 cpp/test/encoding/ts2diff_codec_test.cc | 43 +++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/cpp/src/encoding/ts2diff_decoder.h b/cpp/src/encoding/ts2diff_decoder.h
index 206b7f559..f6fc961b0 100644
--- a/cpp/src/encoding/ts2diff_decoder.h
+++ b/cpp/src/encoding/ts2diff_decoder.h
@@ -952,15 +952,10 @@ class FloatTS2DIFFDecoder : public TS2DIFFDecoder<int32_t> {
 
     int read_batch_float(float* out, int capacity, int& actual,
                          common::ByteStream& in) override {
-        // Reuse SIMD batch decode for int32, then bit-cast to float
-        int32_t* buf = reinterpret_cast<int32_t*>(out);
-        int ret = TS2DIFFDecoder<int32_t>::read_batch_int32(buf, capacity,
-                                                            actual, in);
-        if (ret != common::E_OK) return ret;
-        for (int i = 0; i < actual; ++i) {
-            out[i] = common::int_to_float(buf[i]);
-        }
-        return common::E_OK;
+        // FLOAT TS_2DIFF segments have a scale/overflow prefix before the
+        // integer delta block. The integer batch decoder does not consume
+        // that prefix, so use the segment-aware scalar decoder here.
+        return Decoder::read_batch_float(out, capacity, actual, in);
     }
 
    private:
@@ -989,15 +984,10 @@ class DoubleTS2DIFFDecoder : public TS2DIFFDecoder<int64_t> {
 
     int read_batch_double(double* out, int capacity, int& actual,
                           common::ByteStream& in) override {
-        // Reuse SIMD batch decode for int64, then bit-cast to double
-        int64_t* buf = reinterpret_cast<int64_t*>(out);
-        int ret = TS2DIFFDecoder<int64_t>::read_batch_int64(buf, capacity,
-                                                            actual, in);
-        if (ret != common::E_OK) return ret;
-        for (int i = 0; i < actual; ++i) {
-            out[i] = common::long_to_double(buf[i]);
-        }
-        return common::E_OK;
+        // DOUBLE TS_2DIFF uses the same segment prefix. Bypassing
+        // read_double() misreads that prefix as a block header and can spin
+        // at end-of-input while decoding an otherwise valid page.
+        return Decoder::read_batch_double(out, capacity, actual, in);
     }
 
    private:
diff --git a/cpp/test/encoding/ts2diff_codec_test.cc b/cpp/test/encoding/ts2diff_codec_test.cc
index fb997103c..43c86adef 100644
--- a/cpp/test/encoding/ts2diff_codec_test.cc
+++ b/cpp/test/encoding/ts2diff_codec_test.cc
@@ -187,6 +187,49 @@ TEST_F(FloatDoubleTS2DIFFCodecTest, TestDoubleRoundTrip) {
     EXPECT_FALSE(decoder_double_->has_remaining(out_stream));
 }
 
+TEST_F(FloatDoubleTS2DIFFCodecTest,
+       ReadBatchFloatConsumesPrefixesAcrossSegments) {
+    common::ByteStream out_stream(1024, common::MOD_TS2DIFF_OBJ, false);
+    const int row_num = 300;
+    std::vector<float> expected(row_num);
+    for (int i = 0; i < row_num; ++i) {
+        expected[i] = static_cast<float>(i) * 0.25f + 0.5f;
+        ASSERT_EQ(encoder_float_->encode(expected[i], out_stream),
+                  common::E_OK);
+    }
+    ASSERT_EQ(encoder_float_->flush(out_stream), common::E_OK);
+
+    std::vector<float> actual_values(row_num);
+    int actual = 0;
+    ASSERT_EQ(decoder_float_->read_batch_float(actual_values.data(), row_num,
+                                               actual, out_stream),
+              common::E_OK);
+    ASSERT_EQ(actual, row_num);
+    for (int i = 0; i < row_num; ++i) {
+        EXPECT_FLOAT_EQ(actual_values[i], expected[i]) << "row " << i;
+    }
+    EXPECT_FALSE(decoder_float_->has_remaining(out_stream));
+}
+
+TEST_F(FloatDoubleTS2DIFFCodecTest, ReadBatchDoubleConsumesOverflowPrefix) {
+    common::ByteStream out_stream(1024, common::MOD_TS2DIFF_OBJ, false);
+    const double expected[] = {3.123456768E20, std::nan("")};
+    for (double value : expected) {
+        ASSERT_EQ(encoder_double_->encode(value, out_stream), common::E_OK);
+    }
+    ASSERT_EQ(encoder_double_->flush(out_stream), common::E_OK);
+
+    double actual_values[2] = {};
+    int actual = 0;
+    ASSERT_EQ(decoder_double_->read_batch_double(actual_values, 2, actual,
+                                                 out_stream),
+              common::E_OK);
+    ASSERT_EQ(actual, 2);
+    EXPECT_DOUBLE_EQ(actual_values[0], expected[0]);
+    EXPECT_TRUE(std::isnan(actual_values[1]));
+    EXPECT_FALSE(decoder_double_->has_remaining(out_stream));
+}
+
 TEST_F(TS2DIFFCodecTest, TestIntEncoding1) {
     common::ByteStream out_stream(1024, common::MOD_TS2DIFF_OBJ, false);
     const int row_num = 10000;