diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java index f7aa50a7437f..9eaf441a4bfe 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java @@ -481,6 +481,14 @@ private ReadRowsRequest truncateRequest(ReadRowsRequest request, ByteString last segment.addRowRanges(newRange.build()); } else { // Row is split, remove all read rowKeys and split RowSet at last buffered Row + if (rowRange.getEndKeyCase() == RowRange.EndKeyCase.END_KEY_OPEN + && !rowRange.getEndKeyOpen().isEmpty()) { + ByteString lastKeyWithNull = lastKey.concat(ByteString.copyFrom(new byte[] {0})); + if (ByteStringComparator.INSTANCE.compare(lastKeyWithNull, rowRange.getEndKeyOpen()) + >= 0) { + continue; + } + } segment.addRowRanges(newRange.setStartKeyOpen(lastKey).build()); } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java index bec5a5470cbc..dc756fcbe581 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java @@ -763,6 +763,87 @@ public Void answer(InvocationOnMock invocation) throws Throwable { Mockito.verify(mockCallMetric, Mockito.times(3)).call("ok"); } + /** + * This test ensures that when a range has an open end key that is equal to the start key plus a + * null byte (e.g. [k, k\0)), and the buffer hits the byte limit on key k, truncateRequest + * correctly detects that the range is exhausted and does not create an invalid range (k, k\0). + */ + @Test + public void testReadRangeWithNullByteEndKeyAtByteLimit() throws IOException { + ByteString startKey = ByteString.copyFromUtf8("exact_key"); + ByteString endKey = startKey.concat(ByteString.copyFrom(new byte[] {0})); + RowRange mockRowRange = + RowRange.newBuilder().setStartKeyClosed(startKey).setEndKeyOpen(endKey).build(); + + long segmentByteLimit = DEFAULT_ROW_SIZE / 2; + + byte[] largeMemory = new byte[(int) DEFAULT_ROW_SIZE]; + Row expectedRow = + Row.newBuilder() + .setKey(startKey) + .addFamilies( + Family.newBuilder() + .setName("Family") + .addColumns( + Column.newBuilder() + .setQualifier(ByteString.copyFromUtf8("LargeMemoryRow")) + .addCells( + Cell.newBuilder() + .setValue(ByteString.copyFrom(largeMemory)) + .setTimestampMicros(System.currentTimeMillis()) + .build()) + .build()) + .build()) + .build(); + + List> expectedResults = + ImmutableList.of(ImmutableList.of(expectedRow), ImmutableList.of()); + + ServerStreamingCallable mockCallable = Mockito.mock(ServerStreamingCallable.class); + + StreamController mockController = Mockito.mock(StreamController.class); + doAnswer( + new Answer() { + @Override + public Void answer(InvocationOnMock invocation) throws Throwable { + cancelled.set(true); + return null; + } + }) + .when(mockController) + .cancel(); + + doAnswer(new MultipleAnswer(expectedResults, mockController)) + .when(mockCallable) + .call(any(Query.class), any(ResponseObserver.class), any(ApiCallContext.class)); + when(mockStub.createReadRowsCallable(any(RowAdapter.class))).thenReturn(mockCallable); + ServerStreamingCallable callable = + mockStub.createReadRowsCallable(new BigtableServiceImpl.BigtableRowProtoAdapter()); + when(mockBigtableDataClient.readRowsCallable(any(RowAdapter.class))).thenReturn(callable); + + BigtableService.Reader underTest = + new BigtableServiceImpl.BigtableSegmentReaderImpl( + mockBigtableDataClient, + bigtableDataSettings.getProjectId(), + bigtableDataSettings.getInstanceId(), + TABLE_ID, + RowSet.newBuilder().addRowRanges(mockRowRange).build(), + RowFilter.getDefaultInstance(), + SEGMENT_SIZE, + segmentByteLimit, + mockCallMetric); + + List actualResults = new ArrayList<>(); + Assert.assertTrue(underTest.start()); + do { + actualResults.add(underTest.getCurrentRow()); + } while (underTest.advance()); + + Assert.assertEquals(ImmutableList.of(expectedRow), actualResults); + Mockito.verify(mockCallable, Mockito.times(1)) + .call(any(Query.class), any(ResponseObserver.class), any(ApiCallContext.class)); + } + /** * This test ensures the Exception handling inside of the scanHandler. This test will check if a * StatusRuntimeException was thrown.