Skip to content

Commit 73806e3

Browse files
author
zhilingc
committed
Apply default project to incoming rows without project defined
1 parent 14cae74 commit 73806e3

8 files changed

Lines changed: 85 additions & 2 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
@@ -279,6 +279,7 @@ private ImportOptions getPipelineOptions(
279279
pipelineOptions.setFeatureSetJson(featureSetJsonCompressor.compress(featureSets));
280280
pipelineOptions.setStoreJson(Collections.singletonList(JsonFormat.printer().print(sink)));
281281
pipelineOptions.setProject(projectId);
282+
pipelineOptions.setDefaultFeastProject(Project.DEFAULT_NAME);
282283
pipelineOptions.setUpdate(update);
283284
pipelineOptions.setRunner(DataflowRunner.class);
284285
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: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,12 @@
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("Default feast project to apply to incoming rows that do not specify project in its feature set reference.")
32+
String getDefaultFeastProject();
33+
34+
void setDefaultFeastProject(String defaultProject);
35+
3036
@Required
3137
@Description(
3238
"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 & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,19 @@
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();
33+
String featureSetId = stripVersion(featureRow.getFeatureSet());
34+
featureSetId = applyDefaultProject(featureSetId);
2735
featureRow =
28-
featureRow.toBuilder().setFeatureSet(stripVersion(featureRow.getFeatureSet())).build();
36+
featureRow.toBuilder().setFeatureSet(featureSetId).build();
2937
context.output(featureRow);
3038
}
3139

@@ -34,4 +42,12 @@ private String stripVersion(String featureSetId) {
3442
String[] split = featureSetId.split(":");
3543
return split[0];
3644
}
45+
46+
private String applyDefaultProject(String featureSetId) {
47+
String[] split = featureSetId.split("/");
48+
if (split.length == 1) {
49+
return defaultProject + "/" + featureSetId;
50+
}
51+
return featureSetId;
52+
}
3753
}

‎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,7 @@ public void shouldStripVersions() {
158159
.setCoder(ProtoCoder.of(FeatureRow.class))
159160
.apply(
160161
ProcessAndValidateFeatureRows.newBuilder()
162+
.setDefaultProject("myproject")
161163
.setFailureTag(FAILURE_TAG)
162164
.setSuccessTag(SUCCESS_TAG)
163165
.setFeatureSetSpecs(featureSetSpecs)
@@ -168,6 +170,55 @@ public void shouldStripVersions() {
168170
p.run();
169171
}
170172

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")
212+
.setFailureTag(FAILURE_TAG)
213+
.setSuccessTag(SUCCESS_TAG)
214+
.setFeatureSetSpecs(featureSetSpecs)
215+
.build());
216+
217+
PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected);
218+
219+
p.run();
220+
}
221+
171222
@Test
172223
public void shouldExcludeUnregisteredFields() {
173224
FeatureSetSpec fs1 =
@@ -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)