diff --git a/core/src/main/java/org/apache/iceberg/variants/ShreddedObject.java b/core/src/main/java/org/apache/iceberg/variants/ShreddedObject.java index bd1fbad07547..3391a2378c82 100644 --- a/core/src/main/java/org/apache/iceberg/variants/ShreddedObject.java +++ b/core/src/main/java/org/apache/iceberg/variants/ShreddedObject.java @@ -24,12 +24,10 @@ import java.util.Set; import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; -import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.relocated.com.google.common.collect.Sets; import org.apache.iceberg.util.ByteBuffers; -import org.apache.iceberg.util.SortedMerge; /** * A variant Object that handles full or partial shredding. @@ -143,9 +141,7 @@ public int writeTo(ByteBuffer buffer, int offset) { /** Common state for {@link #size()} and {@link #writeTo(ByteBuffer, int)} */ private static class SerializationState { - private final VariantMetadata metadata; - private final Map unshreddedFields; - private final Map shreddedFields; + private final Entry[] entries; private final int dataSize; private final int numElements; private final boolean isLarge; @@ -157,48 +153,53 @@ private SerializationState( VariantObject unshredded, Map shredded, Set removedFields) { - this.metadata = metadata; - // field ID size is the size needed to store the largest field ID in the data this.fieldIdSize = VariantUtil.sizeOf(metadata.dictionarySize()); - this.shreddedFields = Maps.newHashMap(shredded); + Map sorted = Maps.newTreeMap(); int totalDataSize = 0; - // get the unshredded field names and values as byte buffers - ImmutableMap.Builder unshreddedBuilder = ImmutableMap.builder(); + + for (Map.Entry field : shredded.entrySet()) { + String name = field.getKey(); + int id = metadata.id(name); + Preconditions.checkState(id >= 0, "Invalid metadata, missing: %s", name); + VariantValue value = field.getValue(); + int valueSize = value.sizeInBytes(); + sorted.put(name, Entry.ofShredded(id, value, valueSize)); + totalDataSize += valueSize; + } + if (unshredded instanceof SerializedObject) { - // for serialized objects, use existing buffers instead of materializing values SerializedObject serialized = (SerializedObject) unshredded; for (Map.Entry field : serialized.fields()) { - // if the value is replaced by an unshredded field, don't include it String name = field.getKey(); - boolean replaced = shreddedFields.containsKey(name) || removedFields.contains(name); - if (!replaced) { - ByteBuffer value = serialized.sliceValue(field.getValue()); - unshreddedBuilder.put(name, value); - totalDataSize += value.remaining(); + if (sorted.containsKey(name) || removedFields.contains(name)) { + continue; } + int id = metadata.id(name); + Preconditions.checkState(id >= 0, "Invalid metadata, missing: %s", name); + ByteBuffer value = serialized.sliceValue(field.getValue()); + int valueSize = value.remaining(); + sorted.put(name, Entry.ofBuffer(id, value, valueSize)); + totalDataSize += valueSize; } } else if (unshredded != null) { for (String name : unshredded.fieldNames()) { - boolean replaced = shreddedFields.containsKey(name) || removedFields.contains(name); - if (!replaced) { - shreddedFields.put(name, unshredded.get(name)); + if (sorted.containsKey(name) || removedFields.contains(name)) { + continue; } + int id = metadata.id(name); + Preconditions.checkState(id >= 0, "Invalid metadata, missing: %s", name); + VariantValue value = unshredded.get(name); + int valueSize = value.sizeInBytes(); + sorted.put(name, Entry.ofShredded(id, value, valueSize)); + totalDataSize += valueSize; } } - this.unshreddedFields = unshreddedBuilder.build(); - // duplicates are suppressed when creating unshreddedFields - this.numElements = unshreddedFields.size() + shreddedFields.size(); - // object is large if the number of elements can't be stored in 1 byte + this.entries = sorted.values().toArray(new Entry[0]); + this.numElements = entries.length; this.isLarge = numElements > 0xFF; - - for (VariantValue value : shreddedFields.values()) { - totalDataSize += value.sizeInBytes(); - } - this.dataSize = totalDataSize; - // offset size is the size needed to store the length of the data section this.offsetSize = VariantUtil.sizeOf(totalDataSize); } @@ -211,60 +212,60 @@ private int size() { } private int writeTo(ByteBuffer buffer, int offset) { - int fieldIdListOffset = - offset + 1 /* header size */ + (isLarge ? 4 : 1) /* num elements size */; + int numElementsSize = isLarge ? 4 : 1; + int fieldIdListOffset = offset + 1 /* header size */ + numElementsSize; int offsetListOffset = fieldIdListOffset + (numElements * fieldIdSize); int dataOffset = offsetListOffset + ((1 + numElements) * offsetSize); byte header = VariantUtil.objectHeader(isLarge, fieldIdSize, offsetSize); ByteBuffers.writeByte(buffer, header, offset); - ByteBuffers.writeLittleEndianUnsigned(buffer, numElements, offset + 1, isLarge ? 4 : 1); - - // neither iterable is closeable, so it is okay to use Iterable - Iterable fields = - SortedMerge.of( - () -> unshreddedFields.keySet().stream().sorted().iterator(), - () -> shreddedFields.keySet().stream().sorted().iterator()); + ByteBuffers.writeLittleEndianUnsigned(buffer, numElements, offset + 1, numElementsSize); int nextValueOffset = 0; - int index = 0; - for (String field : fields) { - // write the field ID from the metadata dictionary - int id = metadata.id(field); - Preconditions.checkState(id >= 0, "Invalid metadata, missing: %s", field); + for (int index = 0; index < numElements; index++) { + Entry entry = entries[index]; ByteBuffers.writeLittleEndianUnsigned( - buffer, id, fieldIdListOffset + (index * fieldIdSize), fieldIdSize); - // write the data offset + buffer, entry.id, fieldIdListOffset + (index * fieldIdSize), fieldIdSize); ByteBuffers.writeLittleEndianUnsigned( buffer, nextValueOffset, offsetListOffset + (index * offsetSize), offsetSize); - // copy or serialize the value into the data section - int valueSize; - VariantValue shreddedValue = shreddedFields.get(field); - if (shreddedValue != null) { - valueSize = shreddedValue.writeTo(buffer, dataOffset + nextValueOffset); + if (entry.shredded != null) { + entry.shredded.writeTo(buffer, dataOffset + nextValueOffset); } else { - ByteBuffer unshreddedValue = unshreddedFields.get(field); - buffer.put( - dataOffset + nextValueOffset, - unshreddedValue, - unshreddedValue.position(), - unshreddedValue.remaining()); - valueSize = unshreddedValue.remaining(); + ByteBuffer src = entry.buffer; + buffer.put(dataOffset + nextValueOffset, src, src.position(), src.remaining()); } - - // update tracking - nextValueOffset += valueSize; - index += 1; + nextValueOffset += entry.size; } // write the final size of the data section ByteBuffers.writeLittleEndianUnsigned( - buffer, nextValueOffset, offsetListOffset + (index * offsetSize), offsetSize); + buffer, nextValueOffset, offsetListOffset + (numElements * offsetSize), offsetSize); - // return the total size return (dataOffset - offset) + dataSize; } + + private static final class Entry { + private final int id; + private final VariantValue shredded; + private final ByteBuffer buffer; + private final int size; + + private Entry(int id, VariantValue shredded, ByteBuffer buffer, int size) { + this.id = id; + this.shredded = shredded; + this.buffer = buffer; + this.size = size; + } + + static Entry ofShredded(int id, VariantValue value, int size) { + return new Entry(id, value, null, size); + } + + static Entry ofBuffer(int id, ByteBuffer buffer, int size) { + return new Entry(id, null, buffer, size); + } + } } @Override diff --git a/core/src/test/java/org/apache/iceberg/variants/TestShreddedObject.java b/core/src/test/java/org/apache/iceberg/variants/TestShreddedObject.java index cea1c6922c09..704cc78e0027 100644 --- a/core/src/test/java/org/apache/iceberg/variants/TestShreddedObject.java +++ b/core/src/test/java/org/apache/iceberg/variants/TestShreddedObject.java @@ -280,6 +280,59 @@ public void testRemoveInvalidatesSerializationState() { assertThat(actual.get("b")).as("removed field must not be serialized").isNull(); } + @Test + public void testPutAfterRemoveIsSerialized() { + ShreddedObject partial = createUnshreddedObject(FIELDS); + partial.remove("b"); + partial.put("b", Variants.of("replaced")); + + VariantValue value = roundTripMinimalBuffer(partial, partial.metadata()); + + SerializedObject actual = (SerializedObject) value; + assertThat(actual.numFields()).isEqualTo(3); + assertThat(actual.get("a")).isInstanceOf(VariantPrimitive.class); + assertThat(actual.get("a").asPrimitive().get()).isEqualTo(34); + VariantTestUtil.assertVariantString(actual.get("b"), "replaced"); + assertThat(actual.get("c")).isInstanceOf(VariantPrimitive.class); + assertThat(actual.get("c").asPrimitive().get()).isEqualTo(new BigDecimal("12.21")); + } + + @Test + public void testEmptyObject() { + Set names = Sets.newLinkedHashSet(); + names.add("a"); + names.add("b"); + names.add("c"); + ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(names, false); + VariantMetadata metadata = Variants.metadata(metadataBuffer); + ShreddedObject object = new ShreddedObject(metadata); + + assertThat(object.numFields()).isEqualTo(0); + assertThat(object.sizeInBytes()).isEqualTo(3); + + VariantValue value = roundTripMinimalBuffer(object, metadata); + assertThat(value).isInstanceOf(SerializedObject.class); + assertThat(((SerializedObject) value).numFields()).isEqualTo(0); + } + + @Test + public void testRepeatedWriteToProducesSameBytes() { + ShreddedObject object = createUnshreddedObject(FIELDS); + object.put("c", Variants.ofIsoDate("2024-10-12")); + object.remove("b"); + + int size = object.sizeInBytes(); + ByteBuffer reference = ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN); + object.writeTo(reference, 0); + + for (int i = 0; i < 10; i++) { + ByteBuffer buf = ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN); + int written = object.writeTo(buf, 0); + assertThat(written).as("iteration %d", i).isEqualTo(size); + assertThat(buf).as("iteration %d", i).isEqualTo(reference); + } + } + @ParameterizedTest @ValueSource(ints = {300, 70_000, 16_777_300}) public void testMultiByteOffsets(int len) {