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(); + } }