diff --git a/api/src/test/java/org/apache/iceberg/variants/TestVariantTestUtil.java b/api/src/test/java/org/apache/iceberg/variants/TestVariantTestUtil.java new file mode 100644 index 000000000000..7dede95d105a --- /dev/null +++ b/api/src/test/java/org/apache/iceberg/variants/TestVariantTestUtil.java @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.variants; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Random; +import org.apache.iceberg.util.RandomUtil; +import org.junit.jupiter.api.Test; + +class TestVariantTestUtil { + private static final VariantMetadata EMPTY_METADATA = SerializedMetadata.EMPTY_V1_METADATA; + private static final SerializedPrimitive NULL = SerializedPrimitive.from(new byte[] {0x00}); + private static final SerializedPrimitive TRUE = SerializedPrimitive.from(new byte[] {0b100}); + private static final SerializedPrimitive FALSE = SerializedPrimitive.from(new byte[] {0b1000}); + private static final SerializedPrimitive I34 = SerializedPrimitive.from(new byte[] {0b1100, 34}); + private static final SerializedPrimitive I1234 = + SerializedPrimitive.from(new byte[] {0b10000, (byte) 0xD2, 0x04}); + private static final SerializedShortString A = + SerializedShortString.from(new byte[] {0b101, 'a'}); + private static final SerializedShortString B = + SerializedShortString.from(new byte[] {0b101, 'b'}); + private static final SerializedShortString C = + SerializedShortString.from(new byte[] {0b101, 'c'}); + private static final SerializedShortString ICEBERG = + SerializedShortString.from(new byte[] {0b11101, 'i', 'c', 'e', 'b', 'e', 'r', 'g'}); + + @Test + void serializedPrimitiveSizeMatchesBuffer() { + assertThat(VariantTestUtil.sizeInBytes(NULL)).isEqualTo(NULL.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(TRUE)).isEqualTo(TRUE.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(I34)).isEqualTo(I34.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(I1234)).isEqualTo(I1234.sizeInBytes()); + } + + @Test + void serializedShortStringSizeMatchesBuffer() { + assertThat(VariantTestUtil.sizeInBytes(A)).isEqualTo(A.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(ICEBERG)).isEqualTo(ICEBERG.sizeInBytes()); + } + + @Test + void serializedLongStringSizeMatchesBuffer() { + String longString = RandomUtil.generateString(300, new Random(12345)); + VariantPrimitive longVariant = VariantTestUtil.createString(longString); + assertThat(VariantTestUtil.sizeInBytes(longVariant)).isEqualTo(longVariant.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(longVariant)).isEqualTo(5 + longString.length()); + } + + @Test + void serializedArraySizeMatchesBuffer() { + ByteBuffer buffer = VariantTestUtil.createArray(A, B, C, I34); + SerializedArray array = SerializedArray.from(EMPTY_METADATA, buffer, buffer.get(0)); + + assertThat(VariantTestUtil.sizeInBytes(array)).isEqualTo(array.sizeInBytes()); + assertThat(VariantTestUtil.sizeInBytes(array)).isEqualTo(buffer.remaining()); + } + + @Test + void createArrayMatchesSizeInBytes() { + ByteBuffer buffer = VariantTestUtil.createArray(A, TRUE, ICEBERG, NULL, I1234, FALSE); + SerializedArray array = SerializedArray.from(EMPTY_METADATA, buffer, buffer.get(0)); + + assertThat(array.sizeInBytes()).isEqualTo(buffer.remaining()); + assertThat(VariantTestUtil.sizeInBytes(array)).isEqualTo(buffer.remaining()); + } + + @Test + void nestedArraySizeMatchesBuffer() { + ByteBuffer innerBuffer = VariantTestUtil.createArray(A, B, C); + SerializedArray inner = SerializedArray.from(EMPTY_METADATA, innerBuffer, innerBuffer.get(0)); + ByteBuffer outerBuffer = VariantTestUtil.createArray(inner, I34, ICEBERG); + SerializedArray outer = SerializedArray.from(EMPTY_METADATA, outerBuffer, outerBuffer.get(0)); + + assertThat(VariantTestUtil.sizeInBytes(outer)).isEqualTo(outerBuffer.remaining()); + assertThat(outer.sizeInBytes()).isEqualTo(outerBuffer.remaining()); + } + + @Test + void sizeMatchesForSerializedPrimitivesByType() { + // NULL: 1-byte header + assertThat(VariantTestUtil.sizeInBytes(NULL)).isEqualTo(1); + // BOOLEAN_TRUE / BOOLEAN_FALSE: 1-byte header + assertThat(VariantTestUtil.sizeInBytes(TRUE)).isEqualTo(1); + assertThat(VariantTestUtil.sizeInBytes(FALSE)).isEqualTo(1); + // INT8: 1-byte header + 1-byte value + assertThat(VariantTestUtil.sizeInBytes(I34)).isEqualTo(2); + // INT16: 1-byte header + 2-byte value + assertThat(VariantTestUtil.sizeInBytes(I1234)).isEqualTo(3); + // short string: 1-byte header + utf8 length + assertThat(VariantTestUtil.sizeInBytes(A)).isEqualTo(2); + assertThat(VariantTestUtil.sizeInBytes(ICEBERG)).isEqualTo(1 + "iceberg".length()); + } + + @Test + void sizeMatchesWrittenBytesForArray() { + ByteBuffer arrayBuffer = VariantTestUtil.createArray(A, I34, ICEBERG, NULL, TRUE); + SerializedArray array = SerializedArray.from(EMPTY_METADATA, arrayBuffer, arrayBuffer.get(0)); + + // computed size must match the number of bytes actually written by writeTo + ByteBuffer scratch = + ByteBuffer.allocate(VariantTestUtil.sizeInBytes(array)).order(ByteOrder.LITTLE_ENDIAN); + int written = array.writeTo(scratch, 0); + assertThat(written).isEqualTo(VariantTestUtil.sizeInBytes(array)); + } +} diff --git a/api/src/test/java/org/apache/iceberg/variants/VariantTestUtil.java b/api/src/test/java/org/apache/iceberg/variants/VariantTestUtil.java index 8a1cd515dff6..ff3225969a5c 100644 --- a/api/src/test/java/org/apache/iceberg/variants/VariantTestUtil.java +++ b/api/src/test/java/org/apache/iceberg/variants/VariantTestUtil.java @@ -100,6 +100,58 @@ public static void assertVariantString(VariantValue actual, String expected, int assertThat(VariantUtil.sizeOf(actual.asPrimitive().sizeInBytes())).isEqualTo(offsetSize); } + /** + * Computes the serialized size in bytes of a variant value without allocating an output buffer. + * + *

Pre-serialized values report the size of their backing buffer. Primitives derive their size + * from the physical type (and the encoded byte length for {@code STRING} / {@code BINARY}). + * Arrays are computed recursively from their elements. Non-serialized objects delegate to {@link + * VariantValue#sizeInBytes()} because determining the field id size requires the object's own + * metadata, which is not exposed on the {@link VariantObject} interface. + */ + static int sizeInBytes(VariantValue value) { + if (value instanceof Serialized) { + return ((Serialized) value).buffer().remaining(); + } + + return switch (value.type()) { + case NULL, BOOLEAN_TRUE, BOOLEAN_FALSE -> 1; + case INT8 -> 2; + case INT16 -> 3; + case INT32, DATE, FLOAT -> 5; + case INT64, DOUBLE, TIMESTAMPTZ, TIMESTAMPNTZ, TIMESTAMPTZ_NANOS, TIMESTAMPNTZ_NANOS, TIME -> + 9; + case DECIMAL4 -> 6; + case DECIMAL8 -> 10; + case DECIMAL16 -> 18; + case UUID -> 17; + case STRING -> stringSizeInBytes((String) value.asPrimitive().get()); + case BINARY -> 5 + ((ByteBuffer) value.asPrimitive().get()).remaining(); + case ARRAY -> arraySizeInBytes(value.asArray()); + case OBJECT -> value.sizeInBytes(); + }; + } + + private static int stringSizeInBytes(String value) { + int utf8Length = value.getBytes(StandardCharsets.UTF_8).length; + return (utf8Length <= MAX_SHORT_STRING_LENGTH) ? 1 + utf8Length : 5 + utf8Length; + } + + private static int arraySizeInBytes(VariantArray array) { + int numElements = array.numElements(); + int dataSize = 0; + for (int i = 0; i < numElements; i += 1) { + dataSize += sizeInBytes(array.get(i)); + } + + int numElementsSize = (numElements > 0xFF) ? 4 : 1; + int offsetSize = VariantUtil.sizeOf(dataSize); + return 1 /* header */ + + numElementsSize + + (1 + numElements) * offsetSize /* offset list size */ + + dataSize; + } + private static byte primitiveHeader(int primitiveType) { return (byte) (primitiveType << 2); } @@ -217,7 +269,7 @@ public static ByteBuffer createObject(VariantMetadata metadata, Map field : data.entrySet()) { - dataSize += field.getValue().sizeInBytes(); + dataSize += sizeInBytes(field.getValue()); } // field ID size is the size needed to store the largest field ID in the data @@ -270,8 +322,7 @@ public static ByteBuffer createArray(VariantValue... values) { int dataSize = 0; for (VariantValue value : values) { - // TODO: produce size for every variant without serializing - dataSize += value.sizeInBytes(); + dataSize += sizeInBytes(value); } // offset size is the size needed to store the length of the data section