diff --git a/common-test/src/main/java/feast/common/it/DataGenerator.java b/common-test/src/main/java/feast/common/it/DataGenerator.java index ef31f54..8a0dbb0 100644 --- a/common-test/src/main/java/feast/common/it/DataGenerator.java +++ b/common-test/src/main/java/feast/common/it/DataGenerator.java @@ -249,6 +249,10 @@ public static ValueProto.Value createDoubleValue(double value) { return ValueProto.Value.newBuilder().setDoubleVal(value).build(); } + public static ValueProto.Value createInt32Value(int value) { + return ValueProto.Value.newBuilder().setInt32Val(value).build(); + } + public static ValueProto.Value createInt64Value(long value) { return ValueProto.Value.newBuilder().setInt64Val(value).build(); } diff --git a/serving/pom.xml b/serving/pom.xml index 6eca569..dfcc6e5 100644 --- a/serving/pom.xml +++ b/serving/pom.xml @@ -90,6 +90,12 @@ ${project.version} + + dev.feast + feast-storage-connector-cassandra + ${project.version} + + dev.feast feast-common diff --git a/serving/src/main/java/feast/serving/config/FeastProperties.java b/serving/src/main/java/feast/serving/config/FeastProperties.java index 6794b2d..b9029a0 100644 --- a/serving/src/main/java/feast/serving/config/FeastProperties.java +++ b/serving/src/main/java/feast/serving/config/FeastProperties.java @@ -27,6 +27,7 @@ import feast.common.auth.credentials.CoreAuthenticationProperties; import feast.common.logging.config.LoggingProperties; import feast.storage.connectors.bigtable.retriever.BigTableStoreConfig; +import feast.storage.connectors.cassandra.retriever.CassandraStoreConfig; import feast.storage.connectors.redis.retriever.RedisClusterStoreConfig; import feast.storage.connectors.redis.retriever.RedisStoreConfig; import io.lettuce.core.ReadFrom; @@ -270,7 +271,7 @@ public void setName(String name) { } /** - * Gets the store type. Example are REDIS, REDIS_CLUSTER or BIGTABLE + * Gets the store type. Example are REDIS, REDIS_CLUSTER, BIGTABLE or CASSANDRA * * @return the store type as a String. */ @@ -316,6 +317,13 @@ public BigTableStoreConfig getBigtableConfig() { return new BigTableStoreConfig(this.config.get("project_id"), this.config.get("instance_id")); } + public CassandraStoreConfig getCassandraConfig() { + return new CassandraStoreConfig( + this.config.get("connection_string"), + this.config.get("data_center"), + this.config.get("keyspace")); + } + /** * Sets the store config. Please protos/feast/core/Store.proto for the specific options for each * store. @@ -329,6 +337,7 @@ public void setConfig(Map config) { public enum StoreType { BIGTABLE, + CASSANDRA, REDIS, REDIS_CLUSTER; } diff --git a/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java b/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java index d6ed6db..4c26f4a 100644 --- a/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java +++ b/serving/src/main/java/feast/serving/config/ServingServiceConfigV2.java @@ -16,6 +16,8 @@ */ package feast.serving.config; +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.CqlSessionBuilder; import com.google.cloud.bigtable.data.v2.BigtableDataClient; import com.google.cloud.bigtable.data.v2.BigtableDataSettings; import feast.serving.service.OnlineServingServiceV2; @@ -24,9 +26,15 @@ import feast.storage.api.retriever.OnlineRetrieverV2; import feast.storage.connectors.bigtable.retriever.BigTableOnlineRetriever; import feast.storage.connectors.bigtable.retriever.BigTableStoreConfig; +import feast.storage.connectors.cassandra.retriever.CassandraOnlineRetriever; +import feast.storage.connectors.cassandra.retriever.CassandraStoreConfig; import feast.storage.connectors.redis.retriever.*; import io.opentracing.Tracer; import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; import org.slf4j.Logger; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationContext; @@ -77,6 +85,32 @@ public ServingServiceV2 servingServiceV2( OnlineRetrieverV2 bigtableRetriever = new BigTableOnlineRetriever(bigtableClient); servingService = new OnlineServingServiceV2(bigtableRetriever, specService, tracer); break; + case CASSANDRA: + CassandraStoreConfig config = feastProperties.getActiveStore().getCassandraConfig(); + String connectionString = config.getConnectionString(); + String dataCenter = config.getDataCenter(); + String keySpace = config.getKeySpace(); + + List contactPoints = + Arrays.stream(connectionString.split(",")) + .map(String::trim) + .map(cs -> cs.split(":")) + .map( + hostPort -> { + int port = hostPort.length > 1 ? Integer.parseInt(hostPort[1]) : 9042; + return new InetSocketAddress(hostPort[0], port); + }) + .collect(Collectors.toList()); + + CqlSession session = + new CqlSessionBuilder() + .addContactPoints(contactPoints) + .withLocalDatacenter(dataCenter) + .withKeyspace(keySpace) + .build(); + OnlineRetrieverV2 cassandraRetriever = new CassandraOnlineRetriever(session); + servingService = new OnlineServingServiceV2(cassandraRetriever, specService, tracer); + break; } return servingService; diff --git a/serving/src/main/resources/application.yml b/serving/src/main/resources/application.yml index b23a345..91e27de 100644 --- a/serving/src/main/resources/application.yml +++ b/serving/src/main/resources/application.yml @@ -59,6 +59,12 @@ feast: config: project_id: instance_id: + - name: cassandra + type: CASSANDRA + config: + connection_string: localhost:9042 + data_center: datacenter1 + keyspace: feast tracing: # If true, Feast will provide tracing data (using OpenTracing API) for various RPC method calls # which can be useful to debug performance issues and perform benchmarking diff --git a/serving/src/test/java/feast/serving/it/BaseAuthIT.java b/serving/src/test/java/feast/serving/it/BaseAuthIT.java index ab6e169..d49ac41 100644 --- a/serving/src/test/java/feast/serving/it/BaseAuthIT.java +++ b/serving/src/test/java/feast/serving/it/BaseAuthIT.java @@ -54,6 +54,13 @@ public class BaseAuthIT { static final String BIGTABLE = "bigtable_1"; static final int BIGTABLE_PORT = 8086; + static final String CASSANDRA = "cassandra_1"; + static final int CASSANDRA_PORT = 9042; + static final String CASSANDRA_DATACENTER = "datacenter1"; + static final String CASSANDRA_KEYSPACE = "feast"; + static final String CASSANDRA_SCHEMA_TABLE = "feast_schema_reference"; + static final String CASSANDRA_ENTITY_KEY = "key"; + static final int FEAST_CORE_PORT = 6565; @DynamicPropertySource @@ -72,14 +79,40 @@ static void properties(DynamicPropertyRegistry registry) { } }); registry.add("feast.stores[0].config.port", () -> REDIS_PORT); - registry.add("feast.stores[0].subscriptions[0].name", () -> "*"); - registry.add("feast.stores[0].subscriptions[0].project", () -> "*"); registry.add("feast.stores[1].name", () -> "bigtable"); registry.add("feast.stores[1].type", () -> "BIGTABLE"); registry.add("feast.stores[1].config.project_id", () -> "test-project"); registry.add("feast.stores[1].config.instance_id", () -> "test-instance"); + registry.add("feast.stores[2].name", () -> "cassandra"); + registry.add("feast.stores[2].type", () -> "CASSANDRA"); + registry.add( + "feast.stores[2].config.host", + () -> { + try { + return InetAddress.getLocalHost().getHostAddress(); + } catch (UnknownHostException e) { + e.printStackTrace(); + return ""; + } + }); + + registry.add( + "feast.stores[2].config.connection_string", + () -> { + String hostAddress = ""; + try { + hostAddress = InetAddress.getLocalHost().getHostAddress(); + } catch (UnknownHostException e) { + e.printStackTrace(); + } + + return String.format("%s:%s", hostAddress, CASSANDRA_PORT); + }); + registry.add("feast.stores[2].config.data_center", () -> CASSANDRA_DATACENTER); + registry.add("feast.stores[2].config.keyspace", () -> CASSANDRA_KEYSPACE); + registry.add("feast.core-authentication.options.oauth_url", () -> TOKEN_URL); registry.add("feast.core-authentication.options.grant_type", () -> GRANT_TYPE); registry.add("feast.core-authentication.options.client_id", () -> CLIENT_ID); diff --git a/serving/src/test/java/feast/serving/it/ServingServiceCassandraIT.java b/serving/src/test/java/feast/serving/it/ServingServiceCassandraIT.java new file mode 100644 index 0000000..93ee5f5 --- /dev/null +++ b/serving/src/test/java/feast/serving/it/ServingServiceCassandraIT.java @@ -0,0 +1,728 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2021 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.serving.it; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.cql.PreparedStatement; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.hash.Hashing; +import feast.common.it.DataGenerator; +import feast.common.models.FeatureV2; +import feast.proto.core.EntityProto; +import feast.proto.serving.ServingAPIProto.FeatureReferenceV2; +import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequestV2; +import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesResponse; +import feast.proto.serving.ServingServiceGrpc; +import feast.proto.types.ValueProto; +import io.grpc.ManagedChannel; +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.IOException; +import java.net.InetSocketAddress; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.generic.GenericRecordBuilder; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; +import org.junit.ClassRule; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.test.context.ActiveProfiles; +import org.springframework.test.context.DynamicPropertyRegistry; +import org.springframework.test.context.DynamicPropertySource; +import org.testcontainers.containers.DockerComposeContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +@ActiveProfiles("it") +@SpringBootTest( + webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT, + properties = { + "feast.core-cache-refresh-interval=1", + "feast.active_store=cassandra", + "spring.main.allow-bean-definition-overriding=true" + }) +@Testcontainers +public class ServingServiceCassandraIT extends BaseAuthIT { + + static final Map options = new HashMap<>(); + static CoreSimpleAPIClient coreClient; + static ServingServiceGrpc.ServingServiceBlockingStub servingStub; + + static CqlSession cqlSession; + static final int FEAST_SERVING_PORT = 6570; + + static final FeatureReferenceV2 feature1Reference = + DataGenerator.createFeatureReference("rides", "trip_cost"); + static final FeatureReferenceV2 feature2Reference = + DataGenerator.createFeatureReference("rides", "trip_distance"); + static final FeatureReferenceV2 feature3Reference = + DataGenerator.createFeatureReference("rides", "trip_empty"); + static final FeatureReferenceV2 feature4Reference = + DataGenerator.createFeatureReference("rides", "trip_wrong_type"); + + @ClassRule @Container + public static DockerComposeContainer environment = + new DockerComposeContainer( + new File("src/test/resources/docker-compose/docker-compose-cassandra-it.yml")) + .withExposedService( + CORE, + FEAST_CORE_PORT, + Wait.forLogMessage(".*gRPC Server started.*\\n", 1) + .withStartupTimeout(Duration.ofMinutes(SERVICE_START_MAX_WAIT_TIME_IN_MINUTES))) + .withExposedService(CASSANDRA, CASSANDRA_PORT); + + @DynamicPropertySource + static void initialize(DynamicPropertyRegistry registry) { + registry.add("grpc.server.port", () -> FEAST_SERVING_PORT); + } + + @BeforeAll + static void globalSetup() throws IOException { + coreClient = TestUtils.getApiClientForCore(FEAST_CORE_PORT); + servingStub = TestUtils.getServingServiceStub(false, FEAST_SERVING_PORT, null); + + cqlSession = + CqlSession.builder() + .addContactPoint( + new InetSocketAddress( + environment.getServiceHost("cassandra_1", CASSANDRA_PORT), + environment.getServicePort("cassandra_1", CASSANDRA_PORT))) + .withLocalDatacenter(CASSANDRA_DATACENTER) + .build(); + + /** Feast resource creation Workflow */ + String projectName = "default"; + // Apply Entity (driver_id) + String driverEntityName = "driver_id"; + String driverEntityDescription = "My driver id"; + ValueProto.ValueType.Enum driverEntityType = ValueProto.ValueType.Enum.INT64; + EntityProto.EntitySpecV2 driverEntitySpec = + EntityProto.EntitySpecV2.newBuilder() + .setName(driverEntityName) + .setDescription(driverEntityDescription) + .setValueType(driverEntityType) + .build(); + TestUtils.applyEntity(coreClient, projectName, driverEntitySpec); + + // Apply Entity (merchant_id) + String merchantEntityName = "merchant_id"; + String merchantEntityDescription = "My driver id"; + ValueProto.ValueType.Enum merchantEntityType = ValueProto.ValueType.Enum.INT64; + EntityProto.EntitySpecV2 merchantEntitySpec = + EntityProto.EntitySpecV2.newBuilder() + .setName(merchantEntityName) + .setDescription(merchantEntityDescription) + .setValueType(merchantEntityType) + .build(); + TestUtils.applyEntity(coreClient, projectName, merchantEntitySpec); + + // Apply FeatureTable (rides) + String ridesFeatureTableName = "rides"; + ImmutableList ridesEntities = ImmutableList.of(driverEntityName); + ImmutableMap ridesFeatures = + ImmutableMap.of( + "trip_cost", + ValueProto.ValueType.Enum.INT32, + "trip_distance", + ValueProto.ValueType.Enum.DOUBLE, + "trip_empty", + ValueProto.ValueType.Enum.DOUBLE, + "trip_wrong_type", + ValueProto.ValueType.Enum.STRING); + TestUtils.applyFeatureTable( + coreClient, projectName, ridesFeatureTableName, ridesEntities, ridesFeatures, 7200); + + // Apply FeatureTable (food) + String foodFeatureTableName = "food"; + ImmutableList foodEntities = ImmutableList.of(driverEntityName); + ImmutableMap foodFeatures = + ImmutableMap.of( + "trip_cost", + ValueProto.ValueType.Enum.INT32, + "trip_distance", + ValueProto.ValueType.Enum.DOUBLE); + TestUtils.applyFeatureTable( + coreClient, projectName, foodFeatureTableName, foodEntities, foodFeatures, 7200); + + // Apply FeatureTable (rides_merchant) + String rideMerchantFeatureTableName = "rides_merchant"; + ImmutableList ridesMerchantEntities = + ImmutableList.of(driverEntityName, merchantEntityName); + TestUtils.applyFeatureTable( + coreClient, + projectName, + rideMerchantFeatureTableName, + ridesMerchantEntities, + ridesFeatures, + 7200); + + /** Create Cassandra Tables Workflow */ + String cassandraTableName = String.format("%s__%s", projectName, driverEntityName); + String compoundCassandraTableName = + String.format("%s__%s", projectName, String.join("__", ridesMerchantEntities)); + + cqlSession.execute(String.format("DROP KEYSPACE IF EXISTS %s", CASSANDRA_KEYSPACE)); + cqlSession.execute( + String.format( + "CREATE KEYSPACE %s WITH replication = \n" + + "{'class':'SimpleStrategy','replication_factor':'1'};", + CASSANDRA_KEYSPACE)); + + // Create Cassandra Tables + createCassandraTable(cassandraTableName); + createCassandraTable(compoundCassandraTableName); + + // Add column families + addCassandraTableColumn(cassandraTableName, ridesFeatureTableName); + addCassandraTableColumn(cassandraTableName, foodFeatureTableName); + addCassandraTableColumn(compoundCassandraTableName, rideMerchantFeatureTableName); + + /** Single Entity Ingestion Workflow */ + Schema ftSchema = + SchemaBuilder.record("DriverData") + .namespace(ridesFeatureTableName) + .fields() + .requiredInt(feature1Reference.getName()) + .requiredDouble(feature2Reference.getName()) + .nullableString(feature3Reference.getName(), "null") + .requiredString(feature4Reference.getName()) + .endRecord(); + byte[] schemaReference = + Hashing.murmur3_32().hashBytes(ftSchema.toString().getBytes()).asBytes(); + byte[] schemaKey = createSchemaKey(schemaReference); + + ingestBulk(ridesFeatureTableName, cassandraTableName, ftSchema, 20); + + Schema foodFtSchema = + SchemaBuilder.record("FoodDriverData") + .namespace(foodFeatureTableName) + .fields() + .requiredInt(feature1Reference.getName()) + .requiredDouble(feature2Reference.getName()) + .nullableString(feature3Reference.getName(), "null") + .requiredString(feature4Reference.getName()) + .endRecord(); + byte[] foodSchemaReference = + Hashing.murmur3_32().hashBytes(foodFtSchema.toString().getBytes()).asBytes(); + byte[] foodSchemaKey = createSchemaKey(foodSchemaReference); + + ingestBulk(foodFeatureTableName, cassandraTableName, foodFtSchema, 20); + + /** Compound Entity Ingestion Workflow */ + Schema compoundFtSchema = + SchemaBuilder.record("DriverMerchantData") + .namespace(rideMerchantFeatureTableName) + .fields() + .requiredLong(feature1Reference.getName()) + .requiredDouble(feature2Reference.getName()) + .nullableString(feature3Reference.getName(), "null") + .requiredString(feature4Reference.getName()) + .endRecord(); + byte[] compoundSchemaReference = + Hashing.murmur3_32().hashBytes(compoundFtSchema.toString().getBytes()).asBytes(); + + GenericRecord compoundEntityRecord = + new GenericRecordBuilder(compoundFtSchema) + .set("trip_cost", 10L) + .set("trip_distance", 5.5) + .set("trip_empty", null) + .set("trip_wrong_type", "wrong_type") + .build(); + ValueProto.Value driverEntityValue = ValueProto.Value.newBuilder().setInt64Val(1).build(); + ValueProto.Value merchantEntityValue = ValueProto.Value.newBuilder().setInt64Val(1234).build(); + ImmutableMap compoundEntityMap = + ImmutableMap.of( + driverEntityName, driverEntityValue, merchantEntityName, merchantEntityValue); + GetOnlineFeaturesRequestV2.EntityRow entityRow = + DataGenerator.createCompoundEntityRow(compoundEntityMap, 100); + byte[] compoundEntityFeatureKey = + ridesMerchantEntities.stream() + .map(entity -> DataGenerator.valueToString(entityRow.getFieldsMap().get(entity))) + .collect(Collectors.joining("#")) + .getBytes(); + byte[] compoundEntityFeatureValue = createEntityValue(compoundFtSchema, compoundEntityRecord); + byte[] compoundSchemaKey = createSchemaKey(compoundSchemaReference); + + ingestData( + rideMerchantFeatureTableName, + compoundCassandraTableName, + compoundEntityFeatureKey, + compoundEntityFeatureValue, + compoundSchemaKey); + + /** Schema Ingestion Workflow */ + cqlSession.execute( + String.format( + "CREATE TABLE %s.%s (schema_ref BLOB PRIMARY KEY, avro_schema BLOB);", + CASSANDRA_KEYSPACE, CASSANDRA_SCHEMA_TABLE)); + + ingestSchema(schemaKey, ftSchema); + ingestSchema(foodSchemaKey, foodFtSchema); + ingestSchema(compoundSchemaKey, compoundFtSchema); + + // set up options for call credentials + options.put("oauth_url", TOKEN_URL); + options.put(CLIENT_ID, CLIENT_ID); + options.put(CLIENT_SECRET, CLIENT_SECRET); + options.put("jwkEndpointURI", JWK_URI); + options.put("audience", AUDIENCE); + options.put("grant_type", GRANT_TYPE); + } + + private static byte[] createSchemaKey(byte[] schemaReference) throws IOException { + ByteArrayOutputStream concatOutputStream = new ByteArrayOutputStream(); + concatOutputStream.write(schemaReference); + byte[] schemaKey = concatOutputStream.toByteArray(); + + return schemaKey; + } + + private static byte[] createEntityValue(Schema schema, GenericRecord record) throws IOException { + // Entity-Feature Row + byte[] avroSerializedFeatures = recordToAvro(record, schema); + + ByteArrayOutputStream concatOutputStream = new ByteArrayOutputStream(); + concatOutputStream.write(avroSerializedFeatures); + byte[] entityFeatureValue = concatOutputStream.toByteArray(); + + return entityFeatureValue; + } + + private static void createCassandraTable(String cassandraTableName) { + cqlSession.execute( + String.format( + "CREATE TABLE %s.%s (key BLOB PRIMARY KEY);", CASSANDRA_KEYSPACE, cassandraTableName)); + } + + private static void addCassandraTableColumn(String cassandraTableName, String featureTableName) { + cqlSession.execute( + String.format( + "ALTER TABLE %s.%s ADD (%s BLOB, %s__schema_ref BLOB);", + CASSANDRA_KEYSPACE, cassandraTableName, featureTableName, featureTableName)); + } + + private static void ingestData( + String featureTableName, + String cassandraTableName, + byte[] entityFeatureKey, + byte[] entityFeatureValue, + byte[] schemaKey) { + PreparedStatement statement = + cqlSession.prepare( + String.format( + "INSERT INTO %s.%s (%s, %s__schema_ref, %s) VALUES (?, ?, ?)", + CASSANDRA_KEYSPACE, + cassandraTableName, + CASSANDRA_ENTITY_KEY, + featureTableName, + featureTableName)); + + cqlSession.execute( + statement.bind( + ByteBuffer.wrap(entityFeatureKey), + ByteBuffer.wrap(schemaKey), + ByteBuffer.wrap(entityFeatureValue))); + } + + private static void ingestBulk( + String featureTableName, String cassandraTableName, Schema schema, Integer counts) { + + IntStream.range(0, counts) + .forEach( + i -> { + try { + GenericRecord record = + new GenericRecordBuilder(schema) + .set("trip_cost", i) + .set("trip_distance", (double) i) + .set("trip_empty", null) + .set("trip_wrong_type", "test") + .build(); + byte[] schemaReference = + Hashing.murmur3_32().hashBytes(schema.toString().getBytes()).asBytes(); + + byte[] entityFeatureKey = + String.valueOf(DataGenerator.createInt64Value(i).getInt64Val()).getBytes(); + byte[] entityFeatureValue = createEntityValue(schema, record); + + byte[] schemaKey = createSchemaKey(schemaReference); + ingestData( + featureTableName, + cassandraTableName, + entityFeatureKey, + entityFeatureValue, + schemaKey); + } catch (IOException e) { + e.printStackTrace(); + } + }); + } + + private static void ingestSchema(byte[] schemaKey, Schema schema) { + PreparedStatement schemaStatement = + cqlSession.prepare( + String.format( + "INSERT INTO %s.%s (schema_ref, avro_schema) VALUES (?, ?);", + CASSANDRA_KEYSPACE, CASSANDRA_SCHEMA_TABLE)); + cqlSession.execute( + schemaStatement.bind( + ByteBuffer.wrap(schemaKey), ByteBuffer.wrap(schema.toString().getBytes()))); + } + + private static byte[] recordToAvro(GenericRecord datum, Schema schema) throws IOException { + GenericDatumWriter writer = new GenericDatumWriter<>(schema); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + Encoder encoder = EncoderFactory.get().binaryEncoder(output, null); + writer.write(datum, encoder); + encoder.flush(); + + return output.toByteArray(); + } + + @AfterAll + static void tearDown() { + ((ManagedChannel) servingStub.getChannel()).shutdown(); + } + + @Test + public void shouldRegisterSingleEntityAndGetOnlineFeatures() { + String projectName = "default"; + String entityName = "driver_id"; + ValueProto.Value entityValue = DataGenerator.createInt64Value(1); + + // Instantiate EntityRows + GetOnlineFeaturesRequestV2.EntityRow entityRow = + DataGenerator.createEntityRow(entityName, entityValue, 100); + ImmutableList entityRows = ImmutableList.of(entityRow); + + // Instantiate FeatureReferences + FeatureReferenceV2 featureReference = + DataGenerator.createFeatureReference("rides", "trip_cost"); + FeatureReferenceV2 notFoundFeatureReference = + DataGenerator.createFeatureReference("rides", "trip_transaction"); + + ImmutableList featureReferences = + ImmutableList.of(featureReference, notFoundFeatureReference); + + // Build GetOnlineFeaturesRequestV2 + GetOnlineFeaturesRequestV2 onlineFeatureRequest = + TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows); + GetOnlineFeaturesResponse featureResponse = + servingStub.getOnlineFeaturesV2(onlineFeatureRequest); + + ImmutableMap expectedValueMap = + ImmutableMap.of( + entityName, + entityValue, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(1), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue()); + + ImmutableMap expectedStatusMap = + ImmutableMap.of( + entityName, + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(featureReference), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND); + + GetOnlineFeaturesResponse.FieldValues expectedFieldValues = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap) + .putAllStatuses(expectedStatusMap) + .build(); + ImmutableList expectedFieldValuesList = + ImmutableList.of(expectedFieldValues); + + assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList()); + } + + @Test + public void shouldRegisterCompoundEntityAndGetOnlineFeatures() { + String projectName = "default"; + String driverEntityName = "driver_id"; + String merchantEntityName = "merchant_id"; + ValueProto.Value driverEntityValue = ValueProto.Value.newBuilder().setInt64Val(1).build(); + ValueProto.Value merchantEntityValue = ValueProto.Value.newBuilder().setInt64Val(1234).build(); + + ImmutableMap compoundEntityMap = + ImmutableMap.of( + driverEntityName, driverEntityValue, merchantEntityName, merchantEntityValue); + + // Instantiate EntityRows + GetOnlineFeaturesRequestV2.EntityRow entityRow = + DataGenerator.createCompoundEntityRow(compoundEntityMap, 100); + ImmutableList entityRows = ImmutableList.of(entityRow); + + // Instantiate FeatureReferences + FeatureReferenceV2 featureReference = + DataGenerator.createFeatureReference("rides", "trip_cost"); + FeatureReferenceV2 notFoundFeatureReference = + DataGenerator.createFeatureReference("rides", "trip_transaction"); + + ImmutableList featureReferences = + ImmutableList.of(featureReference, notFoundFeatureReference); + + // Build GetOnlineFeaturesRequestV2 + GetOnlineFeaturesRequestV2 onlineFeatureRequest = + TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows); + GetOnlineFeaturesResponse featureResponse = + servingStub.getOnlineFeaturesV2(onlineFeatureRequest); + + ImmutableMap expectedValueMap = + ImmutableMap.of( + driverEntityName, + driverEntityValue, + merchantEntityName, + merchantEntityValue, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(1), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue()); + + ImmutableMap expectedStatusMap = + ImmutableMap.of( + driverEntityName, + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + merchantEntityName, + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(featureReference), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND); + + GetOnlineFeaturesResponse.FieldValues expectedFieldValues = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap) + .putAllStatuses(expectedStatusMap) + .build(); + ImmutableList expectedFieldValuesList = + ImmutableList.of(expectedFieldValues); + + assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList()); + } + + @Test + public void shouldReturnCorrectRowCountAndOrder() { + // getOnlineFeatures Information + String projectName = "default"; + String entityName = "driver_id"; + ValueProto.Value entityValue1 = ValueProto.Value.newBuilder().setInt64Val(1).build(); + ValueProto.Value entityValue2 = ValueProto.Value.newBuilder().setInt64Val(2).build(); + ValueProto.Value entityValue3 = ValueProto.Value.newBuilder().setInt64Val(3).build(); + ValueProto.Value entityValue4 = ValueProto.Value.newBuilder().setInt64Val(4).build(); + + // Instantiate EntityRows + GetOnlineFeaturesRequestV2.EntityRow entityRow1 = + DataGenerator.createEntityRow(entityName, entityValue1, 100); + GetOnlineFeaturesRequestV2.EntityRow entityRow2 = + DataGenerator.createEntityRow(entityName, entityValue2, 100); + GetOnlineFeaturesRequestV2.EntityRow entityRow3 = + DataGenerator.createEntityRow(entityName, entityValue3, 100); + GetOnlineFeaturesRequestV2.EntityRow entityRow4 = + DataGenerator.createEntityRow(entityName, entityValue4, 100); + ImmutableList entityRows = + ImmutableList.of(entityRow1, entityRow2, entityRow4, entityRow3); + + // Instantiate FeatureReferences + FeatureReferenceV2 featureReference = + DataGenerator.createFeatureReference("rides", "trip_cost"); + FeatureReferenceV2 notFoundFeatureReference = + DataGenerator.createFeatureReference("rides", "trip_transaction"); + FeatureReferenceV2 emptyFeatureReference = + DataGenerator.createFeatureReference("rides", "trip_empty"); + + ImmutableList featureReferences = + ImmutableList.of(featureReference, notFoundFeatureReference, emptyFeatureReference); + + // Build GetOnlineFeaturesRequestV2 + GetOnlineFeaturesRequestV2 onlineFeatureRequest = + TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows); + GetOnlineFeaturesResponse featureResponse = + servingStub.getOnlineFeaturesV2(onlineFeatureRequest); + + ImmutableMap expectedValueMap = + ImmutableMap.of( + entityName, + entityValue1, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(1), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue(), + FeatureV2.getFeatureStringRef(emptyFeatureReference), + DataGenerator.createEmptyValue()); + + ImmutableMap expectedStatusMap = + ImmutableMap.of( + entityName, + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(featureReference), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.NOT_FOUND, + FeatureV2.getFeatureStringRef(emptyFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.NULL_VALUE); + + GetOnlineFeaturesResponse.FieldValues expectedFieldValues = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap) + .putAllStatuses(expectedStatusMap) + .build(); + + ImmutableMap expectedValueMap2 = + ImmutableMap.of( + entityName, + entityValue2, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(2), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue(), + FeatureV2.getFeatureStringRef(emptyFeatureReference), + DataGenerator.createEmptyValue()); + + ImmutableMap expectedValueMap3 = + ImmutableMap.of( + entityName, + entityValue3, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(3), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue(), + FeatureV2.getFeatureStringRef(emptyFeatureReference), + DataGenerator.createEmptyValue()); + + ImmutableMap expectedValueMap4 = + ImmutableMap.of( + entityName, + entityValue4, + FeatureV2.getFeatureStringRef(featureReference), + DataGenerator.createInt32Value(4), + FeatureV2.getFeatureStringRef(notFoundFeatureReference), + DataGenerator.createEmptyValue(), + FeatureV2.getFeatureStringRef(emptyFeatureReference), + DataGenerator.createEmptyValue()); + + GetOnlineFeaturesResponse.FieldValues expectedFieldValues2 = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap2) + .putAllStatuses(expectedStatusMap) + .build(); + GetOnlineFeaturesResponse.FieldValues expectedFieldValues3 = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap3) + .putAllStatuses(expectedStatusMap) + .build(); + GetOnlineFeaturesResponse.FieldValues expectedFieldValues4 = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap4) + .putAllStatuses(expectedStatusMap) + .build(); + ImmutableList expectedFieldValuesList = + ImmutableList.of( + expectedFieldValues, expectedFieldValues2, expectedFieldValues4, expectedFieldValues3); + + assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList()); + } + + @Test + public void shouldReturnFeaturesFromDiffFeatureTable() { + String projectName = "default"; + String entityName = "driver_id"; + ValueProto.Value entityValue = DataGenerator.createInt64Value(1); + + // Instantiate EntityRows + GetOnlineFeaturesRequestV2.EntityRow entityRow = + DataGenerator.createEntityRow(entityName, entityValue, 100); + ImmutableList entityRows = ImmutableList.of(entityRow); + + // Instantiate FeatureReferences + FeatureReferenceV2 rideFeatureReference = + DataGenerator.createFeatureReference("rides", "trip_cost"); + FeatureReferenceV2 rideFeatureReference2 = + DataGenerator.createFeatureReference("rides", "trip_distance"); + FeatureReferenceV2 foodFeatureReference = + DataGenerator.createFeatureReference("food", "trip_cost"); + FeatureReferenceV2 foodFeatureReference2 = + DataGenerator.createFeatureReference("food", "trip_distance"); + + ImmutableList featureReferences = + ImmutableList.of( + rideFeatureReference, + rideFeatureReference2, + foodFeatureReference, + foodFeatureReference2); + + // Build GetOnlineFeaturesRequestV2 + GetOnlineFeaturesRequestV2 onlineFeatureRequest = + TestUtils.createOnlineFeatureRequest(projectName, featureReferences, entityRows); + GetOnlineFeaturesResponse featureResponse = + servingStub.getOnlineFeaturesV2(onlineFeatureRequest); + + ImmutableMap expectedValueMap = + ImmutableMap.of( + entityName, + entityValue, + FeatureV2.getFeatureStringRef(rideFeatureReference), + DataGenerator.createInt32Value(1), + FeatureV2.getFeatureStringRef(rideFeatureReference2), + DataGenerator.createDoubleValue(1.0), + FeatureV2.getFeatureStringRef(foodFeatureReference), + DataGenerator.createInt32Value(1), + FeatureV2.getFeatureStringRef(foodFeatureReference2), + DataGenerator.createDoubleValue(1.0)); + + ImmutableMap expectedStatusMap = + ImmutableMap.of( + entityName, + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(rideFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(rideFeatureReference2), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(foodFeatureReference), + GetOnlineFeaturesResponse.FieldStatus.PRESENT, + FeatureV2.getFeatureStringRef(foodFeatureReference2), + GetOnlineFeaturesResponse.FieldStatus.PRESENT); + + GetOnlineFeaturesResponse.FieldValues expectedFieldValues = + GetOnlineFeaturesResponse.FieldValues.newBuilder() + .putAllFields(expectedValueMap) + .putAllStatuses(expectedStatusMap) + .build(); + ImmutableList expectedFieldValuesList = + ImmutableList.of(expectedFieldValues); + + assertEquals(expectedFieldValuesList, featureResponse.getFieldValuesList()); + } +} diff --git a/serving/src/test/resources/docker-compose/docker-compose-cassandra-it.yml b/serving/src/test/resources/docker-compose/docker-compose-cassandra-it.yml new file mode 100644 index 0000000..15afad0 --- /dev/null +++ b/serving/src/test/resources/docker-compose/docker-compose-cassandra-it.yml @@ -0,0 +1,31 @@ +version: '3' + +services: + core: + image: gcr.io/kf-feast/feast-core:develop + volumes: + - ./core/application-it.yml:/etc/feast/application.yml + environment: + DB_HOST: db + restart: on-failure + depends_on: + - db + ports: + - 6565:6565 + command: + - java + - -jar + - /opt/feast/feast-core.jar + - --spring.config.location=classpath:/application.yml,file:/etc/feast/application.yml + + db: + image: postgres:12-alpine + environment: + POSTGRES_PASSWORD: password + ports: + - "5432:5432" + + cassandra: + image: datastax/cassandra:4.0 + ports: + - "9042:9042" \ No newline at end of file diff --git a/storage/connectors/bigtable/pom.xml b/storage/connectors/bigtable/pom.xml index a53d907..81cd450 100644 --- a/storage/connectors/bigtable/pom.xml +++ b/storage/connectors/bigtable/pom.xml @@ -30,6 +30,12 @@ 1.10.2 + + dev.feast + feast-storage-connector-sstable + ${project.version} + + com.google.guava guava diff --git a/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java b/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java index b9b707e..cf82c14 100644 --- a/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java +++ b/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java @@ -25,10 +25,9 @@ import com.google.protobuf.Timestamp; import feast.proto.serving.ServingAPIProto.FeatureReferenceV2; import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow; -import feast.proto.types.ValueProto; import feast.storage.api.retriever.Feature; import feast.storage.api.retriever.NativeFeature; -import feast.storage.api.retriever.OnlineRetrieverV2; +import feast.storage.connectors.sstable.retriever.SSTableOnlineRetriever; import java.io.IOException; import java.util.*; import java.util.function.Function; @@ -39,7 +38,7 @@ import org.apache.avro.generic.GenericRecord; import org.apache.avro.io.*; -public class BigTableOnlineRetriever implements OnlineRetrieverV2 { +public class BigTableOnlineRetriever implements SSTableOnlineRetriever { private BigtableDataClient client; private BigTableSchemaRegistry schemaRegistry; @@ -49,49 +48,6 @@ public BigTableOnlineRetriever(BigtableDataClient client) { this.schemaRegistry = new BigTableSchemaRegistry(client); } - /** - * Generate name of BigTable table in the form of __ - * - * @param project Name of Feast project - * @param entityNames List of entities used in retrieval call - * @return Name of BigTable table - */ - private String getTableName(String project, List entityNames) { - String tableName = - String.format("%s__%s", project, entityNames.stream().collect(Collectors.joining("__"))); - - return tableName; - } - - /** - * Convert Entity value from Feast valueType to String type. Currently only supports STRING_VAL, - * INT64_VAL, INT32_VAL and BYTES_VAL. - * - * @param v Entity value of Feast valueType - * @return String representation of Entity value - */ - private String valueToString(ValueProto.Value v) { - String stringRepr; - switch (v.getValCase()) { - case STRING_VAL: - stringRepr = v.getStringVal(); - break; - case INT64_VAL: - stringRepr = String.valueOf(v.getInt64Val()); - break; - case INT32_VAL: - stringRepr = String.valueOf(v.getInt32Val()); - break; - case BYTES_VAL: - stringRepr = v.getBytesVal().toString(); - break; - default: - throw new RuntimeException("Type is not supported to be entity"); - } - - return stringRepr; - } - /** * Generate BigTable key in the form of entity values joined by #. * @@ -99,8 +55,8 @@ private String valueToString(ValueProto.Value v) { * @param entityNames List of entities related to feature references in retrieval call * @return BigTable key for retrieval */ - private ByteString convertEntityValueToBigTableKey( - EntityRow entityRow, List entityNames) { + @Override + public ByteString convertEntityValueToKey(EntityRow entityRow, List entityNames) { return ByteString.copyFrom( entityNames.stream() .map(entity -> entityRow.getFieldsMap().get(entity)) @@ -109,118 +65,6 @@ private ByteString convertEntityValueToBigTableKey( .getBytes()); } - /** - * Retrieve BigTable table column families based on FeatureTable names. - * - * @param featureReferences List of feature references of features in retrieval call - * @return List of String of FeatureTable names - */ - private List getColumnFamilies(List featureReferences) { - return featureReferences.stream() - .map(FeatureReferenceV2::getFeatureTable) - .collect(Collectors.toList()); - } - - /** - * AvroRuntimeException is thrown if feature name does not exist in avro schema. Empty Object is - * returned when null is retrieved from BigTable RowCell. - * - * @param tableName Name of BigTable table - * @param value Value of BigTable cell where first 4 bytes represent the schema reference and - * remaining bytes represent avro-serialized features - * @param featureReferences List of feature references - * @param timestamp Timestamp of rowCell - * @return @NativeFeature with retrieved value stored in BigTable RowCell - * @throws IOException - */ - private List decodeFeatures( - String tableName, - ByteString value, - List featureReferences, - BinaryDecoder reusedDecoder, - long timestamp) - throws IOException { - ByteString schemaReferenceBytes = value.substring(0, 4); - byte[] featureValueBytes = value.substring(4).toByteArray(); - - BigTableSchemaRegistry.SchemaReference schemaReference = - new BigTableSchemaRegistry.SchemaReference(tableName, schemaReferenceBytes); - - GenericDatumReader reader = schemaRegistry.getReader(schemaReference); - - reusedDecoder = DecoderFactory.get().binaryDecoder(featureValueBytes, reusedDecoder); - GenericRecord record = reader.read(null, reusedDecoder); - - return featureReferences.stream() - .map( - featureReference -> { - Object featureValue; - try { - featureValue = record.get(featureReference.getName()); - } catch (AvroRuntimeException e) { - // Feature is not found in schema - return null; - } - if (featureValue != null) { - return new NativeFeature( - featureReference, - Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), - featureValue); - } - return new NativeFeature( - featureReference, - Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), - new Object()); - }) - .filter(Objects::nonNull) - .collect(Collectors.toList()); - } - - @Override - public List> getOnlineFeatures( - String project, - List entityRows, - List featureReferences, - List entityNames) { - List columnFamilies = getColumnFamilies(featureReferences); - String tableName = getTableName(project, entityNames); - - List rowKeys = - entityRows.stream() - .map(row -> convertEntityValueToBigTableKey(row, entityNames)) - .collect(Collectors.toList()); - Map rowsFromBigTable = - getFeaturesFromBigTable(tableName, rowKeys, columnFamilies); - List> features = - convertRowToFeature(tableName, rowKeys, rowsFromBigTable, featureReferences); - - return features; - } - - /** - * Retrieve rows for each row entity key by generating BigTable rowQuery with filters based on - * column families. - * - * @param tableName Name of BigTable table - * @param rowKeys List of keys of rows to retrieve - * @param columnFamilies List of FeatureTable names - * @return Map of retrieved features for each rowKey - */ - private Map getFeaturesFromBigTable( - String tableName, List rowKeys, List columnFamilies) { - - Query rowQuery = Query.create(tableName); - Filters.InterleaveFilter familyFilter = Filters.FILTERS.interleave(); - columnFamilies.forEach(cf -> familyFilter.filter(Filters.FILTERS.family().exactMatch(cf))); - - for (ByteString rowKey : rowKeys) { - rowQuery.rowKey(rowKey); - } - - return StreamSupport.stream(client.readRows(rowQuery).spliterator(), false) - .collect(Collectors.toMap(Row::getKey, Function.identity())); - } - /** * Converts rowCell feature value into @NativeFeature type. * @@ -230,7 +74,8 @@ private Map getFeaturesFromBigTable( * @param featureReferences List of feature references * @return List of List of Features associated with respective rowKey */ - private List> convertRowToFeature( + @Override + public List> convertRowToFeature( String tableName, List rowKeys, Map rows, @@ -283,4 +128,84 @@ private List> convertRowToFeature( }) .collect(Collectors.toList()); } + + /** + * Retrieve rows for each row entity key by generating BigTable rowQuery with filters based on + * column families. + * + * @param tableName Name of BigTable table + * @param rowKeys List of keys of rows to retrieve + * @param columnFamilies List of FeatureTable names + * @return Map of retrieved features for each rowKey + */ + @Override + public Map getFeaturesFromSSTable( + String tableName, List rowKeys, List columnFamilies) { + Query rowQuery = Query.create(tableName); + Filters.InterleaveFilter familyFilter = Filters.FILTERS.interleave(); + columnFamilies.forEach(cf -> familyFilter.filter(Filters.FILTERS.family().exactMatch(cf))); + + for (ByteString rowKey : rowKeys) { + rowQuery.rowKey(rowKey); + } + + return StreamSupport.stream(client.readRows(rowQuery).spliterator(), false) + .collect(Collectors.toMap(Row::getKey, Function.identity())); + } + + /** + * AvroRuntimeException is thrown if feature name does not exist in avro schema. Empty Object is + * returned when null is retrieved from BigTable RowCell. + * + * @param tableName Name of BigTable table + * @param value Value of BigTable cell where first 4 bytes represent the schema reference and + * remaining bytes represent avro-serialized features + * @param featureReferences List of feature references + * @param reusedDecoder Decoder for decoding feature values + * @param timestamp Timestamp of rowCell + * @return @NativeFeature with retrieved value stored in BigTable RowCell + * @throws IOException + */ + private List decodeFeatures( + String tableName, + ByteString value, + List featureReferences, + BinaryDecoder reusedDecoder, + long timestamp) + throws IOException { + ByteString schemaReferenceBytes = value.substring(0, 4); + byte[] featureValueBytes = value.substring(4).toByteArray(); + + BigTableSchemaRegistry.SchemaReference schemaReference = + new BigTableSchemaRegistry.SchemaReference(tableName, schemaReferenceBytes); + + GenericDatumReader reader = schemaRegistry.getReader(schemaReference); + + reusedDecoder = DecoderFactory.get().binaryDecoder(featureValueBytes, reusedDecoder); + GenericRecord record = reader.read(null, reusedDecoder); + + return featureReferences.stream() + .map( + featureReference -> { + Object featureValue; + try { + featureValue = record.get(featureReference.getName()); + } catch (AvroRuntimeException e) { + // Feature is not found in schema + return null; + } + if (featureValue != null) { + return new NativeFeature( + featureReference, + Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), + featureValue); + } + return new NativeFeature( + featureReference, + Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), + new Object()); + }) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } } diff --git a/storage/connectors/cassandra/pom.xml b/storage/connectors/cassandra/pom.xml new file mode 100644 index 0000000..32e4d73 --- /dev/null +++ b/storage/connectors/cassandra/pom.xml @@ -0,0 +1,45 @@ + + + + feast-storage-connectors + dev.feast + ${revision} + + + 4.0.0 + feast-storage-connector-cassandra + + + 11 + 11 + + + + + org.apache.avro + avro + 1.10.2 + + + + dev.feast + feast-storage-connector-sstable + ${project.version} + + + + com.datastax.oss + java-driver-core + 4.11.0 + + + + com.datastax.oss + java-driver-query-builder + 4.11.0 + + + + \ No newline at end of file diff --git a/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraOnlineRetriever.java b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraOnlineRetriever.java new file mode 100644 index 0000000..55198e0 --- /dev/null +++ b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraOnlineRetriever.java @@ -0,0 +1,225 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2021 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.cassandra.retriever; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.cql.BoundStatement; +import com.datastax.oss.driver.api.core.cql.Row; +import com.datastax.oss.driver.api.querybuilder.QueryBuilder; +import com.datastax.oss.driver.api.querybuilder.select.Select; +import com.google.protobuf.Timestamp; +import feast.proto.serving.ServingAPIProto.FeatureReferenceV2; +import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow; +import feast.storage.api.retriever.Feature; +import feast.storage.api.retriever.NativeFeature; +import feast.storage.connectors.sstable.retriever.SSTableOnlineRetriever; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.*; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; +import org.apache.avro.AvroRuntimeException; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.BinaryDecoder; +import org.apache.avro.io.DecoderFactory; + +public class CassandraOnlineRetriever implements SSTableOnlineRetriever { + + private final CqlSession session; + private final CassandraSchemaRegistry schemaRegistry; + + private static final String ENTITY_KEY = "key"; + private static final String SCHEMA_REF_SUFFIX = "__schema_ref"; + private static final String EVENT_TIMESTAMP_SUFFIX = "__event_timestamp"; + + public CassandraOnlineRetriever(CqlSession session) { + this.session = session; + this.schemaRegistry = new CassandraSchemaRegistry(session); + } + + /** + * Generate Cassandra key in the form of entity values joined by #. + * + * @param entityRow Single EntityRow representation in feature retrieval call + * @param entityNames List of entities related to feature references in retrieval call + * @return Cassandra key for retrieval + */ + @Override + public ByteBuffer convertEntityValueToKey(EntityRow entityRow, List entityNames) { + return ByteBuffer.wrap( + entityNames.stream() + .map(entity -> entityRow.getFieldsMap().get(entity)) + .map(this::valueToString) + .collect(Collectors.joining("#")) + .getBytes()); + } + + /** + * Converts Cassandra rows into @NativeFeature type. + * + * @param tableName Name of Cassandra table + * @param rowKeys List of keys of rows to retrieve + * @param rows Map of rowKey to Row related to it + * @param featureReferences List of feature references + * @return List of List of Features associated with respective rowKey + */ + @Override + public List> convertRowToFeature( + String tableName, + List rowKeys, + Map rows, + List featureReferences) { + + BinaryDecoder reusedDecoder = DecoderFactory.get().binaryDecoder(new byte[0], null); + + return rowKeys.stream() + .map( + rowKey -> { + if (!rows.containsKey(rowKey)) { + return Collections.emptyList(); + } else { + Row row = rows.get(rowKey); + return featureReferences.stream() + .map(FeatureReferenceV2::getFeatureTable) + .distinct() + .flatMap( + featureTableColumn -> { + ByteBuffer featureValues = row.getByteBuffer(featureTableColumn); + ByteBuffer schemaRefKey = + row.getByteBuffer(featureTableColumn + SCHEMA_REF_SUFFIX); + + // Prevent retrieval of features from incorrect FeatureTable + List localFeatureReferences = + featureReferences.stream() + .filter( + featureReference -> + featureReference + .getFeatureTable() + .equals(featureTableColumn)) + .collect(Collectors.toList()); + + List features; + try { + features = + decodeFeatures( + schemaRefKey, + featureValues, + localFeatureReferences, + reusedDecoder, + row.getLong(featureTableColumn + EVENT_TIMESTAMP_SUFFIX)); + } catch (IOException e) { + throw new RuntimeException("Failed to decode features from Cassandra"); + } + + return features.stream(); + }) + .collect(Collectors.toList()); + } + }) + .collect(Collectors.toList()); + } + + /** + * Retrieve rows for each row entity key by generating Cassandra Query with filters based on + * columns. + * + * @param tableName Name of Cassandra table + * @param rowKeys List of keys of rows to retrieve + * @param columnFamilies List of FeatureTable names + * @return Map of retrieved features for each rowKey + */ + @Override + public Map getFeaturesFromSSTable( + String tableName, List rowKeys, List columnFamilies) { + List schemaRefColumns = + columnFamilies.stream().map(c -> c + SCHEMA_REF_SUFFIX).collect(Collectors.toList()); + Select query = + QueryBuilder.selectFrom(tableName) + .columns(columnFamilies) + .columns(schemaRefColumns) + .column(ENTITY_KEY); + for (String columnFamily : columnFamilies) { + query = query.writeTime(columnFamily).as(columnFamily + EVENT_TIMESTAMP_SUFFIX); + } + query = query.whereColumn(ENTITY_KEY).in(QueryBuilder.bindMarker()); + + BoundStatement statement = session.prepare(query.build()).bind(rowKeys); + + return StreamSupport.stream(session.execute(statement).spliterator(), false) + .collect(Collectors.toMap((Row row) -> row.getByteBuffer(ENTITY_KEY), Function.identity())); + } + + /** + * AvroRuntimeException is thrown if feature name does not exist in avro schema. + * + * @param schemaRefKey Schema reference key + * @param value Value of Cassandra cell where bytes represent avro-serialized features + * @param featureReferences List of feature references + * @param reusedDecoder Decoder for decoding feature values + * @param timestamp Timestamp of rowCell + * @return @NativeFeature with retrieved value stored in Cassandra cell + * @throws IOException + */ + private List decodeFeatures( + ByteBuffer schemaRefKey, + ByteBuffer value, + List featureReferences, + BinaryDecoder reusedDecoder, + long timestamp) + throws IOException { + + if (value == null || schemaRefKey == null) { + return Collections.emptyList(); + } + + CassandraSchemaRegistry.SchemaReference schemaReference = + new CassandraSchemaRegistry.SchemaReference(schemaRefKey); + + // Convert ByteBuffer to ByteArray + byte[] bytesArray = new byte[value.remaining()]; + value.get(bytesArray, 0, bytesArray.length); + GenericDatumReader reader = schemaRegistry.getReader(schemaReference); + reusedDecoder = DecoderFactory.get().binaryDecoder(bytesArray, reusedDecoder); + GenericRecord record = reader.read(null, reusedDecoder); + + return featureReferences.stream() + .map( + featureReference -> { + Object featureValue; + try { + featureValue = record.get(featureReference.getName()); + } catch (AvroRuntimeException e) { + // Feature is not found in schema + return null; + } + if (featureValue != null) { + return new NativeFeature( + featureReference, + Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), + featureValue); + } + return new NativeFeature( + featureReference, + Timestamp.newBuilder().setSeconds(timestamp / 1000).build(), + new Object()); + }) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } +} diff --git a/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraSchemaRegistry.java b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraSchemaRegistry.java new file mode 100644 index 0000000..8001a6f --- /dev/null +++ b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraSchemaRegistry.java @@ -0,0 +1,104 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2021 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.cassandra.retriever; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.cql.BoundStatement; +import com.datastax.oss.driver.api.core.cql.Row; +import com.datastax.oss.driver.api.querybuilder.QueryBuilder; +import com.datastax.oss.driver.api.querybuilder.select.Select; +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.CacheLoader; +import com.google.common.cache.LoadingCache; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericRecord; + +public class CassandraSchemaRegistry { + private final CqlSession session; + private final LoadingCache> cache; + + private static String SCHEMA_REF_TABLE = "feast_schema_reference"; + private static String SCHEMA_REF_COLUMN = "schema_ref"; + private static String SCHEMA_COLUMN = "avro_schema"; + + public static class SchemaReference { + private final ByteBuffer schemaHash; + + public SchemaReference(ByteBuffer schemaHash) { + this.schemaHash = schemaHash; + } + + public ByteBuffer getSchemaHash() { + return schemaHash; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + SchemaReference that = (SchemaReference) o; + return Objects.equals(schemaHash, that.schemaHash); + } + + @Override + public int hashCode() { + return Objects.hash(schemaHash); + } + } + + public CassandraSchemaRegistry(CqlSession session) { + this.session = session; + + CacheLoader> schemaCacheLoader = + CacheLoader.from(this::loadReader); + + cache = CacheBuilder.newBuilder().build(schemaCacheLoader); + } + + public GenericDatumReader getReader(SchemaReference reference) { + GenericDatumReader reader; + try { + reader = this.cache.get(reference); + } catch (ExecutionException | CacheLoader.InvalidCacheLoadException e) { + throw new RuntimeException("Unable to find Schema"); + } + return reader; + } + + private GenericDatumReader loadReader(SchemaReference reference) { + String tableName = String.format("\"%s\"", SCHEMA_REF_TABLE); + Select query = + QueryBuilder.selectFrom(tableName) + .column(SCHEMA_COLUMN) + .whereColumn(SCHEMA_REF_COLUMN) + .isEqualTo(QueryBuilder.bindMarker()); + + BoundStatement statement = session.prepare(query.build()).bind(reference.getSchemaHash()); + + Row row = session.execute(statement).one(); + + Schema schema = + new Schema.Parser() + .parse(StandardCharsets.UTF_8.decode(row.getByteBuffer(SCHEMA_COLUMN)).toString()); + return new GenericDatumReader<>(schema); + } +} diff --git a/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraStoreConfig.java b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraStoreConfig.java new file mode 100644 index 0000000..3ee25df --- /dev/null +++ b/storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraStoreConfig.java @@ -0,0 +1,42 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2021 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.cassandra.retriever; + +public class CassandraStoreConfig { + + private final String connectionString; + private final String dataCenter; + private final String keySpace; + + public CassandraStoreConfig(String connectionString, String dataCenter, String keySpace) { + this.connectionString = connectionString; + this.dataCenter = dataCenter; + this.keySpace = keySpace; + } + + public String getConnectionString() { + return this.connectionString; + } + + public String getDataCenter() { + return this.dataCenter; + } + + public String getKeySpace() { + return this.keySpace; + } +} diff --git a/storage/connectors/pom.xml b/storage/connectors/pom.xml index efa82c5..1c4b75a 100644 --- a/storage/connectors/pom.xml +++ b/storage/connectors/pom.xml @@ -17,6 +17,8 @@ redis bigtable + cassandra + sstable diff --git a/storage/connectors/sstable/pom.xml b/storage/connectors/sstable/pom.xml new file mode 100644 index 0000000..8c7a271 --- /dev/null +++ b/storage/connectors/sstable/pom.xml @@ -0,0 +1,19 @@ + + + + feast-storage-connectors + dev.feast + ${revision} + + + 4.0.0 + feast-storage-connector-sstable + + + 11 + 11 + + + \ No newline at end of file diff --git a/storage/connectors/sstable/src/main/java/feast/storage/connectors/sstable/retriever/SSTableOnlineRetriever.java b/storage/connectors/sstable/src/main/java/feast/storage/connectors/sstable/retriever/SSTableOnlineRetriever.java new file mode 100644 index 0000000..957f0d3 --- /dev/null +++ b/storage/connectors/sstable/src/main/java/feast/storage/connectors/sstable/retriever/SSTableOnlineRetriever.java @@ -0,0 +1,140 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * Copyright 2018-2021 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.sstable.retriever; + +import feast.proto.serving.ServingAPIProto.FeatureReferenceV2; +import feast.proto.serving.ServingAPIProto.GetOnlineFeaturesRequestV2.EntityRow; +import feast.proto.types.ValueProto; +import feast.storage.api.retriever.Feature; +import feast.storage.api.retriever.OnlineRetrieverV2; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +/** + * @param Decoded value type of the partition key + * @param Type of the SSTable row + */ +public interface SSTableOnlineRetriever extends OnlineRetrieverV2 { + + @Override + default List> getOnlineFeatures( + String project, + List entityRows, + List featureReferences, + List entityNames) { + + List columnFamilies = getSSTableColumns(featureReferences); + String tableName = getSSTable(project, entityNames); + + List rowKeys = + entityRows.stream() + .map(row -> convertEntityValueToKey(row, entityNames)) + .collect(Collectors.toList()); + + Map rowsFromSSTable = getFeaturesFromSSTable(tableName, rowKeys, columnFamilies); + + return convertRowToFeature(tableName, rowKeys, rowsFromSSTable, featureReferences); + } + + /** + * Generate SSTable key. + * + * @param entityRow Single EntityRow representation in feature retrieval call + * @param entityNames List of entities related to feature references in retrieval call + * @return SSTable key for retrieval + */ + K convertEntityValueToKey(EntityRow entityRow, List entityNames); + + /** + * Converts SSTable rows into @NativeFeature type. + * + * @param tableName Name of SSTable + * @param rowKeys List of keys of rows to retrieve + * @param rows Map of rowKey to Row related to it + * @param featureReferences List of feature references + * @return List of List of Features associated with respective rowKey + */ + List> convertRowToFeature( + String tableName, + List rowKeys, + Map rows, + List featureReferences); + + /** + * Retrieve rows for each row entity key. + * + * @param tableName Name of SSTable + * @param rowKeys List of keys of rows to retrieve + * @param columnFamilies List of column names + * @return Map of retrieved features for each rowKey + */ + Map getFeaturesFromSSTable(String tableName, List rowKeys, List columnFamilies); + + /** + * Retrieve name of SSTable corresponding to entities in retrieval call + * + * @param project Name of Feast project + * @param entityNames List of entities used in retrieval call + * @return Name of Cassandra table + */ + default String getSSTable(String project, List entityNames) { + return String.format("%s__%s", project, String.join("__", entityNames)); + } + + /** + * Convert Entity value from Feast valueType to String type. Currently only supports STRING_VAL, + * INT64_VAL, INT32_VAL and BYTES_VAL. + * + * @param v Entity value of Feast valueType + * @return String representation of Entity value + */ + default String valueToString(ValueProto.Value v) { + String stringRepr; + switch (v.getValCase()) { + case STRING_VAL: + stringRepr = v.getStringVal(); + break; + case INT64_VAL: + stringRepr = String.valueOf(v.getInt64Val()); + break; + case INT32_VAL: + stringRepr = String.valueOf(v.getInt32Val()); + break; + case BYTES_VAL: + stringRepr = v.getBytesVal().toString(); + break; + default: + throw new RuntimeException("Type is not supported to be entity"); + } + + return stringRepr; + } + + /** + * Retrieve SSTable columns based on Feature references. + * + * @param featureReferences List of feature references in retrieval call + * @return List of String of column names + */ + default List getSSTableColumns(List featureReferences) { + return featureReferences.stream() + .map(FeatureReferenceV2::getFeatureTable) + .distinct() + .collect(Collectors.toList()); + } +}