com.google.bigtable.v2.Mutation Java Examples

/** Tests that when writing to a non-existent table, the write fails. */
@Test
public void testWritingFailsTableDoesNotExist() throws Exception {
  final String table = "TEST-TABLE";

  PCollection<KV<ByteString, Iterable<Mutation>>> emptyInput =
      p.apply(
          Create.empty(
              KvCoder.of(ByteStringCoder.of(), IterableCoder.of(ProtoCoder.of(Mutation.class)))));

  // Exception will be thrown by write.validate() when writeToDynamic is applied.
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage(String.format("Table %s does not exist", table));

  emptyInput.apply("write", defaultWrite.withTableId(table));
  p.run();
}
 

@Override
public CompletionStage<MutateRowResponse> writeRecord(KV<ByteString, Iterable<Mutation>> record)
    throws IOException {
  MutateRowsRequest.Entry request =
      MutateRowsRequest.Entry.newBuilder()
          .setRowKey(record.getKey())
          .addAllMutations(record.getValue())
          .build();

  CompletableFuture<MutateRowResponse> result = new CompletableFuture<>();
  Futures.addCallback(
      new VendoredListenableFutureAdapter<>(bulkMutation.add(request)),
      new FutureCallback<MutateRowResponse>() {
        @Override
        public void onSuccess(MutateRowResponse mutateRowResponse) {
          result.complete(mutateRowResponse);
        }

        @Override
        public void onFailure(Throwable throwable) {
          result.completeExceptionally(throwable);
        }
      },
      directExecutor());
  return result;
}
 

@ProcessElement
public void processElement(ProcessContext c) {
	InputContent i = c.element();
	String jobName = c.getPipelineOptions().getJobName();
	ByteString rowkey = ByteString.copyFromUtf8(jobName + "#" + i.expectedDocumentHash); 
	ByteString value = ByteString.copyFromUtf8(i.text);
	
	Iterable<Mutation> mutations =
		ImmutableList.of(Mutation.newBuilder()
			.setSetCell(
				Mutation.SetCell.newBuilder()
					.setFamilyName(IndexerPipelineUtils.DEAD_LETTER_TABLE_ERR_CF)
					.setColumnQualifier(ByteString.copyFromUtf8("text"))
					.setValue(value)
			)
               .build());
	
	c.output(KV.of(rowkey, mutations));			
}
 

@ProcessElement
public void processElement(ProcessContext c) {
	InputContent i = c.element();
	String jobName = c.getPipelineOptions().getJobName();
	ByteString rowkey = ByteString.copyFromUtf8(jobName + "#" + i.expectedDocumentHash); 
	ByteString value = ByteString.copyFromUtf8(i.text);
	
	Iterable<Mutation> mutations =
		ImmutableList.of(Mutation.newBuilder()
			.setSetCell(
				Mutation.SetCell.newBuilder()
					.setFamilyName(IndexerPipelineUtils.DEAD_LETTER_TABLE_ERR_CF)
					.setColumnQualifier(ByteString.copyFromUtf8("text"))
					.setValue(value)
			)
               .build());
	
	c.output(KV.of(rowkey, mutations));			
}
 

@Override
public BigtableMutation deleteCellsFromColumn(final String columnFamily,
                                              final String columnQualifier,
                                              final Optional<Long> startTimestampMicros,
                                              final Optional<Long> endTimestampMicros) {
  final TimestampRange.Builder timestampRange = TimestampRange.newBuilder();
  startTimestampMicros.ifPresent(timestampRange::setStartTimestampMicros);
  endTimestampMicros.ifPresent(timestampRange::setEndTimestampMicros);

  final Mutation.DeleteFromColumn.Builder deleteFromColumn =
      Mutation.DeleteFromColumn.newBuilder()
          .setFamilyName(columnFamily)
          .setColumnQualifier(ByteString.copyFromUtf8(columnQualifier))
          .setTimeRange(timestampRange);

  mutateRowRequest.addMutations(Mutation.newBuilder().setDeleteFromColumn(deleteFromColumn));
  return this;
}
 

@Test
public void testDeleteRow() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl = (BigtableMutationImpl) this.bigtableMutation.deleteRow();
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(1, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
  final Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
  assertEquals(Mutation.MutationCase.DELETE_FROM_ROW, mutation.getMutationCase());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
  assertEquals(Mutation.DeleteFromColumn.getDefaultInstance(), mutation.getDeleteFromColumn());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@Test
public void testDeleteColumnFamily() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl =
          (BigtableMutationImpl) this.bigtableMutation.deleteColumnFamily("family");
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(1, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
  final Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
  assertEquals(Mutation.MutationCase.DELETE_FROM_FAMILY, mutation.getMutationCase());
  assertEquals("family", mutation.getDeleteFromFamily().getFamilyName());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromColumn.getDefaultInstance(), mutation.getDeleteFromColumn());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@Test
public void testDeleteColumn() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl =
          (BigtableMutationImpl) this.bigtableMutation.deleteColumn("family:qualifier");
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(1, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
  final Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
  assertEquals(Mutation.MutationCase.DELETE_FROM_COLUMN, mutation.getMutationCase());
  assertEquals("family", mutation.getDeleteFromColumn().getFamilyName());
  assertEquals("qualifier", mutation.getDeleteFromColumn().getColumnQualifier().toStringUtf8());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@Test
public void testDeleteColumnFamilyAndQualifier() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl =
          (BigtableMutationImpl) this.bigtableMutation.deleteColumn("family", "qualifier");
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(1, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
  final Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
  assertEquals(Mutation.MutationCase.DELETE_FROM_COLUMN, mutation.getMutationCase());
  assertEquals("family", mutation.getDeleteFromColumn().getFamilyName());
  assertEquals("qualifier", mutation.getDeleteFromColumn().getColumnQualifier().toStringUtf8());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

static void addBigtableMutation(
    KV<ByteString, Iterable<Mutation>> rowMutations,
    String family,
    String qualifier,
    long timestamp,
    String value) {
  SetCell setCell =
      SetCell.newBuilder()
          .setFamilyName(family)
          .setColumnQualifier(toByteString(qualifier))
          .setTimestampMicros(timestamp)
          .setValue(toByteString(value))
          .build();
  Mutation mutation = Mutation.newBuilder().setSetCell(setCell).build();
  ((List) rowMutations.getValue()).add(mutation);
}
 

/**
 * This test ensures that protobuf creation and interactions with {@link BulkMutation} work as
 * expected.
 *
 * @throws IOException
 * @throws InterruptedException
 */
@Test
public void testWrite() throws IOException, InterruptedException {
  BigtableService.Writer underTest =
      new BigtableServiceImpl.BigtableWriterImpl(mockSession, TABLE_NAME);

  Mutation mutation =
      Mutation.newBuilder()
          .setSetCell(SetCell.newBuilder().setFamilyName("Family").build())
          .build();
  ByteString key = ByteString.copyFromUtf8("key");

  SettableFuture<MutateRowResponse> fakeResponse = SettableFuture.create();
  when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(fakeResponse);

  underTest.writeRecord(KV.of(key, ImmutableList.of(mutation)));
  Entry expected =
      MutateRowsRequest.Entry.newBuilder().setRowKey(key).addMutations(mutation).build();
  verify(mockBulkMutation, times(1)).add(expected);

  underTest.close();
  verify(mockBulkMutation, times(1)).flush();
}
 

@Override
public CompletionStage<MutateRowResponse> writeRecord(
    KV<ByteString, Iterable<Mutation>> record) {
  service.verifyTableExists(tableId);
  Map<ByteString, ByteString> table = service.getTable(tableId);
  ByteString key = record.getKey();
  for (Mutation m : record.getValue()) {
    SetCell cell = m.getSetCell();
    if (cell.getValue().isEmpty()) {
      CompletableFuture<MutateRowResponse> result = new CompletableFuture<>();
      result.completeExceptionally(new IOException("cell value missing"));
      return result;
    }
    table.put(key, cell.getValue());
  }
  return CompletableFuture.completedFuture(MutateRowResponse.getDefaultInstance());
}
 

public BigtableWriteException(KV<ByteString, Iterable<Mutation>> record, Throwable cause) {
  super(
      String.format(
          "Error mutating row %s with mutations %s",
          record.getKey().toStringUtf8(), record.getValue()),
      cause);
}
 

@Override
public BigtableMutation deleteColumnFamily(String columnFamily) {
  final Mutation.DeleteFromFamily.Builder deleteFromFamily =
      Mutation.DeleteFromFamily.newBuilder().setFamilyName(columnFamily);
  mutateRowRequest.addMutations(Mutation.newBuilder().setDeleteFromFamily(deleteFromFamily));
  return this;
}
 

@Override
public PCollection<BigtableWriteResult> expand(
    PCollection<KV<ByteString, Iterable<Mutation>>> input) {
  bigtableConfig.validate();

  return input.apply(ParDo.of(new BigtableWriterFn(bigtableConfig)));
}
 

@Test
public void testCoderInference() {
  // Lambdas erase too much type information - use an anonymous class here.
  SerializableFunction<SchemaAndRecord, KV<ByteString, Mutation>> parseFn =
      new SerializableFunction<SchemaAndRecord, KV<ByteString, Mutation>>() {
        @Override
        public KV<ByteString, Mutation> apply(SchemaAndRecord input) {
          return null;
        }
      };

  assertEquals(
      KvCoder.of(ByteStringCoder.of(), ProtoCoder.of(Mutation.class)),
      BigQueryIO.read(parseFn).inferCoder(CoderRegistry.createDefault()));
}
 

/** Helper function to make a single row mutation to be written. */
private static KV<ByteString, Iterable<Mutation>> makeWrite(String key, String value) {
  ByteString rowKey = ByteString.copyFromUtf8(key);
  Iterable<Mutation> mutations =
      ImmutableList.of(
          Mutation.newBuilder()
              .setSetCell(SetCell.newBuilder().setValue(ByteString.copyFromUtf8(value)))
              .build());
  return KV.of(rowKey, mutations);
}
 

/** Tests that when writing to a non-existent table, the write fails. */
@Test
public void testTableCheckIgnoredWhenCanNotAccessConfig() throws Exception {
  PCollection<KV<ByteString, Iterable<Mutation>>> emptyInput =
      p.apply(
          Create.empty(
              KvCoder.of(ByteStringCoder.of(), IterableCoder.of(ProtoCoder.of(Mutation.class)))));

  emptyInput.apply("write", defaultWrite.withTableId(NOT_ACCESSIBLE_VALUE));
  p.run();
}
 

@Override
public BigtableMutation write(final String columnFamily, final String columnQualifier,
                              final ByteString value, final long timestampMicros) {
  final Mutation.SetCell.Builder setCell = Mutation.SetCell.newBuilder()
      .setFamilyName(columnFamily)
      .setColumnQualifier(ByteString.copyFromUtf8(columnQualifier))
      .setValue(value)
      .setTimestampMicros(timestampMicros);
  mutateRowRequest.addMutations(Mutation.newBuilder().setSetCell(setCell));
  return this;
}
 

private Mutation createMutation(String columnFamily, String columnName, ByteString value) {
  SetCell setcell =
      SetCell.newBuilder()
          .setFamilyName(columnFamily)
          .setColumnQualifier(ByteString.copyFrom(Bytes.toBytes(columnName)))
          .setValue(value)
          .build();

  return Mutation.newBuilder().setSetCell(setcell).build();
}
 

@Test
public void applyAvroToBigtableFnSplitLargeRows() throws Exception {
  ValueProvider<Boolean> splitlargeRows = ValueProvider.StaticValueProvider.of(true);
  BigtableRow avroRow1 = createAvroRow("row1");
  addAvroCell(avroRow1, "family1", "column1", 1, "value1");
  addAvroCell(avroRow1, "family1", "column1", 2, "value2");
  addAvroCell(avroRow1, "family1", "column2", 1, "value3");
  addAvroCell(avroRow1, "family2", "column1", 1, "value4");
  BigtableRow avroRow2 = createAvroRow("row2");
  addAvroCell(avroRow2, "family2", "column2", 2, "value2");
  List<BigtableRow> avroRows = ImmutableList.of(avroRow1, avroRow2);

  KV<ByteString, Iterable<Mutation>> rowMutations1 = createBigtableRowMutations("row1");
  addBigtableMutation(rowMutations1, "family1", "column1", 1, "value1");
  addBigtableMutation(rowMutations1, "family1", "column1", 2, "value2");
  KV<ByteString, Iterable<Mutation>> rowMutations2 = createBigtableRowMutations("row1");
  addBigtableMutation(rowMutations2, "family1", "column2", 1, "value3");
  addBigtableMutation(rowMutations2, "family2", "column1", 1, "value4");
  KV<ByteString, Iterable<Mutation>> rowMutations3 = createBigtableRowMutations("row2");
  addBigtableMutation(rowMutations3, "family2", "column2", 2, "value2");
  List<KV<ByteString, Iterable<Mutation>>> expectedBigtableRows =
      ImmutableList.of(rowMutations1, rowMutations2, rowMutations3);

  PCollection<KV<ByteString, Iterable<Mutation>>> bigtableRows =
      pipeline
          .apply("Create", Create.of(avroRows))
          .apply(
              "Transform to Bigtable",
              ParDo.of(AvroToBigtableFn.createWithSplitLargeRows(splitlargeRows, 2)));

  PAssert.that(bigtableRows).containsInAnyOrder(expectedBigtableRows);
  pipeline.run();
}
 

@Test
public void applyAvroToBigtableFn() throws Exception {
  BigtableRow avroRow1 = createAvroRow("row1");
  addAvroCell(avroRow1, "family1", "column1", 1, "value1");
  addAvroCell(avroRow1, "family1", "column1", 2, "value2");
  addAvroCell(avroRow1, "family1", "column2", 1, "value3");
  addAvroCell(avroRow1, "family2", "column1", 1, "value4");
  BigtableRow avroRow2 = createAvroRow("row2");
  addAvroCell(avroRow2, "family2", "column2", 2, "value2");
  List<BigtableRow> avroRows = ImmutableList.of(avroRow1, avroRow2);

  KV<ByteString, Iterable<Mutation>> rowMutations1 = createBigtableRowMutations("row1");
  addBigtableMutation(rowMutations1, "family1", "column1", 1, "value1");
  addBigtableMutation(rowMutations1, "family1", "column1", 2, "value2");
  addBigtableMutation(rowMutations1, "family1", "column2", 1, "value3");
  addBigtableMutation(rowMutations1, "family2", "column1", 1, "value4");
  KV<ByteString, Iterable<Mutation>> rowMutations2 = createBigtableRowMutations("row2");
  addBigtableMutation(rowMutations2, "family2", "column2", 2, "value2");
  List<KV<ByteString, Iterable<Mutation>>> expectedBigtableRows =
      ImmutableList.of(rowMutations1, rowMutations2);

  PCollection<KV<ByteString, Iterable<Mutation>>> bigtableRows =
      pipeline
          .apply("Create", Create.of(avroRows))
          .apply("Transform to Bigtable", ParDo.of(AvroToBigtableFn.create()));

  PAssert.that(bigtableRows).containsInAnyOrder(expectedBigtableRows);
  pipeline.run();
}
 

@Test
public void testDeleteCellsFromColumn() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl = (BigtableMutationImpl) this.bigtableMutation
          .deleteCellsFromColumn("family", "qualifier", Optional.empty(), Optional.empty())
          .deleteCellsFromColumn("family", "qualifier", Optional.of(100L), Optional.of(999L));
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(2, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());

  Mutation mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(0);
  assertEquals(Mutation.MutationCase.DELETE_FROM_COLUMN, mutation.getMutationCase());
  assertEquals("family", mutation.getDeleteFromColumn().getFamilyName());
  assertEquals("qualifier", mutation.getDeleteFromColumn().getColumnQualifier().toStringUtf8());
  assertEquals(TimestampRange.getDefaultInstance(), mutation.getDeleteFromColumn().getTimeRange());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  mutation = bigtableMutationImpl.getMutateRowRequest().getMutations(1);
  assertEquals(Mutation.MutationCase.DELETE_FROM_COLUMN, mutation.getMutationCase());
  assertEquals("family", mutation.getDeleteFromColumn().getFamilyName());
  assertEquals("qualifier", mutation.getDeleteFromColumn().getColumnQualifier().toStringUtf8());
  assertEquals(100L, mutation.getDeleteFromColumn().getTimeRange().getStartTimestampMicros());
  assertEquals(999L, mutation.getDeleteFromColumn().getTimeRange().getEndTimestampMicros());
  assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
  assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
  assertEquals(Mutation.SetCell.getDefaultInstance(), mutation.getSetCell());

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@Test
public void testWriteColumnValue() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl = (BigtableMutationImpl) this.bigtableMutation
          .write("family:qualifier", ByteString.copyFromUtf8("value"))
          .write("family", "qualifier", ByteString.copyFromUtf8("value"));
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(2, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());
  for (Mutation mutation : bigtableMutationImpl.getMutateRowRequest().getMutationsList()) {
    assertEquals(Mutation.MutationCase.SET_CELL, mutation.getMutationCase());
    assertEquals("family", mutation.getSetCell().getFamilyName());
    assertEquals("qualifier", mutation.getSetCell().getColumnQualifier().toStringUtf8());
    assertEquals("value", mutation.getSetCell().getValue().toStringUtf8());
    assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
    assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
    assertEquals(Mutation.DeleteFromColumn.getDefaultInstance(), mutation.getDeleteFromColumn());

    // Check timestamp in last 1 second
    final long tsMillis = TimeUnit.MICROSECONDS.toMillis(mutation.getSetCell().getTimestampMicros());
    final long nowMillis = System.currentTimeMillis();
    assertTrue(tsMillis <= nowMillis && tsMillis > nowMillis - TimeUnit.SECONDS.toMillis(1));
  }

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@Test
public void testWriteColumnAndTimestamp() throws Exception {
  final BigtableMutationImpl bigtableMutationImpl = (BigtableMutationImpl) this.bigtableMutation
          .write("family:qualifier", ByteString.copyFromUtf8("value"), 100L)
          .write("family", "qualifier", ByteString.copyFromUtf8("value"), 100L);
  verifyGetMutateRowRequest();

  bigtableMutationImpl.execute();
  bigtableMutationImpl.executeAsync();

  assertEquals(2, bigtableMutationImpl.getMutateRowRequest().getMutationsCount());

  for (Mutation mutation : bigtableMutationImpl.getMutateRowRequest().getMutationsList()) {
    assertEquals(Mutation.MutationCase.SET_CELL, mutation.getMutationCase());
    assertEquals("family", mutation.getSetCell().getFamilyName());
    assertEquals("qualifier", mutation.getSetCell().getColumnQualifier().toStringUtf8());
    assertEquals("value", mutation.getSetCell().getValue().toStringUtf8());
    assertEquals(100L, mutation.getSetCell().getTimestampMicros());
    assertEquals(Mutation.DeleteFromRow.getDefaultInstance(), mutation.getDeleteFromRow());
    assertEquals(Mutation.DeleteFromFamily.getDefaultInstance(), mutation.getDeleteFromFamily());
    assertEquals(Mutation.DeleteFromColumn.getDefaultInstance(), mutation.getDeleteFromColumn());
  }

  verify(bigtableMock.getMockedDataClient()).mutateRow(bigtableMutation.getMutateRowRequest().build());
  verify(bigtableMock.getMockedDataClient()).mutateRowAsync(bigtableMutation.getMutateRowRequest().build());
  verifyNoMoreInteractions(bigtableMock.getMockedDataClient());
}
 

@ProcessElement
public void processElement(
    @Element BigtableRow row, OutputReceiver<KV<ByteString, Iterable<Mutation>>> out) {
  ByteString key = toByteString(row.getKey());
  // BulkMutation doesn't split rows. Currently, if a single row contains more than 100,000
  // mutations, the service will fail the request.
  ImmutableList.Builder<Mutation> mutations = ImmutableList.builder();
  int cellsProcessed = 0;
  for (BigtableCell cell : row.getCells()) {
    SetCell setCell =
        SetCell.newBuilder()
            .setFamilyName(cell.getFamily().toString())
            .setColumnQualifier(toByteString(cell.getQualifier()))
            .setTimestampMicros(cell.getTimestamp())
            .setValue(toByteString(cell.getValue()))
            .build();

    mutations.add(Mutation.newBuilder().setSetCell(setCell).build());
    cellsProcessed++;

    if (this.splitLargeRows && cellsProcessed % maxMutationsPerRow == 0) {
      // Send a MutateRow request when we have accumulated max mutations per row.
      out.output(KV.of(key, mutations.build()));
      mutations = ImmutableList.builder();
    }
  }

  // Flush any remaining mutations.
  ImmutableList remainingMutations = mutations.build();
  if (!remainingMutations.isEmpty()) {
    out.output(KV.of(key, remainingMutations));
  }
}
 

/**
 * Maps are serialized similarly to {@link BeamRowToBigtableFn#createMapMutations(Row, Field)}
 * collections} however with an additional suffix to denote if the cell is a key or a value. As
 * such a map column in Cassandra named ‘mycolumn’ with two key-value pairs would be expanded into
 * the following column qualifiers: ‘mycolumn[0].key’, ‘mycolumn[1].key’, ‘mycolumn[0].value’,
 * ‘mycolumn[1].value’. Cell values will be serialized with the {@link
 * BeamRowToBigtableFn#primitiveFieldToBytes(TypeName, Object)} primitiveFieldToBytes} method.
 *
 * @param row The row to get the collection from.
 * @param field The field pointing to a collection in the row.
 * @return A set of mutation on the format specified above.
 */
private List<Mutation> createMapMutations(Row row, Field field) {
  List<Mutation> mutations = new ArrayList<>();
  // We use tree-map here to make sure that the keys are sorted.
  // Otherwise we get unpredictable serialization order of the keys in the mutation.
  Map<Object, Object> map = new TreeMap(row.getMap(field.getName()));
  TypeName mapKeyType = field.getType().getMapKeyType().getTypeName();
  TypeName mapValueType = field.getType().getMapValueType().getTypeName();
  Set<Map.Entry<Object, Object>> entries = map.entrySet();
  Iterator<Map.Entry<Object, Object>> iterator = entries.iterator();
  int i = 0;
  while (iterator.hasNext()) {
    Map.Entry<Object, Object> entry = iterator.next();
    // Key
    String keyFieldName = field.getName() + "[" + i + "].key";
    ByteString keyValue = primitiveFieldToBytes(mapKeyType, entry.getKey());
    SetCell keyCell = createCell(defaultColumnFamily.get(), keyFieldName, keyValue);
    mutations.add(Mutation.newBuilder().setSetCell(keyCell).build());

    // Value
    String valueFieldName = field.getName() + "[" + i + "].value";
    ByteString valueValue = primitiveFieldToBytes(mapValueType, entry.getValue());
    SetCell valueCell = createCell(defaultColumnFamily.get(), valueFieldName, valueValue);
    mutations.add(Mutation.newBuilder().setSetCell(valueCell).build());

    i++;
  }
  return mutations;
}
 

/**
 * Method generates a set of Bigtable mutations from the field specified in the input. Example: a
 * collection named ‘mycolumn’ with three values would be expanded into three column qualifiers in
 * inside Cloud Bigtable called, ‘mycolumn[0]’, ‘mycolumn[1]’ ,‘mycolumn[2]’. Cell values will be
 * serialized with the {@link BeamRowToBigtableFn#primitiveFieldToBytes(TypeName, Object)}
 * primitiveFieldToBytes} method.
 *
 * @param row The row to get the collection from.
 * @param field The field pointing to a collection in the row.
 * @return A set of mutation on the format specified above.
 */
private List<Mutation> createCollectionMutations(Row row, Field field) {
  List<Mutation> mutations = new ArrayList<>();
  List<Object> list = new ArrayList<Object>(row.getArray(field.getName()));
  TypeName collectionElementType = field.getType().getCollectionElementType().getTypeName();
  for (int i = 0; i < list.size(); i++) {
    String fieldName = field.getName() + "[" + i + "]";
    ByteString value = primitiveFieldToBytes(collectionElementType, list.get(i));
    SetCell cell = createCell(defaultColumnFamily.get(), fieldName, value);
    mutations.add(Mutation.newBuilder().setSetCell(cell).build());
  }
  return mutations;
}
 

@java.beans.ConstructorProperties({ "mutations" })
public Mutations(final List<Mutation> mutations) {
    this.mutations = mutations;
}
 

@Override
public BigtableMutation deleteRow() {
  mutateRowRequest.addMutations(Mutation.newBuilder()
                                    .setDeleteFromRow(Mutation.DeleteFromRow.newBuilder()));
  return this;
}