Fix FloatEncoder overflow problem

apache · Dec 19, 2024 · 2a23976 · 2a23976
1 parent fc0e9ec
commit 2a23976
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 42 deletions.
diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encoding/decoder/FloatDecoder.java b/java/tsfile/src/main/java/org/apache/tsfile/encoding/decoder/FloatDecoder.java
@@ -24,15 +24,14 @@
 import org.apache.tsfile.exception.encoding.TsFileDecodingException;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.tsfile.utils.Binary;
+import org.apache.tsfile.utils.BitMap;
 import org.apache.tsfile.utils.ReadWriteForEncodingUtils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
 
 /**
  * Decoder for float or double value using rle or two diff. For more info about encoding pattern,
@@ -49,7 +48,7 @@ public class FloatDecoder extends Decoder {
   /** flag that indicates whether we have read maxPointNumber and calculated maxPointValue. */
   private boolean isMaxPointNumberRead;
 
-  private List<Boolean> useMaxPointNumber;
+  private BitMap useMaxPointNumber;
   private int position = 0;
 
   public FloatDecoder(TSEncoding encodingType, TSDataType dataType) {
@@ -107,33 +106,27 @@ public float readFloat(ByteBuffer buffer) {
   public double readDouble(ByteBuffer buffer) {
     readMaxPointValue(buffer);
     long value = decoder.readLong(buffer);
-    return value / maxPointValue;
+    return value / getMaxPointValue();
   }
 
   private double getMaxPointValue() {
     if (useMaxPointNumber == null) {
       return maxPointValue;
     } else {
-      return useMaxPointNumber.get(position) ? maxPointValue : 1;
+      return useMaxPointNumber.isMarked(position) ? maxPointValue : 1;
     }
   }
 
   private void readMaxPointValue(ByteBuffer buffer) {
     if (!isMaxPointNumberRead) {
       int maxPointNumber = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer);
       if (maxPointNumber == Integer.MAX_VALUE) {
-        useMaxPointNumber = new ArrayList<>();
-        while (true) {
-          maxPointNumber = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer);
-          if (maxPointNumber == 1) {
-            useMaxPointNumber.add(true);
-          } else if (maxPointNumber == 0) {
-            useMaxPointNumber.add(false);
-          } else {
-            maxPointValue = Math.pow(10, maxPointNumber);
-            break;
-          }
-        }
+        int size = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer);
+        byte[] tmp = new byte[size / 8 + 1];
+        buffer.get(tmp, 0, size / 8 + 1);
+        useMaxPointNumber = new BitMap(size, tmp);
+        maxPointNumber = ReadWriteForEncodingUtils.readUnsignedVarInt(buffer);
+        maxPointValue = Math.pow(10, maxPointNumber);
       } else if (maxPointNumber <= 0) {
         maxPointValue = 1;
       } else {

diff --git a/java/tsfile/src/main/java/org/apache/tsfile/encoding/encoder/FloatEncoder.java b/java/tsfile/src/main/java/org/apache/tsfile/encoding/encoder/FloatEncoder.java
@@ -22,6 +22,7 @@
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.exception.encoding.TsFileEncodingException;
 import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.utils.BitMap;
 import org.apache.tsfile.utils.ReadWriteForEncodingUtils;
 
 import java.io.ByteArrayOutputStream;
@@ -126,7 +127,13 @@ private int convertFloatToInt(float value) {
   }
 
   private long convertDoubleToLong(double value) {
-    return Math.round(value * maxPointValue);
+    if (value * maxPointValue > Long.MAX_VALUE || value * maxPointValue < Long.MIN_VALUE) {
+      useMaxPointNumber.add(false);
+      return Math.round(value);
+    } else {
+      useMaxPointNumber.add(true);
+      return Math.round(value * maxPointValue);
+    }
   }
 
   @Override
@@ -136,9 +143,14 @@ public void flush(ByteArrayOutputStream out) throws IOException {
       byte[] ba = out.toByteArray();
       out.reset();
       ReadWriteForEncodingUtils.writeUnsignedVarInt(Integer.MAX_VALUE, out);
-      for (boolean b : useMaxPointNumber) {
-        ReadWriteForEncodingUtils.writeUnsignedVarInt(b ? 1 : 0, out);
+      BitMap bitMap = new BitMap(useMaxPointNumber.size());
+      for (int i = 0; i < useMaxPointNumber.size(); i++) {
+        if (useMaxPointNumber.get(i)) {
+          bitMap.mark(i);
+        }
       }
+      ReadWriteForEncodingUtils.writeUnsignedVarInt(useMaxPointNumber.size(), out);
+      out.write(bitMap.getByteArray());
       out.write(ba);
     }
     reset();

diff --git a/java/tsfile/src/test/java/org/apache/tsfile/encoding/decoder/FloatDecoderTest.java b/java/tsfile/src/test/java/org/apache/tsfile/encoding/decoder/FloatDecoderTest.java
@@ -30,7 +30,6 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
@@ -85,25 +84,6 @@ public void setUp() {
   @After
   public void tearDown() {}
 
-  public static void main(String[] args) throws IOException {
-    float value1 = 0.333F;
-    System.out.println(value1);
-    float value = 6.5536403E8F;
-    System.out.println(value);
-    Encoder encoder = new FloatEncoder(TSEncoding.TS_2DIFF, TSDataType.FLOAT, 2);
-    ByteArrayOutputStream baos = new ByteArrayOutputStream();
-    encoder.encode(value1, baos);
-    encoder.encode(value, baos);
-    encoder.flush(baos);
-
-    ByteBuffer buffer = ByteBuffer.wrap(baos.toByteArray());
-    Decoder decoder = new FloatDecoder(TSEncoding.TS_2DIFF, TSDataType.FLOAT);
-    float value_ = decoder.readFloat(buffer);
-    System.out.println(value_);
-    value_ = decoder.readFloat(buffer);
-    System.out.println(value_);
-  }
-
   @Test
   public void testRLEFloat() throws Exception {
     for (int i = 1; i <= 10; i++) {
@@ -223,6 +203,30 @@ private void testDoubleLength(
     }
   }
 
+  @Test
+  public void testBigFloat() throws Exception {
+    float a = 0.333F;
+    float b = 6.5536403E8F;
+    Encoder encoder = new FloatEncoder(TSEncoding.TS_2DIFF, TSDataType.FLOAT, 2);
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    encoder.encode(a, baos);
+    encoder.encode(b, baos);
+    encoder.flush(baos);
+
+    ByteBuffer buffer = ByteBuffer.wrap(baos.toByteArray());
+    Decoder decoder = new FloatDecoder(TSEncoding.TS_2DIFF, TSDataType.FLOAT);
+    assertEquals(roundWithGivenPrecision(a, 2), decoder.readFloat(buffer), delta);
+    assertEquals(roundWithGivenPrecision(b, 2), decoder.readFloat(buffer), delta);
+  }
+
+  public static float roundWithGivenPrecision(float data, int size) {
+    if (size == 0) {
+      return Math.round(data);
+    }
+    return Math.round(data)
+        + Math.round(((data - Math.round(data)) * Math.pow(10, size))) / (float) Math.pow(10, size);
+  }
+
   // private void testDecimalLenght(TSEncoding encoding, List<Double> valueList,
   // int maxPointValue,
   // boolean isDebug, int repeatCount) throws Exception {

diff --git a/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileReadWriteTest.java b/java/tsfile/src/test/java/org/apache/tsfile/write/TsFileReadWriteTest.java
@@ -134,8 +134,13 @@ public void floatTest() throws IOException, WriteProcessException {
 
   public void floatTest(TSEncoding encoding) throws IOException, WriteProcessException {
     writeDataByTSRecord(
-        TSDataType.FLOAT, (i) -> new FloatDataPoint("sensor_1", (float) i), encoding);
-    readData((i, field, delta) -> assertEquals(i, field.getFloatV(), delta));
+        TSDataType.FLOAT,
+        (i) -> new FloatDataPoint("sensor_1", i % 2 == 0 ? 6.55364032E8F : i),
+        encoding);
+    readData(
+        (i, field, delta) ->
+            assertEquals(
+                encoding.toString(), i % 2 == 0 ? 6.55364032E8F : i, field.getFloatV(), delta));
   }
 
   @Test