feat(cache): add control auto cache-count option (#211)

YotpoLtd · Aug 1, 2019 · f3cbf34 · f3cbf34
1 parent 694c46f
commit f3cbf34
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 7 deletions.
diff --git a/config/job_config_sample.yaml b/config/job_config_sample.yaml
@@ -174,4 +174,7 @@ streaming:
   batchMode: true
   # Add any other options supported by the DataStreamWriter
   extraOptions:
-    opt: val
+    opt: val
+
+# Optional: controls caching and counting on each output (default is true)
+cacheCountOnOutput: false
diff --git a/src/main/scala/com/yotpo/metorikku/configuration/job/Configuration.scala b/src/main/scala/com/yotpo/metorikku/configuration/job/Configuration.scala
@@ -16,7 +16,8 @@ case class Configuration(metrics: Option[Seq[String]],
                          var showPreviewLines: Option[Int],
                          var explain: Option[Boolean],
                          var appName: Option[String],
-                         var continueOnFailedStep: Option[Boolean]) {
+                         var continueOnFailedStep: Option[Boolean],
+                         var cacheCountOnOutput: Option[Boolean]) {
 
   require(metrics.isDefined, "metrics files paths are mandatory")
 
@@ -25,6 +26,7 @@ case class Configuration(metrics: Option[Seq[String]],
   explain = Option(explain.getOrElse(false))
   appName = Option(appName.getOrElse("Metorikku"))
   continueOnFailedStep = Option(continueOnFailedStep.getOrElse(false))
+  cacheCountOnOutput = Option(cacheCountOnOutput.getOrElse(true))
 
   def getReaders: Seq[Reader] = inputs.getOrElse(Map()).map {
     case (name, input) => input.getReader(name) }.toSeq

diff --git a/src/main/scala/com/yotpo/metorikku/metric/Metric.scala b/src/main/scala/com/yotpo/metorikku/metric/Metric.scala
@@ -80,10 +80,18 @@ case class Metric(configuration: Configuration, metricDir: File, metricName: Str
                          dataFrameName: String,
                          writer: Writer,
                          outputType: OutputType,
-                         instrumentationProvider: InstrumentationProvider): Unit = {
-    dataFrame.cache()
+                         instrumentationProvider: InstrumentationProvider,
+                         cacheCountOnOutput: Option[Boolean]): Unit = {
+
+    val dataFrameCount = cacheCountOnOutput match {
+      case Some(true) => {
+        dataFrame.cache()
+        dataFrame.count()
+      }
+      case _ => 0
+    }
     val tags = Map("metric" -> metricName, "dataframe" -> dataFrameName, "output_type" -> outputType.toString)
-    instrumentationProvider.count(name="counter", value=dataFrame.count(), tags=tags)
+    instrumentationProvider.count(name="counter", value=dataFrameCount, tags=tags)
     log.info(s"Starting to Write results of ${dataFrameName}")
     try {
       writer.write(dataFrame)
@@ -126,7 +134,7 @@ case class Metric(configuration: Configuration, metricDir: File, metricName: Str
           }
           else {
             writeBatch(dataFrame, dataFrameName, writer,
-              outputConfig.outputType, job.instrumentationClient)
+              outputConfig.outputType, job.instrumentationClient, job.config.cacheCountOnOutput)
           }
         })
 

diff --git a/src/main/scala/com/yotpo/metorikku/test/Tester.scala b/src/main/scala/com/yotpo/metorikku/test/Tester.scala
@@ -48,7 +48,7 @@ case class Tester(config: TesterConfig) {
     val variables = params.variables
     val inputs = getMockFilesFromDir(config.test.mocks, config.basePath)
     Configuration(Option(metrics),inputs, variables, None, None, None, None, None,
-      Option(config.preview > 0), None, None, Option(config.preview), None, None, None)
+      Option(config.preview > 0), None, None, Option(config.preview), None, None, None, None)
   }
 
   private def getMockFilesFromDir(mocks: Option[List[Mock]], testDir: File): Option[Map[String, Input]] = {