From c015543de70486743d37af354474faa212297a8e Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Sat, 3 Apr 2021 15:27:35 +0800 Subject: [PATCH 1/7] Cassandra as alternative online storage Signed-off-by: Khor Shu Heng --- spark/ingestion/pom.xml | 6 + .../scala/feast/ingestion/BasePipeline.scala | 11 ++ .../scala/feast/ingestion/BatchPipeline.scala | 17 ++- .../scala/feast/ingestion/IngestionJob.scala | 3 + .../feast/ingestion/IngestionJobConfig.scala | 7 + .../feast/ingestion/StreamingPipeline.scala | 17 ++- .../bigtable/BigTableSinkRelation.scala | 2 +- .../stores/bigtable/DefaultSource.scala | 2 +- .../cassandra/CassandraSinkRelation.scala | 119 ++++++++++++++++ .../stores/cassandra/DefaultSource.scala | 42 ++++++ .../cassandra/SparkCassandraConfig.scala | 45 ++++++ .../serialization/AvroSerializer.scala | 2 +- .../serialization/Serializer.scala | 2 +- .../ingestion/CassandraIngestionSpec.scala | 130 ++++++++++++++++++ 14 files changed, 393 insertions(+), 12 deletions(-) create mode 100644 spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala create mode 100644 spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/DefaultSource.scala create mode 100644 spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala rename spark/ingestion/src/main/scala/feast/ingestion/stores/{bigtable => }/serialization/AvroSerializer.scala (95%) rename spark/ingestion/src/main/scala/feast/ingestion/stores/{bigtable => }/serialization/Serializer.scala (94%) create mode 100644 spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala diff --git a/spark/ingestion/pom.xml b/spark/ingestion/pom.xml index 003de591..1fc560c5 100644 --- a/spark/ingestion/pom.xml +++ b/spark/ingestion/pom.xml @@ -177,6 +177,12 @@ 2.5.0 + + com.datastax.spark + spark-cassandra-connector_${scala.version} + 3.0.0 + + io.netty netty-all diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala index b971a5ca..cdb224a5 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala @@ -42,6 +42,17 @@ object BasePipeline { conf .set("spark.bigtable.projectId", projectId) .set("spark.bigtable.instanceId", instanceId) + case CassandraConfig(connection, _, properties) => + conf + .set("spark.sql.extensions", "com.datastax.spark.connector.CassandraSparkExtensions") + .set("spark.cassandra.connection.host", connection.host) + .set("spark.cassandra.connection.port", connection.port.toString) + .set("spark.cassandra.output.batch.size.bytes", properties.batchSize.toString) + .set("spark.cassandra.output.concurrent.writes", properties.concurrentWrite.toString) + .set( + s"spark.sql.catalog.feast", + "com.datastax.spark.connector.datasource.CassandraCatalog" + ) } jobConfig.metrics match { diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala index 3dfe2d50..34948759 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala @@ -76,17 +76,26 @@ object BatchPipeline extends BasePipeline { .map(metrics.incrementRead) .filter(rowValidator.allChecks) - validRows.write + val writerWithCommonOptions = validRows.write .format(config.store match { - case _: RedisConfig => "feast.ingestion.stores.redis" - case _: BigTableConfig => "feast.ingestion.stores.bigtable" + case _: RedisConfig => "feast.ingestion.stores.redis" + case _: BigTableConfig => "feast.ingestion.stores.bigtable" + case _: CassandraConfig => "feast.ingestion.stores.cassandra" }) .option("entity_columns", featureTable.entities.map(_.name).mkString(",")) .option("namespace", featureTable.name) .option("project_name", featureTable.project) .option("timestamp_column", config.source.eventTimestampColumn) .option("max_age", config.featureTable.maxAge.getOrElse(0L)) - .save() + + val writer = config.store match { + case storeConfig: CassandraConfig => + writerWithCommonOptions + .option("keyspace", storeConfig.keyspace) + case _ => writerWithCommonOptions + } + + writer.save() config.deadLetterPath foreach { path => projected diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala index 87286156..271823b2 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJob.scala @@ -84,6 +84,9 @@ object IngestionJob { opt[String](name = "bigtable") .action((x, c) => c.copy(store = parseJSON(x).camelizeKeys.extract[BigTableConfig])) + opt[String](name = "cassandra") + .action((x, c) => c.copy(store = parseJSON(x).extract[CassandraConfig])) + opt[String](name = "statsd") .action((x, c) => c.copy(metrics = Some(parseJSON(x).extract[StatsDConfig]))) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala index bd424bd6..bdd873fc 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala @@ -28,6 +28,13 @@ abstract class StoreConfig case class RedisConfig(host: String, port: Int, ssl: Boolean) extends StoreConfig case class BigTableConfig(projectId: String, instanceId: String) extends StoreConfig +case class CassandraConfig( + connection: CassandraConnection, + keyspace: String, + properties: CassandraWriteProperties +) extends StoreConfig +case class CassandraConnection(host: String, port: Int) +case class CassandraWriteProperties(batchSize: Int, concurrentWrite: Int) sealed trait MetricConfig diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala index a533b9e4..e3ac093b 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala @@ -107,20 +107,29 @@ object StreamingPipeline extends BasePipeline with Serializable { implicit val rowEncoder: Encoder[Row] = RowEncoder(rowsAfterValidation.schema) - rowsAfterValidation + val writerWithCommonOptions = rowsAfterValidation .map(metrics.incrementRead) .filter(if (config.doNotIngestInvalidRows) expr("_isValid") else rowValidator.allChecks) .write .format(config.store match { - case _: RedisConfig => "feast.ingestion.stores.redis" - case _: BigTableConfig => "feast.ingestion.stores.bigtable" + case _: RedisConfig => "feast.ingestion.stores.redis" + case _: BigTableConfig => "feast.ingestion.stores.bigtable" + case _: CassandraConfig => "feast.ingestion.stores.cassandra" }) .option("entity_columns", featureTable.entities.map(_.name).mkString(",")) .option("namespace", featureTable.name) .option("project_name", featureTable.project) .option("timestamp_column", config.source.eventTimestampColumn) .option("max_age", config.featureTable.maxAge.getOrElse(0L)) - .save() + + val writer = config.store match { + case storeConfig: CassandraConfig => + writerWithCommonOptions + .option("keyspace", storeConfig.keyspace) + case _ => writerWithCommonOptions + } + + writer.save() config.deadLetterPath match { case Some(path) => diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/BigTableSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/BigTableSinkRelation.scala index a3b7dda4..e76e66ea 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/BigTableSinkRelation.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/BigTableSinkRelation.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.functions.{col, length, struct, udf} import org.apache.spark.sql.{DataFrame, Row, SQLContext} import org.apache.spark.sql.sources.{BaseRelation, InsertableRelation} import org.apache.spark.sql.types.{StringType, StructType} -import feast.ingestion.stores.bigtable.serialization.Serializer +import feast.ingestion.stores.serialization.Serializer import org.apache.hadoop.security.UserGroupInformation class BigTableSinkRelation( diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala index 8b16e097..42014098 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala @@ -19,7 +19,7 @@ package feast.ingestion.stores.bigtable import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode} import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider} import com.google.cloud.bigtable.hbase.BigtableConfiguration -import feast.ingestion.stores.bigtable.serialization.AvroSerializer +import feast.ingestion.stores.serialization.AvroSerializer class DefaultSource extends CreatableRelationProvider { override def createRelation( diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala new file mode 100644 index 00000000..efe412be --- /dev/null +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala @@ -0,0 +1,119 @@ +/* + * 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.ingestion.stores.cassandra + +import feast.ingestion.stores.serialization.Serializer +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.{col, lit, struct, udf} +import org.apache.spark.sql.sources.{BaseRelation, InsertableRelation} +import org.apache.spark.sql.types.{StringType, StructType} +import org.apache.spark.sql.{DataFrame, Row, SQLContext} + +class CassandraSinkRelation( + override val sqlContext: SQLContext, + val serializer: Serializer, + val config: SparkCassandraConfig +) extends BaseRelation + with InsertableRelation + with Serializable { + override def schema: StructType = ??? + + override def insert(data: DataFrame, overwrite: Boolean): Unit = { + + val featureFields = data.schema.fields + .filterNot(f => isSystemColumn(f.name)) + + val featureColumns = featureFields.map(f => col(f.name)) + + val entityColumns = config.entityColumns.map(c => col(c).cast(StringType)) + + val schemaReference = serializer.schemaReference(StructType(featureFields)) + + data + .select( + joinEntityKey(struct(entityColumns: _*)).alias("key"), + serializer.serializeData(struct(featureColumns: _*)).alias(columnName), + col(config.timestampColumn).alias("ts") + ) + .withColumn("schema_ref", lit(schemaReference)) + .writeTo(fullTableReference) + .option("writeTime", "ts") + .append() + } + + def sanitizedForCassandra(expr: String): String = { + expr.replace('-', '_') + } + + val tableName = { + val entities = config.entityColumns.mkString("_") + sanitizedForCassandra(s"${config.projectName}_${entities}") + } + + val keyspace = config.keyspace + + val sparkCatalog = "feast" + + val fullTableReference = s"${sparkCatalog}.${keyspace}.`${tableName}`" + + val columnName = sanitizedForCassandra(config.namespace) + + val schemaTableName = s"${sparkCatalog}.${keyspace}.feast_schema_reference" + + def createTable(): Unit = { + + sqlContext.sql(s""" + |CREATE TABLE IF NOT EXISTS ${fullTableReference} + |(key BINARY, schema_ref BINARY) + |USING cassandra + |PARTITIONED BY (key) + |""".stripMargin) + + sqlContext.sql(s""" + |ALTER TABLE ${fullTableReference} + |ADD COLUMNS (${columnName} BINARY) + |""".stripMargin) + + } + + private def joinEntityKey: UserDefinedFunction = udf { r: Row => + ((0 until r.size)).map(r.getString).mkString("#").getBytes + } + + private def isSystemColumn(name: String) = + (config.entityColumns ++ Seq(config.timestampColumn)).contains(name) + + def saveWriteSchema(data: DataFrame) = { + sqlContext.sql(s""" + |CREATE TABLE IF NOT EXISTS ${schemaTableName} + |(schema_ref BINARY, avro_schema BINARY) + |USING cassandra + |PARTITIONED BY (schema_ref) + |""".stripMargin) + + val featureFields = data.schema.fields + .filterNot(f => isSystemColumn(f.name)) + val featureSchema = StructType(featureFields) + val key = serializer.schemaReference(featureSchema) + val serializedSchema = serializer.serializeSchema(featureSchema).getBytes + + import sqlContext.sparkSession.implicits._ + val schemaData = List((key, serializedSchema)).toDF("schema_ref", "avro_schema") + + schemaData.writeTo(schemaTableName).append() + } +} diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/DefaultSource.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/DefaultSource.scala new file mode 100644 index 00000000..ceab7ec8 --- /dev/null +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/DefaultSource.scala @@ -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.ingestion.stores.cassandra + +import feast.ingestion.stores.serialization.AvroSerializer +import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider} +import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode} + +class DefaultSource extends CreatableRelationProvider { + override def createRelation( + sqlContext: SQLContext, + mode: SaveMode, + parameters: Map[String, String], + data: DataFrame + ): BaseRelation = { + + val rel = + new CassandraSinkRelation( + sqlContext, + new AvroSerializer, + SparkCassandraConfig.parse(parameters) + ) + rel.createTable() + rel.saveWriteSchema(data) + rel.insert(data, overwrite = false) + rel + } +} diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala new file mode 100644 index 00000000..c1fc711b --- /dev/null +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala @@ -0,0 +1,45 @@ +/* + * 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.ingestion.stores.cassandra + +case class SparkCassandraConfig( + keyspace: String, + namespace: String, + projectName: String, + entityColumns: Array[String], + timestampColumn: String, + maxAge: Long +) + +object SparkCassandraConfig { + val KEYSPACE = "keyspace" + val NAMESPACE = "namespace" + val ENTITY_COLUMNS = "entity_columns" + val TS_COLUMN = "timestamp_column" + val PROJECT_NAME = "project_name" + val MAX_AGE = "max_age" + + def parse(parameters: Map[String, String]): SparkCassandraConfig = + SparkCassandraConfig( + keyspace = parameters.getOrElse(KEYSPACE, ""), + namespace = parameters.getOrElse(NAMESPACE, ""), + projectName = parameters.getOrElse(PROJECT_NAME, "default"), + entityColumns = parameters.getOrElse(ENTITY_COLUMNS, "").split(","), + timestampColumn = parameters.getOrElse(TS_COLUMN, "event_timestamp"), + maxAge = parameters.get(MAX_AGE).map(_.toLong).getOrElse(0) + ) +} diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/AvroSerializer.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/AvroSerializer.scala similarity index 95% rename from spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/AvroSerializer.scala rename to spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/AvroSerializer.scala index d2707f4c..1118cf2f 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/AvroSerializer.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/AvroSerializer.scala @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package feast.ingestion.stores.bigtable.serialization +package feast.ingestion.stores.serialization import com.google.common.hash.Hashing import org.apache.spark.sql.Column diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/Serializer.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/Serializer.scala similarity index 94% rename from spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/Serializer.scala rename to spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/Serializer.scala index 575b3834..b6024846 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/serialization/Serializer.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/serialization/Serializer.scala @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package feast.ingestion.stores.bigtable.serialization +package feast.ingestion.stores.serialization import org.apache.spark.sql.Column import org.apache.spark.sql.types.StructType diff --git a/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala new file mode 100644 index 00000000..8958da90 --- /dev/null +++ b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala @@ -0,0 +1,130 @@ +/* + * 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.ingestion + +import com.dimafeng.testcontainers.{ForAllTestContainer, GenericContainer} +import feast.ingestion.helpers.DataHelper.{generateDistinctRows, rowGenerator, storeAsParquet} +import feast.ingestion.helpers.TestRow +import feast.proto.types.ValueProto.ValueType +import org.apache.avro.Schema +import org.apache.avro.generic.{GenericDatumReader, GenericRecord} +import org.apache.avro.io.DecoderFactory +import org.apache.spark.SparkConf +import org.joda.time.DateTime +import org.testcontainers.containers.wait.strategy.Wait + +import java.sql.Timestamp +import java.time.{Duration, Instant} + +class CassandraIngestionSpec extends SparkSpec with ForAllTestContainer { + + override val container = GenericContainer( + "cassandra:3.11.10", + exposedPorts = Seq(9042), + waitStrategy = Wait + .forListeningPort() + .withStartupTimeout(Duration.ofSeconds(120)) + ) + + override def withSparkConfOverrides(conf: SparkConf): SparkConf = conf + .set("spark.sql.extensions", "com.datastax.spark.connector.CassandraSparkExtensions") + .set("spark.cassandra.connection.host", container.host) + .set("spark.cassandra.connection.port", container.mappedPort(9042).toString) + .set( + s"spark.sql.catalog.feast", + "com.datastax.spark.connector.datasource.CassandraCatalog" + ) + .set("spark.cassandra.connection.localDC", "datacenter1") + + trait Scope { + val keyspace = "feast" + + val config = IngestionJobConfig( + featureTable = FeatureTable( + name = "test-fs", + project = "default", + entities = Seq(Field("customer", ValueType.Enum.STRING)), + features = Seq( + Field("feature1", ValueType.Enum.INT32), + Field("feature2", ValueType.Enum.FLOAT) + ) + ), + startTime = DateTime.parse("2020-08-01"), + endTime = DateTime.parse("2020-09-01"), + store = CassandraConfig( + CassandraConnection(container.host, container.mappedPort(9042)), + keyspace, + CassandraWriteProperties(1024, 5) + ) + ) + + val gen = rowGenerator(DateTime.parse("2020-08-01"), DateTime.parse("2020-09-01")) + + } + + def decodeAvroValue(input: Array[Byte], jsonFormatSchema: String): GenericRecord = { + val schema = new Schema.Parser().parse(jsonFormatSchema) + val reader = new GenericDatumReader[Any](schema) + var result: Any = null + + val decoder = DecoderFactory.get().binaryDecoder(input, 0, input.length, null) + result = reader.read(result, decoder) + result.asInstanceOf[GenericRecord] + } + + "Dataset" should "be ingested in cassandra" in new Scope { + sparkSession.sql( + "CREATE DATABASE IF NOT EXISTS feast.feast WITH DBPROPERTIES (class='SimpleStrategy',replication_factor='1')" + ) + val rows = generateDistinctRows(gen, 1000, (_: TestRow).customer).filterNot(_.customer.isEmpty) + val tempPath = storeAsParquet(sparkSession, rows) + val configWithOfflineSource = config.copy( + source = FileSource(tempPath, Map.empty, "eventTimestamp") + ) + + BatchPipeline.createPipeline(sparkSession, configWithOfflineSource) + val avroSchema = sparkSession + .sql( + "SELECT schema_ref, avro_schema FROM feast.feast.feast_schema_reference" + ) + .collect() + .map(row => row.getAs[Array[Byte]](0).toSeq -> new String(row.getAs[Array[Byte]](1))) + .toMap + + val storedRows = sparkSession + .sql( + "SELECT key, test_fs, schema_ref, writeTime(test_fs) FROM feast.feast.default_customer" + ) + .collect() + .map { row => + val record = decodeAvroValue( + row.getAs[Array[Byte]](1), + avroSchema.get(row.getAs[Array[Byte]](2).toSeq).get + ) + + TestRow( + new String(row.getAs[Array[Byte]](0)), + record.get("feature1").asInstanceOf[Integer], + record.get("feature2").asInstanceOf[Float], + Timestamp.from(Instant.ofEpochMilli(row.getLong(3))) + ) + } + + storedRows should contain allElementsOf rows + + } +} From 87b6c0dcea353f12baafc087d49cdccb82c3cf44 Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Mon, 5 Apr 2021 10:51:55 +0800 Subject: [PATCH 2/7] Add TTL Signed-off-by: Khor Shu Heng --- .../stores/cassandra/CassandraSinkRelation.scala | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala index efe412be..0eb5c36d 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala @@ -43,7 +43,7 @@ class CassandraSinkRelation( val schemaReference = serializer.schemaReference(StructType(featureFields)) - data + val writerWithoutTTL = data .select( joinEntityKey(struct(entityColumns: _*)).alias("key"), serializer.serializeData(struct(featureColumns: _*)).alias(columnName), @@ -52,7 +52,13 @@ class CassandraSinkRelation( .withColumn("schema_ref", lit(schemaReference)) .writeTo(fullTableReference) .option("writeTime", "ts") - .append() + + val writer = + if (config.maxAge <= 0) + writerWithoutTTL + else writerWithoutTTL.option("ttl", config.maxAge.toString) + + writer.append() } def sanitizedForCassandra(expr: String): String = { From c6bca90c18008ce97de55f1581311acede65c45a Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Mon, 5 Apr 2021 14:51:15 +0800 Subject: [PATCH 3/7] Add default cassandra write properties Signed-off-by: Khor Shu Heng --- .../src/main/scala/feast/ingestion/IngestionJobConfig.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala index bdd873fc..87150493 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/IngestionJobConfig.scala @@ -31,7 +31,7 @@ case class BigTableConfig(projectId: String, instanceId: String) extends StoreCo case class CassandraConfig( connection: CassandraConnection, keyspace: String, - properties: CassandraWriteProperties + properties: CassandraWriteProperties = CassandraWriteProperties(1024, 5) ) extends StoreConfig case class CassandraConnection(host: String, port: Int) case class CassandraWriteProperties(batchSize: Int, concurrentWrite: Int) From c0e2c9df864ea41ed11fb593cc92ce5ad32b4199 Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Mon, 5 Apr 2021 16:29:22 +0800 Subject: [PATCH 4/7] Use pre-generated schema Signed-off-by: Khor Shu Heng --- .../cassandra/CassandraSinkRelation.scala | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala index 0eb5c36d..145ce269 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala @@ -37,16 +37,16 @@ class CassandraSinkRelation( val featureFields = data.schema.fields .filterNot(f => isSystemColumn(f.name)) - val featureColumns = featureFields.map(f => col(f.name)) + val featureColumns = featureFields.map(f => data(f.name)) - val entityColumns = config.entityColumns.map(c => col(c).cast(StringType)) - - val schemaReference = serializer.schemaReference(StructType(featureFields)) + val entityColumns = config.entityColumns.map(c => data(c).cast(StringType)) + val schema = serializer.convertSchema(StructType(featureFields)) + val schemaReference = serializer.schemaReference(schema) val writerWithoutTTL = data .select( joinEntityKey(struct(entityColumns: _*)).alias("key"), - serializer.serializeData(struct(featureColumns: _*)).alias(columnName), + serializer.serializeData(schema)(struct(featureColumns: _*)).alias(columnName), col(config.timestampColumn).alias("ts") ) .withColumn("schema_ref", lit(schemaReference)) @@ -114,11 +114,12 @@ class CassandraSinkRelation( val featureFields = data.schema.fields .filterNot(f => isSystemColumn(f.name)) val featureSchema = StructType(featureFields) - val key = serializer.schemaReference(featureSchema) - val serializedSchema = serializer.serializeSchema(featureSchema).getBytes + + val schema = serializer.convertSchema(featureSchema) + val key = serializer.schemaReference(schema) import sqlContext.sparkSession.implicits._ - val schemaData = List((key, serializedSchema)).toDF("schema_ref", "avro_schema") + val schemaData = List((key, schema.asInstanceOf[String].getBytes)).toDF("schema_ref", "avro_schema") schemaData.writeTo(schemaTableName).append() } From e5a2e47e78979cdcc10be4a74e69cb19b8df6218 Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Mon, 5 Apr 2021 21:55:47 +0800 Subject: [PATCH 5/7] Use multiple columns for schema ref Signed-off-by: Khor Shu Heng --- .../cassandra/CassandraSinkRelation.scala | 21 +++++++++++-------- .../ingestion/CassandraIngestionSpec.scala | 2 +- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala index 145ce269..ff29fb08 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala @@ -39,8 +39,8 @@ class CassandraSinkRelation( val featureColumns = featureFields.map(f => data(f.name)) - val entityColumns = config.entityColumns.map(c => data(c).cast(StringType)) - val schema = serializer.convertSchema(StructType(featureFields)) + val entityColumns = config.entityColumns.map(c => data(c).cast(StringType)) + val schema = serializer.convertSchema(StructType(featureFields)) val schemaReference = serializer.schemaReference(schema) val writerWithoutTTL = data @@ -49,7 +49,7 @@ class CassandraSinkRelation( serializer.serializeData(schema)(struct(featureColumns: _*)).alias(columnName), col(config.timestampColumn).alias("ts") ) - .withColumn("schema_ref", lit(schemaReference)) + .withColumn(schemaRefColumnName, lit(schemaReference)) .writeTo(fullTableReference) .option("writeTime", "ts") @@ -66,8 +66,8 @@ class CassandraSinkRelation( } val tableName = { - val entities = config.entityColumns.mkString("_") - sanitizedForCassandra(s"${config.projectName}_${entities}") + val entities = config.entityColumns.mkString("__") + sanitizedForCassandra(s"${config.projectName}__${entities}") } val keyspace = config.keyspace @@ -78,20 +78,22 @@ class CassandraSinkRelation( val columnName = sanitizedForCassandra(config.namespace) + val schemaRefColumnName = sanitizedForCassandra(s"${config.namespace}__schema_ref") + val schemaTableName = s"${sparkCatalog}.${keyspace}.feast_schema_reference" def createTable(): Unit = { sqlContext.sql(s""" |CREATE TABLE IF NOT EXISTS ${fullTableReference} - |(key BINARY, schema_ref BINARY) + |(key BINARY) |USING cassandra |PARTITIONED BY (key) |""".stripMargin) sqlContext.sql(s""" |ALTER TABLE ${fullTableReference} - |ADD COLUMNS (${columnName} BINARY) + |ADD COLUMNS (${columnName} BINARY, ${schemaRefColumnName} BINARY) |""".stripMargin) } @@ -113,13 +115,14 @@ class CassandraSinkRelation( val featureFields = data.schema.fields .filterNot(f => isSystemColumn(f.name)) - val featureSchema = StructType(featureFields) + val featureSchema = StructType(featureFields) val schema = serializer.convertSchema(featureSchema) val key = serializer.schemaReference(schema) import sqlContext.sparkSession.implicits._ - val schemaData = List((key, schema.asInstanceOf[String].getBytes)).toDF("schema_ref", "avro_schema") + val schemaData = + List((key, schema.asInstanceOf[String].getBytes)).toDF("schema_ref", "avro_schema") schemaData.writeTo(schemaTableName).append() } diff --git a/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala index 8958da90..26432623 100644 --- a/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala +++ b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala @@ -107,7 +107,7 @@ class CassandraIngestionSpec extends SparkSpec with ForAllTestContainer { val storedRows = sparkSession .sql( - "SELECT key, test_fs, schema_ref, writeTime(test_fs) FROM feast.feast.default_customer" + "SELECT key, test_fs, test_fs__schema_ref, writeTime(test_fs) FROM feast.feast.default__customer" ) .collect() .map { row => From 759814c29610784071fabc18f65a4882616eb6b3 Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Tue, 6 Apr 2021 14:34:44 +0800 Subject: [PATCH 6/7] Move cassandra keyspace config to spark conf Signed-off-by: Khor Shu Heng --- .../main/scala/feast/ingestion/BasePipeline.scala | 3 ++- .../main/scala/feast/ingestion/BatchPipeline.scala | 12 ++---------- .../scala/feast/ingestion/StreamingPipeline.scala | 12 ++---------- .../stores/cassandra/CassandraSinkRelation.scala | 2 +- .../stores/cassandra/SparkCassandraConfig.scala | 3 --- .../feast/ingestion/CassandraIngestionSpec.scala | 6 ++++-- 6 files changed, 11 insertions(+), 27 deletions(-) diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala index cdb224a5..d88e900d 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/BasePipeline.scala @@ -42,7 +42,7 @@ object BasePipeline { conf .set("spark.bigtable.projectId", projectId) .set("spark.bigtable.instanceId", instanceId) - case CassandraConfig(connection, _, properties) => + case CassandraConfig(connection, keyspace, properties) => conf .set("spark.sql.extensions", "com.datastax.spark.connector.CassandraSparkExtensions") .set("spark.cassandra.connection.host", connection.host) @@ -53,6 +53,7 @@ object BasePipeline { s"spark.sql.catalog.feast", "com.datastax.spark.connector.datasource.CassandraCatalog" ) + .set("feast.store.cassandra.keyspace", keyspace) } jobConfig.metrics match { diff --git a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala index 34948759..b9480699 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala @@ -76,7 +76,7 @@ object BatchPipeline extends BasePipeline { .map(metrics.incrementRead) .filter(rowValidator.allChecks) - val writerWithCommonOptions = validRows.write + validRows.write .format(config.store match { case _: RedisConfig => "feast.ingestion.stores.redis" case _: BigTableConfig => "feast.ingestion.stores.bigtable" @@ -87,15 +87,7 @@ object BatchPipeline extends BasePipeline { .option("project_name", featureTable.project) .option("timestamp_column", config.source.eventTimestampColumn) .option("max_age", config.featureTable.maxAge.getOrElse(0L)) - - val writer = config.store match { - case storeConfig: CassandraConfig => - writerWithCommonOptions - .option("keyspace", storeConfig.keyspace) - case _ => writerWithCommonOptions - } - - writer.save() + .save() config.deadLetterPath foreach { path => projected diff --git a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala index e3ac093b..ed0a6fd6 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala @@ -107,7 +107,7 @@ object StreamingPipeline extends BasePipeline with Serializable { implicit val rowEncoder: Encoder[Row] = RowEncoder(rowsAfterValidation.schema) - val writerWithCommonOptions = rowsAfterValidation + rowsAfterValidation .map(metrics.incrementRead) .filter(if (config.doNotIngestInvalidRows) expr("_isValid") else rowValidator.allChecks) .write @@ -121,15 +121,7 @@ object StreamingPipeline extends BasePipeline with Serializable { .option("project_name", featureTable.project) .option("timestamp_column", config.source.eventTimestampColumn) .option("max_age", config.featureTable.maxAge.getOrElse(0L)) - - val writer = config.store match { - case storeConfig: CassandraConfig => - writerWithCommonOptions - .option("keyspace", storeConfig.keyspace) - case _ => writerWithCommonOptions - } - - writer.save() + .save() config.deadLetterPath match { case Some(path) => diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala index ff29fb08..fc40d679 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/CassandraSinkRelation.scala @@ -70,7 +70,7 @@ class CassandraSinkRelation( sanitizedForCassandra(s"${config.projectName}__${entities}") } - val keyspace = config.keyspace + val keyspace = sqlContext.sparkContext.getConf.get("feast.store.cassandra.keyspace") val sparkCatalog = "feast" diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala index c1fc711b..3511aa71 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/stores/cassandra/SparkCassandraConfig.scala @@ -17,7 +17,6 @@ package feast.ingestion.stores.cassandra case class SparkCassandraConfig( - keyspace: String, namespace: String, projectName: String, entityColumns: Array[String], @@ -26,7 +25,6 @@ case class SparkCassandraConfig( ) object SparkCassandraConfig { - val KEYSPACE = "keyspace" val NAMESPACE = "namespace" val ENTITY_COLUMNS = "entity_columns" val TS_COLUMN = "timestamp_column" @@ -35,7 +33,6 @@ object SparkCassandraConfig { def parse(parameters: Map[String, String]): SparkCassandraConfig = SparkCassandraConfig( - keyspace = parameters.getOrElse(KEYSPACE, ""), namespace = parameters.getOrElse(NAMESPACE, ""), projectName = parameters.getOrElse(PROJECT_NAME, "default"), entityColumns = parameters.getOrElse(ENTITY_COLUMNS, "").split(","), diff --git a/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala index 26432623..8e4eb654 100644 --- a/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala +++ b/spark/ingestion/src/test/scala/feast/ingestion/CassandraIngestionSpec.scala @@ -40,6 +40,8 @@ class CassandraIngestionSpec extends SparkSpec with ForAllTestContainer { .withStartupTimeout(Duration.ofSeconds(120)) ) + val keyspace = "feast" + override def withSparkConfOverrides(conf: SparkConf): SparkConf = conf .set("spark.sql.extensions", "com.datastax.spark.connector.CassandraSparkExtensions") .set("spark.cassandra.connection.host", container.host) @@ -49,9 +51,9 @@ class CassandraIngestionSpec extends SparkSpec with ForAllTestContainer { "com.datastax.spark.connector.datasource.CassandraCatalog" ) .set("spark.cassandra.connection.localDC", "datacenter1") + .set("feast.store.cassandra.keyspace", keyspace) trait Scope { - val keyspace = "feast" val config = IngestionJobConfig( featureTable = FeatureTable( @@ -113,7 +115,7 @@ class CassandraIngestionSpec extends SparkSpec with ForAllTestContainer { .map { row => val record = decodeAvroValue( row.getAs[Array[Byte]](1), - avroSchema.get(row.getAs[Array[Byte]](2).toSeq).get + avroSchema(row.getAs[Array[Byte]](2).toSeq) ) TestRow( From 846a3b4333aefb7f793923b64aabe4f69a26ada3 Mon Sep 17 00:00:00 2001 From: Khor Shu Heng Date: Wed, 7 Apr 2021 12:18:42 +0800 Subject: [PATCH 7/7] Expose Cassandra store to python SDK Signed-off-by: Khor Shu Heng --- python/feast_spark/constants.py | 6 ++++++ python/feast_spark/pyspark/abc.py | 18 ++++++++++++++++++ python/feast_spark/pyspark/launcher.py | 3 +++ 3 files changed, 27 insertions(+) diff --git a/python/feast_spark/constants.py b/python/feast_spark/constants.py index 8f6c2496..1ed320a8 100644 --- a/python/feast_spark/constants.py +++ b/python/feast_spark/constants.py @@ -114,6 +114,12 @@ class ConfigOptions(metaclass=ConfigMeta): #: BigTable Instance ID BIGTABLE_INSTANCE: Optional[str] = "" + #: Cassandra host. Can be a comma separated string + CASSANDRA_HOST: Optional[str] = "" + + #: Cassandra port + CASSANDRA_PORT: Optional[str] = "" + #: Enable or disable StatsD STATSD_ENABLED: str = "False" diff --git a/python/feast_spark/pyspark/abc.py b/python/feast_spark/pyspark/abc.py index 2131ac8c..c6bbbf14 100644 --- a/python/feast_spark/pyspark/abc.py +++ b/python/feast_spark/pyspark/abc.py @@ -333,6 +333,8 @@ def __init__( redis_ssl: Optional[bool] = None, bigtable_project: Optional[str] = None, bigtable_instance: Optional[str] = None, + cassandra_host: Optional[str] = None, + cassandra_port: Optional[int] = None, statsd_host: Optional[str] = None, statsd_port: Optional[int] = None, deadletter_path: Optional[str] = None, @@ -347,6 +349,8 @@ def __init__( self._redis_ssl = redis_ssl self._bigtable_project = bigtable_project self._bigtable_instance = bigtable_instance + self._cassandra_host = cassandra_host + self._cassandra_port = cassandra_port self._statsd_host = statsd_host self._statsd_port = statsd_port self._deadletter_path = deadletter_path @@ -361,6 +365,9 @@ def _get_bigtable_config(self): project_id=self._bigtable_project, instance_id=self._bigtable_instance ) + def _get_cassandra_config(self): + return dict(host=self._cassandra_host, port=self._cassandra_port) + def _get_statsd_config(self): return ( dict(host=self._statsd_host, port=self._statsd_port) @@ -394,6 +401,9 @@ def get_arguments(self) -> List[str]: if self._bigtable_project and self._bigtable_instance: args.extend(["--bigtable", json.dumps(self._get_bigtable_config())]) + if self._cassandra_host and self._cassandra_port: + args.extend(["--cassandra", json.dumps(self._get_cassandra_config())]) + if self._get_statsd_config(): args.extend(["--statsd", json.dumps(self._get_statsd_config())]) @@ -427,6 +437,8 @@ def __init__( redis_ssl: Optional[bool], bigtable_project: Optional[str], bigtable_instance: Optional[str], + cassandra_host: Optional[str] = None, + cassandra_port: Optional[int] = None, statsd_host: Optional[str] = None, statsd_port: Optional[int] = None, deadletter_path: Optional[str] = None, @@ -441,6 +453,8 @@ def __init__( redis_ssl, bigtable_project, bigtable_instance, + cassandra_host, + cassandra_port, statsd_host, statsd_port, deadletter_path, @@ -481,6 +495,8 @@ def __init__( redis_ssl: Optional[bool], bigtable_project: Optional[str], bigtable_instance: Optional[str], + cassandra_host: Optional[str] = None, + cassandra_port: Optional[int] = None, statsd_host: Optional[str] = None, statsd_port: Optional[int] = None, deadletter_path: Optional[str] = None, @@ -497,6 +513,8 @@ def __init__( redis_ssl, bigtable_project, bigtable_instance, + cassandra_host, + cassandra_port, statsd_host, statsd_port, deadletter_path, diff --git a/python/feast_spark/pyspark/launcher.py b/python/feast_spark/pyspark/launcher.py index fa8b650b..e86f02d4 100644 --- a/python/feast_spark/pyspark/launcher.py +++ b/python/feast_spark/pyspark/launcher.py @@ -276,6 +276,9 @@ def start_offline_to_online_ingestion( redis_ssl=client.config.getboolean(opt.REDIS_SSL), bigtable_project=client.config.get(opt.BIGTABLE_PROJECT), bigtable_instance=client.config.get(opt.BIGTABLE_INSTANCE), + cassandra_host=client.config.get(opt.CASSANDRA_HOST), + cassandra_port=bool(client.config.get(opt.CASSANDRA_HOST)) + and client.config.getint(opt.CASSANDRA_PORT), statsd_host=( client.config.getboolean(opt.STATSD_ENABLED) and client.config.get(opt.STATSD_HOST)