Skip to content

Commit 86ec4b9

Browse files
luoluoyuyuluoluoyuyu
authored
Add Tablet.serializedSize() and comprehensive size validation tests. (#824)
* Add Tablet.serializedSize() and comprehensive size validation tests. Pre-allocate serialization buffer using exact size estimation, support OBJECT type in tablet serialize/deserialize path, and consolidate serializedSize tests. * Fix serialized size charset and overflow handling * Avoid bitmap copy when sizing tablet serialization
1 parent e7b0d69 commit 86ec4b9

5 files changed

Lines changed: 540 additions & 13 deletions

File tree

‎java/tsfile/src/main/java/org/apache/tsfile/utils/ReadWriteIOUtils.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ public static int write(Map<String, String> map, ByteBuffer buffer) {
185185
if (entry.getKey() == null) {
186186
buffer.putInt(-1);
187187
} else {
188-
bytes = entry.getKey().getBytes();
188+
bytes = entry.getKey().getBytes(TSFileConfig.STRING_CHARSET);
189189
buffer.putInt(bytes.length);
190190
buffer.put(bytes);
191191
length += bytes.length;
@@ -194,7 +194,7 @@ public static int write(Map<String, String> map, ByteBuffer buffer) {
194194
if (entry.getValue() == null) {
195195
buffer.putInt(-1);
196196
} else {
197-
bytes = entry.getValue().getBytes();
197+
bytes = entry.getValue().getBytes(TSFileConfig.STRING_CHARSET);
198198
buffer.putInt(bytes.length);
199199
buffer.put(bytes);
200200
length += bytes.length;
@@ -509,7 +509,7 @@ public static int sizeToWrite(String s) {
509509
if (s == null) {
510510
return INT_LEN;
511511
}
512-
return INT_LEN + s.getBytes().length;
512+
return INT_LEN + s.getBytes(TSFileConfig.STRING_CHARSET).length;
513513
}
514514

515515
/** read a byte var from inputStream. */
@@ -1202,7 +1202,7 @@ public static void writeObject(Object value, DataOutputStream outputStream) {
12021202
outputStream.write(NONE.ordinal());
12031203
} else {
12041204
outputStream.write(STRING.ordinal());
1205-
byte[] bytes = value.toString().getBytes();
1205+
byte[] bytes = value.toString().getBytes(TSFileConfig.STRING_CHARSET);
12061206
outputStream.writeInt(bytes.length);
12071207
outputStream.write(bytes);
12081208
}
@@ -1238,7 +1238,7 @@ public static void writeObject(Object value, ByteBuffer byteBuffer) {
12381238
byteBuffer.putInt(NONE.ordinal());
12391239
} else {
12401240
byteBuffer.putInt(STRING.ordinal());
1241-
byte[] bytes = value.toString().getBytes();
1241+
byte[] bytes = value.toString().getBytes(TSFileConfig.STRING_CHARSET);
12421242
byteBuffer.putInt(bytes.length);
12431243
byteBuffer.put(bytes);
12441244
}
@@ -1271,7 +1271,7 @@ public static Object readObject(ByteBuffer buffer) {
12711271
length = buffer.getInt();
12721272
bytes = new byte[length];
12731273
buffer.get(bytes);
1274-
return new String(bytes);
1274+
return new String(bytes, TSFileConfig.STRING_CHARSET);
12751275
}
12761276
}
12771277

‎java/tsfile/src/main/java/org/apache/tsfile/write/record/Tablet.java‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -748,13 +748,26 @@ private Object createValueColumnOfDataType(TSDataType dataType, int capacity) {
748748

749749
/** Serialize {@link Tablet} */
750750
public ByteBuffer serialize() throws IOException {
751-
try (PublicBAOS byteArrayOutputStream = new PublicBAOS();
751+
final int serializedSize = serializedSize();
752+
try (PublicBAOS byteArrayOutputStream = new PublicBAOS(serializedSize);
752753
DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) {
753754
serialize(outputStream);
754755
return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size());
755756
}
756757
}
757758

759+
/** Return the exact serialized byte size of this tablet. */
760+
public int serializedSize() {
761+
int size = 0;
762+
size = Math.addExact(size, ReadWriteIOUtils.sizeToWrite(insertTargetName));
763+
size = Math.addExact(size, Integer.BYTES);
764+
size = Math.addExact(size, serializedSizeOfMeasurementSchemas());
765+
size = Math.addExact(size, serializedSizeOfTimes());
766+
size = Math.addExact(size, serializedSizeOfBitMaps());
767+
size = Math.addExact(size, serializedSizeOfValues());
768+
return size;
769+
}
770+
758771
public void serialize(DataOutputStream stream) throws IOException {
759772
ReadWriteIOUtils.write(insertTargetName, stream);
760773
ReadWriteIOUtils.write(rowSize, stream);
@@ -764,6 +777,104 @@ public void serialize(DataOutputStream stream) throws IOException {
764777
writeValues(stream);
765778
}
766779

780+
private int serializedSizeOfMeasurementSchemas() {
781+
int size = Byte.BYTES;
782+
if (schemas != null) {
783+
size = Math.addExact(size, Integer.BYTES);
784+
for (int i = 0; i < schemas.size(); i++) {
785+
size = Math.addExact(size, Byte.BYTES);
786+
final IMeasurementSchema schema = schemas.get(i);
787+
if (schema != null) {
788+
size = Math.addExact(size, schema.serializedSize());
789+
size = Math.addExact(size, Byte.BYTES);
790+
}
791+
}
792+
}
793+
return size;
794+
}
795+
796+
private int serializedSizeOfTimes() {
797+
int size = Byte.BYTES;
798+
if (timestamps != null) {
799+
size = Math.addExact(size, Math.multiplyExact(Long.BYTES, rowSize));
800+
}
801+
return size;
802+
}
803+
804+
private int serializedSizeOfBitMaps() {
805+
int size = Byte.BYTES;
806+
if (bitMaps != null) {
807+
final int columnCount = schemas == null ? 0 : schemas.size();
808+
for (int i = 0; i < columnCount; i++) {
809+
if (bitMaps[i] == null || bitMaps[i].isAllUnmarked(rowSize)) {
810+
size = Math.addExact(size, Byte.BYTES);
811+
} else {
812+
size = Math.addExact(size, Byte.BYTES);
813+
size = Math.addExact(size, Integer.BYTES);
814+
size = Math.addExact(size, Integer.BYTES);
815+
size = Math.addExact(size, BitMap.getSizeOfBytes(rowSize));
816+
}
817+
}
818+
}
819+
return size;
820+
}
821+
822+
private int serializedSizeOfValues() {
823+
int size = Byte.BYTES;
824+
if (values != null) {
825+
final int columnCount = schemas == null ? 0 : schemas.size();
826+
for (int i = 0; i < columnCount; i++) {
827+
size = Math.addExact(size, serializedSizeOfColumn(schemas.get(i).getType(), values[i]));
828+
}
829+
}
830+
return size;
831+
}
832+
833+
private int serializedSizeOfColumn(final TSDataType dataType, final Object column) {
834+
int size = Byte.BYTES;
835+
if (column == null) {
836+
return size;
837+
}
838+
switch (dataType) {
839+
case INT32:
840+
return Math.addExact(size, Math.multiplyExact(Integer.BYTES, rowSize));
841+
case DATE:
842+
return Math.addExact(size, Math.multiplyExact(Integer.BYTES, rowSize));
843+
case INT64:
844+
case TIMESTAMP:
845+
return Math.addExact(size, Math.multiplyExact(Long.BYTES, rowSize));
846+
case FLOAT:
847+
return Math.addExact(size, Math.multiplyExact(Float.BYTES, rowSize));
848+
case DOUBLE:
849+
return Math.addExact(size, Math.multiplyExact(Double.BYTES, rowSize));
850+
case BOOLEAN:
851+
return Math.addExact(size, rowSize);
852+
case TEXT:
853+
case STRING:
854+
case BLOB:
855+
case OBJECT:
856+
return Math.addExact(size, serializedSizeOfBinaryValues((Binary[]) column));
857+
default:
858+
throw new UnSupportedDataTypeException(
859+
Messages.format("error.write.type_not_supported", dataType));
860+
}
861+
}
862+
863+
private static int serializedSizeOfBinaryValues(final Binary[] binaryValues, final int rowSize) {
864+
int size = 0;
865+
for (int j = 0; j < rowSize; j++) {
866+
size = Math.addExact(size, Byte.BYTES);
867+
if (binaryValues[j] != null) {
868+
size = Math.addExact(size, ReadWriteIOUtils.sizeToWrite(binaryValues[j]));
869+
}
870+
}
871+
return size;
872+
}
873+
874+
private int serializedSizeOfBinaryValues(final Binary[] binaryValues) {
875+
return serializedSizeOfBinaryValues(binaryValues, rowSize);
876+
}
877+
767878
/** Serialize {@link MeasurementSchema}s */
768879
private void writeMeasurementSchemas(DataOutputStream stream) throws IOException {
769880
ReadWriteIOUtils.write(BytesUtils.boolToByte(schemas != null), stream);

‎java/tsfile/src/main/java/org/apache/tsfile/write/schema/MeasurementSchema.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -319,15 +319,15 @@ public int serializeTo(OutputStream outputStream) throws IOException {
319319
@Override
320320
public int serializedSize() {
321321
int byteLen = 0;
322-
byteLen += ReadWriteIOUtils.sizeToWrite(measurementName);
323-
byteLen += 3 * Byte.BYTES;
322+
byteLen = Math.addExact(byteLen, ReadWriteIOUtils.sizeToWrite(measurementName));
323+
byteLen = Math.addExact(byteLen, 3 * Byte.BYTES);
324324
if (props == null) {
325-
byteLen += Integer.BYTES;
325+
byteLen = Math.addExact(byteLen, Integer.BYTES);
326326
} else {
327-
byteLen += Integer.BYTES;
327+
byteLen = Math.addExact(byteLen, Integer.BYTES);
328328
for (Map.Entry<String, String> entry : props.entrySet()) {
329-
byteLen += ReadWriteIOUtils.sizeToWrite(entry.getKey());
330-
byteLen += ReadWriteIOUtils.sizeToWrite(entry.getValue());
329+
byteLen = Math.addExact(byteLen, ReadWriteIOUtils.sizeToWrite(entry.getKey()));
330+
byteLen = Math.addExact(byteLen, ReadWriteIOUtils.sizeToWrite(entry.getValue()));
331331
}
332332
}
333333

‎java/tsfile/src/test/java/org/apache/tsfile/utils/ReadWriteIOUtilsTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,13 @@ public void mapSerdeTest() {
184184
Assert.assertNotNull(result);
185185
Assert.assertEquals(map, result);
186186

187+
ByteBuffer buffer = ByteBuffer.allocate(DEFAULT_BUFFER_SIZE);
188+
ReadWriteIOUtils.write(map, buffer);
189+
buffer.flip();
190+
result = ReadWriteIOUtils.readMap(buffer);
191+
Assert.assertNotNull(result);
192+
Assert.assertEquals(map, result);
193+
187194
// 7. null
188195
map = null;
189196
byteArrayOutputStream = new ByteArrayOutputStream(DEFAULT_BUFFER_SIZE);

0 commit comments

Comments
 (0)