From d287ce653de129c3b7a4a13597f2be63b3a33b20 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 8 Aug 2026 21:38:22 +0900 Subject: [PATCH] feat(bigtable): expose the row key, mutation count and serialized size on RowMutationEntry RowMutationEntry exposed no accessor for its serialized size, its row key or its mutation count, so a caller that needs any of them had to call the @InternalApi toProto() and build a whole MutateRowsRequest.Entry to read one number. Bounding in-flight memory by bytes is the motivating case: a single row mutation may be megabytes, so bounding by entry count bounds nothing useful. Mutation.byteSize could not simply be exposed - it sums each child mutation's own serialized size and omits the row-key field and the per-mutation tag and length prefixes. Mutation now accumulates the framed size alongside it, and RowMutationEntry.getSerializedSize() adds the row-key field to it. Reads are O(1) and allocate nothing. The counters are maintained on the fromProto factories too, which bypass addMutation, and are transient and recomputed in readObject so an instance restored from an older serialized form is correct rather than reporting zero. numMutations and byteSize are deliberately left alone. --- .../bigtable/data/v2/models/Mutation.java | 34 ++++++ .../data/v2/models/RowMutationEntry.java | 21 ++++ .../data/v2/models/RowMutationEntryTest.java | 110 ++++++++++++++++++ 3 files changed, 165 insertions(+) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java index dc55756241e9..767120af276b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java @@ -17,6 +17,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; +import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.Mutation.AddToCell; import com.google.bigtable.v2.Mutation.DeleteFromColumn; import com.google.bigtable.v2.Mutation.DeleteFromFamily; @@ -28,6 +29,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.primitives.Longs; import com.google.protobuf.ByteString; +import com.google.protobuf.CodedOutputStream; import java.io.IOException; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; @@ -60,6 +62,11 @@ public final class Mutation implements MutationApi, Serializable { private int numMutations; private long byteSize; + private transient int mutationCount; + + /** Tag, length prefix and body of each mutation. Recomputed by readObject. */ + private transient long mutationsSerializedSize; + /** Creates new instance of Mutation object. */ public static Mutation create() { return new Mutation(false); @@ -86,6 +93,7 @@ public static Mutation createUnsafe() { public static Mutation fromProtoUnsafe(List protos) { Mutation mutation = new Mutation(true); mutation.mutations.addAll(protos); + mutation.countAllTowardsSerializedSize(protos); return mutation; } @@ -99,6 +107,7 @@ public static Mutation fromProtoUnsafe(List pro public static Mutation fromProtoUnsafe(Iterable protos) { Mutation mutation = new Mutation(true); mutation.mutations.addAll(protos); + mutation.countAllTowardsSerializedSize(protos); return mutation; } @@ -115,6 +124,7 @@ public static Mutation fromProtoUnsafe(Iterable static Mutation fromProto(List protos) { Mutation mutation = new Mutation(false); mutation.mutations.addAll(protos); + mutation.countAllTowardsSerializedSize(protos); return mutation; } @@ -129,6 +139,7 @@ private void readObject(ObjectInputStream input) throws IOException, ClassNotFou ImmutableList deserialized = (ImmutableList) input.readObject(); this.mutations = ImmutableList.builder().addAll(deserialized); + countAllTowardsSerializedSize(deserialized); } private void writeObject(ObjectOutputStream output) throws IOException { @@ -335,10 +346,33 @@ private void addMutation(com.google.bigtable.v2.Mutation mutation) { numMutations++; byteSize += mutation.getSerializedSize(); + countTowardsSerializedSize(mutation); mutations.add(mutation); } + private void countTowardsSerializedSize(com.google.bigtable.v2.Mutation mutation) { + mutationCount++; + // Every request proto numbers this field below 16, so the tag is one byte in all of them. + mutationsSerializedSize += + CodedOutputStream.computeMessageSize( + MutateRowsRequest.Entry.MUTATIONS_FIELD_NUMBER, mutation); + } + + private void countAllTowardsSerializedSize(Iterable protos) { + for (com.google.bigtable.v2.Mutation proto : protos) { + countTowardsSerializedSize(proto); + } + } + + int getMutationCount() { + return mutationCount; + } + + long getMutationsSerializedSize() { + return mutationsSerializedSize; + } + private static ByteString wrapByteString(String str) { if (str == null) { return null; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntry.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntry.java index 2cd2529d53ef..89493c41462f 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntry.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntry.java @@ -20,6 +20,7 @@ import com.google.bigtable.v2.MutateRowsRequest; import com.google.common.base.Preconditions; import com.google.protobuf.ByteString; +import com.google.protobuf.CodedOutputStream; import java.io.Serializable; import javax.annotation.Nonnull; @@ -200,6 +201,26 @@ public RowMutationEntry mergeToCell( return this; } + /** Returns the row key these mutations apply to. */ + public ByteString getRowKey() { + return key; + } + + /** Returns the number of mutations accumulated in this entry. */ + public int getMutationCount() { + return mutation.getMutationCount(); + } + + /** Returns the serialized size of {@link #toProto()}, without building it. */ + public long getSerializedSize() { + long size = mutation.getMutationsSerializedSize(); + if (!key.isEmpty()) { + // An empty row key is not written at all, so it contributes no tag or length prefix. + size += CodedOutputStream.computeBytesSize(MutateRowsRequest.Entry.ROW_KEY_FIELD_NUMBER, key); + } + return size; + } + @InternalApi public MutateRowsRequest.Entry toProto() { Preconditions.checkArgument( diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntryTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntryTest.java index 4f616d937937..1598d54b7ec0 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntryTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/RowMutationEntryTest.java @@ -165,4 +165,114 @@ public void testWithLongValue() { .setValue(ByteString.copyFrom(Longs.toByteArray(100_000L))) .build()); } + + @Test + public void getRowKeyTest() { + ByteString rowKey = ByteString.copyFromUtf8("row-key"); + + assertThat(RowMutationEntry.create(rowKey).getRowKey()).isEqualTo(rowKey); + assertThat(RowMutationEntry.create("row-key").getRowKey()).isEqualTo(rowKey); + assertThat(RowMutationEntry.create(ByteString.EMPTY).getRowKey()).isEqualTo(ByteString.EMPTY); + } + + @Test + public void getMutationCountTest() { + assertThat(RowMutationEntry.create("row-key").getMutationCount()).isEqualTo(0); + assertThat( + RowMutationEntry.create("row-key") + .setCell("fake-family", "q1", 10_000L, "v1") + .setCell("fake-family", "q2", 10_000L, "v2") + .deleteFamily("fake-family") + .deleteRow() + .getMutationCount()) + .isEqualTo(4); + } + + @Test + public void getSerializedSizeMatchesProtoTest() { + assertSizeMatchesProto(RowMutationEntry.create("row-key")); + assertSizeMatchesProto(RowMutationEntry.create(ByteString.EMPTY)); + assertSizeMatchesProto( + RowMutationEntry.create("row-key").setCell("fake-family", "q", 10_000L, "v")); + assertSizeMatchesProto( + RowMutationEntry.create("row-key").setCell("fake-family", "q", 10_000L, repeat("v", 200))); + assertSizeMatchesProto( + RowMutationEntry.create(ByteString.copyFromUtf8(repeat("k", 5_000))) + .setCell("fake-family", "q", 10_000L, repeat("v", 70_000))); + assertSizeMatchesProto( + RowMutationEntry.create("row-key") + .setCell("fake-family", "q1", 10_000L, "v1") + .setCell("fake-family", "q2", 10_000L, 100_000L) + .deleteCells("fake-family", "q1") + .deleteFamily("fake-family") + .deleteRow()); + } + + @Test + public void getSerializedSizeMatchesProtoForManyMutationsTest() { + RowMutationEntry underTest = RowMutationEntry.create("row-key"); + for (int i = 0; i < 300; i++) { + underTest.setCell("fake-family", "q" + i, 10_000L, "value-" + i); + } + + assertThat(underTest.getMutationCount()).isEqualTo(300); + assertSizeMatchesProto(underTest); + } + + @Test + public void getSerializedSizeFromMutationUnsafeTest() { + com.google.cloud.bigtable.data.v2.models.Mutation mutation = + com.google.cloud.bigtable.data.v2.models.Mutation.fromProtoUnsafe( + ImmutableList.of( + Mutation.newBuilder() + .setSetCell( + Mutation.SetCell.newBuilder() + .setFamilyName("fake-family") + .setColumnQualifier(ByteString.copyFromUtf8("q")) + .setTimestampMicros(10_000L) + .setValue(ByteString.copyFromUtf8(repeat("v", 500)))) + .build(), + Mutation.newBuilder() + .setDeleteFromRow(Mutation.DeleteFromRow.getDefaultInstance()) + .build())); + RowMutationEntry underTest = + RowMutationEntry.createFromMutationUnsafe(ByteString.copyFromUtf8("row-key"), mutation); + + assertThat(underTest.getMutationCount()).isEqualTo(2); + assertSizeMatchesProto(underTest); + } + + @Test + public void getSerializedSizeSurvivesSerializationTest() + throws IOException, ClassNotFoundException { + RowMutationEntry underTest = + RowMutationEntry.create("row-key") + .setCell("fake-family", "q", 10_000L, repeat("v", 500)) + .deleteFamily("fake-family"); + + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + ObjectOutputStream oos = new ObjectOutputStream(bos); + oos.writeObject(underTest); + oos.close(); + ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray())); + RowMutationEntry actual = (RowMutationEntry) ois.readObject(); + + assertThat(actual.getRowKey()).isEqualTo(underTest.getRowKey()); + assertThat(actual.getMutationCount()).isEqualTo(underTest.getMutationCount()); + assertThat(actual.getSerializedSize()).isEqualTo(underTest.getSerializedSize()); + assertSizeMatchesProto(actual); + } + + private static void assertSizeMatchesProto(RowMutationEntry entry) { + assertThat(entry.getSerializedSize()).isEqualTo(entry.toProto().getSerializedSize()); + assertThat(entry.getMutationCount()).isEqualTo(entry.toProto().getMutationsCount()); + } + + private static String repeat(String unit, int times) { + StringBuilder sb = new StringBuilder(unit.length() * times); + for (int i = 0; i < times; i++) { + sb.append(unit); + } + return sb.toString(); + } }