Skip to content

Commit 05a2306

Browse files
khorshuhengkhorshuheng
committed
Add redis key prefix as an option to Redis cluster
1 parent eb150a2 commit 05a2306

9 files changed

Lines changed: 388 additions & 127 deletions

File tree

‎protos/feast/core/Store.proto‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,14 @@ message Store {
134134
int32 max_retries = 3;
135135
// Optional. How often flush data to redis
136136
int32 flush_frequency_seconds = 4;
137+
// Optional. Append a prefix to the Redis Key
138+
string key_prefix = 5;
139+
// Optional. Enable fallback to another key prefix if the original key is not present.
140+
// Useful for migrating key prefix without re-ingestion. Disabled by default.
141+
bool enable_fallback = 6;
142+
// Optional. This would be the fallback prefix to use if enable_fallback is true.
143+
string fallback_prefix = 7;
144+
137145
}
138146

139147
message Subscription {

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

Lines changed: 98 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
*/
1717
package feast.storage.connectors.redis.retriever;
1818

19-
import com.google.protobuf.AbstractMessageLite;
2019
import com.google.protobuf.InvalidProtocolBufferException;
2120
import feast.proto.core.FeatureSetProto.EntitySpec;
2221
import feast.proto.core.FeatureSetProto.FeatureSetSpec;
@@ -27,27 +26,52 @@
2726
import feast.proto.types.ValueProto.Value;
2827
import feast.storage.api.retriever.FeatureSetRequest;
2928
import feast.storage.api.retriever.OnlineRetriever;
29+
import feast.storage.connectors.redis.serializer.RedisKeyPrefixSerializer;
30+
import feast.storage.connectors.redis.serializer.RedisKeySerializer;
3031
import io.grpc.Status;
3132
import io.lettuce.core.RedisURI;
3233
import io.lettuce.core.cluster.RedisClusterClient;
3334
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
3435
import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
3536
import io.lettuce.core.codec.ByteArrayCodec;
36-
import java.util.ArrayList;
37-
import java.util.Arrays;
38-
import java.util.List;
39-
import java.util.Map;
40-
import java.util.Optional;
37+
import java.util.*;
4138
import java.util.concurrent.ExecutionException;
4239
import java.util.stream.Collectors;
40+
import java.util.stream.IntStream;
41+
import javax.annotation.Nullable;
4342

4443
/** Defines a storage retriever */
4544
public class RedisClusterOnlineRetriever implements OnlineRetriever {
4645

4746
private final RedisAdvancedClusterCommands<byte[], byte[]> syncCommands;
47+
private final RedisKeySerializer serializer;
48+
@Nullable private final RedisKeySerializer fallbackSerializer;
4849

49-
private RedisClusterOnlineRetriever(StatefulRedisClusterConnection<byte[], byte[]> connection) {
50-
this.syncCommands = connection.sync();
50+
static class Builder {
51+
private final StatefulRedisClusterConnection<byte[], byte[]> connection;
52+
private final RedisKeySerializer serializer;
53+
@Nullable private RedisKeySerializer fallbackSerializer;
54+
55+
Builder(
56+
StatefulRedisClusterConnection<byte[], byte[]> connection, RedisKeySerializer serializer) {
57+
this.connection = connection;
58+
this.serializer = serializer;
59+
}
60+
61+
Builder withFallbackSerializer(RedisKeySerializer fallbackSerializer) {
62+
this.fallbackSerializer = fallbackSerializer;
63+
return this;
64+
}
65+
66+
RedisClusterOnlineRetriever build() {
67+
return new RedisClusterOnlineRetriever(this);
68+
}
69+
}
70+
71+
private RedisClusterOnlineRetriever(Builder builder) {
72+
this.syncCommands = builder.connection.sync();
73+
this.serializer = builder.serializer;
74+
this.fallbackSerializer = builder.fallbackSerializer;
5175
}
5276

5377
public static OnlineRetriever create(Map<String, String> config) {
@@ -59,15 +83,21 @@ public static OnlineRetriever create(Map<String, String> config) {
5983
return RedisURI.create(hostPortSplit[0], Integer.parseInt(hostPortSplit[1]));
6084
})
6185
.collect(Collectors.toList());
62-
6386
StatefulRedisClusterConnection<byte[], byte[]> connection =
6487
RedisClusterClient.create(redisURIList).connect(new ByteArrayCodec());
6588

66-
return new RedisClusterOnlineRetriever(connection);
67-
}
89+
RedisKeySerializer serializer =
90+
new RedisKeyPrefixSerializer(config.getOrDefault("key_prefix", ""));
91+
92+
Builder builder = new Builder(connection, serializer);
6893

69-
public static OnlineRetriever create(StatefulRedisClusterConnection<byte[], byte[]> connection) {
70-
return new RedisClusterOnlineRetriever(connection);
94+
if (Boolean.parseBoolean(config.getOrDefault("enable_fallback", "false"))) {
95+
RedisKeySerializer fallbackSerializer =
96+
new RedisKeyPrefixSerializer(config.getOrDefault("fallback_prefix", ""));
97+
builder = builder.withFallbackSerializer(fallbackSerializer);
98+
}
99+
100+
return builder.build();
71101
}
72102

73103
/** {@inheritDoc} */
@@ -98,11 +128,9 @@ private List<RedisKey> buildRedisKeys(List<EntityRow> entityRows, FeatureSetSpec
98128
featureSetSpec.getEntitiesList().stream()
99129
.map(EntitySpec::getName)
100130
.collect(Collectors.toList());
101-
List<RedisKey> redisKeys =
102-
entityRows.stream()
103-
.map(row -> makeRedisKey(featureSetRef, featureSetEntityNames, row))
104-
.collect(Collectors.toList());
105-
return redisKeys;
131+
return entityRows.stream()
132+
.map(row -> makeRedisKey(featureSetRef, featureSetEntityNames, row))
133+
.collect(Collectors.toList());
106134
}
107135

108136
/**
@@ -118,9 +146,7 @@ private RedisKey makeRedisKey(
118146
RedisKey.Builder builder = RedisKey.newBuilder().setFeatureSet(featureSet);
119147
Map<String, Value> fieldsMap = entityRow.getFieldsMap();
120148
featureSetEntityNames.sort(String::compareTo);
121-
for (int i = 0; i < featureSetEntityNames.size(); i++) {
122-
String entityName = featureSetEntityNames.get(i);
123-
149+
for (String entityName : featureSetEntityNames) {
124150
if (!fieldsMap.containsKey(entityName)) {
125151
throw Status.INVALID_ARGUMENT
126152
.withDescription(
@@ -180,18 +206,59 @@ private List<byte[]> sendMultiGet(List<RedisKey> keys) {
180206
try {
181207
byte[][] binaryKeys =
182208
keys.stream()
183-
.map(AbstractMessageLite::toByteArray)
209+
.map(serializer::serialize)
184210
.collect(Collectors.toList())
185211
.toArray(new byte[0][0]);
186-
return syncCommands.mget(binaryKeys).stream()
187-
.map(
188-
keyValue -> {
189-
if (keyValue == null) {
190-
return null;
191-
}
192-
return keyValue.getValueOrElse(null);
193-
})
194-
.collect(Collectors.toList());
212+
List<byte[]> redisValues =
213+
syncCommands.mget(binaryKeys).stream()
214+
.map(
215+
keyValue -> {
216+
if (keyValue == null) {
217+
return null;
218+
}
219+
return keyValue.getValueOrElse(null);
220+
})
221+
.collect(Collectors.toList());
222+
223+
List<byte[]> redisValuesWithFallback = redisValues;
224+
if (fallbackSerializer != null) {
225+
List<Integer> indexMissingValue =
226+
IntStream.range(0, keys.size())
227+
.filter(i -> redisValues.get(i) == null)
228+
.boxed()
229+
.collect(Collectors.toList());
230+
231+
byte[][] fallbackBinaryKeys =
232+
indexMissingValue.stream()
233+
.map(i -> fallbackSerializer.serialize(keys.get(i)))
234+
.collect(Collectors.toList())
235+
.toArray(new byte[0][0]);
236+
237+
List<byte[]> fallBackValues =
238+
syncCommands.mget(fallbackBinaryKeys).stream()
239+
.map(
240+
keyValue -> {
241+
if (keyValue == null) {
242+
return null;
243+
}
244+
return keyValue.getValueOrElse(null);
245+
})
246+
.collect(Collectors.toList());
247+
248+
redisValuesWithFallback =
249+
IntStream.range(0, keys.size())
250+
.mapToObj(
251+
i -> {
252+
if (indexMissingValue.contains(i)) {
253+
return fallBackValues.get(indexMissingValue.indexOf(i));
254+
} else {
255+
return redisValues.get(i);
256+
}
257+
})
258+
.collect(Collectors.toList());
259+
}
260+
261+
return redisValuesWithFallback;
195262
} catch (Exception e) {
196263
throw Status.NOT_FOUND
197264
.withDescription("Unable to retrieve feature from Redis")

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/serializer/RedisKeyPrefixSerializer.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
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.serializer;
18+
19+
import feast.proto.storage.RedisProto.RedisKey;
20+
21+
public class RedisKeyPrefixSerializer implements RedisKeySerializer {
22+
23+
private final byte[] prefixBytes;
24+
25+
public RedisKeyPrefixSerializer(String prefix) {
26+
this.prefixBytes = prefix.getBytes();
27+
}
28+
29+
public byte[] serialize(RedisKey redisKey) {
30+
byte[] key = redisKey.toByteArray();
31+
32+
if (prefixBytes.length == 0) {
33+
return key;
34+
}
35+
36+
byte[] keyWithPrefix = new byte[prefixBytes.length + key.length];
37+
System.arraycopy(prefixBytes, 0, keyWithPrefix, 0, prefixBytes.length);
38+
System.arraycopy(key, 0, keyWithPrefix, prefixBytes.length, key.length);
39+
return keyWithPrefix;
40+
}
41+
}

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/serializer/RedisKeyProtoSerializer.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.serializer;
18+
19+
import feast.proto.storage.RedisProto.RedisKey;
20+
21+
public class RedisKeyProtoSerializer implements RedisKeySerializer {
22+
23+
public byte[] serialize(RedisKey redisKey) {
24+
return redisKey.toByteArray();
25+
}
26+
}

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/serializer/RedisKeySerializer.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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.serializer;
18+
19+
import feast.proto.storage.RedisProto.RedisKey;
20+
import java.io.Serializable;
21+
22+
public interface RedisKeySerializer extends Serializable {
23+
24+
byte[] serialize(RedisKey key);
25+
}

‎storage/connectors/redis/src/main/java/feast/storage/connectors/redis/writer/RedisCustomIO.java‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import feast.storage.api.writer.FailedElement;
3232
import feast.storage.api.writer.WriteResult;
3333
import feast.storage.connectors.redis.retriever.FeatureRowDecoder;
34+
import feast.storage.connectors.redis.serializer.RedisKeySerializer;
3435
import java.nio.charset.StandardCharsets;
3536
import java.util.*;
3637
import java.util.HashMap;
@@ -61,23 +62,27 @@ private RedisCustomIO() {}
6162

6263
public static Write write(
6364
RedisIngestionClient redisIngestionClient,
64-
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecs) {
65-
return new Write(redisIngestionClient, featureSetSpecs);
65+
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecs,
66+
RedisKeySerializer serializer) {
67+
return new Write(redisIngestionClient, featureSetSpecs, serializer);
6668
}
6769

6870
/** ServingStoreWrite data to a Redis server. */
6971
public static class Write extends PTransform<PCollection<FeatureRow>, WriteResult> {
7072

7173
private PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecs;
7274
private RedisIngestionClient redisIngestionClient;
75+
private RedisKeySerializer serializer;
7376
private int batchSize;
7477
private Duration flushFrequency;
7578

7679
public Write(
7780
RedisIngestionClient redisIngestionClient,
78-
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecs) {
81+
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecs,
82+
RedisKeySerializer serializer) {
7983
this.redisIngestionClient = redisIngestionClient;
8084
this.featureSetSpecs = featureSetSpecs;
85+
this.serializer = serializer;
8186
}
8287

8388
public Write withBatchSize(int batchSize) {
@@ -108,7 +113,7 @@ public void process(ProcessContext c) {
108113
.apply("ExtractResultValues", Values.create())
109114
.apply("GlobalWindow", Window.<Iterable<FeatureRow>>into(new GlobalWindows()))
110115
.apply(
111-
ParDo.of(new WriteDoFn(redisIngestionClient, featureSetSpecs))
116+
ParDo.of(new WriteDoFn(redisIngestionClient, featureSetSpecs, serializer))
112117
.withOutputTags(successfulInsertsTag, TupleTagList.of(failedInsertsTupleTag))
113118
.withSideInputs(featureSetSpecs));
114119
return WriteResult.in(
@@ -125,13 +130,16 @@ public void process(ProcessContext c) {
125130
*/
126131
public static class WriteDoFn extends BatchDoFnWithRedis<Iterable<FeatureRow>, FeatureRow> {
127132
private final PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecsView;
133+
private final RedisKeySerializer serializer;
128134

129135
WriteDoFn(
130136
RedisIngestionClient redisIngestionClient,
131-
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecsView) {
137+
PCollectionView<Map<String, Iterable<FeatureSetSpec>>> featureSetSpecsView,
138+
RedisKeySerializer serializer) {
132139

133140
super(redisIngestionClient);
134141
this.featureSetSpecsView = featureSetSpecsView;
142+
this.serializer = serializer;
135143
}
136144

137145
private FailedElement toFailedElement(
@@ -230,7 +238,7 @@ public void processElement(ProcessContext context) {
230238
.map(
231239
entry ->
232240
redisIngestionClient
233-
.get(entry.getKey().toByteArray())
241+
.get(serializer.serialize(entry.getKey()))
234242
.thenAccept(
235243
currentValue -> {
236244
FeatureRow newRow = entry.getValue();
@@ -246,7 +254,8 @@ public void processElement(ProcessContext context) {
246254
.map(
247255
row ->
248256
redisIngestionClient.set(
249-
getKey(row, latestSpecs.get(row.getFeatureSet())).toByteArray(),
257+
serializer.serialize(
258+
getKey(row, latestSpecs.get(row.getFeatureSet()))),
250259
getValue(row, latestSpecs.get(row.getFeatureSet()))
251260
.toByteArray()))
252261
.collect(Collectors.toList()));

0 commit comments

Comments
 (0)