fmt

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

online/src/main/scala/ai/chronon/online/ExternalSourceRegistry.scala

@@ -50,8 +50,7 @@ class ExternalSourceRegistry extends Serializable {
   // 1. keys match
   // 2. report missing & extra values
   // 3. schema integrity of returned values
-  def fetchRequests(requests: Seq[Request], context: Context)(implicit
-      ec: ExecutionContext): Future[Seq[Response]] = {
+  def fetchRequests(requests: Seq[Request], context: Context)(implicit ec: ExecutionContext): Future[Seq[Response]] = {
     val startTime = System.currentTimeMillis()
     // we make issue one batch request per external source and flatten out it later
     val responsesByNameF: List[Future[Seq[Response]]] = requests

online/src/main/scala/ai/chronon/online/fetcher/Fetcher.scala

@@ -433,7 +433,7 @@ class Fetcher(val kvStore: KVStore,
 
   def fetchJoinSchema(joinName: String): Try[JoinSchemaResponse] = {
     val startTime = System.currentTimeMillis()
-    val ctx =Metrics.Context(Metrics.Environment.JoinSchemaFetching, join = joinName)
+    val ctx = Metrics.Context(Metrics.Environment.JoinSchemaFetching, join = joinName)
 
     val joinCodecTry = joinCodecCache(joinName)
 
@@ -478,7 +478,6 @@ class Fetcher(val kvStore: KVStore,
                                            joinRequest: Request,
                                            part: ExternalPart) {
 
-
     lazy val context: Metrics.Context =
       Metrics.Context(Metrics.Environment.JoinFetching, join = joinRequest.name, groupBy = part.fullName)
   }

online/src/main/scala/ai/chronon/online/fetcher/GroupByFetcher.scala

@@ -46,7 +46,8 @@ class GroupByFetcher(fetchContext: FetchContext, metadataStore: MetadataStore)
     .map { groupByServingInfo =>
       import ai.chronon.online.metrics
       val context =
-        request.context.getOrElse(metrics.Metrics.Context(metrics.Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
+        request.context.getOrElse(
+          metrics.Metrics.Context(metrics.Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
       context.increment("group_by_request.count")
       var batchKeyBytes: Array[Byte] = null
       var streamingKeyBytes: Array[Byte] = null

online/src/main/scala/ai/chronon/online/fetcher/MetadataStore.scala

@@ -123,7 +123,8 @@ class MetadataStore(fetchContext: FetchContext) {
           }
       },
       { team =>
-       Metrics.Context(environment = "group_by.list.fetch", groupBy = team) }
+        Metrics.Context(environment = "group_by.list.fetch", groupBy = team)
+      }
     )
   }
 
@@ -137,8 +138,10 @@ class MetadataStore(fetchContext: FetchContext) {
             throw e
           }
       },
-      { team => import ai.chronon.online.metrics
-        metrics.Metrics.Context(environment = "join.list.fetch", groupBy = team) }
+      { team =>
+        import ai.chronon.online.metrics
+        metrics.Metrics.Context(environment = "join.list.fetch", groupBy = team)
+      }
     )
   }
 
@@ -167,8 +170,10 @@ class MetadataStore(fetchContext: FetchContext) {
         .distribution(metrics.Metrics.Name.LatencyMillis, System.currentTimeMillis() - startTimeMs)
       result
     },
-    { join => import ai.chronon.online.metrics
-      metrics.Metrics.Context(environment = "join.meta.fetch", join = join) }
+    { join =>
+      import ai.chronon.online.metrics
+      metrics.Metrics.Context(environment = "join.meta.fetch", join = join)
+    }
   )
 
   def putJoinConf(join: Join): Unit = {
@@ -252,8 +257,10 @@ class MetadataStore(fetchContext: FetchContext) {
 
     new TTLCache[String, Try[JoinCodec]](
       codecBuilder,
-      { join: String => import ai.chronon.online.metrics
-        metrics.Metrics.Context(environment = "join.codec.fetch", join = join) },
+      { join: String =>
+        import ai.chronon.online.metrics
+        metrics.Metrics.Context(environment = "join.codec.fetch", join = join)
+      },
       onCreateFunc = onCreateFunc
     )
   }
@@ -374,8 +381,10 @@ class MetadataStore(fetchContext: FetchContext) {
           Success(new GroupByServingInfoParsed(groupByServingInfo, partitionSpec))
         }
       },
-      { gb => import ai.chronon.online.metrics
-        metrics.Metrics.Context(environment = "group_by.serving_info.fetch", groupBy = gb) }
+      { gb =>
+        import ai.chronon.online.metrics
+        metrics.Metrics.Context(environment = "group_by.serving_info.fetch", groupBy = gb)
+      }
     )
 
   def put(

spark/src/main/scala/ai/chronon/spark/streaming/JoinSourceRunner.scala

@@ -30,6 +30,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, Encoder, Encoders, Row, SparkSe
 import org.apache.spark.sql.streaming.{DataStreamWriter, Trigger}
 import org.apache.spark.sql.types.{BooleanType, LongType, StructField, StructType}
 import org.slf4j.{Logger, LoggerFactory}
+import ai.chronon.online.metrics
 
 import java.time.{Instant, ZoneId, ZoneOffset}
 import java.time.format.DateTimeFormatter
@@ -64,11 +65,10 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
     apiImpl: Api)
     extends Serializable {
 
-  import ai.chronon.online.metrics
-
   @transient implicit lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
-  val context: metrics.Metrics.Context = metrics.Metrics.Context(metrics.Metrics.Environment.GroupByStreaming, groupByConf)
+  val context: metrics.Metrics.Context =
+    metrics.Metrics.Context(metrics.Metrics.Environment.GroupByStreaming, groupByConf)
 
   private case class Schemas(leftStreamSchema: StructType,
                              leftSourceSchema: StructType,
@@ -123,7 +123,6 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
     private val streamingDataset: String = groupByConf.streamingDataset
 
     def toPutRequest(input: Row): KVStore.PutRequest = {
-      import ai.chronon.online.metrics
       val keys = keyIndices.map(input.get)
       val values = valueIndices.map(input.get)
 
@@ -230,7 +229,6 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
     val deserialized: Dataset[Mutation] = df
       .as[Array[Byte]]
       .map { arr =>
-        import ai.chronon.online.metrics
         ingressContext.increment(metrics.Metrics.Name.RowCount)
         ingressContext.count(metrics.Metrics.Name.Bytes, arr.length)
         try {
@@ -289,7 +287,6 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
   }
 
   def chainedStreamingQuery: DataStreamWriter[Row] = {
-    import ai.chronon.online.metrics
     val joinSource = groupByConf.streamingSource.get.getJoinSource
     val left = joinSource.join.left
     val topic = TopicInfo.parse(left.topic)
@@ -357,7 +354,6 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
 
           val rowsScala = rows.toScala.toArray
           val requests = rowsScala.map { row =>
-            import ai.chronon.online.metrics
             val keyMap = row.getValuesMap[AnyRef](leftColumns)
             val eventTs = row.get(leftTimeIndex).asInstanceOf[Long]
             context.distribution(metrics.Metrics.Name.LagMillis, System.currentTimeMillis() - eventTs)
@@ -368,12 +364,10 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
           val microBatchTimestamp =
             percentile(rowsScala.map(_.get(leftTimeIndex).asInstanceOf[Long]), timePercentile)
           if (microBatchTimestamp.isDefined) {
-            import ai.chronon.online.metrics
             val microBatchLag = System.currentTimeMillis() - microBatchTimestamp.get
             context.distribution(metrics.Metrics.Name.BatchLagMillis, microBatchLag)
 
             if (minimumQueryDelayMs > 0 && microBatchLag >= 0 && microBatchLag < minimumQueryDelayMs) {
-              import ai.chronon.online.metrics
               val sleepMillis = minimumQueryDelayMs - microBatchLag
               Thread.sleep(sleepMillis)
               context.distribution(metrics.Metrics.Name.QueryDelaySleepMillis, sleepMillis)
@@ -416,7 +410,6 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
 
     def emitRequestMetric(request: PutRequest, context: metrics.Metrics.Context): Unit = {
       request.tsMillis.foreach { ts: Long =>
-        import ai.chronon.online.metrics
         context.distribution(metrics.Metrics.Name.FreshnessMillis, System.currentTimeMillis() - ts)
         context.increment(metrics.Metrics.Name.RowCount)
         context.distribution(metrics.Metrics.Name.ValueBytes, request.valueBytes.length)