diff --git a/.prow/config.yaml b/.prow/config.yaml index 085cfe85423..5b039ff6616 100644 --- a/.prow/config.yaml +++ b/.prow/config.yaml @@ -153,6 +153,19 @@ presubmits: skip_branches: - ^v0\.(3|4)-branch$ + - name: test-end-to-end-redis-cluster + decorate: true + spec: + containers: + - image: maven:3.6-jdk-11 + command: ["infra/scripts/test-end-to-end-redis-cluster.sh"] + resources: + requests: + cpu: "6" + memory: "6144Mi" + skip_branches: + - ^v0\.(3|4)-branch$ + - name: test-end-to-end-java-8 decorate: true always_run: true diff --git a/core/src/main/java/feast/core/model/Store.java b/core/src/main/java/feast/core/model/Store.java index 9dc44bdc73a..debf211ec8d 100644 --- a/core/src/main/java/feast/core/model/Store.java +++ b/core/src/main/java/feast/core/model/Store.java @@ -21,6 +21,7 @@ import feast.core.StoreProto.Store.BigQueryConfig; import feast.core.StoreProto.Store.Builder; import feast.core.StoreProto.Store.CassandraConfig; +import feast.core.StoreProto.Store.RedisClusterConfig; import feast.core.StoreProto.Store.RedisConfig; import feast.core.StoreProto.Store.StoreType; import feast.core.StoreProto.Store.Subscription; @@ -82,6 +83,9 @@ public static Store fromProto(StoreProto.Store storeProto) throws IllegalArgumen case CASSANDRA: config = storeProto.getCassandraConfig().toByteArray(); break; + case REDIS_CLUSTER: + config = storeProto.getRedisClusterConfig().toByteArray(); + break; default: throw new IllegalArgumentException("Invalid store provided"); } @@ -106,6 +110,9 @@ public StoreProto.Store toProto() throws InvalidProtocolBufferException { case CASSANDRA: CassandraConfig cassConfig = CassandraConfig.parseFrom(config); return storeProtoBuilder.setCassandraConfig(cassConfig).build(); + case REDIS_CLUSTER: + RedisClusterConfig redisClusterConfig = RedisClusterConfig.parseFrom(config); + return storeProtoBuilder.setRedisClusterConfig(redisClusterConfig).build(); default: throw new InvalidProtocolBufferException("Invalid store set"); } diff --git a/infra/scripts/setup-redis-cluster.sh b/infra/scripts/setup-redis-cluster.sh new file mode 100755 index 00000000000..a1939705318 --- /dev/null +++ b/infra/scripts/setup-redis-cluster.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash + +apt-get -y install redis-server > /var/log/redis.install.log + +mkdir 7000 7001 7002 7003 7004 7005 +for i in {0..5} ; do +echo "port 700$i +cluster-enabled yes +cluster-config-file nodes-$i.conf +cluster-node-timeout 5000 +appendonly yes" > 700$i/redis.conf +redis-server 700$i/redis.conf --daemonize yes +done +echo yes | redis-cli --cluster create 127.0.0.1:7000 127.0.0.1:7001 \ +127.0.0.1:7002 127.0.0.1:7003 127.0.0.1:7004 127.0.0.1:7005 \ +--cluster-replicas 1 diff --git a/infra/scripts/test-end-to-end-redis-cluster.sh b/infra/scripts/test-end-to-end-redis-cluster.sh new file mode 100755 index 00000000000..7f0d47fc92b --- /dev/null +++ b/infra/scripts/test-end-to-end-redis-cluster.sh @@ -0,0 +1,248 @@ +#!/usr/bin/env bash + +set -e +set -o pipefail + +test -z ${GOOGLE_APPLICATION_CREDENTIALS} && GOOGLE_APPLICATION_CREDENTIALS="/etc/service-account/service-account.json" +test -z ${SKIP_BUILD_JARS} && SKIP_BUILD_JARS="false" +test -z ${GOOGLE_CLOUD_PROJECT} && GOOGLE_CLOUD_PROJECT="kf-feast" +test -z ${TEMP_BUCKET} && TEMP_BUCKET="feast-templocation-kf-feast" +test -z ${JOBS_STAGING_LOCATION} && JOBS_STAGING_LOCATION="gs://${TEMP_BUCKET}/staging-location" +test -z ${JAR_VERSION_SUFFIX} && JAR_VERSION_SUFFIX="-SNAPSHOT" + +echo " +This script will run end-to-end tests for Feast Core and Online Serving. + +1. Install Redis as the store for Feast Online Serving. +2. Install Postgres for persisting Feast metadata. +3. Install Kafka and Zookeeper as the Source in Feast. +4. Install Python 3.7.4, Feast Python SDK and run end-to-end tests from + tests/e2e via pytest. +" + +apt-get -qq update +apt-get -y install wget netcat kafkacat + +echo " +============================================================ +Installing Redis at localhost:6379 +============================================================ +" +# Allow starting serving in this Maven Docker image. Default set to not allowed. +echo "exit 0" > /usr/sbin/policy-rc.d +infra/scripts/setup-redis-cluster.sh +redis-cli -c -p 7000 ping + +echo " +============================================================ +Installing Postgres at localhost:5432 +============================================================ +" +apt-get -y install postgresql > /var/log/postgresql.install.log +service postgresql start +# Initialize with database: 'postgres', user: 'postgres', password: 'password' +cat < /tmp/update-postgres-role.sh +psql -c "ALTER USER postgres PASSWORD 'password';" +EOF +chmod +x /tmp/update-postgres-role.sh +su -s /bin/bash -c /tmp/update-postgres-role.sh postgres +export PGPASSWORD=password +pg_isready + +echo " +============================================================ +Installing Zookeeper at localhost:2181 +Installing Kafka at localhost:9092 +============================================================ +" +wget -qO- https://www-eu.apache.org/dist/kafka/2.3.0/kafka_2.12-2.3.0.tgz | tar xz +mv kafka_2.12-2.3.0/ /tmp/kafka +nohup /tmp/kafka/bin/zookeeper-server-start.sh /tmp/kafka/config/zookeeper.properties &> /var/log/zookeeper.log 2>&1 & +sleep 5 +tail -n10 /var/log/zookeeper.log +nohup /tmp/kafka/bin/kafka-server-start.sh /tmp/kafka/config/server.properties &> /var/log/kafka.log 2>&1 & +sleep 20 +tail -n10 /var/log/kafka.log +kafkacat -b localhost:9092 -L + +if [[ ${SKIP_BUILD_JARS} != "true" ]]; then +echo " +============================================================ +Building jars for Feast +============================================================ +" + +.prow/scripts/download-maven-cache.sh \ + --archive-uri gs://feast-templocation-kf-feast/.m2.2019-10-24.tar \ + --output-dir /root/ + +# Build jars for Feast +mvn --quiet --batch-mode --define skipTests=true clean package + +ls -lh core/target/*jar +ls -lh serving/target/*jar +else + echo "[DEBUG] Skipping building jars" +fi + +echo " +============================================================ +Starting Feast Core +============================================================ +" +# Start Feast Core in background +cat < /tmp/core.application.yml +grpc: + port: 6565 + enable-reflection: true + +feast: + version: 0.3 + jobs: + runner: DirectRunner + options: {} + updates: + timeoutSeconds: 240 + metrics: + enabled: false + + stream: + type: kafka + options: + topic: feast-features + bootstrapServers: localhost:9092 + replicationFactor: 1 + partitions: 1 + +spring: + jpa: + properties.hibernate: + format_sql: true + event.merge.entity_copy_observer: allow + hibernate.naming.physical-strategy=org.hibernate.boot.model.naming: PhysicalNamingStrategyStandardImpl + hibernate.ddl-auto: update + datasource: + url: jdbc:postgresql://localhost:5432/postgres + username: postgres + password: password + +management: + metrics: + export: + simple: + enabled: false + statsd: + enabled: false +EOF + +nohup java -jar core/target/feast-core-*${JAR_VERSION_SUFFIX}.jar \ + --spring.config.location=file:///tmp/core.application.yml \ + &> /var/log/feast-core.log & +sleep 35 +tail -n10 /var/log/feast-core.log +nc -w2 localhost 6565 < /dev/null + +echo " +============================================================ +Starting Feast Online Serving +============================================================ +" +# Start Feast Online Serving in background +cat < /tmp/serving.store.redis.cluster.yml +name: serving +type: REDIS_CLUSTER +redis_cluster_config: + nodes: + - host: localhost + port: 7000 + - host: localhost + port: 7001 + - host: localhost + port: 7002 + - host: localhost + port: 7003 + - host: localhost + port: 7004 + - host: localhost + port: 7005 +subscriptions: + - name: "*" + version: "*" + project: "*" +EOF + +cat < /tmp/serving.online.application.yml +feast: + version: 0.3 + core-host: localhost + core-grpc-port: 6565 + tracing: + enabled: false + store: + config-path: /tmp/serving.store.redis.cluster.yml + redis-pool-max-size: 128 + redis-pool-max-idle: 16 + jobs: + staging-location: ${JOBS_STAGING_LOCATION} + store-type: + store-options: {} + +grpc: + port: 6566 + enable-reflection: true + +spring: + main: + web-environment: false + +EOF + +nohup java -jar serving/target/feast-serving-*${JAR_VERSION_SUFFIX}.jar \ + --spring.config.location=file:///tmp/serving.online.application.yml \ + &> /var/log/feast-serving-online.log & +sleep 15 +tail -n100 /var/log/feast-serving-online.log +nc -w2 localhost 6566 < /dev/null + +echo " +============================================================ +Installing Python 3.7 with Miniconda and Feast SDK +============================================================ +" +# Install Python 3.7 with Miniconda +wget -q https://repo.continuum.io/miniconda/Miniconda3-4.7.12-Linux-x86_64.sh \ + -O /tmp/miniconda.sh +bash /tmp/miniconda.sh -b -p /root/miniconda -f +/root/miniconda/bin/conda init +source ~/.bashrc + +# Install Feast Python SDK and test requirements +pip install -qe sdk/python +pip install -qr tests/e2e/requirements.txt + +echo " +============================================================ +Running end-to-end tests with pytest at 'tests/e2e' +============================================================ +" +# Default artifact location setting in Prow jobs +LOGS_ARTIFACT_PATH=/logs/artifacts + +ORIGINAL_DIR=$(pwd) +cd tests/e2e + +set +e +pytest basic-ingest-redis-serving.py --junitxml=${LOGS_ARTIFACT_PATH}/python-sdk-test-report.xml +TEST_EXIT_CODE=$? + +if [[ ${TEST_EXIT_CODE} != 0 ]]; then + echo "[DEBUG] Printing logs" + ls -ltrh /var/log/feast* + cat /var/log/feast-serving-online.log /var/log/feast-core.log + + echo "[DEBUG] Printing Python packages list" + pip list +fi + +cd ${ORIGINAL_DIR} +exit ${TEST_EXIT_CODE} diff --git a/ingestion/pom.xml b/ingestion/pom.xml index 9386d066bfd..64d5a41f86f 100644 --- a/ingestion/pom.xml +++ b/ingestion/pom.xml @@ -113,6 +113,12 @@ ${project.version} + + dev.feast + feast-storage-connector-redis-cluster + ${project.version} + + dev.feast feast-storage-connector-bigquery diff --git a/ingestion/src/main/java/feast/ingestion/utils/StoreUtil.java b/ingestion/src/main/java/feast/ingestion/utils/StoreUtil.java index b62f83f0f30..566124b3075 100644 --- a/ingestion/src/main/java/feast/ingestion/utils/StoreUtil.java +++ b/ingestion/src/main/java/feast/ingestion/utils/StoreUtil.java @@ -25,6 +25,7 @@ import feast.storage.api.writer.FeatureSink; import feast.storage.connectors.bigquery.writer.BigQueryFeatureSink; import feast.storage.connectors.redis.writer.RedisFeatureSink; +import feast.storage.connectors.rediscluster.writer.RedisClusterFeatureSink; import feast.types.ValueProto.ValueType.Enum; import java.util.HashMap; import java.util.Map; @@ -82,12 +83,14 @@ public static FeatureSink getFeatureSink( Store store, Map featureSetSpecs) { StoreType storeType = store.getType(); switch (storeType) { + case REDIS_CLUSTER: + return RedisClusterFeatureSink.fromConfig(store.getRedisClusterConfig(), featureSetSpecs); case REDIS: return RedisFeatureSink.fromConfig(store.getRedisConfig(), featureSetSpecs); case BIGQUERY: return BigQueryFeatureSink.fromConfig(store.getBigqueryConfig(), featureSetSpecs); default: - throw new RuntimeException(String.format("Store type '{}' is unsupported", storeType)); + throw new RuntimeException(String.format("Store type '%s' is unsupported", storeType)); } } } diff --git a/protos/feast/core/Store.proto b/protos/feast/core/Store.proto index de9af0a99fe..0aa4c8cd420 100644 --- a/protos/feast/core/Store.proto +++ b/protos/feast/core/Store.proto @@ -105,6 +105,8 @@ message Store { // Unsupported in Feast 0.3 CASSANDRA = 3; + + REDIS_CLUSTER = 4; } message RedisConfig { @@ -130,6 +132,13 @@ message Store { int32 port = 2; } + message RedisClusterConfig { + // List of Redis Uri for all the nodes in Redis Cluster, comma separated. Eg. host1:6379, host2:6379 + string connection_string = 1; + int32 initial_backoff_ms = 2; + int32 max_retries = 3; + } + message Subscription { // Name of project that the feature sets belongs to. This can be one of // - [project_name] @@ -172,5 +181,6 @@ message Store { RedisConfig redis_config = 11; BigQueryConfig bigquery_config = 12; CassandraConfig cassandra_config = 13; + RedisClusterConfig redis_cluster_config = 14; } } diff --git a/serving/pom.xml b/serving/pom.xml index bbb694011a3..d3d7ae212fd 100644 --- a/serving/pom.xml +++ b/serving/pom.xml @@ -97,6 +97,12 @@ ${project.version} + + dev.feast + feast-storage-connector-redis-cluster + ${project.version} + + dev.feast feast-storage-connector-bigquery diff --git a/serving/src/main/java/feast/serving/config/FeastProperties.java b/serving/src/main/java/feast/serving/config/FeastProperties.java index bf3387728a7..0c5ec965f31 100644 --- a/serving/src/main/java/feast/serving/config/FeastProperties.java +++ b/serving/src/main/java/feast/serving/config/FeastProperties.java @@ -26,10 +26,7 @@ import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.util.JsonFormat; import feast.core.StoreProto; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; import java.util.stream.Collectors; import javax.validation.constraints.NotBlank; import javax.validation.constraints.Positive; @@ -245,6 +242,11 @@ public StoreProto.Store toProto() // TODO: All of this logic should be moved to the store layer. Only a Map // should be sent to a store and it should do its own validation. switch (StoreProto.Store.StoreType.valueOf(type)) { + case REDIS_CLUSTER: + StoreProto.Store.RedisClusterConfig.Builder redisClusterConfig = + StoreProto.Store.RedisClusterConfig.newBuilder(); + JsonFormat.parser().merge(jsonWriter.writeValueAsString(config), redisClusterConfig); + return storeProtoBuilder.setRedisClusterConfig(redisClusterConfig.build()).build(); case REDIS: StoreProto.Store.RedisConfig.Builder redisConfig = StoreProto.Store.RedisConfig.newBuilder(); diff --git a/serving/src/main/java/feast/serving/config/ServingServiceConfig.java b/serving/src/main/java/feast/serving/config/ServingServiceConfig.java index ec84e6c4fef..a1dbc1db604 100644 --- a/serving/src/main/java/feast/serving/config/ServingServiceConfig.java +++ b/serving/src/main/java/feast/serving/config/ServingServiceConfig.java @@ -29,6 +29,7 @@ import feast.storage.api.retriever.OnlineRetriever; import feast.storage.connectors.bigquery.retriever.BigQueryHistoricalRetriever; import feast.storage.connectors.redis.retriever.RedisOnlineRetriever; +import feast.storage.connectors.rediscluster.retriever.RedisClusterOnlineRetriever; import io.opentracing.Tracer; import java.util.Map; import org.slf4j.Logger; @@ -53,6 +54,10 @@ public ServingService servingService( Map config = store.getConfig(); switch (storeType) { + case REDIS_CLUSTER: + OnlineRetriever redisClusterRetriever = RedisClusterOnlineRetriever.create(config); + servingService = new OnlineServingService(redisClusterRetriever, specService, tracer); + break; case REDIS: OnlineRetriever redisRetriever = RedisOnlineRetriever.create(config); servingService = new OnlineServingService(redisRetriever, specService, tracer); diff --git a/storage/connectors/pom.xml b/storage/connectors/pom.xml index b52668a31a4..b57fe98cd25 100644 --- a/storage/connectors/pom.xml +++ b/storage/connectors/pom.xml @@ -16,6 +16,7 @@ redis + rediscluster bigquery diff --git a/storage/connectors/rediscluster/pom.xml b/storage/connectors/rediscluster/pom.xml new file mode 100644 index 00000000000..5c3cb6e42d3 --- /dev/null +++ b/storage/connectors/rediscluster/pom.xml @@ -0,0 +1,81 @@ + + + + dev.feast + feast-storage-connectors + ${revision} + + + 4.0.0 + feast-storage-connector-redis-cluster + + Feast Storage Connector for Redis Cluster + + + + io.lettuce + lettuce-core + + + + org.apache.commons + commons-lang3 + 3.9 + + + + com.google.auto.value + auto-value-annotations + 1.6.6 + + + + com.google.auto.value + auto-value + 1.6.6 + provided + + + + org.mockito + mockito-core + 2.23.0 + test + + + + org.apache.beam + beam-runners-direct-java + ${org.apache.beam.version} + test + + + + org.hamcrest + hamcrest-core + test + + + + org.hamcrest + hamcrest-library + test + + + + net.ishiis.redis + redis-unit + 1.0.3 + test + + + + junit + junit + 4.12 + test + + + + diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/FeatureRowDecoder.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/FeatureRowDecoder.java new file mode 100644 index 00000000000..d6312c6b6ab --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/FeatureRowDecoder.java @@ -0,0 +1,82 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.retriever; + +import feast.core.FeatureSetProto.FeatureSetSpec; +import feast.core.FeatureSetProto.FeatureSpec; +import feast.types.FeatureRowProto.FeatureRow; +import feast.types.FieldProto.Field; +import java.util.Comparator; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +public class FeatureRowDecoder { + + private final String featureSetRef; + private final FeatureSetSpec spec; + + public FeatureRowDecoder(String featureSetRef, FeatureSetSpec spec) { + this.featureSetRef = featureSetRef; + this.spec = spec; + } + + /** + * Validates if an encoded feature row can be decoded without exception. + * + * @param featureRow Feature row + * @return boolean + */ + public Boolean isEncodingValid(FeatureRow featureRow) { + return featureRow.getFieldsList().size() == spec.getFeaturesList().size(); + } + + /** + * Decoding feature row by repopulating the field names based on the corresponding feature set + * spec. + * + * @param encodedFeatureRow Feature row + * @return boolean + */ + public FeatureRow decode(FeatureRow encodedFeatureRow) { + final List fieldsWithoutName = encodedFeatureRow.getFieldsList(); + + List featureNames = + spec.getFeaturesList().stream() + .sorted(Comparator.comparing(FeatureSpec::getName)) + .map(FeatureSpec::getName) + .collect(Collectors.toList()); + List fields = + IntStream.range(0, featureNames.size()) + .mapToObj( + featureNameIndex -> { + String featureName = featureNames.get(featureNameIndex); + return fieldsWithoutName + .get(featureNameIndex) + .toBuilder() + .setName(featureName) + .build(); + }) + .collect(Collectors.toList()); + return encodedFeatureRow + .toBuilder() + .clearFields() + .setFeatureSet(featureSetRef) + .addAllFields(fields) + .build(); + } +} diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetriever.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetriever.java new file mode 100644 index 00000000000..713b6897b2d --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetriever.java @@ -0,0 +1,226 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.retriever; + +import com.google.protobuf.AbstractMessageLite; +import com.google.protobuf.InvalidProtocolBufferException; +import feast.core.FeatureSetProto.EntitySpec; +import feast.core.FeatureSetProto.FeatureSetSpec; +import feast.serving.ServingAPIProto.FeatureReference; +import feast.serving.ServingAPIProto.GetOnlineFeaturesRequest.EntityRow; +import feast.storage.RedisProto.RedisKey; +import feast.storage.api.retriever.FeatureSetRequest; +import feast.storage.api.retriever.OnlineRetriever; +import feast.types.FeatureRowProto.FeatureRow; +import feast.types.FieldProto.Field; +import feast.types.ValueProto.Value; +import io.grpc.Status; +import io.lettuce.core.RedisURI; +import io.lettuce.core.cluster.RedisClusterClient; +import io.lettuce.core.cluster.api.StatefulRedisClusterConnection; +import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; +import io.lettuce.core.codec.ByteArrayCodec; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; + +public class RedisClusterOnlineRetriever implements OnlineRetriever { + + private final RedisAdvancedClusterCommands syncCommands; + + private RedisClusterOnlineRetriever(StatefulRedisClusterConnection connection) { + this.syncCommands = connection.sync(); + } + + public static OnlineRetriever create(Map config) { + List redisURIList = + Arrays.stream(config.get("connection_string").split(",")) + .map( + hostPort -> { + String[] hostPortSplit = hostPort.trim().split(":"); + return RedisURI.create(hostPortSplit[0], Integer.parseInt(hostPortSplit[1])); + }) + .collect(Collectors.toList()); + + StatefulRedisClusterConnection connection = + RedisClusterClient.create(redisURIList).connect(new ByteArrayCodec()); + + return new RedisClusterOnlineRetriever(connection); + } + + public static OnlineRetriever create(StatefulRedisClusterConnection connection) { + return new RedisClusterOnlineRetriever(connection); + } + + /** + * Gets online features from redis. This method returns a list of {@link FeatureRow}s + * corresponding to each feature set spec. Each feature row in the list then corresponds to an + * {@link EntityRow} provided by the user. + * + * @param entityRows list of entity rows in the feature request + * @param featureSetRequests Map of {@link FeatureSetSpec} to feature references in the request + * tied to that feature set. + * @return List of List of {@link FeatureRow} + */ + @Override + public List> getOnlineFeatures( + List entityRows, List featureSetRequests) { + + List> featureRows = new ArrayList<>(); + for (FeatureSetRequest featureSetRequest : featureSetRequests) { + List redisKeys = buildRedisKeys(entityRows, featureSetRequest.getSpec()); + try { + List featureRowsForFeatureSet = + sendAndProcessMultiGet( + redisKeys, + featureSetRequest.getSpec(), + featureSetRequest.getFeatureReferences().asList()); + featureRows.add(featureRowsForFeatureSet); + } catch (InvalidProtocolBufferException | ExecutionException e) { + throw Status.INTERNAL + .withDescription("Unable to parse protobuf while retrieving feature") + .withCause(e) + .asRuntimeException(); + } + } + return featureRows; + } + + private List buildRedisKeys(List entityRows, FeatureSetSpec featureSetSpec) { + String featureSetRef = generateFeatureSetStringRef(featureSetSpec); + List featureSetEntityNames = + featureSetSpec.getEntitiesList().stream() + .map(EntitySpec::getName) + .collect(Collectors.toList()); + List redisKeys = + entityRows.stream() + .map(row -> makeRedisKey(featureSetRef, featureSetEntityNames, row)) + .collect(Collectors.toList()); + return redisKeys; + } + + /** + * Create {@link RedisKey} + * + * @param featureSet featureSet reference of the feature. E.g. feature_set_1:1 + * @param featureSetEntityNames entity names that belong to the featureSet + * @param entityRow entityRow to build the key from + * @return {@link RedisKey} + */ + private RedisKey makeRedisKey( + String featureSet, List featureSetEntityNames, EntityRow entityRow) { + RedisKey.Builder builder = RedisKey.newBuilder().setFeatureSet(featureSet); + Map fieldsMap = entityRow.getFieldsMap(); + featureSetEntityNames.sort(String::compareTo); + for (int i = 0; i < featureSetEntityNames.size(); i++) { + String entityName = featureSetEntityNames.get(i); + + if (!fieldsMap.containsKey(entityName)) { + throw Status.INVALID_ARGUMENT + .withDescription( + String.format( + "Entity row fields \"%s\" does not contain required entity field \"%s\"", + fieldsMap.keySet().toString(), entityName)) + .asRuntimeException(); + } + + builder.addEntities( + Field.newBuilder().setName(entityName).setValue(fieldsMap.get(entityName))); + } + return builder.build(); + } + + private List sendAndProcessMultiGet( + List redisKeys, + FeatureSetSpec featureSetSpec, + List featureReferences) + throws InvalidProtocolBufferException, ExecutionException { + + List values = sendMultiGet(redisKeys); + List featureRows = new ArrayList<>(); + + FeatureRow.Builder nullFeatureRowBuilder = + FeatureRow.newBuilder().setFeatureSet(generateFeatureSetStringRef(featureSetSpec)); + for (FeatureReference featureReference : featureReferences) { + nullFeatureRowBuilder.addFields(Field.newBuilder().setName(featureReference.getName())); + } + + for (int i = 0; i < values.size(); i++) { + + byte[] value = values.get(i); + if (value == null) { + featureRows.add(nullFeatureRowBuilder.build()); + continue; + } + + FeatureRow featureRow = FeatureRow.parseFrom(value); + String featureSetRef = redisKeys.get(i).getFeatureSet(); + FeatureRowDecoder decoder = new FeatureRowDecoder(featureSetRef, featureSetSpec); + if (decoder.isEncodingValid(featureRow)) { + featureRow = decoder.decode(featureRow); + } else { + featureRows.add(nullFeatureRowBuilder.build()); + continue; + } + + featureRows.add(featureRow); + } + return featureRows; + } + + /** + * Send a list of get request as an mget + * + * @param keys list of {@link RedisKey} + * @return list of {@link FeatureRow} in primitive byte representation for each {@link RedisKey} + */ + private List sendMultiGet(List keys) { + try { + byte[][] binaryKeys = + keys.stream() + .map(AbstractMessageLite::toByteArray) + .collect(Collectors.toList()) + .toArray(new byte[0][0]); + return syncCommands.mget(binaryKeys).stream() + .map( + keyValue -> { + if (keyValue == null) { + return null; + } + return keyValue.getValueOrElse(null); + }) + .collect(Collectors.toList()); + } catch (Exception e) { + throw Status.NOT_FOUND + .withDescription("Unable to retrieve feature from Redis") + .withCause(e) + .asRuntimeException(); + } + } + + // TODO: Refactor this out to common package? + private static String generateFeatureSetStringRef(FeatureSetSpec featureSetSpec) { + String ref = String.format("%s/%s", featureSetSpec.getProject(), featureSetSpec.getName()); + if (featureSetSpec.getVersion() > 0) { + return ref + String.format(":%d", featureSetSpec.getVersion()); + } + return ref; + } +} diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterCustomIO.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterCustomIO.java new file mode 100644 index 00000000000..0a7634c5c5f --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterCustomIO.java @@ -0,0 +1,294 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2019 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.writer; + +import feast.core.FeatureSetProto.EntitySpec; +import feast.core.FeatureSetProto.FeatureSetSpec; +import feast.core.FeatureSetProto.FeatureSpec; +import feast.core.StoreProto.Store.RedisClusterConfig; +import feast.storage.RedisProto.RedisKey; +import feast.storage.RedisProto.RedisKey.Builder; +import feast.storage.api.writer.FailedElement; +import feast.storage.api.writer.WriteResult; +import feast.storage.common.retry.Retriable; +import feast.types.FeatureRowProto.FeatureRow; +import feast.types.FieldProto.Field; +import feast.types.ValueProto; +import io.lettuce.core.RedisException; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.commons.lang3.exception.ExceptionUtils; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class RedisClusterCustomIO { + + private static final int DEFAULT_BATCH_SIZE = 1000; + private static final int DEFAULT_TIMEOUT = 2000; + + private static TupleTag successfulInsertsTag = new TupleTag<>("successfulInserts") {}; + private static TupleTag failedInsertsTupleTag = new TupleTag<>("failedInserts") {}; + + private static final Logger log = LoggerFactory.getLogger(RedisClusterCustomIO.class); + + private RedisClusterCustomIO() {} + + public static Write write( + RedisClusterConfig redisClusterConfig, Map featureSetSpecs) { + return new Write(redisClusterConfig, featureSetSpecs); + } + + /** ServingStoreWrite data to a Redis server. */ + public static class Write extends PTransform, WriteResult> { + + private Map featureSetSpecs; + private RedisClusterConfig redisClusterConfig; + private int batchSize; + private int timeout; + + public Write( + RedisClusterConfig redisClusterConfig, Map featureSetSpecs) { + + this.redisClusterConfig = redisClusterConfig; + this.featureSetSpecs = featureSetSpecs; + } + + public Write withBatchSize(int batchSize) { + this.batchSize = batchSize; + return this; + } + + public Write withTimeout(int timeout) { + this.timeout = timeout; + return this; + } + + @Override + public WriteResult expand(PCollection input) { + PCollectionTuple redisWrite = + input.apply( + ParDo.of(new WriteDoFn(redisClusterConfig, featureSetSpecs)) + .withOutputTags(successfulInsertsTag, TupleTagList.of(failedInsertsTupleTag))); + return WriteResult.in( + input.getPipeline(), + redisWrite.get(successfulInsertsTag), + redisWrite.get(failedInsertsTupleTag)); + } + + public static class WriteDoFn extends DoFn { + + private final List featureRows = new ArrayList<>(); + private Map featureSetSpecs; + private int batchSize = DEFAULT_BATCH_SIZE; + private int timeout = DEFAULT_TIMEOUT; + private RedisIngestionClient redisIngestionClient; + + WriteDoFn(RedisClusterConfig config, Map featureSetSpecs) { + + this.redisIngestionClient = new RedisClusterIngestionClient(config); + this.featureSetSpecs = featureSetSpecs; + } + + public WriteDoFn withBatchSize(int batchSize) { + if (batchSize > 0) { + this.batchSize = batchSize; + } + return this; + } + + public WriteDoFn withTimeout(int timeout) { + if (timeout > 0) { + this.timeout = timeout; + } + return this; + } + + @Setup + public void setup() { + this.redisIngestionClient.setup(); + } + + @StartBundle + public void startBundle() { + try { + redisIngestionClient.connect(); + } catch (RedisException e) { + log.error("Connection to redis cannot be established ", e); + } + featureRows.clear(); + } + + private void executeBatch() throws Exception { + this.redisIngestionClient + .getBackOffExecutor() + .execute( + new Retriable() { + @Override + public void execute() throws ExecutionException, InterruptedException { + if (!redisIngestionClient.isConnected()) { + redisIngestionClient.connect(); + } + featureRows.forEach( + row -> { + redisIngestionClient.set(getKey(row), getValue(row)); + }); + redisIngestionClient.sync(); + } + + @Override + public Boolean isExceptionRetriable(Exception e) { + return e instanceof RedisException; + } + + @Override + public void cleanUpAfterFailure() {} + }); + } + + private FailedElement toFailedElement( + FeatureRow featureRow, Exception exception, String jobName) { + return FailedElement.newBuilder() + .setJobName(jobName) + .setTransformName("RedisClusterCustomIO") + .setPayload(featureRow.toString()) + .setErrorMessage(exception.getMessage()) + .setStackTrace(ExceptionUtils.getStackTrace(exception)) + .build(); + } + + private byte[] getKey(FeatureRow featureRow) { + FeatureSetSpec featureSetSpec = featureSetSpecs.get(featureRow.getFeatureSet()); + List entityNames = + featureSetSpec.getEntitiesList().stream() + .map(EntitySpec::getName) + .sorted() + .collect(Collectors.toList()); + + Map entityFields = new HashMap<>(); + Builder redisKeyBuilder = RedisKey.newBuilder().setFeatureSet(featureRow.getFeatureSet()); + for (Field field : featureRow.getFieldsList()) { + if (entityNames.contains(field.getName())) { + entityFields.putIfAbsent( + field.getName(), + Field.newBuilder().setName(field.getName()).setValue(field.getValue()).build()); + } + } + for (String entityName : entityNames) { + redisKeyBuilder.addEntities(entityFields.get(entityName)); + } + return redisKeyBuilder.build().toByteArray(); + } + + private byte[] getValue(FeatureRow featureRow) { + FeatureSetSpec spec = featureSetSpecs.get(featureRow.getFeatureSet()); + + List featureNames = + spec.getFeaturesList().stream().map(FeatureSpec::getName).collect(Collectors.toList()); + Map fieldValueOnlyMap = + featureRow.getFieldsList().stream() + .filter(field -> featureNames.contains(field.getName())) + .distinct() + .collect( + Collectors.toMap( + Field::getName, + field -> Field.newBuilder().setValue(field.getValue()).build())); + + List values = + featureNames.stream() + .sorted() + .map( + featureName -> + fieldValueOnlyMap.getOrDefault( + featureName, + Field.newBuilder() + .setValue(ValueProto.Value.getDefaultInstance()) + .build())) + .collect(Collectors.toList()); + + return FeatureRow.newBuilder() + .setEventTimestamp(featureRow.getEventTimestamp()) + .addAllFields(values) + .build() + .toByteArray(); + } + + @ProcessElement + public void processElement(ProcessContext context) { + FeatureRow featureRow = context.element(); + featureRows.add(featureRow); + if (featureRows.size() >= batchSize) { + try { + executeBatch(); + featureRows.forEach(row -> context.output(successfulInsertsTag, row)); + featureRows.clear(); + } catch (Exception e) { + featureRows.forEach( + failedMutation -> { + FailedElement failedElement = + toFailedElement(failedMutation, e, context.getPipelineOptions().getJobName()); + context.output(failedInsertsTupleTag, failedElement); + }); + featureRows.clear(); + } + } + } + + @FinishBundle + public void finishBundle(FinishBundleContext context) + throws IOException, InterruptedException { + if (featureRows.size() > 0) { + try { + executeBatch(); + featureRows.forEach( + row -> + context.output( + successfulInsertsTag, row, Instant.now(), GlobalWindow.INSTANCE)); + featureRows.clear(); + } catch (Exception e) { + featureRows.forEach( + failedMutation -> { + FailedElement failedElement = + toFailedElement(failedMutation, e, context.getPipelineOptions().getJobName()); + context.output( + failedInsertsTupleTag, failedElement, Instant.now(), GlobalWindow.INSTANCE); + }); + featureRows.clear(); + } + } + } + + @Teardown + public void teardown() { + redisIngestionClient.shutdown(); + } + } + } +} diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSink.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSink.java new file mode 100644 index 00000000000..c8126c77930 --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSink.java @@ -0,0 +1,75 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.writer; + +import com.google.auto.value.AutoValue; +import feast.core.FeatureSetProto; +import feast.core.StoreProto.Store.RedisClusterConfig; +import feast.storage.api.writer.FeatureSink; +import feast.storage.api.writer.WriteResult; +import feast.types.FeatureRowProto; +import java.util.Map; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PCollection; + +@AutoValue +public abstract class RedisClusterFeatureSink implements FeatureSink { + + /** + * Initialize a {@link RedisClusterFeatureSink.Builder} from a {@link RedisClusterConfig}. + * + * @param redisClusterConfig {@link RedisClusterConfig} + * @param featureSetSpecs + * @return {@link RedisClusterFeatureSink.Builder} + */ + public static FeatureSink fromConfig( + RedisClusterConfig redisClusterConfig, + Map featureSetSpecs) { + return builder() + .setFeatureSetSpecs(featureSetSpecs) + .setRedisClusterConfig(redisClusterConfig) + .build(); + } + + public abstract RedisClusterConfig getRedisClusterConfig(); + + public abstract Map getFeatureSetSpecs(); + + public abstract Builder toBuilder(); + + public static Builder builder() { + return new AutoValue_RedisClusterFeatureSink.Builder(); + } + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setRedisClusterConfig(RedisClusterConfig redisClusterConfig); + + public abstract Builder setFeatureSetSpecs( + Map featureSetSpecs); + + public abstract RedisClusterFeatureSink build(); + } + + @Override + public void prepareWrite(FeatureSetProto.FeatureSet featureSet) {} + + @Override + public PTransform, WriteResult> writer() { + return new RedisClusterCustomIO.Write(getRedisClusterConfig(), getFeatureSetSpecs()); + } +} diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterIngestionClient.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterIngestionClient.java new file mode 100644 index 00000000000..1f395f02e56 --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisClusterIngestionClient.java @@ -0,0 +1,132 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.writer; + +import com.google.common.collect.Lists; +import feast.core.StoreProto; +import feast.storage.common.retry.BackOffExecutor; +import io.lettuce.core.LettuceFutures; +import io.lettuce.core.RedisFuture; +import io.lettuce.core.RedisURI; +import io.lettuce.core.cluster.RedisClusterClient; +import io.lettuce.core.cluster.api.StatefulRedisClusterConnection; +import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands; +import io.lettuce.core.codec.ByteArrayCodec; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.joda.time.Duration; + +public class RedisClusterIngestionClient implements RedisIngestionClient { + + private final BackOffExecutor backOffExecutor; + private final List uriList; + private transient RedisClusterClient clusterClient; + private StatefulRedisClusterConnection connection; + private RedisAdvancedClusterAsyncCommands commands; + private List futures = Lists.newArrayList(); + + public RedisClusterIngestionClient(StoreProto.Store.RedisClusterConfig redisClusterConfig) { + this.uriList = + Arrays.stream(redisClusterConfig.getConnectionString().split(",")) + .map( + hostPort -> { + String[] hostPortSplit = hostPort.trim().split(":"); + return RedisURI.create(hostPortSplit[0], Integer.parseInt(hostPortSplit[1])); + }) + .collect(Collectors.toList()); + + long backoffMs = + redisClusterConfig.getInitialBackoffMs() > 0 ? redisClusterConfig.getInitialBackoffMs() : 1; + this.backOffExecutor = + new BackOffExecutor(redisClusterConfig.getMaxRetries(), Duration.millis(backoffMs)); + this.clusterClient = RedisClusterClient.create(uriList); + } + + @Override + public void setup() { + this.clusterClient = RedisClusterClient.create(this.uriList); + } + + @Override + public BackOffExecutor getBackOffExecutor() { + return this.backOffExecutor; + } + + @Override + public void shutdown() { + this.clusterClient.shutdown(); + } + + @Override + public void connect() { + if (!isConnected()) { + this.connection = clusterClient.connect(new ByteArrayCodec()); + this.commands = connection.async(); + } + } + + @Override + public boolean isConnected() { + return this.connection != null; + } + + @Override + public void sync() { + try { + LettuceFutures.awaitAll(60, TimeUnit.SECONDS, futures.toArray(new RedisFuture[0])); + } finally { + futures.clear(); + } + } + + @Override + public void pexpire(byte[] key, Long expiryMillis) { + futures.add(commands.pexpire(key, expiryMillis)); + } + + @Override + public void append(byte[] key, byte[] value) { + futures.add(commands.append(key, value)); + } + + @Override + public void set(byte[] key, byte[] value) { + futures.add(commands.set(key, value)); + } + + @Override + public void lpush(byte[] key, byte[] value) { + futures.add(commands.lpush(key, value)); + } + + @Override + public void rpush(byte[] key, byte[] value) { + futures.add(commands.rpush(key, value)); + } + + @Override + public void sadd(byte[] key, byte[] value) { + futures.add(commands.sadd(key, value)); + } + + @Override + public void zadd(byte[] key, Long score, byte[] value) { + futures.add(commands.zadd(key, score, value)); + } +} diff --git a/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisIngestionClient.java b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisIngestionClient.java new file mode 100644 index 00000000000..5a0b54e6970 --- /dev/null +++ b/storage/connectors/rediscluster/src/main/java/feast/storage/connectors/rediscluster/writer/RedisIngestionClient.java @@ -0,0 +1,49 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.writer; + +import feast.storage.common.retry.BackOffExecutor; +import java.io.Serializable; + +public interface RedisIngestionClient extends Serializable { + + void setup(); + + BackOffExecutor getBackOffExecutor(); + + void shutdown(); + + void connect(); + + boolean isConnected(); + + void sync(); + + void pexpire(byte[] key, Long expiryMillis); + + void append(byte[] key, byte[] value); + + void set(byte[] key, byte[] value); + + void lpush(byte[] key, byte[] value); + + void rpush(byte[] key, byte[] value); + + void sadd(byte[] key, byte[] value); + + void zadd(byte[] key, Long score, byte[] value); +} diff --git a/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetrieverTest.java b/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetrieverTest.java new file mode 100644 index 00000000000..567e92a3d41 --- /dev/null +++ b/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/retriever/RedisClusterOnlineRetrieverTest.java @@ -0,0 +1,263 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2020 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.retriever; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.mockito.Mockito.when; +import static org.mockito.MockitoAnnotations.initMocks; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.protobuf.AbstractMessageLite; +import com.google.protobuf.Duration; +import com.google.protobuf.Timestamp; +import feast.core.FeatureSetProto.EntitySpec; +import feast.core.FeatureSetProto.FeatureSetSpec; +import feast.core.FeatureSetProto.FeatureSpec; +import feast.serving.ServingAPIProto.FeatureReference; +import feast.serving.ServingAPIProto.GetOnlineFeaturesRequest.EntityRow; +import feast.storage.RedisProto.RedisKey; +import feast.storage.api.retriever.FeatureSetRequest; +import feast.storage.api.retriever.OnlineRetriever; +import feast.types.FeatureRowProto.FeatureRow; +import feast.types.FieldProto.Field; +import feast.types.ValueProto.Value; +import io.lettuce.core.KeyValue; +import io.lettuce.core.cluster.api.StatefulRedisClusterConnection; +import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; +import java.util.List; +import java.util.Optional; +import java.util.stream.Collectors; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; + +public class RedisClusterOnlineRetrieverTest { + + @Mock StatefulRedisClusterConnection connection; + + @Mock RedisAdvancedClusterCommands syncCommands; + + private OnlineRetriever redisClusterOnlineRetriever; + private byte[][] redisKeyList; + + @Before + public void setUp() { + initMocks(this); + when(connection.sync()).thenReturn(syncCommands); + redisClusterOnlineRetriever = RedisClusterOnlineRetriever.create(connection); + redisKeyList = + Lists.newArrayList( + RedisKey.newBuilder() + .setFeatureSet("project/featureSet:1") + .addAllEntities( + Lists.newArrayList( + Field.newBuilder().setName("entity1").setValue(intValue(1)).build(), + Field.newBuilder().setName("entity2").setValue(strValue("a")).build())) + .build(), + RedisKey.newBuilder() + .setFeatureSet("project/featureSet:1") + .addAllEntities( + Lists.newArrayList( + Field.newBuilder().setName("entity1").setValue(intValue(2)).build(), + Field.newBuilder().setName("entity2").setValue(strValue("b")).build())) + .build()) + .stream() + .map(AbstractMessageLite::toByteArray) + .collect(Collectors.toList()) + .toArray(new byte[0][0]); + } + + @Test + public void shouldReturnResponseWithValuesIfKeysPresent() { + FeatureSetRequest featureSetRequest = + FeatureSetRequest.newBuilder() + .setSpec(getFeatureSetSpec()) + .addFeatureReference( + FeatureReference.newBuilder() + .setName("feature1") + .setVersion(1) + .setProject("project") + .build()) + .addFeatureReference( + FeatureReference.newBuilder() + .setName("feature2") + .setVersion(1) + .setProject("project") + .build()) + .build(); + List entityRows = + ImmutableList.of( + EntityRow.newBuilder() + .setEntityTimestamp(Timestamp.newBuilder().setSeconds(100)) + .putFields("entity1", intValue(1)) + .putFields("entity2", strValue("a")) + .build(), + EntityRow.newBuilder() + .setEntityTimestamp(Timestamp.newBuilder().setSeconds(100)) + .putFields("entity1", intValue(2)) + .putFields("entity2", strValue("b")) + .build()); + + List featureRows = + Lists.newArrayList( + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setValue(intValue(1)).build(), + Field.newBuilder().setValue(intValue(1)).build())) + .build(), + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setValue(intValue(2)).build(), + Field.newBuilder().setValue(intValue(2)).build())) + .build()); + + List> featureRowBytes = + featureRows.stream() + .map(x -> KeyValue.from(new byte[1], Optional.of(x.toByteArray()))) + .collect(Collectors.toList()); + + redisClusterOnlineRetriever = RedisClusterOnlineRetriever.create(connection); + when(connection.sync()).thenReturn(syncCommands); + when(syncCommands.mget(redisKeyList)).thenReturn(featureRowBytes); + + List> expected = + List.of( + Lists.newArrayList( + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .setFeatureSet("project/featureSet:1") + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setName("feature1").setValue(intValue(1)).build(), + Field.newBuilder().setName("feature2").setValue(intValue(1)).build())) + .build(), + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .setFeatureSet("project/featureSet:1") + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setName("feature1").setValue(intValue(2)).build(), + Field.newBuilder().setName("feature2").setValue(intValue(2)).build())) + .build())); + + List> actual = + redisClusterOnlineRetriever.getOnlineFeatures(entityRows, List.of(featureSetRequest)); + assertThat(actual, equalTo(expected)); + } + + @Test + public void shouldReturnResponseWithUnsetValuesIfKeysNotPresent() { + FeatureSetRequest featureSetRequest = + FeatureSetRequest.newBuilder() + .setSpec(getFeatureSetSpec()) + .addFeatureReference( + FeatureReference.newBuilder() + .setName("feature1") + .setVersion(1) + .setProject("project") + .build()) + .addFeatureReference( + FeatureReference.newBuilder() + .setName("feature2") + .setVersion(1) + .setProject("project") + .build()) + .build(); + List entityRows = + ImmutableList.of( + EntityRow.newBuilder() + .setEntityTimestamp(Timestamp.newBuilder().setSeconds(100)) + .putFields("entity1", intValue(1)) + .putFields("entity2", strValue("a")) + .build(), + EntityRow.newBuilder() + .setEntityTimestamp(Timestamp.newBuilder().setSeconds(100)) + .putFields("entity1", intValue(2)) + .putFields("entity2", strValue("b")) + .build()); + + List featureRows = + Lists.newArrayList( + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setValue(intValue(1)).build(), + Field.newBuilder().setValue(intValue(1)).build())) + .build()); + + List> featureRowBytes = + featureRows.stream() + .map(x -> KeyValue.from(new byte[1], Optional.of(x.toByteArray()))) + .collect(Collectors.toList()); + featureRowBytes.add(null); + + redisClusterOnlineRetriever = RedisClusterOnlineRetriever.create(connection); + when(connection.sync()).thenReturn(syncCommands); + when(syncCommands.mget(redisKeyList)).thenReturn(featureRowBytes); + + List> expected = + List.of( + Lists.newArrayList( + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(100)) + .setFeatureSet("project/featureSet:1") + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setName("feature1").setValue(intValue(1)).build(), + Field.newBuilder().setName("feature2").setValue(intValue(1)).build())) + .build(), + FeatureRow.newBuilder() + .setFeatureSet("project/featureSet:1") + .addAllFields( + Lists.newArrayList( + Field.newBuilder().setName("feature1").build(), + Field.newBuilder().setName("feature2").build())) + .build())); + + List> actual = + redisClusterOnlineRetriever.getOnlineFeatures(entityRows, List.of(featureSetRequest)); + assertThat(actual, equalTo(expected)); + } + + private Value intValue(int val) { + return Value.newBuilder().setInt64Val(val).build(); + } + + private Value strValue(String val) { + return Value.newBuilder().setStringVal(val).build(); + } + + private FeatureSetSpec getFeatureSetSpec() { + return FeatureSetSpec.newBuilder() + .setProject("project") + .setName("featureSet") + .setVersion(1) + .addEntities(EntitySpec.newBuilder().setName("entity1")) + .addEntities(EntitySpec.newBuilder().setName("entity2")) + .addFeatures(FeatureSpec.newBuilder().setName("feature1")) + .addFeatures(FeatureSpec.newBuilder().setName("feature2")) + .setMaxAge(Duration.newBuilder().setSeconds(30)) // default + .build(); + } +} diff --git a/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSinkTest.java b/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSinkTest.java new file mode 100644 index 00000000000..cc1993636ee --- /dev/null +++ b/storage/connectors/rediscluster/src/test/java/feast/storage/connectors/rediscluster/writer/RedisClusterFeatureSinkTest.java @@ -0,0 +1,506 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2019 The Feast Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package feast.storage.connectors.rediscluster.writer; + +import static feast.storage.common.testing.TestUtil.field; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.Timestamp; +import feast.core.FeatureSetProto.EntitySpec; +import feast.core.FeatureSetProto.FeatureSetSpec; +import feast.core.FeatureSetProto.FeatureSpec; +import feast.core.StoreProto.Store.RedisClusterConfig; +import feast.storage.RedisProto.RedisKey; +import feast.types.FeatureRowProto.FeatureRow; +import feast.types.FieldProto.Field; +import feast.types.ValueProto.Value; +import feast.types.ValueProto.ValueType.Enum; +import io.lettuce.core.RedisURI; +import io.lettuce.core.cluster.RedisClusterClient; +import io.lettuce.core.cluster.api.StatefulRedisClusterConnection; +import io.lettuce.core.cluster.api.sync.RedisClusterCommands; +import io.lettuce.core.codec.ByteArrayCodec; +import java.io.File; +import java.io.IOException; +import java.nio.file.Paths; +import java.util.*; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import net.ishiis.redis.unit.RedisCluster; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; + +public class RedisClusterFeatureSinkTest { + @Rule public transient TestPipeline p = TestPipeline.create(); + + private static String REDIS_CLUSTER_HOST = "localhost"; + private static int REDIS_CLUSTER_PORT1 = 6380; + private static int REDIS_CLUSTER_PORT2 = 6381; + private static int REDIS_CLUSTER_PORT3 = 6382; + private static String CONNECTION_STRING = "localhost:6380,localhost:6381,localhost:6382"; + private RedisCluster redisCluster; + private RedisClusterClient redisClusterClient; + private RedisClusterCommands redisClusterCommands; + + private RedisClusterFeatureSink redisClusterFeatureSink; + + @Before + public void setUp() throws IOException { + redisCluster = new RedisCluster(REDIS_CLUSTER_PORT1, REDIS_CLUSTER_PORT2, REDIS_CLUSTER_PORT3); + redisCluster.start(); + redisClusterClient = + RedisClusterClient.create( + Arrays.asList( + RedisURI.create(REDIS_CLUSTER_HOST, REDIS_CLUSTER_PORT1), + RedisURI.create(REDIS_CLUSTER_HOST, REDIS_CLUSTER_PORT2), + RedisURI.create(REDIS_CLUSTER_HOST, REDIS_CLUSTER_PORT3))); + StatefulRedisClusterConnection connection = + redisClusterClient.connect(new ByteArrayCodec()); + redisClusterCommands = connection.sync(); + redisClusterCommands.setTimeout(java.time.Duration.ofMillis(600000)); + + FeatureSetSpec spec1 = + FeatureSetSpec.newBuilder() + .setName("fs") + .setVersion(1) + .setProject("myproject") + .addEntities(EntitySpec.newBuilder().setName("entity").setValueType(Enum.INT64).build()) + .addFeatures( + FeatureSpec.newBuilder().setName("feature").setValueType(Enum.STRING).build()) + .build(); + + FeatureSetSpec spec2 = + FeatureSetSpec.newBuilder() + .setName("feature_set") + .setProject("myproject") + .setVersion(1) + .addEntities( + EntitySpec.newBuilder() + .setName("entity_id_primary") + .setValueType(Enum.INT32) + .build()) + .addEntities( + EntitySpec.newBuilder() + .setName("entity_id_secondary") + .setValueType(Enum.STRING) + .build()) + .addFeatures( + FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build()) + .addFeatures( + FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build()) + .build(); + + Map specMap = + ImmutableMap.of("myproject/fs:1", spec1, "myproject/feature_set:1", spec2); + RedisClusterConfig redisClusterConfig = + RedisClusterConfig.newBuilder() + .setConnectionString(CONNECTION_STRING) + .setInitialBackoffMs(2000) + .setMaxRetries(4) + .build(); + + redisClusterFeatureSink = + RedisClusterFeatureSink.builder() + .setFeatureSetSpecs(specMap) + .setRedisClusterConfig(redisClusterConfig) + .build(); + } + + static boolean deleteDirectory(File directoryToBeDeleted) { + File[] allContents = directoryToBeDeleted.listFiles(); + if (allContents != null) { + for (File file : allContents) { + deleteDirectory(file); + } + } + return directoryToBeDeleted.delete(); + } + + @After + public void teardown() { + redisClusterClient.shutdown(); + redisCluster.stop(); + deleteDirectory(new File(String.valueOf(Paths.get(System.getProperty("user.dir"), ".redis")))); + } + + @Test + public void shouldWriteToRedis() { + + HashMap kvs = new LinkedHashMap<>(); + kvs.put( + RedisKey.newBuilder() + .setFeatureSet("myproject/fs:1") + .addEntities(field("entity", 1, Enum.INT64)) + .build(), + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.getDefaultInstance()) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("one"))) + .build()); + kvs.put( + RedisKey.newBuilder() + .setFeatureSet("myproject/fs:1") + .addEntities(field("entity", 2, Enum.INT64)) + .build(), + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.getDefaultInstance()) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("two"))) + .build()); + + List featureRows = + ImmutableList.of( + FeatureRow.newBuilder() + .setFeatureSet("myproject/fs:1") + .addFields(field("entity", 1, Enum.INT64)) + .addFields(field("feature", "one", Enum.STRING)) + .build(), + FeatureRow.newBuilder() + .setFeatureSet("myproject/fs:1") + .addFields(field("entity", 2, Enum.INT64)) + .addFields(field("feature", "two", Enum.STRING)) + .build()); + + p.apply(Create.of(featureRows)).apply(redisClusterFeatureSink.writer()); + p.run(); + + kvs.forEach( + (key, value) -> { + byte[] actual = redisClusterCommands.get(key.toByteArray()); + assertThat(actual, equalTo(value.toByteArray())); + }); + } + + @Test(timeout = 15000) + public void shouldRetryFailConnection() throws InterruptedException { + HashMap kvs = new LinkedHashMap<>(); + kvs.put( + RedisKey.newBuilder() + .setFeatureSet("myproject/fs:1") + .addEntities(field("entity", 1, Enum.INT64)) + .build(), + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.getDefaultInstance()) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("one"))) + .build()); + + List featureRows = + ImmutableList.of( + FeatureRow.newBuilder() + .setFeatureSet("myproject/fs:1") + .addFields(field("entity", 1, Enum.INT64)) + .addFields(field("feature", "one", Enum.STRING)) + .build()); + + PCollection failedElementCount = + p.apply(Create.of(featureRows)) + .apply(redisClusterFeatureSink.writer()) + .getFailedInserts() + .apply(Count.globally()); + + redisCluster.stop(); + final ScheduledThreadPoolExecutor redisRestartExecutor = new ScheduledThreadPoolExecutor(1); + ScheduledFuture scheduledRedisRestart = + redisRestartExecutor.schedule( + () -> { + redisCluster.start(); + }, + 3, + TimeUnit.SECONDS); + + PAssert.that(failedElementCount).containsInAnyOrder(0L); + p.run(); + scheduledRedisRestart.cancel(true); + + kvs.forEach( + (key, value) -> { + byte[] actual = redisClusterCommands.get(key.toByteArray()); + assertThat(actual, equalTo(value.toByteArray())); + }); + } + + @Test + public void shouldProduceFailedElementIfRetryExceeded() { + RedisClusterConfig redisClusterConfig = + RedisClusterConfig.newBuilder() + .setConnectionString(CONNECTION_STRING) + .setInitialBackoffMs(2000) + .setMaxRetries(1) + .build(); + + FeatureSetSpec spec1 = + FeatureSetSpec.newBuilder() + .setName("fs") + .setVersion(1) + .setProject("myproject") + .addEntities(EntitySpec.newBuilder().setName("entity").setValueType(Enum.INT64).build()) + .addFeatures( + FeatureSpec.newBuilder().setName("feature").setValueType(Enum.STRING).build()) + .build(); + Map specMap = ImmutableMap.of("myproject/fs:1", spec1); + redisClusterFeatureSink = + RedisClusterFeatureSink.builder() + .setFeatureSetSpecs(specMap) + .setRedisClusterConfig(redisClusterConfig) + .build(); + redisCluster.stop(); + + List featureRows = + ImmutableList.of( + FeatureRow.newBuilder() + .setFeatureSet("myproject/fs:1") + .addFields(field("entity", 1, Enum.INT64)) + .addFields(field("feature", "one", Enum.STRING)) + .build()); + + PCollection failedElementCount = + p.apply(Create.of(featureRows)) + .apply(redisClusterFeatureSink.writer()) + .getFailedInserts() + .apply(Count.globally()); + + PAssert.that(failedElementCount).containsInAnyOrder(1L); + p.run(); + } + + @Test + public void shouldConvertRowWithDuplicateEntitiesToValidKey() { + + FeatureRow offendingRow = + FeatureRow.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addFields( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(2))) + .addFields( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .addFields( + Field.newBuilder() + .setName("feature_1") + .setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields( + Field.newBuilder() + .setName("feature_2") + .setValue(Value.newBuilder().setInt64Val(1001))) + .build(); + + RedisKey expectedKey = + RedisKey.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .addEntities( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addEntities( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .build(); + + FeatureRow expectedValue = + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setInt64Val(1001))) + .build(); + + p.apply(Create.of(offendingRow)).apply(redisClusterFeatureSink.writer()); + + p.run(); + + byte[] actual = redisClusterCommands.get(expectedKey.toByteArray()); + assertThat(actual, equalTo(expectedValue.toByteArray())); + } + + @Test + public void shouldConvertRowWithOutOfOrderFieldsToValidKey() { + FeatureRow offendingRow = + FeatureRow.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .addFields( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addFields( + Field.newBuilder() + .setName("feature_2") + .setValue(Value.newBuilder().setInt64Val(1001))) + .addFields( + Field.newBuilder() + .setName("feature_1") + .setValue(Value.newBuilder().setStringVal("strValue1"))) + .build(); + + RedisKey expectedKey = + RedisKey.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .addEntities( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addEntities( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .build(); + + List expectedFields = + Arrays.asList( + Field.newBuilder().setValue(Value.newBuilder().setStringVal("strValue1")).build(), + Field.newBuilder().setValue(Value.newBuilder().setInt64Val(1001)).build()); + FeatureRow expectedValue = + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addAllFields(expectedFields) + .build(); + + p.apply(Create.of(offendingRow)).apply(redisClusterFeatureSink.writer()); + + p.run(); + + byte[] actual = redisClusterCommands.get(expectedKey.toByteArray()); + assertThat(actual, equalTo(expectedValue.toByteArray())); + } + + @Test + public void shouldMergeDuplicateFeatureFields() { + FeatureRow featureRowWithDuplicatedFeatureFields = + FeatureRow.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addFields( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .addFields( + Field.newBuilder() + .setName("feature_1") + .setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields( + Field.newBuilder() + .setName("feature_1") + .setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields( + Field.newBuilder() + .setName("feature_2") + .setValue(Value.newBuilder().setInt64Val(1001))) + .build(); + + RedisKey expectedKey = + RedisKey.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .addEntities( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addEntities( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .build(); + + FeatureRow expectedValue = + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setInt64Val(1001))) + .build(); + + p.apply(Create.of(featureRowWithDuplicatedFeatureFields)) + .apply(redisClusterFeatureSink.writer()); + + p.run(); + + byte[] actual = redisClusterCommands.get(expectedKey.toByteArray()); + assertThat(actual, equalTo(expectedValue.toByteArray())); + } + + @Test + public void shouldPopulateMissingFeatureValuesWithDefaultInstance() { + FeatureRow featureRowWithDuplicatedFeatureFields = + FeatureRow.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addFields( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .addFields( + Field.newBuilder() + .setName("feature_1") + .setValue(Value.newBuilder().setStringVal("strValue1"))) + .build(); + + RedisKey expectedKey = + RedisKey.newBuilder() + .setFeatureSet("myproject/feature_set:1") + .addEntities( + Field.newBuilder() + .setName("entity_id_primary") + .setValue(Value.newBuilder().setInt32Val(1))) + .addEntities( + Field.newBuilder() + .setName("entity_id_secondary") + .setValue(Value.newBuilder().setStringVal("a"))) + .build(); + + FeatureRow expectedValue = + FeatureRow.newBuilder() + .setEventTimestamp(Timestamp.newBuilder().setSeconds(10)) + .addFields(Field.newBuilder().setValue(Value.newBuilder().setStringVal("strValue1"))) + .addFields(Field.newBuilder().setValue(Value.getDefaultInstance())) + .build(); + + p.apply(Create.of(featureRowWithDuplicatedFeatureFields)) + .apply(redisClusterFeatureSink.writer()); + + p.run(); + + byte[] actual = redisClusterCommands.get(expectedKey.toByteArray()); + assertThat(actual, equalTo(expectedValue.toByteArray())); + } +}