WIP

airbnb · Dec 6, 2023 · 3695355 · 3695355
1 parent b322934
commit 3695355
Show file tree

Hide file tree

Showing 9 changed files with 24 additions and 17 deletions.
diff --git a/api/src/main/scala/ai/chronon/api/Extensions.scala b/api/src/main/scala/ai/chronon/api/Extensions.scala
@@ -784,7 +784,7 @@ object Extensions {
   }
 
   implicit class JoinOps(val join: Join) extends Serializable {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     // all keys as they should appear in left that are being used on right
     def leftKeyCols: Array[String] = {
       join.joinParts.toScala

diff --git a/spark/src/main/scala/ai/chronon/spark/Analyzer.scala b/spark/src/main/scala/ai/chronon/spark/Analyzer.scala
@@ -336,7 +336,8 @@ class Analyzer(tableUtils: TableUtils,
       logger.info(s"----- Schema validation completed. Found ${keysWithError.size} errors")
       val keyErrorSet: Set[(String, String)] = keysWithError.toSet
       logger.info(keyErrorSet.map { case (key, errorMsg) => s"$key => $errorMsg" }.mkString("\n"))
-      logger.info(s"---- Table permission check completed. Found permission errors in ${noAccessTables.size} tables ----")
+      logger.info(
+        s"---- Table permission check completed. Found permission errors in ${noAccessTables.size} tables ----")
       logger.info(noAccessTables.mkString("\n"))
       logger.info(s"---- Data availability check completed. Found issue in ${dataAvailabilityErrors.size} tables ----")
       dataAvailabilityErrors.foreach(error =>

diff --git a/spark/src/main/scala/ai/chronon/spark/Driver.scala b/spark/src/main/scala/ai/chronon/spark/Driver.scala
@@ -214,7 +214,7 @@ object Driver {
   }
 
   object JoinBackfill {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     class Args
         extends Subcommand("join")
         with OfflineSubcommand
@@ -256,12 +256,13 @@ object Driver {
       }
 
       df.show(numRows = 3, truncate = 0, vertical = true)
-      logger.info(s"\nShowing three rows of output above.\nQuery table `${args.joinConf.metaData.outputTable}` for more.\n")
+      logger.info(
+        s"\nShowing three rows of output above.\nQuery table `${args.joinConf.metaData.outputTable}` for more.\n")
     }
   }
 
   object GroupByBackfill {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     class Args
         extends Subcommand("group-by-backfill")
         with OfflineSubcommand
@@ -525,7 +526,7 @@ object Driver {
   }
 
   object FetcherCli {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
 
     class Args extends Subcommand("fetch") with OnlineSubcommand {
       val keyJson: ScallopOption[String] = opt[String](required = false, descr = "json of the keys to fetch")
@@ -565,7 +566,8 @@ object Driver {
         )
           series.get(keyMap("statsKey").asInstanceOf[String])
         else series
-      logger.info(s"--- [FETCHED RESULT] ---\n${objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(toPrint)}")
+      logger.info(
+        s"--- [FETCHED RESULT] ---\n${objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(toPrint)}")
     }
 
     def run(args: Args): Unit = {
@@ -644,7 +646,7 @@ object Driver {
   }
 
   object MetadataUploader {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     class Args extends Subcommand("metadata-upload") with OnlineSubcommand {
       val confPath: ScallopOption[String] =
         opt[String](required = true, descr = "Path to the Chronon config file or directory")
@@ -689,7 +691,7 @@ object Driver {
   }
 
   object GroupByStreaming {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     def dataStream(session: SparkSession, host: String, topic: String): DataFrame = {
       TopicChecker.topicShouldExist(topic, host)
       session.streams.addListener(new StreamingQueryListener() {
@@ -713,7 +715,7 @@ object Driver {
     }
 
     class Args extends Subcommand("group-by-streaming") with OnlineSubcommand {
-  private val logger = LoggerFactory.getLogger(getClass)
+      private val logger = LoggerFactory.getLogger(getClass)
       val confPath: ScallopOption[String] = opt[String](required = true, descr = "path to groupBy conf")
       val DEFAULT_LAG_MILLIS = 2000 // 2seconds
       val kafkaBootstrap: ScallopOption[String] =

diff --git a/spark/src/main/scala/ai/chronon/spark/Extensions.scala b/spark/src/main/scala/ai/chronon/spark/Extensions.scala
@@ -84,7 +84,7 @@ object Extensions {
   }
 
   implicit class DataframeOps(df: DataFrame) {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     private implicit val tableUtils: TableUtils = TableUtils(df.sparkSession)
 
     // This is safe to call on dataframes that are un-shuffled from their disk sources -

diff --git a/spark/src/main/scala/ai/chronon/spark/JoinBase.scala b/spark/src/main/scala/ai/chronon/spark/JoinBase.scala
@@ -167,7 +167,8 @@ abstract class JoinBase(joinConf: api.Join,
         }
       } catch {
         case e: Exception =>
-          logger.info(s"Error while processing groupBy: ${joinConf.metaData.name}/${joinPart.groupBy.getMetaData.getName}")
+          logger.info(
+            s"Error while processing groupBy: ${joinConf.metaData.name}/${joinPart.groupBy.getMetaData.getName}")
           throw e
       }
       if (tableUtils.tableExists(partTable)) {
@@ -193,7 +194,8 @@ abstract class JoinBase(joinConf: api.Join,
     val rowCount = leftDfWithStats.get.count
     val unfilledRange = leftDfWithStats.get.partitionRange
 
-    logger.info(s"\nBackfill is required for ${joinPart.groupBy.metaData.name} for $rowCount rows on range $unfilledRange")
+    logger.info(
+      s"\nBackfill is required for ${joinPart.groupBy.metaData.name} for $rowCount rows on range $unfilledRange")
     val rightBloomMap =
       JoinUtils.genBloomFilterIfNeeded(leftDf,
                                        joinPart,

diff --git a/spark/src/main/scala/ai/chronon/spark/LabelJoin.scala b/spark/src/main/scala/ai/chronon/spark/LabelJoin.scala
@@ -94,7 +94,8 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) {
       labelTable
     } else {
       // creating final join view with feature join output table
-      logger.info(s"Joining label table : ${outputLabelTable} with joined output table : ${joinConf.metaData.outputTable}")
+      logger.info(
+        s"Joining label table : ${outputLabelTable} with joined output table : ${joinConf.metaData.outputTable}")
       JoinUtils.createOrReplaceView(
         joinConf.metaData.outputFinalView,
         leftTable = joinConf.metaData.outputTable,

diff --git a/spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala b/spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala
@@ -71,7 +71,8 @@ class CompareJob(
 
     // Save the comparison table
     logger.info("Saving comparison output..")
-    logger.info(s"Comparison schema ${compareDf.schema.fields.map(sb => (sb.name, sb.dataType)).toMap.mkString("\n - ")}")
+    logger.info(
+      s"Comparison schema ${compareDf.schema.fields.map(sb => (sb.name, sb.dataType)).toMap.mkString("\n - ")}")
     tableUtils.insertUnPartitioned(compareDf, comparisonTableName, tableProps, saveMode = SaveMode.Overwrite)
 
     // Save the metrics table

diff --git a/spark/src/main/scala/ai/chronon/spark/streaming/JoinSourceRunner.scala b/spark/src/main/scala/ai/chronon/spark/streaming/JoinSourceRunner.scala
@@ -107,7 +107,7 @@ class JoinSourceRunner(groupByConf: api.GroupBy, conf: Map[String, String] = Map
   private val microBatchIntervalMillis: Int = getProp("batch_interval_millis", "1000").toInt
 
   private case class PutRequestHelper(inputSchema: StructType) extends Serializable {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     private val keyIndices: Array[Int] = keyColumns.map(inputSchema.fieldIndex)
     private val valueIndices: Array[Int] = valueColumns.map(inputSchema.fieldIndex)
     private val tsIndex: Int = inputSchema.fieldIndex(eventTimeColumn)

diff --git a/spark/src/main/scala/ai/chronon/spark/streaming/TopicChecker.scala b/spark/src/main/scala/ai/chronon/spark/streaming/TopicChecker.scala
@@ -86,7 +86,7 @@ object TopicChecker {
   }
 
   class Args(arguments: Seq[String]) extends ScallopConf(arguments) {
-  private val logger = LoggerFactory.getLogger(getClass)
+    private val logger = LoggerFactory.getLogger(getClass)
     val conf: ScallopOption[String] = opt[String](descr = "Conf to pull topic and bootstrap server information")
     val bootstrap: ScallopOption[String] = opt[String](descr = "Kafka bootstrap server in host:port format")
     val topic: ScallopOption[String] = opt[String](descr = "kafka topic to check metadata for")