BQ sink produces sample of successful inserts (#875)

Oleksii Moskalenko · woop · web-flow · commit 2327b299ee4d · 2020-07-13T19:28:05.000+08:00
* feature row batch produces sample

* lint

Co-authored-by: Willem Pienaar &lt;git@willem.co&gt;
diff --git a/storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/compression/FeatureRowsBatch.java b/storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/compression/FeatureRowsBatch.java
@@ -207,48 +207,58 @@ public static FeatureRowsBatch fromRow(Row row) {
     return new FeatureRowsBatch(row.getSchema(), row.getValues());
   }
 
-  public Iterator<FeatureRowProto.FeatureRow> getFeatureRows() {
+  private FeatureRowProto.FeatureRow restoreFeatureRow(int rowIdx) {
     int timestampColumnIdx = schema.indexOf("eventTimestamp");
     int ingestionIdColumnIdx = schema.indexOf("ingestionId");
 
-    return IntStream.range(0, ((List<Object>) values.get(0)).size())
+    return FeatureRowProto.FeatureRow.newBuilder()
+        .setFeatureSet(getFeatureSetReference())
+        .setEventTimestamp(
+            Timestamp.newBuilder()
+                .setSeconds((long) (((List<Object>) values.get(timestampColumnIdx)).get(rowIdx)))
+                .build())
+        .setIngestionId((String) (((List<Object>) values.get(ingestionIdColumnIdx)).get(rowIdx)))
+        .addAllFields(
+            schema.getFieldNames().stream()
+                .map(
+                    fieldName -> {
+                      if (SERVICE_FIELDS.contains(fieldName)) {
+                        return null;
+                      }
+                      int fieldIdx = schema.indexOf(fieldName);
+
+                      return FieldProto.Field.newBuilder()
+                          .setName(schema.getField(fieldIdx).getName())
+                          .setValue(
+                              objectToProtoValue(
+                                  ((List<Object>) values.get(fieldIdx)).get(rowIdx),
+                                  schemaToProtoTypes.get(
+                                      schema
+                                          .getField(fieldIdx)
+                                          .getType()
+                                          .getCollectionElementType())))
+                          .build();
+                    })
+                .filter(Objects::nonNull)
+                .collect(Collectors.toList()))
+        .build();
+  }
+
+  public Iterator<FeatureRowProto.FeatureRow> getFeatureRows() {
+    int featureCount = ((List<Object>) values.get(0)).size();
+
+    return IntStream.range(0, featureCount).parallel().mapToObj(this::restoreFeatureRow).iterator();
+  }
+
+  public Iterator<FeatureRowProto.FeatureRow> getFeatureRowsSample(int maxCount) {
+    int featureCount = ((List<Object>) values.get(0)).size();
+    Random rd = new Random(42);
+
+    return IntStream.range(0, featureCount)
+        .filter(idx -> rd.nextInt(featureCount) < maxCount)
         .parallel()
-        .mapToObj(
-            rowIdx ->
-                FeatureRowProto.FeatureRow.newBuilder()
-                    .setFeatureSet(getFeatureSetReference())
-                    .setEventTimestamp(
-                        Timestamp.newBuilder()
-                            .setSeconds(
-                                (long)
-                                    (((List<Object>) values.get(timestampColumnIdx)).get(rowIdx)))
-                            .build())
-                    .setIngestionId(
-                        (String) (((List<Object>) values.get(ingestionIdColumnIdx)).get(rowIdx)))
-                    .addAllFields(
-                        schema.getFieldNames().stream()
-                            .map(
-                                fieldName -> {
-                                  if (SERVICE_FIELDS.contains(fieldName)) {
-                                    return null;
-                                  }
-                                  int fieldIdx = schema.indexOf(fieldName);
-
-                                  return FieldProto.Field.newBuilder()
-                                      .setName(schema.getField(fieldIdx).getName())
-                                      .setValue(
-                                          objectToProtoValue(
-                                              ((List<Object>) values.get(fieldIdx)).get(rowIdx),
-                                              schemaToProtoTypes.get(
-                                                  schema
-                                                      .getField(fieldIdx)
-                                                      .getType()
-                                                      .getCollectionElementType())))
-                                      .build();
-                                })
-                            .filter(Objects::nonNull)
-                            .collect(Collectors.toList()))
-                    .build())
+        .mapToObj(this::restoreFeatureRow)
+        .limit(maxCount)
         .iterator();
   }
 
diff --git a/storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/writer/BigQueryWrite.java b/storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/writer/BigQueryWrite.java
@@ -55,6 +55,7 @@ public class BigQueryWrite extends PTransform<PCollection<FeatureRow>, WriteResu
   private static final Duration BIGQUERY_JOB_MAX_EXPECTING_RESULT_TIME = Duration.standardHours(1);
   private static final int BIGQUERY_MAX_JOB_RETRIES = 20;
   private static final int DEFAULT_COMPACTION_BATCH_SIZE = 10000;
+  private static final int MAX_SUCCESSFUL_OUTPUTS_PER_DESTINATION = 1000;
 
   private DatasetId destination;
   private PCollectionView<Map<String, Iterable<TableSchema>>> schemas;
@@ -63,6 +64,7 @@ public class BigQueryWrite extends PTransform<PCollection<FeatureRow>, WriteResu
   private Duration expectingResultTime = BIGQUERY_JOB_MAX_EXPECTING_RESULT_TIME;
   private BigQueryServices testServices;
   private int compactionBatchSize = DEFAULT_COMPACTION_BATCH_SIZE;
+  private int maxSuccessfulOutputs = MAX_SUCCESSFUL_OUTPUTS_PER_DESTINATION;
 
   public BigQueryWrite(
       DatasetId destination, PCollectionView<Map<String, Iterable<TableSchema>>> schemas) {
@@ -90,6 +92,11 @@ public BigQueryWrite withCompactionBatchSize(int batchSize) {
     return this;
   }
 
+  public BigQueryWrite withMaxSuccessfulOutputs(int maxSuccessfulOutputs) {
+    this.maxSuccessfulOutputs = maxSuccessfulOutputs;
+    return this;
+  }
+
   /**
    * BigQuery writer 1. choose destination based on featureSetName {@link
    * FeatureDynamicDestinations} 2. dynamically pull destination's schema from schemas' view 3.
@@ -225,7 +232,10 @@ public void process(ProcessContext c) {
 
                 result
                     .getAll(inputTag)
-                    .forEach(rows -> rows.getFeatureRows().forEachRemaining(c::output));
+                    .forEach(
+                        rows ->
+                            rows.getFeatureRowsSample(maxSuccessfulOutputs)
+                                .forEachRemaining(c::output));
               }
             }));
   }
diff --git a/storage/connectors/bigquery/src/test/java/feast/storage/connectors/bigquery/writer/BigQuerySinkTest.java b/storage/connectors/bigquery/src/test/java/feast/storage/connectors/bigquery/writer/BigQuerySinkTest.java
@@ -21,6 +21,7 @@
 import static feast.storage.connectors.bigquery.writer.FeatureSetSpecToTableSchema.*;
 import static org.hamcrest.CoreMatchers.*;
 import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.lessThan;
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
 import static org.mockito.MockitoAnnotations.initMocks;
@@ -498,6 +499,30 @@ public void featureRowCompressShouldPackAndUnpackSuccessfully() {
     p.run();
   }
 
+  @Test
+  public void featureRowBatchShouldSampleOnRestore() {
+    List<FeatureRow> stream =
+        IntStream.range(0, 1000)
+            .mapToObj(i -> generateRow("project/fs"))
+            .collect(Collectors.toList());
+
+    PCollection<Long> result =
+        p.apply(Create.of(stream))
+            .apply("KV", ParDo.of(new ExtractKV()))
+            .apply(new CompactFeatureRows(1000))
+            .apply(ParDo.of(new FlatMapWithSample(100)))
+            .apply(Count.globally());
+
+    PAssert.that(result)
+        .satisfies(
+            r -> {
+              // sample size is within bound of required size
+              assertThat(Math.abs(r.iterator().next() - 100), lessThan(5L));
+              return null;
+            });
+    p.run();
+  }
+
   private List<FeatureRow> dropNullFeature(List<FeatureRow> input) {
     return input.stream()
         .map(
@@ -549,4 +574,17 @@ public Table answer(InvocationOnMock invocationOnMock) throws Throwable {
       return FakeTable.create(mock(BigQuery.class), tableId, tableDefinition);
     }
   }
+
+  private static class FlatMapWithSample extends DoFn<KV<String, FeatureRowsBatch>, FeatureRow> {
+    private int sampleSize;
+
+    FlatMapWithSample(int sampleSize) {
+      this.sampleSize = sampleSize;
+    }
+
+    @ProcessElement
+    public void process(ProcessContext c) {
+      c.element().getValue().getFeatureRowsSample(sampleSize).forEachRemaining(c::output);
+    }
+  }
 }
