Use RedisKeyV2 as key serializer and java murmur implementation in Redis Sink (feast-dev#1064)

pyalex · web-flow · commit 0d5cb8df2b2b · 2020-10-19T19:58:38.000+08:00
diff --git a/.github/workflows/complete.yml b/.github/workflows/complete.yml
@@ -143,9 +143,31 @@ jobs:
           path: load-test-output/
 
   tests-docker-compose:
-    needs: build-push-docker-images
+    needs:
+    - build-push-docker-images
+    - publish-ingestion-jar
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - name: Test docker compose
         run: ./infra/scripts/test-docker-compose.sh ${GITHUB_SHA}
+
+  publish-ingestion-jar:
+    runs-on: [self-hosted]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: GoogleCloudPlatform/github-actions/setup-gcloud@master
+        with:
+          version: '290.0.1'
+          export_default_credentials: true
+      - uses: actions/setup-java@v1
+        with:
+          java-version: '11'
+      - uses: stCarolas/setup-maven@v3
+        with:
+          maven-version: 3.6.3
+      - name: build-jar
+        run: make build-java-no-tests REVISION=${GITHUB_SHA}
+      - name: copy to gs
+        run: gsutil cp ./spark/ingestion/target/feast-ingestion-spark-${GITHUB_SHA}.jar gs://feast-jobs/spark/ingestion/
+
diff --git a/.github/workflows/master_only.yml b/.github/workflows/master_only.yml
@@ -62,3 +62,28 @@ jobs:
               docker push gcr.io/kf-feast/feast-${{ matrix.component }}:latest
             fi
           fi
+
+  publish-ingestion-jar:
+    runs-on: [ self-hosted ]
+    env:
+      PUBLISH_BUCKET: feast-jobs
+    steps:
+      - uses: actions/checkout@v2
+      - uses: GoogleCloudPlatform/github-actions/setup-gcloud@master
+        with:
+          version: '290.0.1'
+          export_default_credentials: true
+      - uses: actions/setup-java@v1
+        with:
+          java-version: '11'
+      - uses: stCarolas/setup-maven@v3
+        with:
+          maven-version: 3.6.3
+      - name: build-jar
+        run: |
+          SEMVER_REGEX='^v[0-9]+\.[0-9]+\.[0-9]+(-([0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*))?$'
+          if echo "${RELEASE_VERSION}" | grep -P "$SEMVER_REGEX" &>/dev/null ; then
+            VERSION=${RELEASE_VERSION:1}
+            make build-java-no-tests REVISION=${VERSION}
+            gsutil cp ./spark/ingestion/target/feast-ingestion-spark-${VERSION}.jar gs://${PUBLISH_BUCKET}/spark/ingestion/
+          fi
diff --git a/Makefile b/Makefile
@@ -55,6 +55,9 @@ test-java-with-coverage:
 build-java:
 	mvn clean verify
 
+build-java-no-tests:
+	mvn --no-transfer-progress -Dmaven.javadoc.skip=true -Dgpg.skip -DskipUTs=true -Drevision=${REVISION} clean package
+
 # Python SDK
 
 install-python-ci-dependencies:
diff --git a/infra/docker/jupyter/Dockerfile b/infra/docker/jupyter/Dockerfile
@@ -1,4 +1,4 @@
-FROM jupyter/minimal-notebook:619e9cc2fc07
+FROM jupyter/pyspark-notebook:ae5f7e104dd5
 
 USER root
 WORKDIR /feast
diff --git a/infra/scripts/test-docker-compose.sh b/infra/scripts/test-docker-compose.sh
@@ -63,4 +63,5 @@ export FEAST_ONLINE_SERVING_CONTAINER_IP_ADDRESS=$(docker inspect -f '{{range .N
 ${PROJECT_ROOT_DIR}/infra/scripts/wait-for-it.sh ${FEAST_ONLINE_SERVING_CONTAINER_IP_ADDRESS}:6566 --timeout=120
 
 # Run e2e tests for Redis
-docker exec feast_jupyter_1 bash -c 'cd /feast/tests/e2e && pytest *.py -m "not bq" --core_url core:6565 --serving_url=online_serving:6566 --kafka_brokers=kafka:9092'
+docker exec -e FEAST_VERSION=${FEAST_VERSION} feast_jupyter_1 bash \
+-c 'cd /feast/tests/e2e && unset GOOGLE_APPLICATION_CREDENTIALS && pytest *.py -m "not bq" --ingestion-jar gs://feast-jobs/spark/ingestion/feast-ingestion-spark-${FEAST_VERSION}.jar --redis-url redis:6379 --core_url core:6565 --serving_url=online_serving:6566 --kafka_brokers=kafka:9092'
diff --git a/infra/scripts/test-end-to-end-redis-cluster.sh b/infra/scripts/test-end-to-end-redis-cluster.sh
@@ -103,7 +103,8 @@ cd tests/e2e
 
 set +e
 CORE_NO=$(nproc --all)
-pytest *.py -n ${CORE_NO} --dist=loadscope --junitxml=${LOGS_ARTIFACT_PATH}/python-sdk-test-report.xml
+pytest *.py -n ${CORE_NO} --redis-url localhost:7000 \
+  --dist=loadscope --junitxml=${LOGS_ARTIFACT_PATH}/python-sdk-test-report.xml
 TEST_EXIT_CODE=$?
 
 if [[ ${TEST_EXIT_CODE} != 0 ]]; then
diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py
@@ -889,7 +889,16 @@ def _get_feature_tables_from_feature_refs(
     def start_offline_to_online_ingestion(
         self, feature_table: FeatureTable, start: datetime, end: datetime,
     ) -> SparkJob:
-        return start_offline_to_online_ingestion(feature_table, start, end, self)  # type: ignore
+        """
+
+        Launch Ingestion Job from Batch Source to Online Store for given featureTable
+
+        :param feature_table: FeatureTable which will be ingested
+        :param start: lower datetime boundary
+        :param end: upper datetime boundary
+        :return: Spark Job Proxy object
+        """
+        return start_offline_to_online_ingestion(feature_table, start, end, self)
 
     def start_stream_to_online_ingestion(
         self, feature_table: FeatureTable, extra_jars: Optional[List[str]] = None,
diff --git a/sdk/python/feast/config.py b/sdk/python/feast/config.py
@@ -47,8 +47,7 @@ def _init_config(path: str):
     config_dir = os.path.dirname(path)
     config_dir = config_dir.rstrip("/") + "/"
 
-    if not os.path.exists(os.path.dirname(config_dir)):
-        os.makedirs(os.path.dirname(config_dir))
+    os.makedirs(os.path.dirname(config_dir), exist_ok=True)
 
     # Create the configuration file itself
     config = ConfigParser(defaults=DEFAULTS)
diff --git a/sdk/python/feast/constants.py b/sdk/python/feast/constants.py
@@ -124,4 +124,7 @@ class AuthProvider(Enum):
     CONFIG_AUTH_PROVIDER: "google",
     CONFIG_SPARK_LAUNCHER: "dataproc",
     CONFIG_SPARK_INGESTION_JOB_JAR: "gs://feast-jobs/feast-ingestion-spark-0.8-SNAPSHOT.jar",
+    CONFIG_REDIS_HOST: "localhost",
+    CONFIG_REDIS_PORT: "6379",
+    CONFIG_REDIS_SSL: "False",
 }
diff --git a/sdk/python/feast/staging/storage_client.py b/sdk/python/feast/staging/storage_client.py
@@ -22,6 +22,8 @@
 from typing.io import IO
 from urllib.parse import ParseResult
 
+from google.auth.exceptions import DefaultCredentialsError
+
 GS = "gs"
 S3 = "s3"
 LOCAL_FILE = "file"
@@ -73,7 +75,10 @@ def __init__(self):
                 "Install package google-cloud-storage==1.20.* for gcs staging support"
                 "run ```pip install google-cloud-storage==1.20.*```"
             )
-        self.gcs_client = storage.Client(project=None)
+        try:
+            self.gcs_client = storage.Client(project=None)
+        except DefaultCredentialsError:
+            self.gcs_client = storage.Client.create_anonymous_client()
 
     def download_file(self, uri: ParseResult) -> IO[bytes]:
         """
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/HashTypePersistence.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/HashTypePersistence.scala
@@ -19,11 +19,13 @@ package feast.ingestion.stores.redis
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.types._
 import redis.clients.jedis.{Pipeline, Response}
+import java.nio.charset.StandardCharsets
+
+import com.google.common.hash.Hashing
 
 import scala.jdk.CollectionConverters._
 import com.google.protobuf.Timestamp
 import feast.ingestion.utils.TypeConversion
-import scala.util.hashing.MurmurHash3
 
 /**
   * Use Redis hash type as storage layout. Every feature is stored as separate entry in Hash.
@@ -71,27 +73,27 @@ class HashTypePersistence(config: SparkRedisConfig) extends Persistence with Ser
 
   def encodeKey(key: String): Array[Byte] = {
     val fullFeatureReference = s"${config.namespace}:$key"
-    MurmurHash3.stringHash(fullFeatureReference).toHexString.getBytes
+    Hashing.murmur3_32.hashString(fullFeatureReference, StandardCharsets.UTF_8).asBytes()
   }
 
   def save(
       pipeline: Pipeline,
-      key: String,
+      key: Array[Byte],
       value: Map[Array[Byte], Array[Byte]],
       ttl: Int
   ): Unit = {
-    pipeline.hset(key.getBytes, value.asJava)
+    pipeline.hset(key, value.asJava)
     if (ttl > 0) {
       pipeline.expire(key, ttl)
     }
   }
 
   def getTimestamp(
       pipeline: Pipeline,
-      key: String,
+      key: Array[Byte],
       timestampField: String
   ): Response[Array[Byte]] = {
-    pipeline.hget(key.getBytes(), timestampField.getBytes)
+    pipeline.hget(key, timestampField.getBytes)
   }
 
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/Persistence.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/Persistence.scala
@@ -28,14 +28,14 @@ trait Persistence {
 
   def save(
       pipeline: Pipeline,
-      key: String,
+      key: Array[Byte],
       value: Map[Array[Byte], Array[Byte]],
       ttl: Int
   ): Unit
 
   def getTimestamp(
       pipeline: Pipeline,
-      key: String,
+      key: Array[Byte],
       timestampField: String
   ): Response[Array[Byte]]
 }
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/RedisSinkRelation.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/redis/RedisSinkRelation.scala
@@ -17,16 +17,21 @@
 package feast.ingestion.stores.redis
 
 import com.google.protobuf.Timestamp
-import com.redislabs.provider.redis.{ReadWriteConfig, RedisConfig, RedisEndpoint}
-import com.redislabs.provider.redis.rdd.Keys.groupKeysByNode
+import com.redislabs.provider.redis.{ReadWriteConfig, RedisConfig, RedisEndpoint, RedisNode}
+import redis.clients.jedis.util.JedisClusterCRC16
 import com.redislabs.provider.redis.util.PipelineUtils.{foreachWithPipeline, mapWithPipeline}
+import feast.ingestion.utils.TypeConversion
 import org.apache.spark.SparkEnv
 import org.apache.spark.metrics.source.RedisSinkMetricSource
 import org.apache.spark.sql.sources.{BaseRelation, InsertableRelation}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
+import collection.JavaConverters._
+import feast.proto.storage.RedisProto.RedisKeyV2
+import feast.proto.types.ValueProto
+
 /**
   * High-level writer to Redis. Relies on `Persistence` implementation for actual storage layout.
   * Here we define general flow:
@@ -65,14 +70,14 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
       // grouped iterator to only allocate memory for a portion of rows
       partition.grouped(config.iteratorGroupingSize).foreach { batch =>
         // group by key and keep only latest row per each key
-        val rowsWithKey: Map[String, Row] =
+        val rowsWithKey: Map[RedisKeyV2, Row] =
           compactRowsToLatestTimestamp(batch.map(row => dataKeyId(row) -> row)).toMap
 
         groupKeysByNode(redisConfig.hosts, rowsWithKey.keysIterator).foreach { case (node, keys) =>
           val conn = node.connect()
           // retrieve latest stored timestamp per key
           val timestamps = mapWithPipeline(conn, keys) { (pipeline, key) =>
-            persistence.getTimestamp(pipeline, key, timestampField)
+            persistence.getTimestamp(pipeline, key.toByteArray, timestampField)
           }
 
           val timestampByKey = timestamps
@@ -82,7 +87,7 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
                 .map(Timestamp.parseFrom)
                 .map(t => new java.sql.Timestamp(t.getSeconds * 1000))
             )
-            .zip(rowsWithKey.keys)
+            .zip(keys)
             .map(_.swap)
             .toMap
 
@@ -102,7 +107,7 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
                 }
 
                 val encodedRow = persistence.encodeRow(config.entityColumns, timestampField, row)
-                persistence.save(pipeline, key, encodedRow, ttl = 0)
+                persistence.save(pipeline, key.toByteArray, encodedRow, ttl = 0)
             }
           }
           conn.close()
@@ -111,19 +116,30 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
     }
   }
 
-  private def compactRowsToLatestTimestamp(rows: Seq[(String, Row)]) = rows
+  private def compactRowsToLatestTimestamp(rows: Seq[(RedisKeyV2, Row)]) = rows
     .groupBy(_._1)
     .values
     .map(_.maxBy(_._2.getAs[java.sql.Timestamp](config.timestampColumn).getTime))
 
   /**
     * Key is built from entities columns values with prefix of entities columns names.
     */
-  private def dataKeyId(row: Row): String = {
-    val sortedEntities = config.entityColumns.sorted
-    val entityKey      = sortedEntities.map(row.getAs[Any]).map(_.toString).mkString(":")
-    val entityPrefix   = sortedEntities.mkString("_")
-    s"${config.projectName}_${entityPrefix}:$entityKey"
+  private def dataKeyId(row: Row): RedisKeyV2 = {
+    val types = row.schema.fields.map(f => (f.name, f.dataType)).toMap
+
+    val sortedEntities = config.entityColumns.sorted.toSeq
+    val entityValues = sortedEntities
+      .map(entity => (row.getAs[Any](entity), types(entity)))
+      .map { case (value, v_type) =>
+        TypeConversion.sqlTypeToProtoValue(value, v_type).asInstanceOf[ValueProto.Value]
+      }
+
+    RedisKeyV2
+      .newBuilder()
+      .setProject(config.projectName)
+      .addAllEntityNames(sortedEntities.asJava)
+      .addAllEntityValues(entityValues.asJava)
+      .build
   }
 
   private def timestampField: String = {
@@ -135,4 +151,22 @@ class RedisSinkRelation(override val sqlContext: SQLContext, config: SparkRedisC
       case Seq(head) => Some(head.asInstanceOf[RedisSinkMetricSource])
       case _         => None
     }
+
+  private def groupKeysByNode(
+      nodes: Array[RedisNode],
+      keys: Iterator[RedisKeyV2]
+  ): Iterator[(RedisNode, Array[RedisKeyV2])] = {
+    keys
+      .map(key => (getMasterNode(nodes, key), key))
+      .toArray
+      .groupBy(_._1)
+      .map(x => (x._1, x._2.map(_._2)))
+      .iterator
+  }
+
+  private def getMasterNode(nodes: Array[RedisNode], key: RedisKeyV2): RedisNode = {
+    val slot = JedisClusterCRC16.getSlot(key.toByteArray)
+
+    nodes.filter { node => node.startSlot <= slot && node.endSlot >= slot }.filter(_.idx == 0)(0)
+  }
 }
diff --git a/spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala b/spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala
@@ -23,11 +23,11 @@ import org.apache.spark.SparkConf
 import org.joda.time.{DateTime, Seconds}
 import org.scalacheck._
 import org.scalatest._
-
 import redis.clients.jedis.Jedis
-
 import feast.ingestion.helpers.RedisStorageHelper._
 import feast.ingestion.helpers.DataHelper._
+import feast.proto.storage.RedisProto.RedisKeyV2
+import feast.proto.types.ValueProto
 
 case class Row(customer: String, feature1: Int, feature2: Float, eventTimestamp: java.sql.Timestamp)
 
@@ -54,8 +54,13 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
       } yield Row(customer, feature1, feature2, new java.sql.Timestamp(eventTimestamp.getMillis))
 
     def encodeEntityKey(row: Row, featureTable: FeatureTable): Array[Byte] = {
-      val entityPrefix = featureTable.entities.map(_.name).mkString("_")
-      s"${featureTable.project}_${entityPrefix}:${row.customer}".getBytes
+      RedisKeyV2
+        .newBuilder()
+        .setProject(featureTable.project)
+        .addAllEntityNames(featureTable.entities.map(_.name).sorted.asJava)
+        .addEntityValues(ValueProto.Value.newBuilder().setStringVal(row.customer))
+        .build
+        .toByteArray
     }
 
     def groupByEntity(row: Row) =
diff --git a/spark/ingestion/src/test/scala/feast/ingestion/StreamingPipelineIT.scala b/spark/ingestion/src/test/scala/feast/ingestion/StreamingPipelineIT.scala
diff --git a/spark/ingestion/src/test/scala/feast/ingestion/helpers/RedisStorageHelper.scala b/spark/ingestion/src/test/scala/feast/ingestion/helpers/RedisStorageHelper.scala
diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
diff --git a/tests/e2e/test_online_features.py b/tests/e2e/test_online_features.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-FROM jupyter/minimal-notebook:619e9cc2fc07`
	`1`	`+FROM jupyter/pyspark-notebook:ae5f7e104dd5`
`2`	`2`
`3`	`3`	`USER root`
`4`	`4`	`WORKDIR /feast`
Original file line number	Diff line number	Diff line change
`@@ -124,4 +124,7 @@ class AuthProvider(Enum):`
`124`	`124`	`CONFIG_AUTH_PROVIDER: "google",`
`125`	`125`	`CONFIG_SPARK_LAUNCHER: "dataproc",`
`126`	`126`	`CONFIG_SPARK_INGESTION_JOB_JAR: "gs://feast-jobs/feast-ingestion-spark-0.8-SNAPSHOT.jar",`
	`127`	`+ CONFIG_REDIS_HOST: "localhost",`
	`128`	`+ CONFIG_REDIS_PORT: "6379",`
	`129`	`+ CONFIG_REDIS_SSL: "False",`
`127`	`130`	`}`