[WIP] Branch versioning logic - RepoIndex

Author: nikhil-zlai

orchestration/README.md

@@ -270,9 +270,21 @@ Make a new copy of the conf & update the name (file name & metadata.name) -
 This needs to be followed by changing references in the downstream nodes - 
 all tables and nodes downstream will have the branch suffix.
 
+```
+# 1. cli sends file_hash_map to remote
+
+local_file_map = repo.compiled.file_hash_map
+remote_file_map = remote.file_map
+deleted = remote_file_map - local_file_map
+added = local_file_map - remote_file_map
+updated = [k in intersect(local_file_map, remote_file_map)] 
+# 2. remote marks the changed files it needs
+
+(node, lineage_hash) => 
+```
+
 ### Merging changes into `main` 
 
 - Deletes should actually trigger asset and pipeline clean up.
 - Updates should trigger asset renaming
 - Adds can work as is - we are not suffixing adds.
-

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

@@ -0,0 +1,184 @@
+package ai.chronon.orchestration
+
+import ai.chronon.orchestration.RepoIndex._
+import ai.chronon.orchestration.RepoTypes._
+import ai.chronon.orchestration.utils.CollectionExtensions.IteratorExtensions
+import ai.chronon.orchestration.utils.SequenceMap
+
+import scala.collection.mutable
+
+class RepoIndex[T](proc: ConfProcessor[T]) {
+
+  // first pass updates
+  private val branchToFileHash: TriMap[Branch, Name, FileHash] = mutable.Map.empty
+  private val fileHashToContent: TriMap[Name, FileHash, NodeContent[T]] = mutable.Map.empty
+
+  // second pass updates
+  private val branchVersionIndex: TriMap[Branch, Name, Version] = mutable.Map.empty
+  private val versionSequencer: SequenceMap[Name, GlobalHash] = new SequenceMap[Name, GlobalHash]
+
+  def addNodes(fileHashes: mutable.Map[Name, FileHash],
+               nodes: 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 globalHashes = mutable.Map.empty[Name, GlobalHash]
+
+    // memoizes into globalHashes and recursively computes global hash from parents
+    def computeGlobalHash(name: Name): GlobalHash = {
+
+      if (globalHashes.contains(name)) return globalHashes(name)
+
+      val fileHash = fileHashes(name)
+      val content = getContents(name, fileHash)
+
+      val localHash = content.localData.localHash
+      val parents = content.localData.inputs
+
+      // recursively compute parent hashes
+      val parentHashes = parents
+        .map { parent =>
+          val parentHash = globalHashes.getOrElse(parent, computeGlobalHash(parent)).hash
+          s"${parent.name}:$parentHash"
+
+        }
+        .mkString(",")
+
+      // combine parent hashcode with local hash
+      val codeString = s"node=${name.name}:$localHash|parents=$parentHashes"
+      val globalHash = GlobalHash(codeString.hashCode().toHexString)
+
+      globalHashes.update(name, globalHash)
+      globalHash
+    }
+
+    val newVersions = mutable.Map.empty[Name, Version]
+
+    fileHashes.foreach {
+      case (name, _) =>
+        val globalHash = computeGlobalHash(name)
+
+        val versionIndex = versionSequencer.potentialIndex(name, globalHash)
+        newVersions.update(name, Version("v" + versionIndex.toString))
+    }
+
+    val existingVersions = branchVersionIndex.getOrElse(branch, mutable.Map.empty)
+    val mainVersions = branchVersionIndex.getOrElse(Branch.main, mutable.Map.empty)
+
+    val versionUpdates = VersionUpdate.join(newVersions, existingVersions, mainVersions)
+    VersionUpdate.print(versionUpdates)
+
+    if (!dryRun) {
+
+      newContents.foreach {
+        case (name, (fileHash, content)) => update(fileHashToContent, name, fileHash, content)
+      }
+
+      branchToFileHash.update(branch, fileHashes)
+      branchVersionIndex.update(branch, newVersions)
+
+    }
+
+    versionUpdates
+  }
+
+  // returns the contents of the files not present in the index
+  def diff(incomingFileHashes: mutable.Map[Name, FileHash]): Seq[Name] = {
+
+    incomingFileHashes
+      .filter {
+        case (name, incomingHash) =>
+          val fileHashMap = fileHashToContent.get(name)
+
+          lazy val nameAbsentInIndex = fileHashMap.isEmpty
+          lazy val fileHashAbsentForName = !fileHashMap.get.contains(incomingHash)
+
+          nameAbsentInIndex || fileHashAbsentForName
+
+      }
+      .keys
+      .toSeq
+  }
+
+  def pruneBranch(branch: Branch): Unit = {
+
+    branchToFileHash.remove(branch)
+
+    pruneContents()
+  }
+
+  private def pruneContents(): Unit = {
+
+    // collect unique hashes per name from every branch
+    val validHashes: mutable.Map[Name, mutable.HashSet[FileHash]] = innerValues(branchToFileHash)
+
+    fileHashToContent.retain {
+      case (name, fileHashMap) =>
+        fileHashMap.retain {
+
+          case (fileHash, _) =>
+            validHashes.get(name) match {
+              case None         => false // no branch has this name
+              case Some(hashes) => hashes.contains(fileHash) // this branch has this fileHash
+            }
+
+        }
+
+        fileHashMap.nonEmpty
+    }
+  }
+
+  def addFiles(fileHashes: mutable.Map[Name, FileHash], updatedFiles: Map[String, String], branch: Branch): Unit = {
+
+    val nodes: Seq[T] = updatedFiles.iterator.flatMap {
+      case (name, content) =>
+        proc.parse(name, FileContent(content))
+    }.distinct
+
+    addNodes(fileHashes, nodes, branch)
+  }
+
+}
+
+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]] = {
+    val result = mutable.Map.empty[K2, mutable.HashSet[V]]
+    map.values.foreach { innerMap =>
+      innerMap.foreach {
+        case (k2, v) =>
+          result.getOrElseUpdate(k2, mutable.HashSet.empty).add(v)
+      }
+    }
+    result
+  }
+
+}