Search in sources :

Example 1 with FlatRow

use of com.google.cloud.bigtable.grpc.scanner.FlatRow in project beam by apache.

the class BigtableServiceImplTest method testReadByteLimitBuffer.

/**
 * This test checks that the buffer will stop filling up once the byte limit is reached. It will
 * cancel the ScanHandler after reached the limit. This test completes one fill and contains one
 * Row after the first buffer has been completed. The test cheaks the current available memory in
 * the JVM and uses a percent of it to mock the original behavior. Range: [b00000, b00010)
 *
 * @throws IOException
 */
@Test
public void testReadByteLimitBuffer() throws IOException {
    long segmentByteLimit = DEFAULT_ROW_SIZE * (SEGMENT_SIZE / 2);
    int numOfRowsInsideBuffer = (int) (segmentByteLimit / DEFAULT_ROW_SIZE) + 1;
    RowRange mockRowRange = generateRowRange(generateByteString(DEFAULT_PREFIX, 0), generateByteString(DEFAULT_PREFIX, SEGMENT_SIZE));
    OngoingStubbing<?> stub = when(mockBigtableDataClient.readFlatRows(any(ReadRowsRequest.class), any()));
    List<List<FlatRow>> expectedResults = ImmutableList.of(generateLargeSegmentResult(DEFAULT_PREFIX, 0, numOfRowsInsideBuffer), generateSegmentResult(DEFAULT_PREFIX, numOfRowsInsideBuffer, SEGMENT_SIZE - numOfRowsInsideBuffer), ImmutableList.of());
    expectRowResults(stub, expectedResults);
    BigtableService.Reader underTest = new BigtableServiceImpl.BigtableSegmentReaderImpl(mockSession, TABLE_ID, RowSet.newBuilder().addRowRanges(mockRowRange).build(), SEGMENT_SIZE, segmentByteLimit, RowFilter.getDefaultInstance(), mockCallMetric);
    List<Row> actualResults = new ArrayList<>();
    Assert.assertTrue(underTest.start());
    do {
        actualResults.add(underTest.getCurrentRow());
    } while (underTest.advance());
    Assert.assertEquals(expectedResults.stream().flatMap(Collection::stream).map(i -> FlatRowConverter.convert(i)).collect(Collectors.toList()), actualResults);
    underTest.close();
    Mockito.verify(mockCallMetric, Mockito.times(3)).call("ok");
}
Also used : RowRange(com.google.bigtable.v2.RowRange) ArrayList(java.util.ArrayList) List(java.util.List) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) ArrayList(java.util.ArrayList) Row(com.google.bigtable.v2.Row) FlatRow(com.google.cloud.bigtable.grpc.scanner.FlatRow) Test(org.junit.Test)

Example 2 with FlatRow

use of com.google.cloud.bigtable.grpc.scanner.FlatRow in project beam by apache.

the class BigtableServiceImplTest method testReadSingleRangeBelowSegmentLimit.

/**
 * This test ensures that protobuf creation and interactions with {@link BigtableDataClient} work
 * as expected. This test checks that a single row is returned from the future.
 *
 * @throws IOException
 */
@Test
public void testReadSingleRangeBelowSegmentLimit() throws Exception {
    RowSet.Builder ranges = RowSet.newBuilder();
    ranges.addRowRanges(generateRowRange(generateByteString(DEFAULT_PREFIX, 0), generateByteString(DEFAULT_PREFIX, 1)));
    FlatRow expectedRow = FlatRow.newBuilder().withRowKey(ByteString.copyFromUtf8("a")).build();
    when(mockBigtableDataClient.readFlatRows(any(ReadRowsRequest.class), any())).thenAnswer(mockReadRowsAnswer(Arrays.asList(expectedRow)));
    BigtableService.Reader underTest = new BigtableServiceImpl.BigtableSegmentReaderImpl(mockSession, TABLE_ID, ranges.build(), SEGMENT_SIZE, DEFAULT_BYTE_SEGMENT_SIZE, RowFilter.getDefaultInstance(), mockCallMetric);
    underTest.start();
    Assert.assertEquals(FlatRowConverter.convert(expectedRow), underTest.getCurrentRow());
    Assert.assertFalse(underTest.advance());
    underTest.close();
    Mockito.verify(mockCallMetric, Mockito.times(2)).call("ok");
}
Also used : RowSet(com.google.bigtable.v2.RowSet) ReadRowsRequest(com.google.bigtable.v2.ReadRowsRequest) FlatRow(com.google.cloud.bigtable.grpc.scanner.FlatRow) Test(org.junit.Test)

Aggregations

FlatRow (com.google.cloud.bigtable.grpc.scanner.FlatRow)2 Test (org.junit.Test)2 ReadRowsRequest (com.google.bigtable.v2.ReadRowsRequest)1 Row (com.google.bigtable.v2.Row)1 RowRange (com.google.bigtable.v2.RowRange)1 RowSet (com.google.bigtable.v2.RowSet)1 ArrayList (java.util.ArrayList)1 List (java.util.List)1 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)1