RedisClusterOnlineRetriever implements OnlineRetriever { @Override public List<Optional<FeatureRow>> getOnlineFeatures( List<EntityRow> entityRows, FeatureSetRequest featureSetRequest) { FeatureSetSpec featureSetSpec = featureSetRequest.getSpec(); List<RedisKey> redisKeys = buildRedisKeys(entityRows, featureSetSpec); FeatureRowDecoder decoder = new FeatureRowDecoder(generateFeatureSetStringRef(featureSetSpec), featureSetSpec); List<Optional<FeatureRow>> featureRows = new ArrayList<>(); try { featureRows = getFeaturesFromRedis(redisKeys, decoder); } catch (InvalidProtocolBufferException | ExecutionException e) { throw Status.INTERNAL .withDescription("Unable to parse protobuf while retrieving feature") .withCause(e) .asRuntimeException(); } return featureRows; } private  RedisClusterOnlineRetriever(Builder builder); static OnlineRetriever create(Map<String, String> config); @Override List<Optional<FeatureRow>> getOnlineFeatures(
      List<EntityRow> entityRows, FeatureSetRequest featureSetRequest);  }
@Test public void shouldReturnResponseWithValuesIfKeysPresent() { byte[] serializedKey1 = serializer.serialize(redisKeys.get(0)); byte[] serializedKey2 = serializer.serialize(redisKeys.get(1)); KeyValue<byte[], byte[]> keyValue1 = KeyValue.from(serializedKey1, Optional.of(featureRows.get(0).toByteArray())); KeyValue<byte[], byte[]> keyValue2 = KeyValue.from(serializedKey2, Optional.of(featureRows.get(1).toByteArray())); List<KeyValue<byte[], byte[]>> featureRowBytes = Lists.newArrayList(keyValue1, keyValue2); OnlineRetriever redisClusterOnlineRetriever = new RedisClusterOnlineRetriever.Builder(connection, serializer) .withFallbackSerializer(fallbackSerializer) .build(); when(syncCommands.mget(serializedKey1, serializedKey2)).thenReturn(featureRowBytes); List<Optional<FeatureRow>> expected = Lists.newArrayList( Optional.of( FeatureRow.newBuilder() .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) .setFeatureSet("project/featureSet") .addAllFields( Lists.newArrayList( Field.newBuilder().setName("feature1").setValue(intValue(1)).build(), Field.newBuilder().setName("feature2").setValue(intValue(1)).build())) .build()), Optional.of( FeatureRow.newBuilder() .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) .setFeatureSet("project/featureSet") .addAllFields( Lists.newArrayList( Field.newBuilder().setName("feature1").setValue(intValue(2)).build(), Field.newBuilder().setName("feature2").setValue(intValue(2)).build())) .build())); List<Optional<FeatureRow>> actual = redisClusterOnlineRetriever.getOnlineFeatures(entityRows, featureSetRequest); assertThat(actual, equalTo(expected)); verify(syncCommands, never()).mget(); }
@Test public void shouldReturnNullIfKeysNotPresent() { byte[] serializedKey1 = serializer.serialize(redisKeys.get(0)); byte[] serializedKey2 = serializer.serialize(redisKeys.get(1)); KeyValue<byte[], byte[]> keyValue1 = KeyValue.from(serializedKey1, Optional.of(featureRows.get(0).toByteArray())); KeyValue<byte[], byte[]> keyValue2 = KeyValue.empty(serializedKey2); List<KeyValue<byte[], byte[]>> featureRowBytes = Lists.newArrayList(keyValue1, keyValue2); OnlineRetriever redisClusterOnlineRetriever = new RedisClusterOnlineRetriever.Builder(connection, serializer).build(); when(syncCommands.mget(serializedKey1, serializedKey2)).thenReturn(featureRowBytes); List<Optional<FeatureRow>> expected = Lists.newArrayList( Optional.of( FeatureRow.newBuilder() .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) .setFeatureSet("project/featureSet") .addAllFields( Lists.newArrayList( Field.newBuilder().setName("feature1").setValue(intValue(1)).build(), Field.newBuilder().setName("feature2").setValue(intValue(1)).build())) .build()), Optional.empty()); List<Optional<FeatureRow>> actual = redisClusterOnlineRetriever.getOnlineFeatures(entityRows, featureSetRequest); assertThat(actual, equalTo(expected)); }
@Test public void shouldUseFallbackIfAvailable() { byte[] serializedKey1 = serializer.serialize(redisKeys.get(0)); byte[] serializedKey2 = serializer.serialize(redisKeys.get(1)); byte[] fallbackSerializedKey2 = fallbackSerializer.serialize(redisKeys.get(1)); KeyValue<byte[], byte[]> keyValue1 = KeyValue.from(serializedKey1, Optional.of(featureRows.get(0).toByteArray())); KeyValue<byte[], byte[]> keyValue2 = KeyValue.empty(serializedKey2); KeyValue<byte[], byte[]> fallbackKeyValue2 = KeyValue.from(serializedKey2, Optional.of(featureRows.get(1).toByteArray())); List<KeyValue<byte[], byte[]>> featureRowBytes = Lists.newArrayList(keyValue1, keyValue2); List<KeyValue<byte[], byte[]>> fallbackFeatureRowBytes = Lists.newArrayList(fallbackKeyValue2); OnlineRetriever redisClusterOnlineRetriever = new RedisClusterOnlineRetriever.Builder(connection, serializer) .withFallbackSerializer(fallbackSerializer) .build(); when(syncCommands.mget(serializedKey1, serializedKey2)).thenReturn(featureRowBytes); when(syncCommands.mget(fallbackSerializedKey2)).thenReturn(fallbackFeatureRowBytes); List<Optional<FeatureRow>> expected = Lists.newArrayList( Optional.of( FeatureRow.newBuilder() .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) .setFeatureSet("project/featureSet") .addAllFields( Lists.newArrayList( Field.newBuilder().setName("feature1").setValue(intValue(1)).build(), Field.newBuilder().setName("feature2").setValue(intValue(1)).build())) .build()), Optional.of( FeatureRow.newBuilder() .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) .setFeatureSet("project/featureSet") .addAllFields( Lists.newArrayList( Field.newBuilder().setName("feature1").setValue(intValue(2)).build(), Field.newBuilder().setName("feature2").setValue(intValue(2)).build())) .build())); List<Optional<FeatureRow>> actual = redisClusterOnlineRetriever.getOnlineFeatures(entityRows, featureSetRequest); assertThat(actual, equalTo(expected)); }