Skip to content

Commit 19c8bf3

Browse files
author
Chen Zhiling
committed
Storage refactor API and docstring tweaks (#569)
* API and docstring tweaks * Fix javadoc linting errors * Apply spotless * Fix javadoc formatting * Drop result from HistoricalRetrievalResult constructors
1 parent 36ca802 commit 19c8bf3

File tree

17 files changed

+223
-132
lines changed

17 files changed

+223
-132
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public static PipelineResult runPipeline(ImportOptions options) throws IOExcepti
137137

138138
// Step 3. Write FeatureRow to the corresponding Store.
139139
WriteResult writeFeatureRows =
140-
validatedRows.get(FEATURE_ROW_OUT).apply("WriteFeatureRowToStore", featureSink.write());
140+
validatedRows.get(FEATURE_ROW_OUT).apply("WriteFeatureRowToStore", featureSink.writer());
141141

142142
// Step 4. Write FailedElements to a dead letter table in BigQuery.
143143
if (options.getDeadLetterTableSpec() != null) {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,9 @@ public static void publishFeatureRowsToKafka(
162162
* Create a Feature Row with random value according to the FeatureSetSpec
163163
*
164164
* <p>See {@link #createRandomFeatureRow(FeatureSetSpec, int)}
165+
*
166+
* @param featureSetSpec {@link FeatureSetSpec}
167+
* @return {@link FeatureRow}
165168
*/
166169
public static FeatureRow createRandomFeatureRow(FeatureSetSpec featureSetSpec) {
167170
ThreadLocalRandom random = ThreadLocalRandom.current();

‎serving/src/main/java/feast/serving/configuration/ServingServiceConfig.java‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,11 @@
2525
import feast.core.StoreProto.Store.RedisConfig;
2626
import feast.core.StoreProto.Store.Subscription;
2727
import feast.serving.FeastProperties;
28-
import feast.serving.service.BatchServingService;
29-
import feast.serving.service.JobService;
30-
import feast.serving.service.NoopJobService;
31-
import feast.serving.service.OnlineServingService;
32-
import feast.serving.service.ServingService;
28+
import feast.serving.service.*;
3329
import feast.serving.specs.CachedSpecService;
34-
import feast.storage.api.retrieval.BatchRetriever;
30+
import feast.storage.api.retrieval.HistoricalRetriever;
3531
import feast.storage.api.retrieval.OnlineRetriever;
36-
import feast.storage.connectors.bigquery.retrieval.BigQueryBatchRetriever;
32+
import feast.storage.connectors.bigquery.retrieval.BigQueryHistoricalRetriever;
3733
import feast.storage.connectors.redis.retrieval.RedisOnlineRetriever;
3834
import io.opentracing.Tracer;
3935
import java.util.Map;
@@ -109,8 +105,8 @@ public ServingService servingService(
109105
"Unable to instantiate jobService for BigQuery store.");
110106
}
111107

112-
BatchRetriever bqRetriever =
113-
BigQueryBatchRetriever.builder()
108+
HistoricalRetriever bqRetriever =
109+
BigQueryHistoricalRetriever.builder()
114110
.setBigquery(bigquery)
115111
.setDatasetId(bqConfig.getDatasetId())
116112
.setProjectId(bqConfig.getProjectId())
@@ -121,7 +117,7 @@ public ServingService servingService(
121117
.setStorage(storage)
122118
.build();
123119

124-
servingService = new BatchServingService(bqRetriever, specService, jobService);
120+
servingService = new HistoricalServingService(bqRetriever, specService, jobService);
125121
break;
126122
case CASSANDRA:
127123
case UNRECOGNIZED:

serving/src/main/java/feast/serving/service/BatchServingService.java renamed to serving/src/main/java/feast/serving/service/HistoricalServingService.java

Lines changed: 51 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,25 +18,29 @@
1818

1919
import feast.serving.ServingAPIProto;
2020
import feast.serving.ServingAPIProto.*;
21+
import feast.serving.ServingAPIProto.Job.Builder;
2122
import feast.serving.specs.CachedSpecService;
22-
import feast.storage.api.retrieval.BatchRetriever;
2323
import feast.storage.api.retrieval.FeatureSetRequest;
24-
import feast.storage.connectors.bigquery.retrieval.BigQueryBatchRetriever;
24+
import feast.storage.api.retrieval.HistoricalRetrievalResult;
25+
import feast.storage.api.retrieval.HistoricalRetriever;
26+
import feast.storage.connectors.bigquery.retrieval.BigQueryHistoricalRetriever;
2527
import io.grpc.Status;
2628
import java.util.List;
2729
import java.util.Optional;
30+
import java.util.UUID;
2831
import org.slf4j.Logger;
2932

30-
public class BatchServingService implements ServingService {
33+
public class HistoricalServingService implements ServingService {
3134

32-
private static final Logger log = org.slf4j.LoggerFactory.getLogger(BatchServingService.class);
35+
private static final Logger log =
36+
org.slf4j.LoggerFactory.getLogger(HistoricalServingService.class);
3337

34-
private final BatchRetriever retriever;
38+
private final HistoricalRetriever retriever;
3539
private final CachedSpecService specService;
3640
private final JobService jobService;
3741

38-
public BatchServingService(
39-
BatchRetriever retriever, CachedSpecService specService, JobService jobService) {
42+
public HistoricalServingService(
43+
HistoricalRetriever retriever, CachedSpecService specService, JobService jobService) {
4044
this.retriever = retriever;
4145
this.specService = specService;
4246
this.jobService = jobService;
@@ -47,10 +51,11 @@ public BatchServingService(
4751
public GetFeastServingInfoResponse getFeastServingInfo(
4852
GetFeastServingInfoRequest getFeastServingInfoRequest) {
4953
try {
50-
BigQueryBatchRetriever bigQueryBatchRetriever = (BigQueryBatchRetriever) retriever;
54+
BigQueryHistoricalRetriever bigQueryHistoricalRetriever =
55+
(BigQueryHistoricalRetriever) retriever;
5156
return GetFeastServingInfoResponse.newBuilder()
5257
.setType(FeastServingType.FEAST_SERVING_TYPE_BATCH)
53-
.setJobStagingLocation(bigQueryBatchRetriever.jobStagingLocation())
58+
.setJobStagingLocation(bigQueryHistoricalRetriever.jobStagingLocation())
5459
.build();
5560
} catch (Exception e) {
5661
return GetFeastServingInfoResponse.newBuilder()
@@ -70,9 +75,28 @@ public GetOnlineFeaturesResponse getOnlineFeatures(GetOnlineFeaturesRequest getF
7075
public GetBatchFeaturesResponse getBatchFeatures(GetBatchFeaturesRequest getFeaturesRequest) {
7176
List<FeatureSetRequest> featureSetRequests =
7277
specService.getFeatureSets(getFeaturesRequest.getFeaturesList());
73-
Job feastJob = retriever.getBatchFeatures(getFeaturesRequest, featureSetRequests);
74-
jobService.upsert(feastJob);
75-
return GetBatchFeaturesResponse.newBuilder().setJob(feastJob).build();
78+
String retrievalId = UUID.randomUUID().toString();
79+
Job runningJob =
80+
Job.newBuilder()
81+
.setId(retrievalId)
82+
.setType(JobType.JOB_TYPE_DOWNLOAD)
83+
.setStatus(JobStatus.JOB_STATUS_RUNNING)
84+
.build();
85+
jobService.upsert(runningJob);
86+
Thread thread =
87+
new Thread(
88+
new Runnable() {
89+
@Override
90+
public void run() {
91+
HistoricalRetrievalResult result =
92+
retriever.getHistoricalFeatures(
93+
retrievalId, getFeaturesRequest.getDatasetSource(), featureSetRequests);
94+
jobService.upsert(resultToJob(result));
95+
}
96+
});
97+
thread.start();
98+
99+
return GetBatchFeaturesResponse.newBuilder().setJob(runningJob).build();
76100
}
77101

78102
/** {@inheritDoc} */
@@ -86,4 +110,19 @@ public GetJobResponse getJob(GetJobRequest getJobRequest) {
86110
}
87111
return GetJobResponse.newBuilder().setJob(job.get()).build();
88112
}
113+
114+
private Job resultToJob(HistoricalRetrievalResult result) {
115+
Builder builder =
116+
Job.newBuilder()
117+
.setId(result.getId())
118+
.setType(JobType.JOB_TYPE_DOWNLOAD)
119+
.setStatus(result.getStatus());
120+
if (result.hasError()) {
121+
return builder.setError(result.getError()).build();
122+
}
123+
return builder
124+
.addAllFileUris(result.getFileUris())
125+
.setDataFormat(result.getDataFormat())
126+
.build();
127+
}
89128
}

‎storage/api/src/main/java/feast/storage/api/retrieval/HistoricalRetrievalResult.java‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.storage.api.retrieval;
18+
19+
import com.google.auto.value.AutoValue;
20+
import feast.serving.ServingAPIProto.DataFormat;
21+
import feast.serving.ServingAPIProto.JobStatus;
22+
import java.io.Serializable;
23+
import java.util.List;
24+
import javax.annotation.Nullable;
25+
26+
/** Result of a historical feature retrieval request. */
27+
@AutoValue
28+
public abstract class HistoricalRetrievalResult implements Serializable {
29+
30+
public abstract String getId();
31+
32+
public abstract JobStatus getStatus();
33+
34+
@Nullable
35+
public abstract String getError();
36+
37+
@Nullable
38+
public abstract List<String> getFileUris();
39+
40+
@Nullable
41+
public abstract DataFormat getDataFormat();
42+
43+
/**
44+
* Instantiates a {@link HistoricalRetrievalResult} indicating that the retrieval was a failure,
45+
* together with its associated error.
46+
*
47+
* @param id retrieval id identifying the retrieval request.
48+
* @param error error that occurred
49+
* @return {@link HistoricalRetrievalResult}
50+
*/
51+
public static HistoricalRetrievalResult error(String id, Exception error) {
52+
return newBuilder()
53+
.setId(id)
54+
.setStatus(JobStatus.JOB_STATUS_DONE)
55+
.setError(error.getMessage())
56+
.build();
57+
}
58+
59+
/**
60+
* Instantiates a {@link HistoricalRetrievalResult} indicating that the retrieval was a success,
61+
* together with the location of the output.
62+
*
63+
* @param id retrieval id identifying the retrieval request
64+
* @param fileUris list of output file URIs
65+
* @param dataFormat data format of the output files
66+
* @return
67+
*/
68+
public static HistoricalRetrievalResult success(
69+
String id, List<String> fileUris, DataFormat dataFormat) {
70+
return newBuilder()
71+
.setId(id)
72+
.setStatus(JobStatus.JOB_STATUS_DONE)
73+
.setFileUris(fileUris)
74+
.setDataFormat(dataFormat)
75+
.build();
76+
}
77+
78+
static Builder newBuilder() {
79+
return new AutoValue_HistoricalRetrievalResult.Builder();
80+
}
81+
82+
@AutoValue.Builder
83+
abstract static class Builder {
84+
abstract Builder setId(String id);
85+
86+
abstract Builder setStatus(JobStatus jobStatus);
87+
88+
abstract Builder setError(String error);
89+
90+
abstract Builder setFileUris(List<String> fileUris);
91+
92+
abstract Builder setDataFormat(DataFormat dataFormat);
93+
94+
abstract HistoricalRetrievalResult build();
95+
}
96+
97+
public boolean hasError() {
98+
return getError() != null;
99+
}
100+
}

storage/api/src/main/java/feast/storage/api/retrieval/BatchRetriever.java renamed to storage/api/src/main/java/feast/storage/api/retrieval/HistoricalRetriever.java

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,26 @@
1616
*/
1717
package feast.storage.api.retrieval;
1818

19-
import feast.serving.ServingAPIProto;
19+
import feast.serving.ServingAPIProto.DatasetSource;
2020
import java.util.List;
2121

22-
/** Interface for implementing user defined retrieval functionality from Batch/historical stores. */
23-
public interface BatchRetriever {
22+
/**
23+
* A historical retriever is a feature retriever that retrieves feature data corresponding to
24+
* provided entities over a given period of time.
25+
*/
26+
public interface HistoricalRetriever {
2427

2528
/**
2629
* Get all features corresponding to the provided batch features request.
2730
*
28-
* @param request {@link ServingAPIProto.GetBatchFeaturesRequest} containing requested features
29-
* and file containing entity columns.
31+
* @param retrievalId String that uniquely identifies this retrieval request.
32+
* @param datasetSource {@link DatasetSource} containing source to load the dataset containing
33+
* entity columns.
3034
* @param featureSetRequests List of {@link FeatureSetRequest} to feature references in the
3135
* request tied to that feature set.
32-
* @return {@link ServingAPIProto.Job} if successful, contains the location of the results, else
33-
* contains the error to be returned to the user.
36+
* @return {@link HistoricalRetrievalResult} if successful, contains the location of the results,
37+
* else contains the error to be returned to the user.
3438
*/
35-
ServingAPIProto.Job getBatchFeatures(
36-
ServingAPIProto.GetBatchFeaturesRequest request, List<FeatureSetRequest> featureSetRequests);
39+
HistoricalRetrievalResult getHistoricalFeatures(
40+
String retrievalId, DatasetSource datasetSource, List<FeatureSetRequest> featureSetRequests);
3741
}

‎storage/api/src/main/java/feast/storage/api/retrieval/OnlineRetriever.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@
2020
import feast.types.FeatureRowProto.FeatureRow;
2121
import java.util.List;
2222

23-
/** Interface for implementing user defined retrieval functionality from Online stores. */
23+
/**
24+
* An online retriever is a feature retriever that retrieves the latest feature data corresponding
25+
* to provided entities.
26+
*/
2427
public interface OnlineRetriever {
2528

2629
/**
@@ -29,7 +32,8 @@ public interface OnlineRetriever {
2932
* @param entityRows list of entity rows in the feature request
3033
* @param featureSetRequests List of {@link FeatureSetRequest} to feature references in the
3134
* request tied to that feature set.
32-
* @return list of {@link OnlineRetrieverResponse} for each entity row
35+
* @return list of lists of {@link FeatureRow}s corresponding to each feature set request and
36+
* entity row.
3337
*/
3438
List<List<FeatureRow>> getOnlineFeatures(
3539
List<EntityRow> entityRows, List<FeatureSetRequest> featureSetRequests);

‎storage/api/src/main/java/feast/storage/api/retrieval/OnlineRetrieverResponse.java‎

Lines changed: 0 additions & 44 deletions
This file was deleted.

‎storage/api/src/main/java/feast/storage/api/write/FeatureSink.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,5 +50,5 @@ public interface FeatureSink extends Serializable {
5050
*
5151
* @return {@link PTransform}
5252
*/
53-
PTransform<PCollection<FeatureRow>, WriteResult> write();
53+
PTransform<PCollection<FeatureRow>, WriteResult> writer();
5454
}

‎storage/api/src/main/java/feast/storage/api/write/WriteResult.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,15 @@ public final class WriteResult implements Serializable, POutput {
3434
private static TupleTag<FeatureRow> successfulInsertsTag = new TupleTag<>("successfulInserts");
3535
private static TupleTag<FailedElement> failedInsertsTupleTag = new TupleTag<>("failedInserts");
3636

37-
/** Creates a {@link WriteResult} in the given {@link Pipeline}. */
37+
/**
38+
* Creates a {@link WriteResult} in the given {@link Pipeline}.
39+
*
40+
* @param pipeline {@link Pipeline}
41+
* @param successfulInserts {@link PCollection} of {@link FeatureRow}s successfully inserted into
42+
* the store
43+
* @param failedInserts {@link PCollection} of {@link FailedElement}s
44+
* @return {@link WriteResult}
45+
*/
3846
public static WriteResult in(
3947
Pipeline pipeline,
4048
PCollection<FeatureRow> successfulInserts,

0 commit comments

Comments
 (0)