feat(bigtable): expose the row key, mutation count and serialized size on RowMutationEntry

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<Mutation>, 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<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
     mutation.mutations.addAll(protos);
+    mutation.countAllTowardsSerializedSize(protos);
     return mutation;
   }
 
@@ -99,6 +107,7 @@ public static Mutation fromProtoUnsafe(List<com.google.bigtable.v2.Mutation> pro
   public static Mutation fromProtoUnsafe(Iterable<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
     mutation.mutations.addAll(protos);
+    mutation.countAllTowardsSerializedSize(protos);
     return mutation;
   }
 
@@ -115,6 +124,7 @@ public static Mutation fromProtoUnsafe(Iterable<com.google.bigtable.v2.Mutation>
   static Mutation fromProto(List<com.google.bigtable.v2.Mutation> 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<com.google.bigtable.v2.Mutation> deserialized =
         (ImmutableList<com.google.bigtable.v2.Mutation>) input.readObject();
     this.mutations = ImmutableList.<com.google.bigtable.v2.Mutation>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<com.google.bigtable.v2.Mutation> 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;

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(

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