Skip to content

Commit aacce71

Browse files
terryyylimterryyylim
committed
Refactor duplicated online retriever code
Signed-off-by: Terence <terencelimxp@gmail.com>
1 parent 68973c9 commit aacce71

6 files changed

Lines changed: 209 additions & 249 deletions

File tree

‎serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,7 @@
2323
import feast.serving.service.ServingServiceV2;
2424
import feast.serving.specs.CachedSpecService;
2525
import feast.storage.api.retriever.OnlineRetrieverV2;
26-
import feast.storage.connectors.redis.retriever.RedisClusterOnlineRetrieverV2;
27-
import feast.storage.connectors.redis.retriever.RedisOnlineRetrieverV2;
26+
import feast.storage.connectors.redis.retriever.*;
2827
import io.opentracing.Tracer;
2928
import java.util.Map;
3029
import org.slf4j.Logger;
@@ -46,11 +45,13 @@ public ServingServiceV2 servingServiceV2(
4645

4746
switch (storeType) {
4847
case REDIS_CLUSTER:
49-
OnlineRetrieverV2 redisClusterRetriever = RedisClusterOnlineRetrieverV2.create(config);
48+
RedisClientWrapper redisClusterClient = RedisClusterClient.create(config);
49+
OnlineRetrieverV2 redisClusterRetriever = new OnlineRetriever(redisClusterClient);
5050
servingService = new OnlineServingServiceV2(redisClusterRetriever, specService, tracer);
5151
break;
5252
case REDIS:
53-
OnlineRetrieverV2 redisRetriever = RedisOnlineRetrieverV2.create(config);
53+
RedisClientWrapper redisClient = RedisClient.create(config);
54+
OnlineRetrieverV2 redisRetriever = new OnlineRetriever(redisClient);
5455
servingService = new OnlineServingServiceV2(redisRetriever, specService, tracer);
5556
break;
5657
case CASSANDRA:

storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/RedisOnlineRetrieverV2.java renamed to storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/OnlineRetriever.java

Lines changed: 15 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -18,66 +18,46 @@
1818

1919
import com.google.common.collect.Lists;
2020
import com.google.protobuf.InvalidProtocolBufferException;
21-
import feast.proto.serving.ServingAPIProto.FeatureReferenceV2;
22-
import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow;
23-
import feast.proto.storage.RedisProto.RedisKeyV2;
21+
import feast.proto.serving.ServingAPIProto;
22+
import feast.proto.storage.RedisProto;
2423
import feast.storage.api.retriever.Feature;
2524
import feast.storage.api.retriever.OnlineRetrieverV2;
2625
import feast.storage.connectors.redis.common.RedisHashDecoder;
2726
import feast.storage.connectors.redis.common.RedisKeyGenerator;
2827
import io.grpc.Status;
2928
import io.lettuce.core.KeyValue;
30-
import io.lettuce.core.RedisClient;
3129
import io.lettuce.core.RedisFuture;
32-
import io.lettuce.core.RedisURI;
33-
import io.lettuce.core.api.StatefulRedisConnection;
34-
import io.lettuce.core.api.async.RedisAsyncCommands;
35-
import io.lettuce.core.codec.ByteArrayCodec;
3630
import java.util.*;
3731
import java.util.concurrent.ExecutionException;
3832
import java.util.stream.Collectors;
3933

40-
public class RedisOnlineRetrieverV2 implements OnlineRetrieverV2 {
34+
public class OnlineRetriever implements OnlineRetrieverV2 {
4135

4236
private static final String timestampPrefix = "_ts";
43-
private final RedisAsyncCommands<byte[], byte[]> asyncCommands;
37+
RedisClientWrapper redisClientWrapper;
4438

45-
private RedisOnlineRetrieverV2(StatefulRedisConnection<byte[], byte[]> connection) {
46-
this.asyncCommands = connection.async();
47-
48-
// Disable auto-flushing
49-
this.asyncCommands.setAutoFlushCommands(false);
50-
}
51-
52-
public static OnlineRetrieverV2 create(Map<String, String> config) {
53-
54-
StatefulRedisConnection<byte[], byte[]> connection =
55-
RedisClient.create(
56-
RedisURI.create(config.get("host"), Integer.parseInt(config.get("port"))))
57-
.connect(new ByteArrayCodec());
58-
59-
return new RedisOnlineRetrieverV2(connection);
60-
}
61-
62-
public static OnlineRetrieverV2 create(StatefulRedisConnection<byte[], byte[]> connection) {
63-
return new RedisOnlineRetrieverV2(connection);
39+
public OnlineRetriever(RedisClientWrapper redisClientWrapper) {
40+
this.redisClientWrapper = redisClientWrapper;
6441
}
6542

6643
@Override
6744
public List<List<Optional<Feature>>> getOnlineFeatures(
68-
String project, List<EntityRow> entityRows, List<FeatureReferenceV2> featureReferences) {
45+
String project,
46+
List<ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow> entityRows,
47+
List<ServingAPIProto.FeatureReferenceV2> featureReferences) {
6948

70-
List<RedisKeyV2> redisKeys = RedisKeyGenerator.buildRedisKeys(project, entityRows);
49+
List<RedisProto.RedisKeyV2> redisKeys = RedisKeyGenerator.buildRedisKeys(project, entityRows);
7150
List<List<Optional<Feature>>> features = getFeaturesFromRedis(redisKeys, featureReferences);
7251

7352
return features;
7453
}
7554

7655
private List<List<Optional<Feature>>> getFeaturesFromRedis(
77-
List<RedisKeyV2> redisKeys, List<FeatureReferenceV2> featureReferences) {
56+
List<RedisProto.RedisKeyV2> redisKeys,
57+
List<ServingAPIProto.FeatureReferenceV2> featureReferences) {
7858
List<List<Optional<Feature>>> features = new ArrayList<>();
7959
// To decode bytes back to Feature Reference
80-
Map<String, FeatureReferenceV2> byteToFeatureReferenceMap = new HashMap<>();
60+
Map<String, ServingAPIProto.FeatureReferenceV2> byteToFeatureReferenceMap = new HashMap<>();
8161

8262
// Serialize using proto
8363
List<byte[]> binaryRedisKeys =
@@ -106,11 +86,11 @@ private List<List<Optional<Feature>>> getFeaturesFromRedis(
10686
byte[][] featureReferenceWithTsByteArrays =
10787
featureReferenceWithTsByteList.toArray(new byte[0][]);
10888
// Access redis keys and extract features
109-
futures.add(asyncCommands.hmget(binaryRedisKey, featureReferenceWithTsByteArrays));
89+
futures.add(redisClientWrapper.hmget(binaryRedisKey, featureReferenceWithTsByteArrays));
11090
}
11191

11292
// Write all commands to the transport layer
113-
asyncCommands.flushCommands();
93+
redisClientWrapper.flushCommands();
11494

11595
futures.forEach(
11696
future -> {

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/RedisClient.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
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.connectors.redis.retriever;
18+
19+
import io.lettuce.core.KeyValue;
20+
import io.lettuce.core.RedisFuture;
21+
import io.lettuce.core.RedisURI;
22+
import io.lettuce.core.api.StatefulRedisConnection;
23+
import io.lettuce.core.api.async.RedisAsyncCommands;
24+
import io.lettuce.core.codec.ByteArrayCodec;
25+
import java.util.List;
26+
import java.util.Map;
27+
28+
public class RedisClient implements RedisClientWrapper {
29+
30+
public final RedisAsyncCommands<byte[], byte[]> asyncCommands;
31+
32+
@Override
33+
public RedisFuture<List<KeyValue<byte[], byte[]>>> hmget(byte[] key, byte[]... fields) {
34+
return asyncCommands.hmget(key, fields);
35+
}
36+
37+
@Override
38+
public void flushCommands() {
39+
asyncCommands.flushCommands();
40+
}
41+
42+
private RedisClient(StatefulRedisConnection<byte[], byte[]> connection) {
43+
this.asyncCommands = connection.async();
44+
45+
// Disable auto-flushing
46+
this.asyncCommands.setAutoFlushCommands(false);
47+
}
48+
49+
public static RedisClientWrapper create(Map<String, String> config) {
50+
StatefulRedisConnection<byte[], byte[]> connection =
51+
io.lettuce.core.RedisClient.create(
52+
RedisURI.create(config.get("host"), Integer.parseInt(config.get("port"))))
53+
.connect(new ByteArrayCodec());
54+
55+
return new RedisClient(connection);
56+
}
57+
}

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/RedisClientWrapper.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
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.connectors.redis.retriever;
18+
19+
import io.lettuce.core.*;
20+
import java.util.List;
21+
22+
public interface RedisClientWrapper {
23+
RedisFuture<List<KeyValue<byte[], byte[]>>> hmget(byte[] key, byte[]... fields);
24+
25+
void flushCommands();
26+
}

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/retriever/RedisClusterClient.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
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.connectors.redis.retriever;
18+
19+
import feast.storage.connectors.redis.serializer.RedisKeyPrefixSerializerV2;
20+
import feast.storage.connectors.redis.serializer.RedisKeySerializerV2;
21+
import io.lettuce.core.KeyValue;
22+
import io.lettuce.core.RedisFuture;
23+
import io.lettuce.core.RedisURI;
24+
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
25+
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
26+
import io.lettuce.core.codec.ByteArrayCodec;
27+
import java.util.Arrays;
28+
import java.util.List;
29+
import java.util.Map;
30+
import java.util.stream.Collectors;
31+
import javax.annotation.Nullable;
32+
33+
public class RedisClusterClient implements RedisClientWrapper {
34+
35+
public final RedisAdvancedClusterAsyncCommands<byte[], byte[]> asyncCommands;
36+
public final RedisKeySerializerV2 serializer;
37+
@Nullable public final RedisKeySerializerV2 fallbackSerializer;
38+
39+
@Override
40+
public RedisFuture<List<KeyValue<byte[], byte[]>>> hmget(byte[] key, byte[]... fields) {
41+
return asyncCommands.hmget(key, fields);
42+
}
43+
44+
@Override
45+
public void flushCommands() {
46+
asyncCommands.flushCommands();
47+
}
48+
49+
static class Builder {
50+
private final StatefulRedisClusterConnection<byte[], byte[]> connection;
51+
private final RedisKeySerializerV2 serializer;
52+
@Nullable private RedisKeySerializerV2 fallbackSerializer;
53+
54+
Builder(
55+
StatefulRedisClusterConnection<byte[], byte[]> connection,
56+
RedisKeySerializerV2 serializer) {
57+
this.connection = connection;
58+
this.serializer = serializer;
59+
}
60+
61+
Builder withFallbackSerializer(RedisKeySerializerV2 fallbackSerializer) {
62+
this.fallbackSerializer = fallbackSerializer;
63+
return this;
64+
}
65+
66+
RedisClusterClient build() {
67+
return new RedisClusterClient(this);
68+
}
69+
}
70+
71+
private RedisClusterClient(Builder builder) {
72+
this.asyncCommands = builder.connection.async();
73+
this.serializer = builder.serializer;
74+
this.fallbackSerializer = builder.fallbackSerializer;
75+
76+
// Disable auto-flushing
77+
this.asyncCommands.setAutoFlushCommands(false);
78+
}
79+
80+
public static RedisClientWrapper create(Map<String, String> config) {
81+
List<RedisURI> redisURIList =
82+
Arrays.stream(config.get("connection_string").split(","))
83+
.map(
84+
hostPort -> {
85+
String[] hostPortSplit = hostPort.trim().split(":");
86+
return RedisURI.create(hostPortSplit[0], Integer.parseInt(hostPortSplit[1]));
87+
})
88+
.collect(Collectors.toList());
89+
StatefulRedisClusterConnection<byte[], byte[]> connection =
90+
io.lettuce.core.cluster.RedisClusterClient.create(redisURIList)
91+
.connect(new ByteArrayCodec());
92+
93+
RedisKeySerializerV2 serializer =
94+
new RedisKeyPrefixSerializerV2(config.getOrDefault("key_prefix", ""));
95+
96+
Builder builder = new Builder(connection, serializer);
97+
98+
if (Boolean.parseBoolean(config.getOrDefault("enable_fallback", "false"))) {
99+
RedisKeySerializerV2 fallbackSerializer =
100+
new RedisKeyPrefixSerializerV2(config.getOrDefault("fallback_prefix", ""));
101+
builder = builder.withFallbackSerializer(fallbackSerializer);
102+
}
103+
104+
return builder.build();
105+
}
106+
}

0 commit comments

Comments
 (0)