Skip to content

Commit 372bd69

Browse files
pyalexpyalex
committed
tests passing
Signed-off-by: Oleksii Moskalenko <moskalenko.alexey@gmail.com>
1 parent 158a887 commit 372bd69

8 files changed

Lines changed: 92 additions & 33 deletions

File tree

‎spark/ingestion/pom.xml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,14 +189,14 @@
189189
<dependency>
190190
<groupId>com.dimafeng</groupId>
191191
<artifactId>testcontainers-scala-scalatest_${scala.version}</artifactId>
192-
<version>0.38.6</version>
192+
<version>0.38.8</version>
193193
<scope>test</scope>
194194
</dependency>
195195

196196
<dependency>
197197
<groupId>com.dimafeng</groupId>
198198
<artifactId>testcontainers-scala-kafka_${scala.version}</artifactId>
199-
<version>0.38.6</version>
199+
<version>0.38.8</version>
200200
<scope>test</scope>
201201
</dependency>
202202

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ object BatchPipeline extends BasePipeline {
3939
val projection =
4040
inputProjection(config.source, featureTable.features, featureTable.entities)
4141
val rowValidator = new RowValidator(featureTable, config.source.eventTimestampColumn)
42+
val metrics = new IngestionPipelineMetrics
4243

4344
val input = config.source match {
4445
case source: BQSource =>
@@ -68,7 +69,7 @@ object BatchPipeline extends BasePipeline {
6869
}
6970

7071
val validRows = projected
71-
.mapPartitions(IngestionPipelineMetrics.incrementRead)
72+
.mapPartitions(metrics.incrementRead)
7273
.filter(rowValidator.allChecks)
7374

7475
validRows.write
@@ -84,7 +85,7 @@ object BatchPipeline extends BasePipeline {
8485
case Some(path) =>
8586
projected
8687
.filter(!rowValidator.allChecks)
87-
.mapPartitions(IngestionPipelineMetrics.incrementDeadletters)
88+
.mapPartitions(metrics.incrementDeadLetters)
8889
.write
8990
.format("parquet")
9091
.mode(SaveMode.Append)

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@ object StreamingPipeline extends BasePipeline with Serializable {
5656
val featureTable = config.featureTable
5757
val projection =
5858
inputProjection(config.source, featureTable.features, featureTable.entities)
59-
val rowValidator = new RowValidator(featureTable, config.source.eventTimestampColumn)
60-
59+
val rowValidator = new RowValidator(featureTable, config.source.eventTimestampColumn)
60+
val metrics = new IngestionPipelineMetrics
6161
val validationUDF = createValidationUDF(sparkSession, config)
6262

6363
val input = config.source match {
@@ -107,7 +107,7 @@ object StreamingPipeline extends BasePipeline with Serializable {
107107
implicit def rowEncoder: Encoder[Row] = RowEncoder(rowsAfterValidation.schema)
108108

109109
rowsAfterValidation
110-
.mapPartitions(IngestionPipelineMetrics.incrementRead)
110+
.mapPartitions(metrics.incrementRead)
111111
.filter(if (config.doNotIngestInvalidRows) expr("_isValid") else rowValidator.allChecks)
112112
.write
113113
.format("feast.ingestion.stores.redis")
@@ -120,10 +120,9 @@ object StreamingPipeline extends BasePipeline with Serializable {
120120

121121
config.deadLetterPath match {
122122
case Some(path) =>
123-
124123
rowsAfterValidation
125124
.filter("!_isValid")
126-
.mapPartitions(IngestionPipelineMetrics.incrementDeadletters)
125+
.mapPartitions(metrics.incrementDeadLetters)
127126
.write
128127
.format("parquet")
129128
.mode(SaveMode.Append)

‎spark/ingestion/src/main/scala/feast/ingestion/metrics/IngestionPipelineMetrics.scala‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,36 +18,39 @@ package feast.ingestion.metrics
1818

1919
import org.apache.spark.SparkEnv
2020
import org.apache.spark.metrics.source.IngestionPipelineMetricSource
21+
import org.apache.spark.sql.Row
2122

22-
object IngestionPipelineMetrics {
23-
def incrementDeadletters[A](rowIterator: Iterator[A]): Iterator[A] = {
23+
class IngestionPipelineMetrics extends Serializable {
24+
25+
def incrementDeadLetters(rowIterator: Iterator[Row]): Iterator[Row] = {
26+
val materialized = rowIterator.toArray
2427
if (metricSource.nonEmpty)
25-
metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc(rowIterator.length)
28+
metricSource.get.METRIC_DEADLETTER_ROWS_INSERTED.inc(materialized.length)
2629

27-
rowIterator
30+
materialized.toIterator
2831
}
2932

30-
def incrementRead[A](rowIterator: Iterator[A]): Iterator[A] = {
33+
def incrementRead(rowIterator: Iterator[Row]): Iterator[Row] = {
34+
val materialized = rowIterator.toArray
3135
if (metricSource.nonEmpty)
32-
metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc(rowIterator.length)
36+
metricSource.get.METRIC_ROWS_READ_FROM_SOURCE.inc(materialized.length)
3337

34-
rowIterator
38+
materialized.toIterator
3539
}
3640

3741
private lazy val metricSource: Option[IngestionPipelineMetricSource] = {
38-
this.synchronized {
39-
if (
40-
SparkEnv.get.metricsSystem
41-
.getSourcesByName(IngestionPipelineMetricSource.sourceName)
42-
.isEmpty
43-
) {
44-
SparkEnv.get.metricsSystem.registerSource(new IngestionPipelineMetricSource)
42+
val metricsSystem = SparkEnv.get.metricsSystem
43+
IngestionPipelineMetricsLock.synchronized {
44+
if (metricsSystem.getSourcesByName(IngestionPipelineMetricSource.sourceName).isEmpty) {
45+
metricsSystem.registerSource(new IngestionPipelineMetricSource)
4546
}
4647
}
4748

48-
SparkEnv.get.metricsSystem.getSourcesByName(IngestionPipelineMetricSource.sourceName) match {
49+
metricsSystem.getSourcesByName(IngestionPipelineMetricSource.sourceName) match {
4950
case Seq(head) => Some(head.asInstanceOf[IngestionPipelineMetricSource])
5051
case _ => None
5152
}
5253
}
5354
}
55+
56+
private object IngestionPipelineMetricsLock

‎spark/ingestion/src/main/scala/feast/ingestion/metrics/StatsdReporterWithTags.scala‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,8 +125,13 @@ class StatsdReporterWithTags(
125125
private def reportGauge(name: String, gauge: Gauge[_])(implicit socket: DatagramSocket): Unit =
126126
formatAny(gauge.getValue).foreach(v => send(fullName(name), v, GAUGE))
127127

128-
private def reportCounter(name: String, counter: Counter)(implicit socket: DatagramSocket): Unit =
129-
send(fullName(name), format(counter.getCount), COUNTER)
128+
private def reportCounter(name: String, counter: Counter)(implicit
129+
socket: DatagramSocket
130+
): Unit = {
131+
val snapshot = counter.getCount
132+
send(fullName(name), format(snapshot), COUNTER)
133+
counter.dec(snapshot) // reset counter
134+
}
130135

131136
private def reportHistogram(name: String, histogram: Histogram)(implicit
132137
socket: DatagramSocket

‎spark/ingestion/src/test/scala/feast/ingestion/BatchPipelineIT.scala‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import collection.JavaConverters._
2323
import com.dimafeng.testcontainers.{ForAllTestContainer, GenericContainer}
2424
import com.google.protobuf.util.Timestamps
2525
import feast.proto.types.ValueProto.ValueType
26-
import org.apache.spark.SparkConf
26+
import org.apache.spark.{SparkConf, SparkEnv}
2727
import org.joda.time.{DateTime, Seconds}
2828
import org.scalacheck._
2929
import org.scalatest._
@@ -46,25 +46,28 @@ case class TestRow(
4646
class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
4747

4848
override val container = GenericContainer("redis:6.0.8", exposedPorts = Seq(6379))
49+
val statsDStub = new StatsDStub
4950

5051
override def withSparkConfOverrides(conf: SparkConf): SparkConf = conf
5152
.set("spark.redis.host", container.host)
5253
.set("spark.redis.port", container.mappedPort(6379).toString)
54+
.set("spark.metrics.conf.*.sink.statsd.port", statsDStub.port.toString)
5355

5456
trait Scope {
5557
val jedis = new Jedis("localhost", container.mappedPort(6379))
5658
jedis.flushAll()
5759

60+
statsDStub.receivedMetrics // clean the buffer
61+
5862
implicit def testRowEncoder: Encoder[TestRow] = ExpressionEncoder()
59-
val statsDStub = new StatsDStub
6063

6164
def rowGenerator(start: DateTime, end: DateTime, customerGen: Option[Gen[String]] = None) =
6265
for {
6366
customer <- customerGen.getOrElse(Gen.asciiPrintableStr)
6467
feature1 <- Gen.choose(0, 100)
6568
feature2 <- Gen.choose[Float](0, 1)
6669
eventTimestamp <- Gen
67-
.choose(0, Seconds.secondsBetween(start, end).getSeconds)
70+
.choose(0, Seconds.secondsBetween(start, end).getSeconds - 1)
6871
.map(start.withMillisOfSecond(0).plusSeconds)
6972
} yield TestRow(
7073
customer,
@@ -98,7 +101,7 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
98101
),
99102
startTime = DateTime.parse("2020-08-01"),
100103
endTime = DateTime.parse("2020-09-01"),
101-
metrics = Some(StatsDConfig(host="localhost", port=statsDStub.port))
104+
metrics = Some(StatsDConfig(host = "localhost", port = statsDStub.port))
102105
)
103106
}
104107

@@ -129,6 +132,14 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
129132
keyTTL shouldEqual -1
130133

131134
})
135+
136+
SparkEnv.get.metricsSystem.report()
137+
statsDStub.receivedMetrics should contain.allElementsOf(
138+
Map(
139+
"driver.ingestion_pipeline.read_from_source_count" -> rows.length,
140+
"driver.redis_sink.feature_row_ingested_count" -> rows.length
141+
)
142+
)
132143
}
133144

134145
"Parquet source file" should "be ingested in redis with expiry time equal to the largest of (event_timestamp + max_age) for" +
@@ -466,6 +477,13 @@ class BatchPipelineIT extends SparkSpec with ForAllTestContainer {
466477
.toString
467478
)
468479
.count() should be(rows.length)
480+
481+
SparkEnv.get.metricsSystem.report()
482+
statsDStub.receivedMetrics should contain.allElementsOf(
483+
Map(
484+
"driver.ingestion_pipeline.deadletter_count" -> rows.length
485+
)
486+
)
469487
}
470488

471489
"Columns from source" should "be mapped according to configuration" in new Scope {

‎spark/ingestion/src/test/scala/feast/ingestion/SparkSpec.scala‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,10 @@ class SparkSpec extends UnitSpec with BeforeAndAfter {
3636
"org.apache.spark.metrics.sink.StatsdSinkWithTags"
3737
)
3838
.set("spark.metrics.conf.*.sink.statsd.host", "localhost")
39-
.set("spark.metrics.conf.*.sink.statsd.port", "8125")
39+
.set("spark.metrics.conf.*.sink.statsd.period", "999") // disable scheduled reporting
40+
.set("spark.metrics.conf.*.sink.statsd.unit", "minutes")
41+
.set("spark.metrics.labels", "job_id=test")
42+
.set("spark.metrics.namespace", "")
4043
.set("spark.sql.legacy.allowUntypedScalaUDF", "true")
4144
.set("spark.sql.execution.arrow.maxRecordsPerBatch", "50000")
4245

‎spark/ingestion/src/test/scala/feast/ingestion/metrics/StatsDStub.scala‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,19 @@
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+
*/
117
package feast.ingestion.metrics
218

319
import java.net.{DatagramPacket, DatagramSocket, SocketTimeoutException}
@@ -12,11 +28,11 @@ class StatsDStub {
1228

1329
def receive: Array[String] = {
1430
val messages: ArrayBuffer[String] = ArrayBuffer()
15-
var finished = false
31+
var finished = false
1632

1733
do {
1834
val buf = new Array[Byte](65535)
19-
val p = new DatagramPacket(buf, buf.length)
35+
val p = new DatagramPacket(buf, buf.length)
2036
try {
2137
socket.receive(p)
2238
} catch {
@@ -28,4 +44,18 @@ class StatsDStub {
2844

2945
messages.toArray
3046
}
47+
48+
private val metricLine = """(.+):(.+)\|(.+)#(.+)""".r
49+
50+
def receivedMetrics: Map[String, Float] = {
51+
receive
52+
.flatMap {
53+
case metricLine(name, value, type_, tags) =>
54+
Seq(name -> value.toFloat)
55+
case s: String =>
56+
Seq()
57+
}
58+
.groupBy(_._1)
59+
.mapValues(_.map(_._2).sum)
60+
}
3161
}

0 commit comments

Comments
 (0)