ProcessAndValidateFeatureRows extends PTransform<PCollection<FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_ProcessAndValidateFeatureRows.Builder(); }  abstract PCollectionView<Map<String, Iterable<FeatureSetProto.FeatureSetSpec>>> getFeatureSetSpecs(); abstract String getDefaultProject(); abstract TupleTag<FeatureRow> getSuccessTag(); abstract TupleTag<FailedElement> getFailureTag(); static Builder newBuilder(); @Override PCollectionTuple expand(PCollection<FeatureRow> input);  }
@Test public void shouldWriteSuccessAndFailureTagsCorrectly() { FeatureSetSpec fs1 = FeatureSetSpec.newBuilder() .setName("feature_set") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) .build(); FeatureSetSpec fs2 = FeatureSetSpec.newBuilder() .setName("feature_set_2") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) .build(); Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>(); featureSetSpecs.put("myproject/feature_set", fs1); featureSetSpecs.put("myproject/feature_set_2", fs2); List<FeatureRow> input = new ArrayList<>(); List<FeatureRow> expected = new ArrayList<>(); for (FeatureSetSpec featureSetSpec : featureSetSpecs.values()) { FeatureRow randomRow = TestUtil.createRandomFeatureRow(featureSetSpec); input.add(randomRow); expected.add(randomRow); } FeatureRow invalidRow = FeatureRow.newBuilder() .setFeatureSet(getFeatureSetStringRef(fs1)) .addFields( Field.newBuilder() .setName("feature_1") .setValue(Value.newBuilder().setBoolVal(false).build()) .build()) .build(); input.add(invalidRow); PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView = p.apply("StaticSpecs", Create.of(featureSetSpecs)).apply(View.asMultimap()); PCollectionTuple output = p.apply(Create.of(input)) .setCoder(ProtoCoder.of(FeatureRow.class)) .apply( ProcessAndValidateFeatureRows.newBuilder() .setDefaultProject("myproject") .setFailureTag(FAILURE_TAG) .setSuccessTag(SUCCESS_TAG) .setFeatureSetSpecs(specsView) .build()); PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected); PAssert.that(output.get(FAILURE_TAG).apply(Count.globally())).containsInAnyOrder(1L); p.run(); }
@Test public void shouldStripVersions() { FeatureSetSpec fs1 = FeatureSetSpec.newBuilder() .setName("feature_set") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) .build(); Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>(); featureSetSpecs.put("myproject/feature_set", fs1); List<FeatureRow> input = new ArrayList<>(); List<FeatureRow> expected = new ArrayList<>(); FeatureRow randomRow = TestUtil.createRandomFeatureRow(fs1); expected.add(randomRow); randomRow = randomRow.toBuilder().setFeatureSet("myproject/feature_set:1").build(); input.add(randomRow); PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView = p.apply("StaticSpecs", Create.of(featureSetSpecs)).apply(View.asMultimap()); PCollectionTuple output = p.apply(Create.of(input)) .setCoder(ProtoCoder.of(FeatureRow.class)) .apply( ProcessAndValidateFeatureRows.newBuilder() .setDefaultProject("myproject") .setFailureTag(FAILURE_TAG) .setSuccessTag(SUCCESS_TAG) .setFeatureSetSpecs(specsView) .build()); PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected); p.run(); }
@Test public void shouldApplyDefaultProject() { FeatureSetSpec fs1 = FeatureSetSpec.newBuilder() .setName("feature_set") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) .build(); Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>(); featureSetSpecs.put("myproject/feature_set", fs1); List<FeatureRow> input = new ArrayList<>(); List<FeatureRow> expected = new ArrayList<>(); FeatureRow randomRow = TestUtil.createRandomFeatureRow(fs1); expected.add(randomRow); randomRow = randomRow.toBuilder().setFeatureSet("feature_set").build(); input.add(randomRow); PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView = p.apply("StaticSpecs", Create.of(featureSetSpecs)).apply(View.asMultimap()); PCollectionTuple output = p.apply(Create.of(input)) .setCoder(ProtoCoder.of(FeatureRow.class)) .apply( ProcessAndValidateFeatureRows.newBuilder() .setDefaultProject("myproject") .setFailureTag(FAILURE_TAG) .setSuccessTag(SUCCESS_TAG) .setFeatureSetSpecs(specsView) .build()); PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected); p.run(); }
@Test public void shouldExcludeUnregisteredFields() { FeatureSetSpec fs1 = FeatureSetSpec.newBuilder() .setName("feature_set") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) .build(); Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>(); featureSetSpecs.put("myproject/feature_set", fs1); List<FeatureRow> input = new ArrayList<>(); List<FeatureRow> expected = new ArrayList<>(); FeatureRow randomRow = TestUtil.createRandomFeatureRow(fs1); expected.add(randomRow); input.add( randomRow .toBuilder() .addFields( Field.newBuilder() .setName("extra") .setValue(Value.newBuilder().setStringVal("hello"))) .build()); PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView = p.apply("StaticSpecs", Create.of(featureSetSpecs)).apply(View.asMultimap()); PCollectionTuple output = p.apply(Create.of(input)) .setCoder(ProtoCoder.of(FeatureRow.class)) .apply( ProcessAndValidateFeatureRows.newBuilder() .setDefaultProject("myproject") .setFailureTag(FAILURE_TAG) .setSuccessTag(SUCCESS_TAG) .setFeatureSetSpecs(specsView) .build()); PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected); p.run(); }