Skip to content

Use bzip2 compressed feature set json as pipeline option #466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
feast-ci-bot merged 3 commits into from
Feb 14, 2020
Merged

Use bzip2 compressed feature set json as pipeline option #466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7128a0a
Use bzip2 compressed feature set json as pipeline option
Feb 7, 2020
0fe5654
Make decompressor and compressor more generic and extensible
Feb 12, 2020
813b4a7
Avoid code duplication in test
Feb 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 16 additions & 18 deletions core/src/main/java/feast/core/job/dataflow/DataflowJobManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,23 +22,20 @@
import com.google.common.base.Strings;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.google.protobuf.util.JsonFormat.Printer;
import feast.core.FeatureSetProto;
import feast.core.SourceProto;
import feast.core.StoreProto;
import feast.core.config.FeastProperties.MetricsProperties;
import feast.core.exception.JobExecutionException;
import feast.core.job.JobManager;
import feast.core.job.Runner;
import feast.core.model.FeatureSet;
import feast.core.model.Job;
import feast.core.model.JobStatus;
import feast.core.model.Project;
import feast.core.model.Source;
import feast.core.model.Store;
import feast.core.job.option.FeatureSetJsonByteConverter;
import feast.core.model.*;
import feast.core.util.TypeConversion;
import feast.ingestion.ImportJob;
import feast.ingestion.options.BZip2Compressor;
import feast.ingestion.options.ImportOptions;
import feast.ingestion.options.OptionCompressor;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
Expand Down Expand Up @@ -93,7 +90,8 @@ public Job startJob(Job job) {
} catch (InvalidProtocolBufferException e) {
log.error(e.getMessage());
throw new IllegalArgumentException(
String.format("DataflowJobManager failed to START job with id '%s' because the job"
String.format(
"DataflowJobManager failed to START job with id '%s' because the job"
+ "has an invalid spec. Please check the FeatureSet, Source and Store specs. Actual error message: %s",
job.getId(), e.getMessage()));
}
Expand All @@ -112,12 +110,13 @@ public Job updateJob(Job job) {
for (FeatureSet featureSet : job.getFeatureSets()) {
featureSetProtos.add(featureSet.toProto());
}
return submitDataflowJob(job.getId(), featureSetProtos, job.getSource().toProto(),
job.getStore().toProto(), true);
return submitDataflowJob(
job.getId(), featureSetProtos, job.getSource().toProto(), job.getStore().toProto(), true);
} catch (InvalidProtocolBufferException e) {
log.error(e.getMessage());
throw new IllegalArgumentException(
String.format("DataflowJobManager failed to UPDATE job with id '%s' because the job"
String.format(
"DataflowJobManager failed to UPDATE job with id '%s' because the job"
+ "has an invalid spec. Please check the FeatureSet, Source and Store specs. Actual error message: %s",
job.getId(), e.getMessage()));
}
Expand Down Expand Up @@ -221,13 +220,12 @@ private ImportOptions getPipelineOptions(
throws IOException {
String[] args = TypeConversion.convertMapToArgs(defaultOptions);
ImportOptions pipelineOptions = PipelineOptionsFactory.fromArgs(args).as(ImportOptions.class);
Printer printer = JsonFormat.printer();
List<String> featureSetsJson = new ArrayList<>();
for (FeatureSetProto.FeatureSet featureSet : featureSets) {
featureSetsJson.add(printer.print(featureSet.getSpec()));
}
pipelineOptions.setFeatureSetJson(featureSetsJson);
pipelineOptions.setStoreJson(Collections.singletonList(printer.print(sink)));

OptionCompressor<List<FeatureSetProto.FeatureSet>> featureSetJsonCompressor =
new BZip2Compressor<>(new FeatureSetJsonByteConverter());

pipelineOptions.setFeatureSetJson(featureSetJsonCompressor.compress(featureSets));
pipelineOptions.setStoreJson(Collections.singletonList(JsonFormat.printer().print(sink)));
pipelineOptions.setProject(projectId);
pipelineOptions.setUpdate(update);
pipelineOptions.setRunner(DataflowRunner.class);
Expand Down
21 changes: 10 additions & 11 deletions core/src/main/java/feast/core/job/direct/DirectRunnerJobManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,21 +17,22 @@
package feast.core.job.direct;

import com.google.common.base.Strings;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.google.protobuf.util.JsonFormat.Printer;
import feast.core.FeatureSetProto;
import feast.core.StoreProto;
import feast.core.config.FeastProperties.MetricsProperties;
import feast.core.exception.JobExecutionException;
import feast.core.job.JobManager;
import feast.core.job.Runner;
import feast.core.job.option.FeatureSetJsonByteConverter;
import feast.core.model.FeatureSet;
import feast.core.model.Job;
import feast.core.model.JobStatus;
import feast.core.util.TypeConversion;
import feast.ingestion.ImportJob;
import feast.ingestion.options.BZip2Compressor;
import feast.ingestion.options.ImportOptions;
import feast.ingestion.options.OptionCompressor;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
Expand Down Expand Up @@ -92,17 +93,15 @@ public Job startJob(Job job) {
}

private ImportOptions getPipelineOptions(
List<FeatureSetProto.FeatureSet> featureSets, StoreProto.Store sink)
throws InvalidProtocolBufferException {
List<FeatureSetProto.FeatureSet> featureSets, StoreProto.Store sink) throws IOException {
String[] args = TypeConversion.convertMapToArgs(defaultOptions);
ImportOptions pipelineOptions = PipelineOptionsFactory.fromArgs(args).as(ImportOptions.class);
Printer printer = JsonFormat.printer();
List<String> featureSetsJson = new ArrayList<>();
for (FeatureSetProto.FeatureSet featureSet : featureSets) {
featureSetsJson.add(printer.print(featureSet.getSpec()));
}
pipelineOptions.setFeatureSetJson(featureSetsJson);
pipelineOptions.setStoreJson(Collections.singletonList(printer.print(sink)));

OptionCompressor<List<FeatureSetProto.FeatureSet>> featureSetJsonCompressor =
new BZip2Compressor<>(new FeatureSetJsonByteConverter());

pipelineOptions.setFeatureSetJson(featureSetJsonCompressor.compress(featureSets));
pipelineOptions.setStoreJson(Collections.singletonList(JsonFormat.printer().print(sink)));
pipelineOptions.setRunner(DirectRunner.class);
pipelineOptions.setProject(""); // set to default value to satisfy validation
if (metrics.isEnabled()) {
Expand Down
47 changes: 47 additions & 0 deletions core/src/main/java/feast/core/job/option/FeatureSetJsonByteConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/*
* SPDX-License-Identifier: Apache-2.0
* Copyright 2018-2020 The Feast Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package feast.core.job.option;

import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import feast.core.FeatureSetProto;
import feast.ingestion.options.OptionByteConverter;
import java.util.ArrayList;
import java.util.List;

public class FeatureSetJsonByteConverter
implements OptionByteConverter<List<FeatureSetProto.FeatureSet>> {

/**
* Convert list of feature sets to json strings joined by new line, represented as byte arrays
*
* @param featureSets List of feature set protobufs
* @return Byte array representation of the json strings
* @throws InvalidProtocolBufferException
*/
@Override
public byte[] toByte(List<FeatureSetProto.FeatureSet> featureSets)
throws InvalidProtocolBufferException {
JsonFormat.Printer printer =
JsonFormat.printer().omittingInsignificantWhitespace().printingEnumsAsInts();
List<String> featureSetsJson = new ArrayList<>();
for (FeatureSetProto.FeatureSet featureSet : featureSets) {
featureSetsJson.add(printer.print(featureSet.getSpec()));
}
return String.join("\n", featureSetsJson).getBytes();
}
}
16 changes: 8 additions & 8 deletions core/src/main/java/feast/core/model/FeatureSet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,8 +264,8 @@ private void setEntitySpecFields(EntitySpec.Builder entitySpecBuilder, Field ent
if (entityField.getPresence() != null) {
entitySpecBuilder.setPresence(FeaturePresence.parseFrom(entityField.getPresence()));
} else if (entityField.getGroupPresence() != null) {
entitySpecBuilder
.setGroupPresence(FeaturePresenceWithinGroup.parseFrom(entityField.getGroupPresence()));
entitySpecBuilder.setGroupPresence(
FeaturePresenceWithinGroup.parseFrom(entityField.getGroupPresence()));
}

if (entityField.getShape() != null) {
Expand Down Expand Up @@ -298,8 +298,8 @@ private void setEntitySpecFields(EntitySpec.Builder entitySpecBuilder, Field ent
} else if (entityField.getTimeDomain() != null) {
entitySpecBuilder.setTimeDomain(TimeDomain.parseFrom(entityField.getTimeDomain()));
} else if (entityField.getTimeOfDayDomain() != null) {
entitySpecBuilder
.setTimeOfDayDomain(TimeOfDayDomain.parseFrom(entityField.getTimeOfDayDomain()));
entitySpecBuilder.setTimeOfDayDomain(
TimeOfDayDomain.parseFrom(entityField.getTimeOfDayDomain()));
}
}

Expand All @@ -314,8 +314,8 @@ private void setFeatureSpecFields(FeatureSpec.Builder featureSpecBuilder, Field
if (featureField.getPresence() != null) {
featureSpecBuilder.setPresence(FeaturePresence.parseFrom(featureField.getPresence()));
} else if (featureField.getGroupPresence() != null) {
featureSpecBuilder
.setGroupPresence(FeaturePresenceWithinGroup.parseFrom(featureField.getGroupPresence()));
featureSpecBuilder.setGroupPresence(
FeaturePresenceWithinGroup.parseFrom(featureField.getGroupPresence()));
}

if (featureField.getShape() != null) {
Expand Down Expand Up @@ -348,8 +348,8 @@ private void setFeatureSpecFields(FeatureSpec.Builder featureSpecBuilder, Field
} else if (featureField.getTimeDomain() != null) {
featureSpecBuilder.setTimeDomain(TimeDomain.parseFrom(featureField.getTimeDomain()));
} else if (featureField.getTimeOfDayDomain() != null) {
featureSpecBuilder
.setTimeOfDayDomain(TimeOfDayDomain.parseFrom(featureField.getTimeOfDayDomain()));
featureSpecBuilder.setTimeOfDayDomain(
TimeOfDayDomain.parseFrom(featureField.getTimeOfDayDomain()));
}
}

Expand Down
3 changes: 1 addition & 2 deletions core/src/main/java/feast/core/model/Field.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,8 +71,7 @@ public class Field {
private byte[] timeDomain;
private byte[] timeOfDayDomain;

public Field() {
}
public Field() {}

public Field(String name, ValueType.Enum type) {
this.name = name;
Expand Down
9 changes: 4 additions & 5 deletions core/src/main/java/feast/core/service/SpecService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,8 +167,7 @@ public ListFeatureSetsResponse listFeatureSets(ListFeatureSetsRequest.Filter fil
checkValidCharactersAllowAsterisk(name, "featureSetName");
checkValidCharactersAllowAsterisk(project, "projectName");

List<FeatureSet> featureSets = new ArrayList<FeatureSet>() {
};
List<FeatureSet> featureSets = new ArrayList<FeatureSet>() {};

if (project.equals("*")) {
// Matching all projects
Expand Down Expand Up @@ -277,9 +276,9 @@ public ListStoresResponse listStores(ListStoresRequest.Filter filter) {
* Creates or updates a feature set in the repository. If there is a change in the feature set
* schema, then the feature set version will be incremented.
*
* <p>This function is idempotent. If no changes are detected in the incoming featureSet's
* schema, this method will update the incoming featureSet spec with the latest version stored in
* the repository, and return that.
* <p>This function is idempotent. If no changes are detected in the incoming featureSet's schema,
* this method will update the incoming featureSet spec with the latest version stored in the
* repository, and return that.
*
* @param newFeatureSet Feature set that will be created or updated.
*/
Expand Down
36 changes: 24 additions & 12 deletions core/src/test/java/feast/core/job/dataflow/DataflowJobManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,7 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.*;
import static org.mockito.MockitoAnnotations.initMocks;

import com.google.api.services.dataflow.Dataflow;
Expand All @@ -44,14 +40,15 @@
import feast.core.config.FeastProperties.MetricsProperties;
import feast.core.exception.JobExecutionException;
import feast.core.job.Runner;
import feast.core.model.FeatureSet;
import feast.core.model.Job;
import feast.core.model.JobStatus;
import feast.core.model.Source;
import feast.core.model.Store;
import feast.core.job.option.FeatureSetJsonByteConverter;
import feast.core.model.*;
import feast.ingestion.options.BZip2Compressor;
import feast.ingestion.options.ImportOptions;
import feast.ingestion.options.OptionCompressor;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.beam.runners.dataflow.DataflowPipelineJob;
import org.apache.beam.runners.dataflow.DataflowRunner;
Expand Down Expand Up @@ -131,8 +128,11 @@ public void shouldStartJobWithCorrectPipelineOptions() throws IOException {
expectedPipelineOptions.setAppName("DataflowJobManager");
expectedPipelineOptions.setJobName(jobName);
expectedPipelineOptions.setStoreJson(Lists.newArrayList(printer.print(store)));

OptionCompressor<List<FeatureSetProto.FeatureSet>> featureSetJsonCompressor =
new BZip2Compressor<>(new FeatureSetJsonByteConverter());
expectedPipelineOptions.setFeatureSetJson(
Lists.newArrayList(printer.print(featureSet.getSpec())));
featureSetJsonCompressor.compress(Collections.singletonList(featureSet)));

ArgumentCaptor<ImportOptions> captor = ArgumentCaptor.forClass(ImportOptions.class);

Expand Down Expand Up @@ -170,7 +170,19 @@ public void shouldStartJobWithCorrectPipelineOptions() throws IOException {
// Assume the files that are staged are correct
expectedPipelineOptions.setFilesToStage(actualPipelineOptions.getFilesToStage());

assertThat(actualPipelineOptions.toString(), equalTo(expectedPipelineOptions.toString()));
assertThat(
actualPipelineOptions.getFeatureSetJson(),
equalTo(expectedPipelineOptions.getFeatureSetJson()));
assertThat(
actualPipelineOptions.getDeadLetterTableSpec(),
equalTo(expectedPipelineOptions.getDeadLetterTableSpec()));
assertThat(
actualPipelineOptions.getStatsdHost(), equalTo(expectedPipelineOptions.getStatsdHost()));
assertThat(
actualPipelineOptions.getMetricsExporterType(),
equalTo(expectedPipelineOptions.getMetricsExporterType()));
assertThat(
actualPipelineOptions.getStoreJson(), equalTo(expectedPipelineOptions.getStoreJson()));
assertThat(actual.getExtId(), equalTo(expectedExtJobId));
}

Expand Down
25 changes: 23 additions & 2 deletions core/src/test/java/feast/core/job/direct/DirectRunnerJobManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,19 @@
import feast.core.StoreProto.Store.Subscription;
import feast.core.config.FeastProperties.MetricsProperties;
import feast.core.job.Runner;
import feast.core.job.option.FeatureSetJsonByteConverter;
import feast.core.model.FeatureSet;
import feast.core.model.Job;
import feast.core.model.JobStatus;
import feast.core.model.Source;
import feast.core.model.Store;
import feast.ingestion.options.BZip2Compressor;
import feast.ingestion.options.ImportOptions;
import feast.ingestion.options.OptionCompressor;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.beam.runners.direct.DirectRunner;
import org.apache.beam.sdk.PipelineResult;
Expand Down Expand Up @@ -121,8 +126,11 @@ public void shouldStartDirectJobAndRegisterPipelineResult() throws IOException {
expectedPipelineOptions.setProject("");
expectedPipelineOptions.setStoreJson(Lists.newArrayList(printer.print(store)));
expectedPipelineOptions.setProject("");

OptionCompressor<List<FeatureSetProto.FeatureSet>> featureSetJsonCompressor =
new BZip2Compressor<>(new FeatureSetJsonByteConverter());
expectedPipelineOptions.setFeatureSetJson(
Lists.newArrayList(printer.print(featureSet.getSpec())));
featureSetJsonCompressor.compress(Collections.singletonList(featureSet)));

String expectedJobId = "feast-job-0";
ArgumentCaptor<ImportOptions> pipelineOptionsCaptor =
Expand Down Expand Up @@ -150,7 +158,20 @@ public void shouldStartDirectJobAndRegisterPipelineResult() throws IOException {
expectedPipelineOptions.setOptionsId(
actualPipelineOptions.getOptionsId()); // avoid comparing this value

assertThat(actualPipelineOptions.toString(), equalTo(expectedPipelineOptions.toString()));
assertThat(
actualPipelineOptions.getFeatureSetJson(),
equalTo(expectedPipelineOptions.getFeatureSetJson()));
assertThat(
actualPipelineOptions.getDeadLetterTableSpec(),
equalTo(expectedPipelineOptions.getDeadLetterTableSpec()));
assertThat(
actualPipelineOptions.getStatsdHost(), equalTo(expectedPipelineOptions.getStatsdHost()));
assertThat(
actualPipelineOptions.getMetricsExporterType(),
equalTo(expectedPipelineOptions.getMetricsExporterType()));
assertThat(
actualPipelineOptions.getStoreJson(), equalTo(expectedPipelineOptions.getStoreJson()));

assertThat(jobStarted.getPipelineResult(), equalTo(mockPipelineResult));
assertThat(jobStarted.getJobId(), equalTo(expectedJobId));
assertThat(actual.getExtId(), equalTo(expectedJobId));
Expand Down
Loading