Skip to content

Commit 810f59e

Browse files
wooppyalex
woop
authored and
pyalex
committed
feature row batch produces sample
1 parent 4c1b152 commit 810f59e

3 files changed

Lines changed: 84 additions & 22 deletions

File tree

‎storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/compression/FeatureRowsBatch.java‎

Lines changed: 35 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -176,30 +176,44 @@ public static FeatureRowsBatch fromRow(Row row) {
176176
return new FeatureRowsBatch(row.getSchema(), row.getValues());
177177
}
178178

179+
private FeatureRowProto.FeatureRow restoreFeatureRow(int rowIdx) {
180+
return FeatureRowProto.FeatureRow.newBuilder()
181+
.setFeatureSet(getFeatureSetReference())
182+
.addAllFields(
183+
IntStream.range(0, schema.getFieldCount())
184+
.mapToObj(
185+
fieldIdx ->
186+
FieldProto.Field.newBuilder()
187+
.setName(schema.getField(fieldIdx).getName())
188+
.setValue(
189+
objectToProtoValue(
190+
((List<Object>) values.get(fieldIdx)).get(rowIdx),
191+
schemaToProtoTypes.get(
192+
schema
193+
.getField(fieldIdx)
194+
.getType()
195+
.getCollectionElementType())))
196+
.build())
197+
.collect(Collectors.toList()))
198+
.build();
199+
}
200+
179201
public Iterator<FeatureRowProto.FeatureRow> getFeatureRows() {
180202
return IntStream.range(0, ((List<Object>) values.get(0)).size())
181203
.parallel()
182-
.mapToObj(
183-
rowIdx ->
184-
FeatureRowProto.FeatureRow.newBuilder()
185-
.setFeatureSet(getFeatureSetReference())
186-
.addAllFields(
187-
IntStream.range(0, schema.getFieldCount())
188-
.mapToObj(
189-
fieldIdx ->
190-
FieldProto.Field.newBuilder()
191-
.setName(schema.getField(fieldIdx).getName())
192-
.setValue(
193-
objectToProtoValue(
194-
((List<Object>) values.get(fieldIdx)).get(rowIdx),
195-
schemaToProtoTypes.get(
196-
schema
197-
.getField(fieldIdx)
198-
.getType()
199-
.getCollectionElementType())))
200-
.build())
201-
.collect(Collectors.toList()))
202-
.build())
204+
.mapToObj(this::restoreFeatureRow)
205+
.iterator();
206+
}
207+
208+
public Iterator<FeatureRowProto.FeatureRow> getFeatureRowsSample(int maxCount) {
209+
int featureCount = ((List<Object>) values.get(0)).size();
210+
Random rd = new Random(42);
211+
212+
return IntStream.range(0, featureCount)
213+
.filter(idx -> rd.nextInt(featureCount) < maxCount)
214+
.parallel()
215+
.mapToObj(this::restoreFeatureRow)
216+
.limit(maxCount)
203217
.iterator();
204218
}
205219

‎storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/writer/BigQueryWrite.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ public class BigQueryWrite extends PTransform<PCollection<FeatureRow>, WriteResu
5555
private static final Duration BIGQUERY_JOB_MAX_EXPECTING_RESULT_TIME = Duration.standardHours(1);
5656
private static final int BIGQUERY_MAX_JOB_RETRIES = 20;
5757
private static final int DEFAULT_COMPACTION_BATCH_SIZE = 10000;
58+
private static final int MAX_SUCCESSFUL_OUTPUTS_PER_DESTINATION = 1000;
5859

5960
private DatasetId destination;
6061
private PCollectionView<Map<String, Iterable<TableSchema>>> schemas;
@@ -63,6 +64,7 @@ public class BigQueryWrite extends PTransform<PCollection<FeatureRow>, WriteResu
6364
private Duration expectingResultTime = BIGQUERY_JOB_MAX_EXPECTING_RESULT_TIME;
6465
private BigQueryServices testServices;
6566
private int compactionBatchSize = DEFAULT_COMPACTION_BATCH_SIZE;
67+
private int maxSuccessfulOutputs = MAX_SUCCESSFUL_OUTPUTS_PER_DESTINATION;
6668

6769
public BigQueryWrite(
6870
DatasetId destination, PCollectionView<Map<String, Iterable<TableSchema>>> schemas) {
@@ -90,6 +92,11 @@ public BigQueryWrite withCompactionBatchSize(int batchSize) {
9092
return this;
9193
}
9294

95+
public BigQueryWrite withMaxSuccessfulOutputs(int maxSuccessfulOutputs) {
96+
this.maxSuccessfulOutputs = maxSuccessfulOutputs;
97+
return this;
98+
}
99+
93100
/**
94101
* BigQuery writer 1. choose destination based on featureSetName {@link
95102
* FeatureDynamicDestinations} 2. dynamically pull destination's schema from schemas' view 3.
@@ -225,7 +232,10 @@ public void process(ProcessContext c) {
225232

226233
result
227234
.getAll(inputTag)
228-
.forEach(rows -> rows.getFeatureRows().forEachRemaining(c::output));
235+
.forEach(
236+
rows ->
237+
rows.getFeatureRowsSample(maxSuccessfulOutputs)
238+
.forEachRemaining(c::output));
229239
}
230240
}));
231241
}

‎storage/connectors/bigquery/src/test/java/feast/storage/connectors/bigquery/writer/BigQuerySinkTest.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static feast.storage.connectors.bigquery.writer.FeatureSetSpecToTableSchema.*;
2222
import static org.hamcrest.CoreMatchers.*;
2323
import static org.hamcrest.Matchers.containsInAnyOrder;
24+
import static org.hamcrest.Matchers.lessThan;
2425
import static org.junit.Assert.*;
2526
import static org.mockito.Mockito.*;
2627
import static org.mockito.MockitoAnnotations.initMocks;
@@ -493,6 +494,30 @@ public void featureRowCompressShouldPackAndUnpackSuccessfully() {
493494
p.run();
494495
}
495496

497+
@Test
498+
public void featureRowBatchShouldSampleOnRestore() {
499+
List<FeatureRow> stream =
500+
IntStream.range(0, 1000)
501+
.mapToObj(i -> generateRow("project/fs"))
502+
.collect(Collectors.toList());
503+
504+
PCollection<Long> result =
505+
p.apply(Create.of(stream))
506+
.apply("KV", ParDo.of(new ExtractKV()))
507+
.apply(new CompactFeatureRows(1000))
508+
.apply(ParDo.of(new FlatMapWithSample(100)))
509+
.apply(Count.globally());
510+
511+
PAssert.that(result)
512+
.satisfies(
513+
r -> {
514+
// sample size is within bound of required size
515+
assertThat(Math.abs(r.iterator().next() - 100), lessThan(5L));
516+
return null;
517+
});
518+
p.run();
519+
}
520+
496521
private List<FeatureRow> dropNullFeature(List<FeatureRow> input) {
497522
return input.stream()
498523
.map(
@@ -540,4 +565,17 @@ public Table answer(InvocationOnMock invocationOnMock) throws Throwable {
540565
return FakeTable.create(mock(BigQuery.class), tableId, tableDefinition);
541566
}
542567
}
568+
569+
private static class FlatMapWithSample extends DoFn<KV<String, FeatureRowsBatch>, FeatureRow> {
570+
private int sampleSize;
571+
572+
FlatMapWithSample(int sampleSize) {
573+
this.sampleSize = sampleSize;
574+
}
575+
576+
@ProcessElement
577+
public void process(ProcessContext c) {
578+
c.element().getValue().getFeatureRowsSample(sampleSize).forEachRemaining(c::output);
579+
}
580+
}
543581
}

0 commit comments

Comments
 (0)