Examples with BatchableMutationFilterFn org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn used on opensource projects

Example 1 with BatchableMutationFilterFn

use of org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn in project beam by apache.

the class SpannerIOWriteTest method testBatchableMutationFilterFn_size.

@Test
public void testBatchableMutationFilterFn_size() {
    Mutation all = Mutation.delete("test", KeySet.all());
    Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
    Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
    MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too big.
    g(m(1L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range) };
    long mutationSize = MutationSizeEstimator.sizeOf(m(1L));
    BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, mutationSize * 3, 1000, 1000);
    DoFn<MutationGroup, MutationGroup>.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
    when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
    // Capture the outputs.
    doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
    doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
    // Process all elements.
    for (MutationGroup m : mutationGroups) {
        when(mockProcessContext.element()).thenReturn(m);
        testFn.processElement(mockProcessContext);
    }
    // Verify captured batchable elements.
    assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
    // Verify captured unbatchable mutations
    Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
    assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too big.
    g(m(1L), m(3L), m(4L), m(5L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range)));
}

Example 2 with BatchableMutationFilterFn

use of org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn in project beam by apache.

the class SpannerIOWriteTest method testBatchableMutationFilterFn_rows.

@Test
public void testBatchableMutationFilterFn_rows() {
    Mutation all = Mutation.delete("test", KeySet.all());
    Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
    Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
    MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too many rows.
    g(m(1L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range) };
    BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, 1000, 1000, 3);
    BatchableMutationFilterFn.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
    when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
    // Capture the outputs.
    doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
    doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
    // Process all elements.
    for (MutationGroup m : mutationGroups) {
        when(mockProcessContext.element()).thenReturn(m);
        testFn.processElement(mockProcessContext);
    }
    // Verify captured batchable elements.
    assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
    // Verify captured unbatchable mutations
    Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
    assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too many rows.
    g(m(1L), m(3L), m(4L), m(5L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range)));
}

Example 3 with BatchableMutationFilterFn

use of org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn in project beam by apache.

the class SpannerIOWriteTest method testBatchableMutationFilterFn_batchingDisabled.

@Test
public void testBatchableMutationFilterFn_batchingDisabled() {
    MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L)), g(del(1L)), g(del(5L, 6L)) };
    BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, 0, 0, 0);
    BatchableMutationFilterFn.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
    when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
    // Capture the outputs.
    doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
    doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
    // Process all elements.
    for (MutationGroup m : mutationGroups) {
        when(mockProcessContext.element()).thenReturn(m);
        testFn.processElement(mockProcessContext);
    }
    // Verify captured batchable elements.
    assertTrue(mutationGroupCaptor.getAllValues().isEmpty());
    // Verify captured unbatchable mutations
    Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
    assertThat(unbatchableMutations, containsInAnyOrder(mutationGroups));
}

Example 4 with BatchableMutationFilterFn

use of org.apache.beam.sdk.io.gcp.spanner.SpannerIO.BatchableMutationFilterFn in project beam by apache.

the class SpannerIOWriteTest method testBatchableMutationFilterFn_cells.

@Test
public void testBatchableMutationFilterFn_cells() {
    Mutation all = Mutation.delete("test", KeySet.all());
    Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
    Mutation range = Mutation.delete("test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));
    MutationGroup[] mutationGroups = new MutationGroup[] { g(m(1L)), g(m(2L), m(3L)), // not batchable - too big.
    g(m(2L), m(3L), m(4L), m(5L)), g(del(1L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range) };
    BatchableMutationFilterFn testFn = new BatchableMutationFilterFn(null, null, 10000000, 3 * CELLS_PER_KEY, 1000);
    BatchableMutationFilterFn.ProcessContext mockProcessContext = Mockito.mock(ProcessContext.class);
    when(mockProcessContext.sideInput(any())).thenReturn(getSchema());
    // Capture the outputs.
    doNothing().when(mockProcessContext).output(mutationGroupCaptor.capture());
    doNothing().when(mockProcessContext).output(any(), mutationGroupListCaptor.capture());
    // Process all elements.
    for (MutationGroup m : mutationGroups) {
        when(mockProcessContext.element()).thenReturn(m);
        testFn.processElement(mockProcessContext);
    }
    // Verify captured batchable elements.
    assertThat(mutationGroupCaptor.getAllValues(), containsInAnyOrder(g(m(1L)), g(m(2L), m(3L)), g(del(1L))));
    // Verify captured unbatchable mutations
    Iterable<MutationGroup> unbatchableMutations = Iterables.concat(mutationGroupListCaptor.getAllValues());
    assertThat(unbatchableMutations, containsInAnyOrder(// not batchable - too big.
    g(m(2L), m(3L), m(4L), m(5L)), // not point delete.
    g(del(5L, 6L)), g(all), g(prefix), g(range)));
}