Skip to content

Commit d8459e0

Browse files
author
Chen Zhiling
authored
Apply default project to rows without project during ingestion (#701)
* Apply default project to incoming rows without project defined * Apply spotless * Fix broken test
1 parent c1a5c43 commit d8459e0

9 files changed

Lines changed: 87 additions & 3 deletions

File tree

‎core/src/main/java/feast/core/job/dataflow/DataflowJobManager.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,7 @@ private ImportOptions getPipelineOptions(
280280
pipelineOptions.setFeatureSetJson(featureSetJsonCompressor.compress(featureSets));
281281
pipelineOptions.setStoreJson(Collections.singletonList(JsonFormat.printer().print(sink)));
282282
pipelineOptions.setProject(projectId);
283+
pipelineOptions.setDefaultFeastProject(Project.DEFAULT_NAME);
283284
pipelineOptions.setUpdate(update);
284285
pipelineOptions.setRunner(DataflowRunner.class);
285286
pipelineOptions.setJobName(jobName);

‎core/src/main/java/feast/core/job/direct/DirectRunnerJobManager.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import feast.core.model.FeatureSet;
2727
import feast.core.model.Job;
2828
import feast.core.model.JobStatus;
29+
import feast.core.model.Project;
2930
import feast.core.util.TypeConversion;
3031
import feast.ingestion.ImportJob;
3132
import feast.ingestion.options.BZip2Compressor;
@@ -105,6 +106,7 @@ private ImportOptions getPipelineOptions(
105106
pipelineOptions.setJobName(jobName);
106107
pipelineOptions.setStoreJson(Collections.singletonList(JsonFormat.printer().print(sink)));
107108
pipelineOptions.setRunner(DirectRunner.class);
109+
pipelineOptions.setDefaultFeastProject(Project.DEFAULT_NAME);
108110
pipelineOptions.setProject(""); // set to default value to satisfy validation
109111
if (metrics.isEnabled()) {
110112
pipelineOptions.setMetricsExporterType(metrics.getType());

‎core/src/main/java/feast/core/model/Project.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
@Entity
3535
@Table(name = "projects")
3636
public class Project {
37+
public static final String DEFAULT_NAME = "default";
3738

3839
// Name of the project
3940
@Id

‎ingestion/src/main/java/feast/ingestion/ImportJob.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ public static PipelineResult runPipeline(ImportOptions options) throws IOExcepti
130130
.get(FEATURE_ROW_OUT)
131131
.apply(
132132
ProcessAndValidateFeatureRows.newBuilder()
133+
.setDefaultProject(options.getDefaultFeastProject())
133134
.setFeatureSetSpecs(featureSetSpecsByKey)
134135
.setSuccessTag(FEATURE_ROW_OUT)
135136
.setFailureTag(DEADLETTER_OUT)

‎ingestion/src/main/java/feast/ingestion/options/ImportOptions.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,13 @@
2727
/** Options passed to Beam to influence the job's execution environment */
2828
public interface ImportOptions extends PipelineOptions, DataflowPipelineOptions, DirectOptions {
2929

30+
@Required
31+
@Description(
32+
"Default feast project to apply to incoming rows that do not specify project in its feature set reference.")
33+
String getDefaultFeastProject();
34+
35+
void setDefaultFeastProject(String defaultProject);
36+
3037
@Required
3138
@Description(
3239
"JSON string representation of the FeatureSet that the import job will process, in BZip2 binary format."

‎ingestion/src/main/java/feast/ingestion/transform/ProcessAndValidateFeatureRows.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ public abstract class ProcessAndValidateFeatureRows
3939

4040
public abstract Map<String, FeatureSetProto.FeatureSetSpec> getFeatureSetSpecs();
4141

42+
public abstract String getDefaultProject();
43+
4244
public abstract TupleTag<FeatureRow> getSuccessTag();
4345

4446
public abstract TupleTag<FailedElement> getFailureTag();
@@ -53,6 +55,8 @@ public abstract static class Builder {
5355
public abstract Builder setFeatureSetSpecs(
5456
Map<String, FeatureSetProto.FeatureSetSpec> featureSets);
5557

58+
public abstract Builder setDefaultProject(String defaultProject);
59+
5660
public abstract Builder setSuccessTag(TupleTag<FeatureRow> successTag);
5761

5862
public abstract Builder setFailureTag(TupleTag<FailedElement> failureTag);
@@ -69,7 +73,7 @@ public PCollectionTuple expand(PCollection<FeatureRow> input) {
6973
.collect(Collectors.toMap(Pair::getLeft, Pair::getRight));
7074

7175
return input
72-
.apply("ProcessFeatureRows", ParDo.of(new ProcessFeatureRowDoFn()))
76+
.apply("ProcessFeatureRows", ParDo.of(new ProcessFeatureRowDoFn(getDefaultProject())))
7377
.apply(
7478
"ValidateFeatureRows",
7579
ParDo.of(

‎ingestion/src/main/java/feast/ingestion/transform/fn/ProcessFeatureRowDoFn.java‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,18 @@
2121

2222
public class ProcessFeatureRowDoFn extends DoFn<FeatureRow, FeatureRow> {
2323

24+
private String defaultProject;
25+
26+
public ProcessFeatureRowDoFn(String defaultProject) {
27+
this.defaultProject = defaultProject;
28+
}
29+
2430
@ProcessElement
2531
public void processElement(ProcessContext context) {
2632
FeatureRow featureRow = context.element();
27-
featureRow =
28-
featureRow.toBuilder().setFeatureSet(stripVersion(featureRow.getFeatureSet())).build();
33+
String featureSetId = stripVersion(featureRow.getFeatureSet());
34+
featureSetId = applyDefaultProject(featureSetId);
35+
featureRow = featureRow.toBuilder().setFeatureSet(featureSetId).build();
2936
context.output(featureRow);
3037
}
3138

@@ -34,4 +41,12 @@ private String stripVersion(String featureSetId) {
3441
String[] split = featureSetId.split(":");
3542
return split[0];
3643
}
44+
45+
private String applyDefaultProject(String featureSetId) {
46+
String[] split = featureSetId.split("/");
47+
if (split.length == 1) {
48+
return defaultProject + "/" + featureSetId;
49+
}
50+
return featureSetId;
51+
}
3752
}

‎ingestion/src/test/java/feast/ingestion/ImportJobTest.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,7 @@ public void runPipeline_ShouldWriteToRedisCorrectlyGivenValidSpecAndFeatureRow()
176176
});
177177
options.setFeatureSetJson(compressor.compress(spec));
178178
options.setStoreJson(Collections.singletonList(JsonFormat.printer().print(redis)));
179+
options.setDefaultFeastProject("myproject");
179180
options.setProject("");
180181
options.setBlockOnRun(false);
181182

‎ingestion/src/test/java/feast/ingestion/transform/ProcessAndValidateFeatureRowsTest.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ public void shouldWriteSuccessAndFailureTagsCorrectly() {
109109
.setCoder(ProtoCoder.of(FeatureRow.class))
110110
.apply(
111111
ProcessAndValidateFeatureRows.newBuilder()
112+
.setDefaultProject("myproject")
112113
.setFailureTag(FAILURE_TAG)
113114
.setSuccessTag(SUCCESS_TAG)
114115
.setFeatureSetSpecs(featureSetSpecs)
@@ -158,6 +159,56 @@ public void shouldStripVersions() {
158159
.setCoder(ProtoCoder.of(FeatureRow.class))
159160
.apply(
160161
ProcessAndValidateFeatureRows.newBuilder()
162+
.setDefaultProject("myproject")
163+
.setFailureTag(FAILURE_TAG)
164+
.setSuccessTag(SUCCESS_TAG)
165+
.setFeatureSetSpecs(featureSetSpecs)
166+
.build());
167+
168+
PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected);
169+
170+
p.run();
171+
}
172+
173+
@Test
174+
public void shouldApplyDefaultProject() {
175+
FeatureSetSpec fs1 =
176+
FeatureSetSpec.newBuilder()
177+
.setName("feature_set")
178+
.setProject("myproject")
179+
.addEntities(
180+
EntitySpec.newBuilder()
181+
.setName("entity_id_primary")
182+
.setValueType(Enum.INT32)
183+
.build())
184+
.addEntities(
185+
EntitySpec.newBuilder()
186+
.setName("entity_id_secondary")
187+
.setValueType(Enum.STRING)
188+
.build())
189+
.addFeatures(
190+
FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build())
191+
.addFeatures(
192+
FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build())
193+
.build();
194+
195+
Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>();
196+
featureSetSpecs.put("myproject/feature_set", fs1);
197+
198+
List<FeatureRow> input = new ArrayList<>();
199+
List<FeatureRow> expected = new ArrayList<>();
200+
201+
FeatureRow randomRow = TestUtil.createRandomFeatureRow(fs1);
202+
expected.add(randomRow);
203+
randomRow = randomRow.toBuilder().setFeatureSet("feature_set").build();
204+
input.add(randomRow);
205+
206+
PCollectionTuple output =
207+
p.apply(Create.of(input))
208+
.setCoder(ProtoCoder.of(FeatureRow.class))
209+
.apply(
210+
ProcessAndValidateFeatureRows.newBuilder()
211+
.setDefaultProject("myproject")
161212
.setFailureTag(FAILURE_TAG)
162213
.setSuccessTag(SUCCESS_TAG)
163214
.setFeatureSetSpecs(featureSetSpecs)
@@ -212,6 +263,7 @@ public void shouldExcludeUnregisteredFields() {
212263
.setCoder(ProtoCoder.of(FeatureRow.class))
213264
.apply(
214265
ProcessAndValidateFeatureRows.newBuilder()
266+
.setDefaultProject("myproject")
215267
.setFailureTag(FAILURE_TAG)
216268
.setSuccessTag(SUCCESS_TAG)
217269
.setFeatureSetSpecs(featureSets)

0 commit comments

Comments
 (0)