feat:large-row-skip-in-bigtable | added experimental options to skip …

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java

@@ -631,6 +631,24 @@ public Read withMaxBufferElementCount(@Nullable Integer maxBufferElementCount) {
           .build();
     }
 
+    /**
+     * Returns a new {@link BigtableIO.Read} that will skip the large rows while reading. This
+     * function will switch the base BigtableIO.Reader class to using the
+     * BigtableReaderWithExperimentalOptions. If
+     *
+     * <p>Does not modify this object.
+     *
+     * <p>This is incompatible with withMaxBufferElementCount()
+     */
+    public Read withExperimentalSkipLargeRows(@Nullable Boolean skipLargeRows) {
+      BigtableReadOptions bigtableReadOptions = getBigtableReadOptions();
+
+      return toBuilder()
+          .setBigtableReadOptions(
+              bigtableReadOptions.toBuilder().setExperimentalSkipLargeRows(skipLargeRows).build())
+          .build();
+    }
+
     /**
      * Returns a new {@link BigtableIO.Read} that will read only rows in the specified range.
      *

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadOptions.java

@@ -37,6 +37,9 @@
 })
 abstract class BigtableReadOptions implements Serializable {
 
+  /** Returns the experimental skip large rows option value. */
+  abstract @Nullable Boolean getExperimentalSkipLargeRows();
+
   /** Returns the table id. */
   abstract ValueProvider<String> getTableId();
 
@@ -69,6 +72,8 @@ static BigtableReadOptions.Builder builder() {
   @AutoValue.Builder
   abstract static class Builder {
 
+    abstract Builder setExperimentalSkipLargeRows(Boolean experimentalSkipLargeRows);
+
     abstract Builder setTableId(ValueProvider<String> tableId);
 
     abstract Builder setRowFilter(ValueProvider<RowFilter> rowFilter);
@@ -86,6 +91,10 @@ abstract static class Builder {
     abstract BigtableReadOptions build();
   }
 
+  BigtableReadOptions setExperimentalSkipLargeRows(@Nullable Boolean experimentalSkipLargeRows) {
+    return toBuilder().setExperimentalSkipLargeRows(experimentalSkipLargeRows).build();
+  }
+
   BigtableReadOptions setMaxBufferElementCount(@Nullable Integer maxBufferElementCount) {
     return toBuilder().setMaxBufferElementCount(maxBufferElementCount).build();
   }
@@ -116,7 +125,10 @@ void populateDisplayData(DisplayData.Builder builder) {
                 .withLabel("Read Attempt Timeout"))
         .addIfNotNull(
             DisplayData.item("operationTimeout", getOperationTimeout())
-                .withLabel("Read Operation Timeout"));
+                .withLabel("Read Operation Timeout"))
+        .addIfNotNull(
+            DisplayData.item("experimentalSkipLargeRows", getExperimentalSkipLargeRows())
+                .withLabel("Skip Large Rows"));
   }
 
   void validate() {

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java

@@ -142,6 +142,7 @@ static class BigtableReaderImpl implements Reader {
     private ServerStream<Row> stream;
 
     private boolean exhausted;
+    private BigtableReadOptions bigtableReadOptions;
 
     @VisibleForTesting
     BigtableReaderImpl(
@@ -150,13 +151,15 @@ static class BigtableReaderImpl implements Reader {
         String instanceId,
         String tableId,
         List<ByteKeyRange> ranges,
-        @Nullable RowFilter rowFilter) {
+        @Nullable RowFilter rowFilter,
+        BigtableReadOptions bigtableReadOptions) {
       this.client = client;
       this.projectId = projectId;
       this.instanceId = instanceId;
       this.tableId = tableId;
       this.ranges = ranges;
       this.rowFilter = rowFilter;
+      this.bigtableReadOptions = bigtableReadOptions;
     }
 
     @Override
@@ -174,10 +177,18 @@ public boolean start() throws IOException {
         query.filter(Filters.FILTERS.fromProto(rowFilter));
       }
       try {
-        stream =
-            client
-                .readRowsCallable(new BigtableRowProtoAdapter())
-                .call(query, GrpcCallContext.createDefault());
+        if (bigtableReadOptions != null
+            && Boolean.TRUE.equals(bigtableReadOptions.getExperimentalSkipLargeRows())) {
+          stream =
+              client
+                  .skipLargeRowsCallable(new BigtableRowProtoAdapter())
+                  .call(query, GrpcCallContext.createDefault());
+        } else {
+          stream =
+              client
+                  .readRowsCallable(new BigtableRowProtoAdapter())
+                  .call(query, GrpcCallContext.createDefault());
+        }
         results = stream.iterator();
         serviceCallMetric.call("ok");
       } catch (StatusRuntimeException e) {
@@ -667,7 +678,8 @@ public Reader createReader(BigtableSource source) throws IOException {
           instanceId,
           source.getTableId().get(),
           source.getRanges(),
-          source.getRowFilter());
+          source.getRowFilter(),
+          source.getReadOptions());
     }
   }
 

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIOTest.java

@@ -231,6 +231,27 @@ public void testReadBuildsCorrectly() {
     assertEquals(PORT_CONFIGURATOR, read.getBigtableConfig().getBigtableOptionsConfigurator());
   }
 
+  @Test
+  public void testReadWithSkippingLargeRows() {
+    BigtableIO.Read read =
+        BigtableIO.read()
+            .withBigtableOptions(BIGTABLE_OPTIONS)
+            .withExperimentalSkipLargeRows(true)
+            .withTableId("table")
+            .withInstanceId("instance")
+            .withProjectId("project")
+            .withAppProfileId("app-profile")
+            .withBigtableOptionsConfigurator(PORT_CONFIGURATOR);
+    assertEquals("options_project", read.getBigtableOptions().getProjectId());
+    assertEquals("options_instance", read.getBigtableOptions().getInstanceId());
+    assertEquals("instance", read.getBigtableConfig().getInstanceId().get());
+    assertEquals("project", read.getBigtableConfig().getProjectId().get());
+    assertEquals("app-profile", read.getBigtableConfig().getAppProfileId().get());
+    assertEquals("table", read.getTableId());
+    assertEquals(true, read.getBigtableReadOptions().getExperimentalSkipLargeRows());
+    assertEquals(PORT_CONFIGURATOR, read.getBigtableConfig().getBigtableOptionsConfigurator());
+  }
+
   @Test
   public void testReadValidationFailsMissingTable() {
     BigtableIO.Read read = BigtableIO.read().withBigtableOptions(BIGTABLE_OPTIONS);

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java

@@ -119,6 +119,27 @@ public void testE2EBigtableRead() {
     checkLineageSourceMetric(r, tableId);
   }
 
+  @Test
+  public void testE2EBigtableReadWithSkippingLargeRows() {
+    BigtableOptions.Builder bigtableOptionsBuilder =
+        new BigtableOptions.Builder().setProjectId(project).setInstanceId(options.getInstanceId());
+
+    final String tableId = "BigtableReadTest";
+    final long numRows = 1000L;
+
+    Pipeline p = Pipeline.create(options);
+    PCollection<Long> count =
+        p.apply(
+                BigtableIO.read()
+                    .withBigtableOptions(bigtableOptionsBuilder)
+                    .withTableId(tableId)
+                    .withExperimentalSkipLargeRows(true))
+            .apply(Count.globally());
+    PAssert.thatSingleton(count).isEqualTo(numRows);
+    PipelineResult r = p.run();
+    checkLineageSourceMetric(r, tableId);
+  }
+
   @Test
   public void testE2EBigtableSegmentRead() {
     tableAdminClient.createTable(CreateTableRequest.of(tableId).addFamily(COLUMN_FAMILY_NAME));

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImplTest.java

@@ -179,7 +179,8 @@ public void testRead() throws IOException {
             bigtableDataSettings.getInstanceId(),
             mockBigtableSource.getTableId().get(),
             mockBigtableSource.getRanges(),
-            null);
+            null,
+            BigtableReadOptions.builder().setTableId(mockBigtableSource.getTableId()).build());
 
     underTest.start();
     Assert.assertEquals(expectedRow, underTest.getCurrentRow());
@@ -188,6 +189,69 @@ public void testRead() throws IOException {
     verifyMetricWasSet("google.bigtable.v2.ReadRows", "ok", 1);
   }
 
+  /**
+   * This test ensures that protobuf creation and interactions with {@link BigtableDataClient} work
+   * as expected.
+   *
+   * @throws IOException
+   * @throws InterruptedException
+   */
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testReadWithSkippingLargeRows() throws IOException {
+    ByteKey start = ByteKey.copyFrom("a".getBytes(StandardCharsets.UTF_8));
+    ByteKey end = ByteKey.copyFrom("d".getBytes(StandardCharsets.UTF_8));
+    when(mockBigtableSource.getRanges()).thenReturn(Arrays.asList(ByteKeyRange.of(start, end)));
+    when(mockBigtableSource.getTableId()).thenReturn(StaticValueProvider.of(TABLE_ID));
+
+    Row expectedRowA = Row.newBuilder().setKey(ByteString.copyFromUtf8("a")).build();
+    //    b is large row that's skipped
+    //    Row expectedRowB = Row.newBuilder().setKey(ByteString.copyFromUtf8("c")).build();
+    Row expectedRowC = Row.newBuilder().setKey(ByteString.copyFromUtf8("c")).build();
+
+    // Set up iterator to be returned by ServerStream.iterator()
+    Iterator<Row> mockIterator = Mockito.mock(Iterator.class);
+    when(mockIterator.next()).thenReturn(expectedRowA).thenReturn(expectedRowC).thenReturn(null);
+    when(mockIterator.hasNext()).thenReturn(true).thenReturn(true).thenReturn(false);
+    // Set up ServerStream to be returned by callable.call(Query)
+    ServerStream<Row> mockRows = Mockito.mock(ServerStream.class);
+    when(mockRows.iterator()).thenReturn(mockIterator);
+    // Set up Callable to be returned by stub.createReadRowsCallable()
+    ServerStreamingCallable<Query, Row> mockCallable = Mockito.mock(ServerStreamingCallable.class);
+    when(mockCallable.call(
+            any(com.google.cloud.bigtable.data.v2.models.Query.class), any(ApiCallContext.class)))
+        .thenReturn(mockRows);
+    when(mockStub.createSkipLargeRowsCallable(any(RowAdapter.class))).thenReturn(mockCallable);
+    ServerStreamingCallable<Query, Row> callable =
+        mockStub.createSkipLargeRowsCallable(new BigtableServiceImpl.BigtableRowProtoAdapter());
+
+    when(mockBigtableDataClient.skipLargeRowsCallable(any(RowAdapter.class))).thenReturn(callable);
+
+    BigtableService.Reader underTest =
+        new BigtableServiceImpl.BigtableReaderImpl(
+            mockBigtableDataClient,
+            bigtableDataSettings.getProjectId(),
+            bigtableDataSettings.getInstanceId(),
+            mockBigtableSource.getTableId().get(),
+            mockBigtableSource.getRanges(),
+            null,
+            BigtableReadOptions.builder()
+                .setExperimentalSkipLargeRows(true)
+                .setTableId(mockBigtableSource.getTableId())
+                .build());
+
+    int rowCount = 0;
+    underTest.start();
+    Assert.assertEquals(expectedRowA, underTest.getCurrentRow());
+    rowCount++;
+    underTest.advance();
+    Assert.assertEquals(expectedRowC, underTest.getCurrentRow());
+    rowCount++;
+    Assert.assertFalse(underTest.advance());
+    Assert.assertEquals(2, rowCount);
+    verify(mockIterator, times(2)).next();
+  }
+
   /**
    * 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.