Skip to content

Commit b04022a

Browse files
author
zhilingc
committed
Fix batch retrieval, use templating
1 parent ca1e7c6 commit b04022a

4 files changed

Lines changed: 137 additions & 63 deletions

File tree

‎serving/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,12 @@
165165
<version>${protobufVersion}</version>
166166
</dependency>
167167

168+
<dependency>
169+
<groupId>io.pebbletemplates</groupId>
170+
<artifactId>pebble</artifactId>
171+
<version>3.1.0</version>
172+
</dependency>
173+
168174
<!--compile 'redis.clients:jedis:2.9.0'-->
169175
<dependency>
170176
<groupId>redis.clients</groupId>

‎serving/src/main/java/feast/serving/service/BigQueryServingService.java‎

Lines changed: 85 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
11
package feast.serving.service;
22

3+
import static feast.serving.util.BigQueryUtil.getTimestampLimitQuery;
4+
5+
import com.google.api.services.bigquery.model.TableReference;
36
import com.google.cloud.bigquery.BigQuery;
7+
import com.google.cloud.bigquery.BigQuery.TableField;
8+
import com.google.cloud.bigquery.BigQuery.TableOption;
49
import com.google.cloud.bigquery.BigQueryException;
10+
import com.google.cloud.bigquery.Dataset;
511
import com.google.cloud.bigquery.DatasetId;
612
import com.google.cloud.bigquery.ExtractJobConfiguration;
713
import com.google.cloud.bigquery.Field;
14+
import com.google.cloud.bigquery.FieldValueList;
815
import com.google.cloud.bigquery.FormatOptions;
916
import com.google.cloud.bigquery.Job;
1017
import com.google.cloud.bigquery.JobInfo;
@@ -15,6 +22,7 @@
1522
import com.google.cloud.bigquery.TableDefinition;
1623
import com.google.cloud.bigquery.TableId;
1724
import com.google.cloud.bigquery.TableInfo;
25+
import com.google.cloud.bigquery.TableResult;
1826
import com.google.cloud.storage.Blob;
1927
import com.google.cloud.storage.Storage;
2028
import com.google.cloud.storage.Storage.BlobListOption;
@@ -36,6 +44,7 @@
3644
import feast.serving.ServingAPIProto.JobType;
3745
import feast.serving.util.BigQueryUtil;
3846
import io.grpc.Status;
47+
import java.io.IOException;
3948
import java.util.ArrayList;
4049
import java.util.List;
4150
import java.util.Optional;
@@ -46,6 +55,8 @@
4655
@Slf4j
4756
public class BigQueryServingService implements ServingService {
4857

58+
private static final Long TABLE_EXPIRATION_TIME = 172800000L;
59+
4960
private final BigQuery bigquery;
5061
private final String projectId;
5162
private final String datasetId;
@@ -112,19 +123,34 @@ public GetBatchFeaturesResponse getBatchFeatures(GetBatchFeaturesRequest getFeat
112123
}
113124

114125
Table entityTable = loadEntities(getFeaturesRequest.getDatasetSource());
126+
String entityTableName = entityTable.getTableId().getTable();
127+
//TODO: add expiration to temp tables
128+
// entityTable = entityTable.toBuilder().setExpirationTime(TABLE_EXPIRATION_TIME).build();
129+
// entityTable.update(TableOption.fields(TableField.EXPIRATION_TIME));
130+
FieldValueList timestampLimits = getTimestampLimits(entityTableName);
131+
115132
Schema entityTableSchema = entityTable.getDefinition().getSchema();
116133
List<String> entityNames = entityTableSchema.getFields().stream()
117134
.map(Field::getName)
118135
.filter(name -> !name.equals("event_timestamp"))
119136
.collect(Collectors.toList());
120137

121-
final String query =
122-
BigQueryUtil.createQuery(
123-
getFeaturesRequest.getFeatureSetsList(),
124-
featureSetSpecs,
125-
entityNames,
126-
datasetId, entityTable.getFriendlyName());
127-
log.debug("Running BigQuery query: {}", query);
138+
String query;
139+
try {
140+
query =
141+
BigQueryUtil.createQuery(
142+
getFeaturesRequest.getFeatureSetsList(),
143+
featureSetSpecs,
144+
entityNames,
145+
projectId,
146+
datasetId,
147+
entityTableName,
148+
timestampLimits.get("min").getStringValue(),
149+
timestampLimits.get("max").getStringValue());
150+
log.info("Running BigQuery query: {}", query);
151+
} catch (IOException e) {
152+
throw new RuntimeException("Unable to generate query for batch retrieval");
153+
}
128154

129155
String feastJobId = UUID.randomUUID().toString();
130156
ServingAPIProto.Job feastJob =
@@ -217,6 +243,32 @@ public GetBatchFeaturesResponse getBatchFeatures(GetBatchFeaturesRequest getFeat
217243
return GetBatchFeaturesResponse.newBuilder().setJob(feastJob).build();
218244
}
219245

246+
private FieldValueList getTimestampLimits(String entityTableName) {
247+
QueryJobConfiguration getTimestampLimitsQuery = QueryJobConfiguration
248+
.newBuilder(getTimestampLimitQuery(projectId, datasetId, entityTableName))
249+
.setDefaultDataset(DatasetId.of(projectId, datasetId)).build();
250+
try {
251+
Job job = bigquery
252+
.create(JobInfo.of(getTimestampLimitsQuery));
253+
TableResult getTimestampLimitsQueryResult = job
254+
.waitFor()
255+
.getQueryResults();
256+
FieldValueList result = null;
257+
for (FieldValueList fields : getTimestampLimitsQueryResult.getValues()) {
258+
result = fields;
259+
}
260+
if (result == null || result.get("min").isNull() || result.get("max").isNull()) {
261+
throw new RuntimeException("query returned insufficient values");
262+
}
263+
return result;
264+
} catch (InterruptedException e) {
265+
throw Status.INTERNAL
266+
.withDescription("Unable to extract min and max timestamps from query")
267+
.withCause(e)
268+
.asRuntimeException();
269+
}
270+
}
271+
220272
/**
221273
* {@inheritDoc}
222274
*/
@@ -234,22 +286,30 @@ public GetJobResponse getJob(GetJobRequest getJobRequest) {
234286
private Table loadEntities(DatasetSource datasetSource) {
235287
switch (datasetSource.getDatasetSourceCase()) {
236288
case FILE_SOURCE:
237-
String tableName = generateTemporaryTableName();
238-
TableId tableId = TableId.of(projectId, datasetId, tableName);
239-
// Currently only avro supported
240-
if (datasetSource.getFileSource().getDataFormat() != DataFormat.DATA_FORMAT_AVRO) {
241-
throw Status.INVALID_ARGUMENT
242-
.withDescription("Invalid file format, only avro supported")
243-
.asRuntimeException();
244-
}
245-
LoadJobConfiguration loadJobConfiguration = LoadJobConfiguration.of(tableId,
246-
datasetSource.getFileSource().getFileUrisList(),
247-
FormatOptions.avro());
248-
Job job = bigquery.create(JobInfo.of(loadJobConfiguration));
249289
try {
290+
String tableName = generateTemporaryTableName();
291+
log.info("Loading entity dataset to table {}.{}.{}", projectId, datasetId, tableName);
292+
TableId tableId = TableId.of(projectId, datasetId, tableName);
293+
// Currently only avro supported
294+
if (datasetSource.getFileSource().getDataFormat() != DataFormat.DATA_FORMAT_AVRO) {
295+
throw Status.INVALID_ARGUMENT
296+
.withDescription("Invalid file format, only avro supported")
297+
.asRuntimeException();
298+
}
299+
LoadJobConfiguration loadJobConfiguration = LoadJobConfiguration.of(tableId,
300+
datasetSource.getFileSource().getFileUrisList(),
301+
FormatOptions.avro());
302+
loadJobConfiguration = loadJobConfiguration.toBuilder()
303+
.setUseAvroLogicalTypes(true)
304+
.build();
305+
Job job = bigquery.create(JobInfo.of(loadJobConfiguration));
250306
job.waitFor();
251-
return bigquery.getTable(tableId);
252-
} catch (InterruptedException e) {
307+
Table entityTable = bigquery.getTable(tableId);
308+
if (!entityTable.exists()) {
309+
throw new RuntimeException("Unable to create entity dataset table");
310+
}
311+
return entityTable;
312+
} catch (Exception e) {
253313
throw Status.INTERNAL
254314
.withDescription("Failed to load entity dataset into store")
255315
.withCause(e)
@@ -264,8 +324,9 @@ private Table loadEntities(DatasetSource datasetSource) {
264324
}
265325

266326
private String generateTemporaryTableName() {
267-
String source = String.format("feast_serving_%d", System.currentTimeMillis());
268-
UUID uuid = UUID.fromString(source);
269-
return uuid.toString().replaceAll("-", "_");
327+
String source = String.format("feastserving%d", System.currentTimeMillis());
328+
String guid = UUID.nameUUIDFromBytes(source.getBytes()).toString();
329+
String suffix = guid.substring(0, Math.min(guid.length(), 10)).replaceAll("-", "");
330+
return String.format("temp_%s", suffix);
270331
}
271332
}

‎serving/src/main/java/feast/serving/service/CachedSpecService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ private Store readConfig(Path path) {
136136
try {
137137
List<String> fileContents = Files.readAllLines(path);
138138
String yaml = fileContents.stream().reduce("", (l1, l2) -> l1 + "\n" + l2);
139+
log.info("loaded store config at {}: \n{}", path.toString(), yaml);
139140
return yamlToStoreProto(yaml);
140141
} catch (IOException e) {
141142
throw new RuntimeException(

‎serving/src/main/java/feast/serving/util/BigQueryUtil.java‎

Lines changed: 45 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,39 @@
11
package feast.serving.util;
22

33
import com.google.protobuf.Duration;
4-
import feast.core.FeatureSetProto.EntitySpec;
4+
import com.mitchellbosecke.pebble.PebbleEngine;
5+
import com.mitchellbosecke.pebble.template.PebbleTemplate;
56
import feast.core.FeatureSetProto.FeatureSetSpec;
67
import feast.serving.ServingAPIProto.GetBatchFeaturesRequest.FeatureSet;
8+
import java.io.IOException;
9+
import java.io.StringWriter;
10+
import java.io.Writer;
711
import java.util.ArrayList;
12+
import java.util.HashMap;
813
import java.util.List;
14+
import java.util.Map;
915

1016
public class BigQueryUtil {
1117

18+
private static final PebbleEngine engine = new PebbleEngine.Builder().build();
19+
private static final String FEATURESET_TEMPLATE_NAME = "templates/bq_featureset_query.sql";
20+
21+
public static String getTimestampLimitQuery(String projectId, String datasetId,
22+
String leftTableName) {
23+
return String.format(
24+
"SELECT DATETIME(MAX(event_timestamp)) as max, DATETIME(MIN(event_timestamp)) as min FROM `%s.%s.%s`",
25+
projectId, datasetId, leftTableName);
26+
}
27+
1228
public static String createQuery(
1329
List<FeatureSet> featureSets,
1430
List<FeatureSetSpec> featureSetSpecs,
1531
List<String> entityNames,
32+
String projectId,
1633
String bigqueryDataset,
17-
String leftTableName) {
34+
String leftTableName,
35+
String minTimestamp,
36+
String maxTimestamp) throws IOException {
1837
if (featureSets == null
1938
|| featureSetSpecs == null
2039
|| entityNames == null
@@ -33,10 +52,14 @@ public static String createQuery(
3352
featureSets.get(i),
3453
featureSetSpecs.get(i),
3554
entityNames,
55+
projectId,
3656
bigqueryDataset,
37-
leftTableName));
57+
leftTableName,
58+
minTimestamp,
59+
maxTimestamp));
3860
}
3961

62+
// TODO: templatize this as well
4063
if (featureSetQueries.size() > 1) {
4164
StringBuilder selectColumns = new StringBuilder("SELECT ");
4265
for (int i = 0; i < featureSets.size(); i++) {
@@ -70,44 +93,27 @@ private static String createQueryForFeatureSet(
7093
FeatureSet featureSet,
7194
FeatureSetSpec featureSetSpec,
7295
List<String> entityNames,
96+
String projectId,
7397
String datasetName,
74-
String leftTableName) {
98+
String leftTableName,
99+
String minTimestamp,
100+
String maxTimestamp) throws IOException {
101+
102+
PebbleTemplate template = engine.getTemplate(FEATURESET_TEMPLATE_NAME);
103+
Map<String, Object> context = new HashMap<>();
104+
context.put("entityNames", entityNames);
105+
context.put("featureSet", featureSet);
106+
context.put("featureSetSpec", featureSetSpec);
107+
context.put("maxAge", getMaxAge(featureSet, featureSetSpec).getSeconds());
108+
context.put("minTimestamp", minTimestamp);
109+
context.put("maxTimestamp", maxTimestamp);
110+
context.put("leftTableName", String.format("%s.%s.%s", projectId, datasetName, leftTableName));
111+
context.put("rightTableName",
112+
String.format("%s.%s.%s", projectId, datasetName, getTableName(featureSet)));
75113

76-
String joinedEntities = entityNames.stream().map(s -> String.format("joined.%s, ", s))
77-
.reduce(String::concat).get();
78-
String selectEntities = entityNames.stream().map(s -> String.format("l.%s, ", s))
79-
.reduce(String::concat).get();
80-
String joinedFeatures = featureSet.getFeatureNamesList().stream()
81-
.map(s -> String.format("joined.%s_%s, ", getTableName(featureSet), s))
82-
.reduce(String::concat).get();
83-
String selectFeatures = featureSet.getFeatureNamesList().stream()
84-
.map(s -> String.format("r.%s as %s_%s, ", s, getTableName(featureSet), s))
85-
.reduce(String::concat).get();
86-
String joinConditions = featureSetSpec.getEntitiesList().stream()
87-
.map(EntitySpec::getName)
88-
.map(s -> String.format("AND l.%s = r.%s ", s, s))
89-
.reduce(String::concat).get();
90-
91-
StringBuilder queryBuilder = new StringBuilder();
92-
queryBuilder.append(String.format(
93-
"SELECT * FROM (SELECT joined.event_timestamp, %s %s ROW_NUMBER() ",
94-
joinedEntities, joinedFeatures));
95-
queryBuilder.append(String.format(
96-
"OVER ( PARTITION BY %s joined.event_timestamp ORDER BY joined.r_event_timestamp DESC) rank ",
97-
joinedEntities));
98-
queryBuilder.append(String.format(
99-
"FROM (SELECT %s l.event_timestamp, %s r.event_timestamp AS r_event_timestamp ",
100-
selectEntities, selectFeatures));
101-
queryBuilder.append(String.format(
102-
"FROM %s.%s AS l LEFT OUTER JOIN %s.%s AS r ",
103-
datasetName, leftTableName, datasetName, getTableName(featureSet)));
104-
queryBuilder.append(String.format("ON l.event_timestamp >= r.event_timestamp \n"
105-
+ "AND Timestamp_sub(l.event_timestamp, interval %d second) < r.event_timestamp ",
106-
getMaxAge(featureSet, featureSetSpec).getSeconds()));
107-
queryBuilder
108-
.append(String.format("%s) AS joined) AS reduce WHERE reduce.rank = 1",
109-
joinConditions));
110-
return queryBuilder.toString();
114+
Writer writer = new StringWriter();
115+
template.evaluate(writer, context);
116+
return writer.toString();
111117
}
112118

113119
private static Duration getMaxAge(FeatureSet featureSet, FeatureSetSpec featureSetSpec) {

0 commit comments

Comments
 (0)