Skip to content

Commit 71e74ca

Browse files
terryyylimterryyylim
committed
Add compound entity IT
Signed-off-by: Terence Lim <terencelimxp@gmail.com>
1 parent 2ce92a3 commit 71e74ca

2 files changed

Lines changed: 203 additions & 16 deletions

File tree

‎common-test/src/main/java/feast/common/it/DataGenerator.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,28 @@ public static Triple<String, String, Boolean> getDefaultSubscription() {
5353
return defaultSubscription;
5454
}
5555

56+
public static String valueToString(ValueProto.Value v) {
57+
String stringRepr;
58+
switch (v.getValCase()) {
59+
case STRING_VAL:
60+
stringRepr = v.getStringVal();
61+
break;
62+
case INT64_VAL:
63+
stringRepr = String.valueOf(v.getInt64Val());
64+
break;
65+
case INT32_VAL:
66+
stringRepr = String.valueOf(v.getInt32Val());
67+
break;
68+
case BYTES_VAL:
69+
stringRepr = v.getBytesVal().toString();
70+
break;
71+
default:
72+
throw new RuntimeException("Type is not supported to be entity");
73+
}
74+
75+
return stringRepr;
76+
}
77+
5678
public static StoreProto.Store getDefaultStore() {
5779
return defaultStore;
5880
}
@@ -247,6 +269,18 @@ public static ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow createEntityR
247269
.build();
248270
}
249271

272+
public static ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow createCompoundEntityRow(
273+
ImmutableMap<String, ValueProto.Value> entityNameValues, long seconds) {
274+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow.Builder entityRow =
275+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow.newBuilder()
276+
.setTimestamp(Timestamp.newBuilder().setSeconds(seconds));
277+
278+
entityNameValues.entrySet().stream()
279+
.forEach(entry -> entityRow.putFields(entry.getKey(), entry.getValue()));
280+
281+
return entityRow.build();
282+
}
283+
250284
public static DataSource createKinesisDataSourceSpec(
251285
String region, String streamName, String classPath, String timestampColumn) {
252286
return DataSource.newBuilder()

‎serving/src/test/java/feast/serving/it/ServingServiceBigTableIT.java‎

Lines changed: 169 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.google.api.gax.rpc.TransportChannelProvider;
2626
import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient;
2727
import com.google.cloud.bigtable.admin.v2.models.CreateTableRequest;
28-
import com.google.cloud.bigtable.admin.v2.models.Table;
2928
import com.google.cloud.bigtable.admin.v2.stub.BigtableTableAdminStubSettings;
3029
import com.google.cloud.bigtable.admin.v2.stub.EnhancedBigtableTableAdminStub;
3130
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
@@ -51,6 +50,7 @@
5150
import java.util.HashMap;
5251
import java.util.List;
5352
import java.util.Map;
53+
import java.util.stream.Collectors;
5454
import org.apache.avro.Schema;
5555
import org.apache.avro.SchemaBuilder;
5656
import org.apache.avro.generic.GenericDatumWriter;
@@ -88,7 +88,7 @@ public class ServingServiceBigTableIT extends BaseAuthIT {
8888
static CoreSimpleAPIClient coreClient;
8989
static ServingServiceGrpc.ServingServiceBlockingStub servingStub;
9090

91-
static final int FEAST_SERVING_PORT = 6568;
91+
static final int FEAST_SERVING_PORT = 6569;
9292

9393
static final String PROJECT_ID = "test-project";
9494
static final String INSTANCE_ID = "test-instance";
@@ -125,6 +125,15 @@ static void globalSetup() throws IOException {
125125
.setInstanceId(INSTANCE_ID)
126126
.build());
127127

128+
String endpoint =
129+
environment.getServiceHost("bigtable_1", BIGTABLE_PORT)
130+
+ ":"
131+
+ environment.getServicePort("bigtable_1", BIGTABLE_PORT);
132+
channel = ManagedChannelBuilder.forTarget(endpoint).usePlaintext().build();
133+
TransportChannelProvider channelProvider =
134+
FixedTransportChannelProvider.create(GrpcTransportChannel.create(channel));
135+
NoCredentialsProvider credentialsProvider = NoCredentialsProvider.create();
136+
128137
String projectName = "default";
129138
// Apply Entity (driver_id)
130139
String driverEntityName = "driver_id";
@@ -187,6 +196,8 @@ static void globalSetup() throws IOException {
187196
ImmutableList<String> columnFamilies = ImmutableList.of(featureTableName, metadataColumnFamily);
188197
String emptyQualifier = "";
189198

199+
createTable(channelProvider, credentialsProvider, btTableName, columnFamilies);
200+
190201
Schema ftSchema =
191202
SchemaBuilder.record("DriverData")
192203
.namespace(featureTableName)
@@ -209,6 +220,7 @@ static void globalSetup() throws IOException {
209220
.build();
210221
byte[] avroSerializedFeatures = recordToAvro(record, ftSchema);
211222

223+
// Single Entity Key
212224
byte[] entityFeatureKey =
213225
String.valueOf(DataGenerator.createInt64Value(1).getInt64Val()).getBytes();
214226

@@ -218,24 +230,12 @@ static void globalSetup() throws IOException {
218230
entityFeatureOutputStream.write(avroSerializedFeatures);
219231
byte[] entityFeatureValue = entityFeatureOutputStream.toByteArray();
220232

221-
// SchemaKey
233+
// Single Entity SchemaKey
222234
ByteArrayOutputStream concatOutputStream = new ByteArrayOutputStream();
223235
concatOutputStream.write("schema#".getBytes());
224236
concatOutputStream.write(schemaReference);
225237
byte[] schemaKey = concatOutputStream.toByteArray();
226238

227-
String endpoint =
228-
environment.getServiceHost("bigtable_1", BIGTABLE_PORT)
229-
+ ":"
230-
+ environment.getServicePort("bigtable_1", BIGTABLE_PORT);
231-
// ManagedChannel channel = ManagedChannelBuilder.forTarget(endpoint).usePlaintext().build();
232-
channel = ManagedChannelBuilder.forTarget(endpoint).usePlaintext().build();
233-
TransportChannelProvider channelProvider =
234-
FixedTransportChannelProvider.create(GrpcTransportChannel.create(channel));
235-
NoCredentialsProvider credentialsProvider = NoCredentialsProvider.create();
236-
237-
createTable(channelProvider, credentialsProvider, btTableName, columnFamilies);
238-
239239
// Update Entity-Feature Row
240240
client.mutateRow(
241241
RowMutation.create(btTableName, ByteString.copyFrom(entityFeatureKey))
@@ -252,6 +252,93 @@ static void globalSetup() throws IOException {
252252
ByteString.copyFrom("avro".getBytes()),
253253
ByteString.copyFrom(ftSchema.toString().getBytes())));
254254

255+
// Compound Entity Key
256+
// Apply FeatureTable (rides_merchant)
257+
String rideMerchantFeatureTableName = "rides_merchant";
258+
ImmutableList<String> ridesMerchantEntities =
259+
ImmutableList.of(driverEntityName, merchantEntityName);
260+
ImmutableList<String> compoundColumnFamilies =
261+
ImmutableList.of(rideMerchantFeatureTableName, metadataColumnFamily);
262+
263+
TestUtils.applyFeatureTable(
264+
coreClient,
265+
projectName,
266+
rideMerchantFeatureTableName,
267+
ridesMerchantEntities,
268+
ridesFeatures,
269+
7200);
270+
271+
String compoundBtTableName =
272+
String.format(
273+
"%s__%s",
274+
projectName, ridesMerchantEntities.stream().collect(Collectors.joining("__")));
275+
ValueProto.Value driverEntityValue = ValueProto.Value.newBuilder().setInt64Val(1).build();
276+
ValueProto.Value merchantEntityValue = ValueProto.Value.newBuilder().setInt64Val(1234).build();
277+
ImmutableMap<String, ValueProto.Value> compoundEntityMap =
278+
ImmutableMap.of(
279+
driverEntityName, driverEntityValue, merchantEntityName, merchantEntityValue);
280+
281+
createTable(channelProvider, credentialsProvider, compoundBtTableName, compoundColumnFamilies);
282+
283+
// Instantiate EntityRows
284+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow entityRow =
285+
DataGenerator.createCompoundEntityRow(compoundEntityMap, 100);
286+
byte[] compoundEntityFeatureKey =
287+
ridesMerchantEntities.stream()
288+
.map(entity -> DataGenerator.valueToString(entityRow.getFieldsMap().get(entity)))
289+
.collect(Collectors.joining("#"))
290+
.getBytes();
291+
292+
Schema compoundFtSchema =
293+
SchemaBuilder.record("DriverMerchantData")
294+
.namespace(rideMerchantFeatureTableName)
295+
.fields()
296+
.requiredInt(feature1Reference.getName())
297+
.requiredDouble(feature2Reference.getName())
298+
.nullableString(feature3Reference.getName(), "null")
299+
.requiredString(feature4Reference.getName())
300+
.endRecord();
301+
byte[] compoundSchemaReference =
302+
Hashing.murmur3_32().hashBytes(compoundFtSchema.toString().getBytes()).asBytes();
303+
304+
// Entity-Feature Row
305+
GenericRecord compoundEntityRecord =
306+
new GenericRecordBuilder(ftSchema)
307+
.set("trip_cost", 10)
308+
.set("trip_distance", 5.5)
309+
.set("trip_empty", null)
310+
.set("trip_wrong_type", "wrong_type")
311+
.build();
312+
byte[] compoundAvroSerializedFeatures = recordToAvro(compoundEntityRecord, compoundFtSchema);
313+
314+
// Compound Entity SchemaKey
315+
ByteArrayOutputStream compoundEntityFeatureOutputStream = new ByteArrayOutputStream();
316+
compoundEntityFeatureOutputStream.write(schemaReference);
317+
compoundEntityFeatureOutputStream.write("".getBytes());
318+
compoundEntityFeatureOutputStream.write(compoundAvroSerializedFeatures);
319+
byte[] compoundEntityFeatureValue = compoundEntityFeatureOutputStream.toByteArray();
320+
321+
ByteArrayOutputStream compoundConcatOutputStream = new ByteArrayOutputStream();
322+
compoundConcatOutputStream.write("schema#".getBytes());
323+
compoundConcatOutputStream.write(compoundSchemaReference);
324+
byte[] compoundSchemaKey = concatOutputStream.toByteArray();
325+
326+
// Update Compound Entity-Feature Row
327+
client.mutateRow(
328+
RowMutation.create(compoundBtTableName, ByteString.copyFrom(compoundEntityFeatureKey))
329+
.setCell(
330+
rideMerchantFeatureTableName,
331+
ByteString.copyFrom(emptyQualifier.getBytes()),
332+
ByteString.copyFrom(compoundEntityFeatureValue)));
333+
334+
// Update Schema Row
335+
client.mutateRow(
336+
RowMutation.create(compoundBtTableName, ByteString.copyFrom(compoundSchemaKey))
337+
.setCell(
338+
metadataColumnFamily,
339+
ByteString.copyFrom("avro".getBytes()),
340+
ByteString.copyFrom(compoundFtSchema.toString().getBytes())));
341+
255342
// set up options for call credentials
256343
options.put("oauth_url", TOKEN_URL);
257344
options.put(CLIENT_ID, CLIENT_ID);
@@ -286,7 +373,7 @@ private static void createTable(
286373
for (String columnFamily : columnFamilies) {
287374
createTableRequest.addFamily(columnFamily);
288375
}
289-
Table table = client.createTable(createTableRequest);
376+
client.createTable(createTableRequest);
290377
}
291378
}
292379

@@ -357,6 +444,72 @@ public void shouldRegisterAndGetOnlineFeaturesWithNotFound() {
357444
assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList());
358445
}
359446

447+
@Test
448+
public void shouldRegisterCompoundEntityAndGetOnlineFeatures() {
449+
String projectName = "default";
450+
String driverEntityName = "driver_id";
451+
String merchantEntityName = "merchant_id";
452+
ValueProto.Value driverEntityValue = ValueProto.Value.newBuilder().setInt64Val(1).build();
453+
ValueProto.Value merchantEntityValue = ValueProto.Value.newBuilder().setInt64Val(1234).build();
454+
455+
ImmutableMap<String, ValueProto.Value> compoundEntityMap =
456+
ImmutableMap.of(
457+
driverEntityName, driverEntityValue, merchantEntityName, merchantEntityValue);
458+
459+
// Instantiate EntityRows
460+
ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow entityRow =
461+
DataGenerator.createCompoundEntityRow(compoundEntityMap, 100);
462+
ImmutableList<ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow> entityRows =
463+
ImmutableList.of(entityRow);
464+
465+
// Instantiate FeatureReferences
466+
ServingAPIProto.FeatureReferenceV2 featureReference =
467+
DataGenerator.createFeatureReference("rides", "trip_cost");
468+
ServingAPIProto.FeatureReferenceV2 notFoundFeatureReference =
469+
DataGenerator.createFeatureReference("rides", "trip_transaction");
470+
471+
ImmutableList<ServingAPIProto.FeatureReferenceV2> featureReferences =
472+
ImmutableList.of(featureReference, notFoundFeatureReference);
473+
474+
// Build GetOnlineFeaturesRequestV2
475+
ServingAPIProto.GetOnlineFeaturesRequestV2 onlineFeatureRequest =
476+
TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows);
477+
ServingAPIProto.GetOnlineFeaturesResponse featureResponse =
478+
servingStub.getOnlineFeaturesV2(onlineFeatureRequest);
479+
480+
ImmutableMap<String, ValueProto.Value> expectedValueMap =
481+
ImmutableMap.of(
482+
driverEntityName,
483+
driverEntityValue,
484+
merchantEntityName,
485+
merchantEntityValue,
486+
FeatureV2.getFeatureStringRef(featureReference),
487+
DataGenerator.createInt64Value(5),
488+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
489+
DataGenerator.createEmptyValue());
490+
491+
ImmutableMap<String, ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus> expectedStatusMap =
492+
ImmutableMap.of(
493+
driverEntityName,
494+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
495+
merchantEntityName,
496+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
497+
FeatureV2.getFeatureStringRef(featureReference),
498+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.PRESENT,
499+
FeatureV2.getFeatureStringRef(notFoundFeatureReference),
500+
ServingAPIProto.GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND);
501+
502+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues expectedFieldValues =
503+
ServingAPIProto.GetOnlineFeaturesResponse.FieldValues.newBuilder()
504+
.putAllFields(expectedValueMap)
505+
.putAllStatuses(expectedStatusMap)
506+
.build();
507+
ImmutableList<ServingAPIProto.GetOnlineFeaturesResponse.FieldValues> expectedFieldValuesList =
508+
ImmutableList.of(expectedFieldValues);
509+
510+
assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList());
511+
}
512+
360513
@TestConfiguration
361514
public static class TestConfig {
362515
@Bean

0 commit comments

Comments
 (0)