Skip to content

Commit 6b59ec0

Browse files
pyalexpyalex
committed
all types test
Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 9e028a3 commit 6b59ec0

12 files changed

Lines changed: 422 additions & 115 deletions

File tree

‎spark/ingestion/pom.xml‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
<spark.version>2.4.7</spark.version>
3737
<scala-maven-plugin.version>4.4.0</scala-maven-plugin.version>
3838
<maven-assembly-plugin.version>3.3.0</maven-assembly-plugin.version>
39+
<protobuf.version>3.12.2</protobuf.version>
3940
</properties>
4041

4142
<dependencies>
@@ -54,7 +55,13 @@
5455
<dependency>
5556
<groupId>com.google.protobuf</groupId>
5657
<artifactId>protobuf-java</artifactId>
57-
<version>3.12.2</version>
58+
<version>${protobuf.version}</version>
59+
</dependency>
60+
61+
<dependency>
62+
<groupId>com.gojek</groupId>
63+
<artifactId>stencil</artifactId>
64+
<version>4.2.0</version>
5865
</dependency>
5966

6067
<dependency>
@@ -100,7 +107,12 @@
100107
<groupId>org.apache.spark</groupId>
101108
<artifactId>spark-sql-kafka-0-10_${scala.version}</artifactId>
102109
<version>${spark.version}</version>
103-
<scope>provided</scope>
110+
</dependency>
111+
112+
<dependency>
113+
<groupId>org.apache.kafka</groupId>
114+
<artifactId>kafka-clients</artifactId>
115+
<version>2.5.1</version>
104116
</dependency>
105117

106118
<dependency>
@@ -173,13 +185,6 @@
173185
<version>0.38.3</version>
174186
<scope>test</scope>
175187
</dependency>
176-
177-
<dependency>
178-
<groupId>org.apache.kafka</groupId>
179-
<artifactId>kafka-clients</artifactId>
180-
<version>2.5.1</version>
181-
<scope>test</scope>
182-
</dependency>
183188
</dependencies>
184189

185190

@@ -302,8 +307,8 @@
302307
<goal>run</goal>
303308
</goals>
304309
<configuration>
305-
<protocArtifact>com.google.protobuf:protoc:3.12.2</protocArtifact>
306-
<protocVersion>3.12.2</protocVersion>
310+
<protocArtifact>com.google.protobuf:protoc:${protobuf.version}</protocArtifact>
311+
<protocVersion>${protobuf.version}</protocVersion>
307312
<inputDirectories>
308313
<include>src/test/proto</include>
309314
</inputDirectories>

‎spark/ingestion/src/main/scala/feast/ingestion/BatchPipeline.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ object BatchPipeline extends BasePipeline {
3535
val featureTable = config.featureTable
3636
val projection =
3737
inputProjection(config.source, featureTable.features, featureTable.entities)
38-
val validator = new RowValidator(featureTable)
38+
val validator = new RowValidator(featureTable, config.source.timestampColumn)
3939

4040
val input = config.source match {
4141
case source: BQSource =>

‎spark/ingestion/src/main/scala/feast/ingestion/StreamingPipeline.scala‎

Lines changed: 10 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,13 @@
1616
*/
1717
package feast.ingestion
1818

19-
import java.sql
20-
21-
import com.google.protobuf.Descriptors.{Descriptor, EnumValueDescriptor, FieldDescriptor}
22-
import com.google.protobuf.{AbstractMessage, ByteString, GeneratedMessageV3, Parser, Timestamp}
19+
import feast.ingestion.registry.proto.ProtoRegistryFactory
2320
import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
2421
import org.apache.spark.sql.functions.udf
25-
import org.apache.spark.sql.types._
26-
import com.google.protobuf.Descriptors.FieldDescriptor.JavaType._
27-
import feast.ingestion.BatchPipeline.inputProjection
2822
import feast.ingestion.utils.ProtoReflection
2923
import feast.ingestion.validation.RowValidator
3024
import org.apache.spark.sql.streaming.StreamingQuery
3125

32-
import scala.collection.convert.ImplicitConversions._
33-
import collection.JavaConverters._
34-
3526
object StreamingPipeline extends BasePipeline with Serializable {
3627
override def createPipeline(
3728
sparkSession: SparkSession,
@@ -42,15 +33,10 @@ object StreamingPipeline extends BasePipeline with Serializable {
4233
val featureTable = config.featureTable
4334
val projection =
4435
inputProjection(config.source, featureTable.features, featureTable.entities)
45-
val validator = new RowValidator(featureTable)
36+
val validator = new RowValidator(featureTable, config.source.timestampColumn)
4637

47-
val defaultInstance = defaultInstanceFromProtoClass(
48-
config.source.asInstanceOf[StreamingSource].classpath
49-
)
50-
val protoParser = udf(
51-
ProtoReflection.createMessageParser(defaultInstance),
52-
inferSchemaFromProto(defaultInstance)
53-
)
38+
val messageParser =
39+
protoParser(sparkSession, config.source.asInstanceOf[StreamingSource].classpath)
5440

5541
val input = config.source match {
5642
case source: KafkaSource =>
@@ -62,7 +48,7 @@ object StreamingPipeline extends BasePipeline with Serializable {
6248
}
6349

6450
val projected = input
65-
.withColumn("features", protoParser($"value"))
51+
.withColumn("features", messageParser($"value"))
6652
.select("features.*")
6753
.select(projection: _*)
6854

@@ -105,17 +91,11 @@ object StreamingPipeline extends BasePipeline with Serializable {
10591
Some(query)
10692
}
10793

108-
private def defaultInstanceFromProtoClass(className: String): GeneratedMessageV3 =
109-
Class
110-
.forName(className, true, getClass.getClassLoader)
111-
.asInstanceOf[Class[GeneratedMessageV3]]
112-
.getMethod("getDefaultInstance")
113-
.invoke(null)
114-
.asInstanceOf[GeneratedMessageV3]
94+
private def protoParser(sparkSession: SparkSession, className: String) = {
95+
val protoRegistry = ProtoRegistryFactory.resolveProtoRegistry(sparkSession)
11596

116-
private def inferSchemaFromProto(defaultInstance: GeneratedMessageV3) =
117-
StructType(
118-
defaultInstance.getDescriptorForType.getFields.flatMap(ProtoReflection.structFieldFor)
119-
)
97+
val parser: Array[Byte] => Row = ProtoReflection.createMessageParser(protoRegistry, className)
12098

99+
udf(parser, ProtoReflection.inferSchema(protoRegistry.getProtoDescriptor(className)))
100+
}
121101
}

‎spark/ingestion/src/main/scala/feast/ingestion/registry/proto/LocalProtoRegistry.scala‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion.registry.proto
18+
import java.io.{IOException, ObjectInputStream}
19+
20+
import com.google.protobuf.Descriptors.Descriptor
21+
22+
import collection.mutable
23+
import scala.util.control.NonFatal
24+
25+
class LocalProtoRegistry extends ProtoRegistry {
26+
@transient
27+
private var cache: mutable.Map[String, Descriptor] = mutable.Map.empty
28+
29+
@throws(classOf[IOException])
30+
private def readObject(ois: ObjectInputStream): Unit = {
31+
try {
32+
ois.defaultReadObject()
33+
cache = mutable.Map.empty
34+
} catch {
35+
case NonFatal(e) =>
36+
throw new IOException(e)
37+
}
38+
}
39+
40+
override def getProtoDescriptor(className: String): Descriptor = {
41+
if (!cache.contains(className)) {
42+
cache(className) = Class
43+
.forName(className, true, getClass.getClassLoader)
44+
.getMethod("getDescriptor")
45+
.invoke(null)
46+
.asInstanceOf[Descriptor]
47+
}
48+
49+
cache(className)
50+
}
51+
}

‎spark/ingestion/src/main/scala/feast/ingestion/registry/proto/ProtoRegistry.scala‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion.registry.proto
18+
19+
import com.google.protobuf.Descriptors.Descriptor
20+
21+
trait ProtoRegistry extends Serializable {
22+
def getProtoDescriptor(className: String): Descriptor
23+
}

‎spark/ingestion/src/main/scala/feast/ingestion/registry/proto/ProtoRegistryFactory.scala‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion.registry.proto
18+
19+
import org.apache.spark.sql.SparkSession
20+
21+
object ProtoRegistryFactory {
22+
val CONFIG_PREFIX = "feast.ingestion.registry.proto"
23+
val PROTO_REGISTRY_KIND = s"$CONFIG_PREFIX.kind"
24+
val DEFAULT_KIND = "local"
25+
26+
def resolveProtoRegistry(sparkSession: SparkSession): ProtoRegistry = {
27+
val config = sparkSession.sparkContext.getConf
28+
val kind = config.get(PROTO_REGISTRY_KIND, DEFAULT_KIND)
29+
val properties = config.getAllWithPrefix(CONFIG_PREFIX).toMap
30+
protoRegistry(kind, properties)
31+
}
32+
33+
private def protoRegistry(name: String, properties: Map[String, String]): ProtoRegistry =
34+
name match {
35+
case "local" => new LocalProtoRegistry
36+
case "stencil" => new StencilProtoRegistry(properties("url"))
37+
}
38+
}

‎spark/ingestion/src/main/scala/feast/ingestion/registry/proto/StencilProtoRegistry.scala‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright 2018-2020 The Feast Authors
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* https://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package feast.ingestion.registry.proto
18+
import java.util.Collections
19+
20+
import com.google.protobuf.Descriptors
21+
import com.gojek.de.stencil.StencilClientFactory
22+
import com.gojek.de.stencil.client.StencilClient
23+
24+
class StencilProtoRegistry(val url: String) extends ProtoRegistry {
25+
26+
val stencilClient: StencilClient =
27+
StencilClientFactory.getClient(url, Collections.emptyMap[String, String])
28+
29+
override def getProtoDescriptor(className: String): Descriptors.Descriptor = {
30+
stencilClient.get(className)
31+
}
32+
}

0 commit comments

Comments
 (0)