working tests

Author: nikhil-zlai

api/py/ai/chronon/utils.py

@@ -262,6 +262,7 @@ def join_part_output_table_name(join, jp, full_name: bool = False):
     def partOutputTable(jp: JoinPart): String = (Seq(join.metaData.outputTable) ++ Option(jp.prefix) :+
       jp.groupBy.metaData.cleanName).mkString("_")
     """
+    print(join)
     if not join.metaData.name and isinstance(join, api.Join):
         __set_name(join, api.Join, "joins")
     return "_".join(

orchestration/src/main/scala/ai/chronon/orchestration/RepoIndex.scala

@@ -4,11 +4,12 @@ import ai.chronon.orchestration.RepoIndex._
 import ai.chronon.orchestration.RepoTypes._
 import ai.chronon.orchestration.utils.CollectionExtensions.IteratorExtensions
 import ai.chronon.orchestration.utils.SequenceMap
+import ai.chronon.orchestration.utils.StringExtensions.StringOps
 import org.apache.logging.log4j.scala.Logging
 
 import scala.collection.mutable
 
-class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
+class RepoIndex[T >: Null](proc: ConfProcessor[T]) extends Logging {
 
   // first pass updates
   private val branchToFileHash: TriMap[Branch, Name, FileHash] = mutable.Map.empty
@@ -19,30 +20,14 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
   private val versionSequencer: SequenceMap[Name, GlobalHash] = new SequenceMap[Name, GlobalHash]
 
   def addNodes(fileHashes: mutable.Map[Name, FileHash],
-               nodes: Seq[T],
+               newNodes: Seq[T],
                branch: Branch,
                dryRun: Boolean = true): Seq[VersionUpdate] = {
 
-    val newContents = nodes.map { node =>
-      val data = proc.toLocalData(node)
-      val nodeContent = NodeContent(data, node)
-
-      require(data.fileHash == fileHashes(data.name), s"File hash mismatch for ${data.name}")
-
-      data.name -> (data.fileHash -> nodeContent)
-
-    }.toMap
-
-    def getContents(name: Name, fileHash: FileHash): NodeContent[T] = {
-
-      val incomingContents = newContents.get(name).map(_._2)
-
-      lazy val existingContents = fileHashToContent
-        .get(name)
-        .flatMap(_.get(fileHash))
-
-      incomingContents.orElse(existingContents).get
-    }
+    val newContents = buildContentMap(proc, newNodes, fileHashes)
+    val enrichedFileHashes = newContents.map {
+      case (name, content) => name -> content.localData.fileHash
+    } ++ fileHashes
 
     val globalHashes = mutable.Map.empty[Name, GlobalHash]
 
@@ -51,8 +36,27 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
 
       if (globalHashes.contains(name)) return globalHashes(name)
 
-      val fileHash = fileHashes(name)
-      val content = getContents(name, fileHash)
+      val fileHash = enrichedFileHashes.get(name) match {
+        case Some(hash) => hash
+
+        // this could be an artifact related to unchanged files on the branch
+        // we reach out to content index
+        // artifacts are just names with no content - so there should be just one entry
+        case None =>
+          val hashToContent = fileHashToContent(name)
+
+          require(hashToContent.nonEmpty, s"Expected 1 entry for artifact $name, found none")
+          require(hashToContent.size == 1, s"Expected 1 entry for artifact $name, found ${hashToContent.size}")
+
+          hashToContent.head._1
+      }
+
+      val content = if (newContents.contains(name)) {
+        newContents(name)
+      } else {
+        // fetch
+        fileHashToContent(name)(fileHash)
+      }
 
       val localHash = content.localData.localHash
       val parents = content.localData.inputs
@@ -70,7 +74,7 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
 
       logger.info(s"codeString: $codeString")
 
-      val globalHash = GlobalHash(codeString.hashCode().toHexString)
+      val globalHash = GlobalHash(codeString.md5)
 
       globalHashes.update(name, globalHash)
       globalHash
@@ -90,12 +94,12 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
     val mainVersions = branchVersionIndex.getOrElse(Branch.main, mutable.Map.empty)
 
     val versionUpdates = VersionUpdate.join(newVersions, existingVersions, mainVersions)
-    VersionUpdate.print(versionUpdates)
 
     if (!dryRun) {
 
+      logger.info("Not a dry run! Inserting new nodes into the index into branch: " + branch.name)
       newContents.foreach {
-        case (name, (fileHash, content)) => update(fileHashToContent, name, fileHash, content)
+        case (name, content) => update(fileHashToContent, name, content.localData.fileHash, content)
       }
 
       val newVersions = globalHashes.map {
@@ -105,7 +109,7 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
           name -> version
       }
 
-      branchToFileHash.update(branch, fileHashes)
+      branchToFileHash.update(branch, enrichedFileHashes)
       branchVersionIndex.update(branch, newVersions)
 
     }
@@ -141,7 +145,7 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
   private def pruneContents(): Unit = {
 
     // collect unique hashes per name from every branch
-    val validHashes: mutable.Map[Name, mutable.HashSet[FileHash]] = innerValues(branchToFileHash)
+    val validHashes: mutable.Map[Name, mutable.HashSet[FileHash]] = innerKeyToValueSet(branchToFileHash)
 
     fileHashToContent.retain {
       case (name, fileHashMap) =>
@@ -173,14 +177,12 @@ class RepoIndex[T](proc: ConfProcessor[T]) extends Logging {
 
 object RepoIndex {
 
-  private case class NodeContent[T](localData: LocalData, conf: T)
-
   private type TriMap[K1, K2, V] = mutable.Map[K1, mutable.Map[K2, V]]
 
   private def update[K1, K2, V](map: TriMap[K1, K2, V], k1: K1, k2: K2, v: V): Unit =
     map.getOrElseUpdate(k1, mutable.Map.empty).update(k2, v)
 
-  private def innerValues[K1, K2, V](map: TriMap[K1, K2, V]): mutable.Map[K2, mutable.HashSet[V]] = {
+  private def innerKeyToValueSet[K1, K2, V](map: TriMap[K1, K2, V]): mutable.Map[K2, mutable.HashSet[V]] = {
     val result = mutable.Map.empty[K2, mutable.HashSet[V]]
     map.values.foreach { innerMap =>
       innerMap.foreach {
@@ -191,4 +193,54 @@ object RepoIndex {
     result
   }
 
+  /**
+    * Takes data from repo parser and builds a local index for the repo parser
+    * We treat inputs and outputs that are not present in FileHashes as artifacts
+    * For these artifacts we create additional entries in the result
+    */
+  def buildContentMap[T >: Null](proc: ConfProcessor[T],
+                                 nodes: Seq[T],
+                                 fileHashes: mutable.Map[Name, FileHash]): mutable.Map[Name, NodeContent[T]] = {
+
+    val contentMap = mutable.Map.empty[Name, NodeContent[T]]
+
+    // first pass - update non-artifact contents
+    for (
+      node <- nodes;
+      nodeContent <- proc.nodeContents(node)
+    ) {
+
+      val name = nodeContent.localData.name
+      contentMap.update(name, nodeContent)
+
+      def updateContents(artifactName: Name, isOutput: Boolean): Unit = {
+
+        // artifacts are not present in file hashes
+        if (fileHashes.contains(artifactName)) return
+
+        val existingParents = if (contentMap.contains(artifactName)) {
+          contentMap(artifactName).localData.inputs
+        } else {
+          Seq.empty
+        }
+
+        val newParents = if (isOutput) Seq(name) else Seq.empty
+
+        val parents = (existingParents ++ newParents).distinct
+
+        val artifactData = LocalData.forArtifact(artifactName, parents)
+        val artifactContent = NodeContent[T](artifactData, null)
+
+        contentMap.update(artifactName, artifactContent)
+
+      }
+
+      nodeContent.localData.outputs.foreach { output => updateContents(output, isOutput = true) }
+      nodeContent.localData.inputs.foreach { input => updateContents(input, isOutput = false) }
+
+    }
+
+    contentMap
+  }
+
 }

orchestration/src/main/scala/ai/chronon/orchestration/RepoTypes.scala

@@ -1,33 +1,96 @@
 package ai.chronon.orchestration
 
+import ai.chronon.orchestration.utils.StringExtensions.StringOps
+
+/**
+  * Types relevant to the orchestration layer.
+  * It is very easy to get raw strings mixed up in the indexing logic.
+  * So we guard them using case classes. For that reason we have a lot of case classes here.
+  */
 object RepoTypes {
 
+  /**
+    * name of the node
+    * example: group_bys.<team>.<file>.<var>, joins.<team>.<file>.<var>, staging_queries.<team>.<file>.<var>
+    * and also table.<namespace>.<name> - adding a dummy node for the table makes the code easier to write
+    */
   case class Name(name: String)
 
   case class Branch(name: String)
 
+  object Branch {
+    val main: Branch = Branch("main")
+  }
+
+  /**
+    * Take the file content string and hashes it
+    * Whenever this changes the cli will upload the file into the index.
+    */
   case class FileHash(hash: String)
 
+  /**
+    * Local hash represents the computation defined in the file.
+    * In chronon api, anything field other than metadata is
+    * considered to impact the computation & consequently the output.
+    */
   case class LocalHash(hash: String)
 
+  /**
+    * Global hash represents the computation defined in the file and all its dependencies.
+    * We recursively scan all the parents of the node to compute the global hash.
+    *
+    * `global_hash(node) = hash(node.local_hash + node.parents.map(global_hash))`
+    */
   case class GlobalHash(hash: String)
 
+  /**
+    * Local data represents relevant information for lineage tracking
+    * that can be computed by parsing a file in *isolation*.
+    */
   case class LocalData(name: Name, fileHash: FileHash, localHash: LocalHash, inputs: Seq[Name], outputs: Seq[Name])
 
-  object Branch {
-    val main: Branch = Branch("main")
+  object LocalData {
+
+    def forArtifact(name: Name, parents: Seq[Name]): LocalData = {
+
+      val nameHash = name.name.md5
+
+      LocalData(
+        name,
+        FileHash(nameHash),
+        LocalHash(nameHash),
+        inputs = parents,
+        outputs = Seq.empty
+      )
+    }
   }
 
+  /**
+    * Node content represents the actual data that is stored in the index.
+    * It is a combination of local data and the actual data that is stored in the index.
+    */
   case class Version(name: String)
 
+  /**
+    * Content of the compiled file
+    * Currently TSimpleJsonProtocol serialized StagingQuery, Join, GroupBy thrift objects.
+    * Python compile.py will serialize user's python code into these objects and
+    * the [[RepoParser]] will pick them up and sync into [[RepoIndex]].
+    */
   case class FileContent(content: String) {
-    def hash: FileHash = FileHash(content.hashCode().toHexString)
+    def hash: FileHash = FileHash(content.md5)
   }
 
+  case class NodeContent[T](localData: LocalData, conf: T)
+
   case class Table(name: String)
 
+  /**
+    * To make the code testable, we parameterize the Config with `T`
+    * You can see how this is used in [[RepoIndexSpec]]
+    */
   trait ConfProcessor[T] {
-    def toLocalData(t: T): LocalData
+    def nodeContents(t: T): Seq[NodeContent[T]]
     def parse(name: String, fileContent: FileContent): Seq[T]
   }
 }

orchestration/src/main/scala/ai/chronon/orchestration/VersionUpdate.scala

@@ -7,8 +7,6 @@ import scala.collection.mutable
 
 case class VersionUpdate(name: Name, previous: Option[Version], next: Option[Version], main: Option[Version]) {
 
-  private def isChanged: Boolean = next != main || next != previous
-
   private def toRow: Seq[String] =
     Seq(
       name.name,
@@ -30,7 +28,6 @@ object VersionUpdate {
       .map { name =>
         VersionUpdate(name, previous.get(name), next.get(name), main.get(name))
       }
-      .filter(_.isChanged)
       .toSeq
       .sortBy(_.name.name)
   }

orchestration/src/main/scala/ai/chronon/orchestration/utils/StringExtensions.scala

@@ -0,0 +1,22 @@
+package ai.chronon.orchestration.utils
+
+import ai.chronon.api.Constants
+
+import java.security.MessageDigest
+
+object StringExtensions {
+
+  lazy val digester: ThreadLocal[MessageDigest] = new ThreadLocal[MessageDigest]() {
+    override def initialValue(): MessageDigest = MessageDigest.getInstance("MD5")
+  }
+
+  implicit class StringOps(s: String) {
+    def md5: String =
+      digester
+        .get()
+        .digest(s.getBytes(Constants.UTF8))
+        .map("%02x".format(_))
+        .mkString
+        .take(8)
+  }
+}