Skip to content

Commit 95ca756

Browse files
wooppyalex
woop
authored and
pyalex
committed
feature row batch produces sample
1 parent 558453a commit 95ca756

3 files changed

Lines changed: 103 additions & 39 deletions

File tree

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

Lines changed: 54 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -207,48 +207,64 @@ public static FeatureRowsBatch fromRow(Row row) {
207207
return new FeatureRowsBatch(row.getSchema(), row.getValues());
208208
}
209209

210-
public Iterator<FeatureRowProto.FeatureRow> getFeatureRows() {
210+
private FeatureRowProto.FeatureRow restoreFeatureRow(int rowIdx) {
211211
int timestampColumnIdx = schema.indexOf("eventTimestamp");
212212
int ingestionIdColumnIdx = schema.indexOf("ingestionId");
213213

214-
return IntStream.range(0, ((List<Object>) values.get(0)).size())
214+
return FeatureRowProto.FeatureRow.newBuilder()
215+
.setFeatureSet(getFeatureSetReference())
216+
.setEventTimestamp(
217+
Timestamp.newBuilder()
218+
.setSeconds(
219+
(long)
220+
(((List<Object>) values.get(timestampColumnIdx)).get(rowIdx)))
221+
.build())
222+
.setIngestionId(
223+
(String) (((List<Object>) values.get(ingestionIdColumnIdx)).get(rowIdx)))
224+
.addAllFields(
225+
schema.getFieldNames().stream()
226+
.map(
227+
fieldName -> {
228+
if (SERVICE_FIELDS.contains(fieldName)) {
229+
return null;
230+
}
231+
int fieldIdx = schema.indexOf(fieldName);
232+
233+
return FieldProto.Field.newBuilder()
234+
.setName(schema.getField(fieldIdx).getName())
235+
.setValue(
236+
objectToProtoValue(
237+
((List<Object>) values.get(fieldIdx)).get(rowIdx),
238+
schemaToProtoTypes.get(
239+
schema
240+
.getField(fieldIdx)
241+
.getType()
242+
.getCollectionElementType())))
243+
.build();
244+
})
245+
.filter(Objects::nonNull)
246+
.collect(Collectors.toList()))
247+
.build();
248+
}
249+
250+
public Iterator<FeatureRowProto.FeatureRow> getFeatureRows() {
251+
int featureCount = ((List<Object>) values.get(0)).size();
252+
253+
return IntStream.range(0, featureCount)
254+
.parallel()
255+
.mapToObj(this::restoreFeatureRow)
256+
.iterator();
257+
}
258+
259+
public Iterator<FeatureRowProto.FeatureRow> getFeatureRowsSample(int maxCount) {
260+
int featureCount = ((List<Object>) values.get(0)).size();
261+
Random rd = new Random(42);
262+
263+
return IntStream.range(0, featureCount)
264+
.filter(idx -> rd.nextInt(featureCount) < maxCount)
215265
.parallel()
216-
.mapToObj(
217-
rowIdx ->
218-
FeatureRowProto.FeatureRow.newBuilder()
219-
.setFeatureSet(getFeatureSetReference())
220-
.setEventTimestamp(
221-
Timestamp.newBuilder()
222-
.setSeconds(
223-
(long)
224-
(((List<Object>) values.get(timestampColumnIdx)).get(rowIdx)))
225-
.build())
226-
.setIngestionId(
227-
(String) (((List<Object>) values.get(ingestionIdColumnIdx)).get(rowIdx)))
228-
.addAllFields(
229-
schema.getFieldNames().stream()
230-
.map(
231-
fieldName -> {
232-
if (SERVICE_FIELDS.contains(fieldName)) {
233-
return null;
234-
}
235-
int fieldIdx = schema.indexOf(fieldName);
236-
237-
return FieldProto.Field.newBuilder()
238-
.setName(schema.getField(fieldIdx).getName())
239-
.setValue(
240-
objectToProtoValue(
241-
((List<Object>) values.get(fieldIdx)).get(rowIdx),
242-
schemaToProtoTypes.get(
243-
schema
244-
.getField(fieldIdx)
245-
.getType()
246-
.getCollectionElementType())))
247-
.build();
248-
})
249-
.filter(Objects::nonNull)
250-
.collect(Collectors.toList()))
251-
.build())
266+
.mapToObj(this::restoreFeatureRow)
267+
.limit(maxCount)
252268
.iterator();
253269
}
254270

‎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;
@@ -498,6 +499,30 @@ public void featureRowCompressShouldPackAndUnpackSuccessfully() {
498499
p.run();
499500
}
500501

502+
@Test
503+
public void featureRowBatchShouldSampleOnRestore() {
504+
List<FeatureRow> stream =
505+
IntStream.range(0, 1000)
506+
.mapToObj(i -> generateRow("project/fs"))
507+
.collect(Collectors.toList());
508+
509+
PCollection<Long> result =
510+
p.apply(Create.of(stream))
511+
.apply("KV", ParDo.of(new ExtractKV()))
512+
.apply(new CompactFeatureRows(1000))
513+
.apply(ParDo.of(new FlatMapWithSample(100)))
514+
.apply(Count.globally());
515+
516+
PAssert.that(result)
517+
.satisfies(
518+
r -> {
519+
// sample size is within bound of required size
520+
assertThat(Math.abs(r.iterator().next() - 100), lessThan(5L));
521+
return null;
522+
});
523+
p.run();
524+
}
525+
501526
private List<FeatureRow> dropNullFeature(List<FeatureRow> input) {
502527
return input.stream()
503528
.map(
@@ -549,4 +574,17 @@ public Table answer(InvocationOnMock invocationOnMock) throws Throwable {
549574
return FakeTable.create(mock(BigQuery.class), tableId, tableDefinition);
550575
}
551576
}
577+
578+
private static class FlatMapWithSample extends DoFn<KV<String, FeatureRowsBatch>, FeatureRow> {
579+
private int sampleSize;
580+
581+
FlatMapWithSample(int sampleSize) {
582+
this.sampleSize = sampleSize;
583+
}
584+
585+
@ProcessElement
586+
public void process(ProcessContext c) {
587+
c.element().getValue().getFeatureRowsSample(sampleSize).forEachRemaining(c::output);
588+
}
589+
}
552590
}

0 commit comments

Comments
 (0)