modify Bitmap

apache · Nov 20, 2024 · 6db0a1d · 6db0a1d
1 parent a27323c
commit 6db0a1d
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 3 deletions.
diff --git a/java/common/src/main/java/org/apache/tsfile/utils/BitMap.java b/java/common/src/main/java/org/apache/tsfile/utils/BitMap.java
@@ -100,6 +100,24 @@ public boolean isAllUnmarked() {
     return true;
   }
 
+  // whether all bits in the range are unmarked
+  public boolean isAllUnmarked(int rangeSize) {
+    int j;
+    for (j = 0; j < rangeSize / Byte.SIZE; j++) {
+      if (bits[j] != (byte) 0) {
+        return false;
+      }
+    }
+    int remainingBits = rangeSize % Byte.SIZE;
+    if (remainingBits > 0) {
+      byte mask = (byte) (0xFF >> (Byte.SIZE - remainingBits));
+      if ((bits[rangeSize / Byte.SIZE] & mask) != 0) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /** whether all bits are one, i.e., all are Null */
   public boolean isAllMarked() {
     int j;
@@ -191,4 +209,12 @@ public BitMap getRegion(int positionOffset, int length) {
     copyOfRange(this, positionOffset, newBitMap, 0, length);
     return newBitMap;
   }
+
+  public int getTruncatedSize(int size) {
+    return size / Byte.SIZE + (size % Byte.SIZE == 0 ? 0 : 1);
+  }
+
+  public byte[] getTruncatedByteArray(int size) {
+    return Arrays.copyOf(this.bits, getTruncatedSize(size));
+  }
 }
diff --git a/java/tsfile/src/main/java/org/apache/tsfile/write/record/Tablet.java b/java/tsfile/src/main/java/org/apache/tsfile/write/record/Tablet.java
@@ -598,12 +598,12 @@ private void writeBitMaps(DataOutputStream stream) throws IOException {
     if (bitMaps != null) {
       int size = (schemas == null ? 0 : schemas.size());
       for (int i = 0; i < size; i++) {
-        if (bitMaps[i] == null) {
+        if (bitMaps[i] == null || bitMaps[i].isAllUnmarked(rowSize)) {
           ReadWriteIOUtils.write(BytesUtils.boolToByte(false), stream);
         } else {
           ReadWriteIOUtils.write(BytesUtils.boolToByte(true), stream);
-          ReadWriteIOUtils.write(bitMaps[i].getSize(), stream);
-          ReadWriteIOUtils.write(new Binary(bitMaps[i].getByteArray()), stream);
+          ReadWriteIOUtils.write(bitMaps[i].getTruncatedSize(rowSize), stream);
+          ReadWriteIOUtils.write(new Binary(bitMaps[i].getTruncatedByteArray(rowSize)), stream);
         }
       }
     }

diff --git a/java/tsfile/src/test/java/org/apache/tsfile/utils/BitMapTest.java b/java/tsfile/src/test/java/org/apache/tsfile/utils/BitMapTest.java
@@ -20,6 +20,7 @@
 
 import org.junit.Test;
 
+import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -64,4 +65,43 @@ public void testInitFromBytes() {
       assertEquals(bitmap1.isMarked(i), bitmap2.isMarked(i));
     }
   }
+
+  @Test
+  public void testIsAllUnmarkedInRange() {
+    BitMap bitMap = new BitMap(16);
+    assertTrue(bitMap.isAllUnmarked(6));
+    assertTrue(bitMap.isAllUnmarked(8));
+    assertTrue(bitMap.isAllUnmarked(9));
+    assertTrue(bitMap.isAllUnmarked(16));
+
+    bitMap.mark(3);
+    assertTrue(bitMap.isAllUnmarked(2));
+    assertTrue(bitMap.isAllUnmarked(3));
+    assertFalse(bitMap.isAllUnmarked(4));
+    assertFalse(bitMap.isAllUnmarked(16));
+    bitMap.unmark(3);
+
+    bitMap.mark(9);
+    assertTrue(bitMap.isAllUnmarked(9));
+    assertFalse(bitMap.isAllUnmarked(10));
+  }
+
+  @Test
+  public void testGetTruncatedByteArray() {
+    BitMap bitMap = new BitMap(16);
+    assertArrayEquals(new byte[2], bitMap.getTruncatedByteArray(13));
+    assertArrayEquals(new byte[2], bitMap.getTruncatedByteArray(16));
+
+    bitMap.mark(3);
+    byte[] truncatedArray = bitMap.getTruncatedByteArray(12);
+    assertEquals(2, truncatedArray.length);
+
+    assertEquals((byte) 0b00001000, truncatedArray[0]);
+    assertEquals((byte) 0b00000000, truncatedArray[1]);
+
+    truncatedArray = bitMap.getTruncatedByteArray(8);
+    assertEquals(1, truncatedArray.length);
+
+    assertEquals((byte) 0b00001000, truncatedArray[0]);
+  }
 }