WIP - introducing confNode

Author: nikhil-zlai

api/src/main/scala/ai/chronon/api/planner/GroupByOfflinePlanner.scala

@@ -0,0 +1,45 @@
+package ai.chronon.api.planner
+import ai.chronon.api.Extensions.GroupByOps
+import ai.chronon.api.{DataModel, GroupBy, MetaData, PartitionSpec, TableDependency, ThriftJsonCodec}
+import ai.chronon.planner.{ConfPlan, GroupByBackfillNode, Node, NodeContent}
+
+import scala.util.Try
+
+class GroupByPlanner(groupBy: GroupBy)(implicit outputPartitionSpec: PartitionSpec)
+    extends Planner[GroupBy](groupBy)(outputPartitionSpec) {
+
+  private def tableDeps: Seq[TableDependency] = TableDependencies.fromGroupBy(groupBy)
+
+  private def effectiveStepDays: Int = {
+    val defaultStepDays = if (groupBy.dataModel == DataModel.EVENTS) 15 else 1
+    val configuredStepDaysOpt = Option(groupBy.metaData.executionInfo).flatMap(e => Option(e.stepDays))
+    configuredStepDaysOpt.getOrElse(defaultStepDays)
+  }
+
+  // execInfo can be heavy - and we don't want to duplicate it
+  private def eraseExecutionInfo: GroupBy = {
+    val result = groupBy.deepCopy()
+    result.metaData.unsetExecutionInfo()
+    result
+  }
+
+  def backfillNode: Node = {
+
+    val metaData = MetaDataUtils.layer(groupBy.metaData,
+      "backfill",
+      groupBy.metaData.name + "/backfill",
+      tableDeps,
+      Some(effectiveStepDays))
+
+    val node = new GroupByBackfillNode().setGroupBy(eraseExecutionInfo)
+
+    toNode(metaData, _.setGroupByBackfill(node), eraseExecutionInfo)
+  }
+
+  def streamingNode: Option[Node] = groupBy.streamingSource.map { _ =>
+
+  }
+
+  override def buildPlan: ConfPlan = ???
+}
+

api/src/main/scala/ai/chronon/api/planner/GroupByPlanner.scala

@@ -3,18 +3,33 @@ package ai.chronon.api.planner
 import ai.chronon.api.Extensions.{GroupByOps, MetadataOps, SourceOps, StringOps}
 import ai.chronon.api.ScalaJavaConversions.{IterableOps, IteratorOps}
 import ai.chronon.api._
-import ai.chronon.api.planner.JoinOfflinePlanner._
-import ai.chronon.api.planner.GroupByOfflinePlanner._
 import ai.chronon.planner._
 
-import scala.collection.mutable
-import scala.language.{implicitConversions, reflectiveCalls}
 import scala.collection.Seq
+import scala.language.{implicitConversions, reflectiveCalls}
 
-class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec)
+class JoinPlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec)
     extends Planner[Join](join)(outputPartitionSpec) {
 
-  val leftSourceNode: SourceWithFilterNode = {
+  // will mutate the join in place - use on deepCopy-ied objects only
+  private def unsetNestedMetadata(join: Join): Unit = {
+    join.unsetMetaData()
+    Option(join.joinParts).foreach(_.iterator().toScala.foreach(_.groupBy.unsetMetaData()))
+    Option(join.labelParts).foreach(_.labels.iterator().toScala.foreach(_.groupBy.unsetMetaData()))
+    join.unsetOnlineExternalParts()
+  }
+
+  private def joinWithoutExecutionInfo: Join = {
+    val copied = join.deepCopy()
+    copied.metaData.unsetExecutionInfo()
+    Option(copied.joinParts).foreach(_.iterator().toScala.foreach(_.groupBy.metaData.unsetExecutionInfo()))
+    Option(copied.labelParts).foreach(_.labels.iterator().toScala.foreach(_.groupBy.metaData.unsetExecutionInfo()))
+    copied.unsetOnlineExternalParts()
+    copied
+  }
+
+  val leftSourceNode: Node = {
+
     val left = join.left
     val result = new SourceWithFilterNode()
       .setSource(left)
@@ -34,12 +49,12 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
       stepDays = Some(1)
     )
 
-    result.setMetaData(metaData)
+    toNode(metaData, _.setSourceWithFilter(result), result)
   }
 
-  val bootstrapNodeOpt: Option[JoinBootstrapNode] = Option(join.bootstrapParts).map { bootstrapParts =>
+  private val bootstrapNodeOpt: Option[Node] = Option(join.bootstrapParts).map { bootstrapParts =>
     val result = new JoinBootstrapNode()
-      .setJoin(join)
+      .setJoin(joinWithoutExecutionInfo)
 
     // bootstrap tables are unfortunately unique to the join - can't be re-used if a new join part is added
     val bootstrapNodeName = join.metaData.name + "/boostrap"
@@ -56,12 +71,25 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
       stepDays = Some(1)
     )
 
-    result.setMetaData(metaData)
+    val content = new NodeContent()
+    content.setJoinBootstrap(result)
+
+    val copy = result.deepCopy()
+    unsetNestedMetadata(copy.join)
+
+    toNode(metaData, _.setJoinBootstrap(result), copy)
+  }
+
+  private def copyAndEraseExecutionInfo(joinPart: JoinPart): JoinPart = {
+    val copy = joinPart.deepCopy()
+    copy.groupBy.metaData.unsetExecutionInfo()
+    copy
   }
 
-  private def buildJoinPartNode(joinPart: JoinPart): JoinPartNode = {
+  private def buildJoinPartNode(joinPart: JoinPart): Node = {
+
     val result = new JoinPartNode()
-      .setJoinPart(joinPart)
+      .setJoinPart(copyAndEraseExecutionInfo(joinPart))
       .setLeftDataModel(join.left.dataModel)
       .setLeftSourceTable(leftSourceNode.metaData.outputTable)
 
@@ -76,7 +104,7 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
       .map(_.stepDays)
       .getOrElse(joinPart.groupBy.dataModel match {
         case DataModel.ENTITIES => 1
-        case DataModel.EVENTS   => 15
+        case DataModel.EVENTS => 15
       })
 
     // pull conf params from the groupBy metadata, but use the join namespace to write to.
@@ -90,12 +118,15 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
       )
       .setOutputNamespace(join.metaData.outputNamespace)
 
-    result.setMetaData(metaData)
+    val copy = result.deepCopy()
+    copy.joinPart.groupBy.unsetMetaData()
+
+    toNode(metaData, _.setJoinPart(result), copy)
   }
 
-  private val joinPartNodes: Seq[JoinPartNode] = join.joinParts.toScala.map { buildJoinPartNode }.toSeq
+  private val joinPartNodes: Seq[Node] = join.joinParts.toScala.map { buildJoinPartNode }.toSeq
 
-  val mergeNode: JoinMergeNode = {
+  val mergeNode: Node = {
     val result = new JoinMergeNode()
       .setJoin(join)
 
@@ -124,10 +155,16 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
         stepDays = Some(1)
       )
 
-    result.setMetaData(metaData)
+    val copy = result.deepCopy()
+    unsetNestedMetadata(copy.join)
+    copy.join.unsetDerivations()
+    copy.join.unsetLabelParts()
+
+    toNode(metaData, _.setJoinMerge(result), copy)
   }
 
-  private val derivationNodeOpt: Option[JoinDerivationNode] = Option(join.derivations).map { _ =>
+  private val derivationNodeOpt: Option[Node] = Option(join.derivations).map { _ =>
+
     val result = new JoinDerivationNode()
       .setJoin(join)
 
@@ -142,7 +179,11 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
         stepDays = Some(1)
       )
 
-    result.setMetaData(metaData)
+    val copy = result.deepCopy()
+    unsetNestedMetadata(copy.join)
+    copy.join.unsetLabelParts()
+
+    toNode(metaData, _.setJoinDerivation(result), copy)
   }
 
   // these need us to additionally (groupBy backfill) generate the snapshot tables
@@ -156,7 +197,7 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
     )
     .getOrElse(Array.empty)
 
-  private val labelJoinNodeOpt: Option[LabelJoinNode] = Option(join.labelParts).map { labelParts =>
+  private val labelJoinNodeOpt: Option[Node] = Option(join.labelParts).map { labelParts =>
     val result = new LabelJoinNode()
       .setJoin(join)
 
@@ -177,112 +218,46 @@ class JoinOfflinePlanner(join: Join)(implicit outputPartitionSpec: PartitionSpec
         stepDays = Some(1)
       )
 
-    result.setMetaData(metaData)
-  }
+    val copy = result.deepCopy()
+    unsetNestedMetadata(result.join)
 
-  override def offlineNodes: Seq[PlanNode] = {
-    val result: mutable.ArrayBuffer[PlanNode] = mutable.ArrayBuffer.empty[PlanNode]
+    toNode(metaData, _.setLabelJoin(result), copy)
+  }
 
-    result.append(leftSourceNode)
-    bootstrapNodeOpt.foreach(bn => result.append(bn))
-    joinPartNodes.foreach(jpn => result.append(jpn))
+  def offlineNodes: Seq[Node] = {
+
+    Seq(leftSourceNode) ++
+      bootstrapNodeOpt ++
+      joinPartNodes ++
+      Seq(mergeNode) ++
+      derivationNodeOpt ++
+      snapshotLabelParts.map(_.groupBy).map{
+        gb => new GroupByPlanner(gb).backfillNode
+      } ++
+      labelJoinNodeOpt
+  }
 
-    result.append(mergeNode)
-    derivationNodeOpt.foreach(dn => result.append(dn))
-    snapshotLabelParts.foreach(lp => result.append(lp.groupBy))
-    labelJoinNodeOpt.foreach(ljn => result.append(ljn))
 
-    result
+  override def onlineNodes: Seq[Node] = {
+    // depends on all
   }
 
-  override def onlineNodes: Seq[PlanNode] = ???
+
+  override def buildPlan: ConfPlan = {
+
+    val result = new ConfPlan()
+    result.setNodes()
+  }
 }
 
-object JoinOfflinePlanner {
+object JoinPlanner {
 
+  // will mutate the join in place - use on deepCopy-ied objects only
   private def unsetNestedMetadata(join: Join): Unit = {
     join.unsetMetaData()
     Option(join.joinParts).foreach(_.iterator().toScala.foreach(_.groupBy.unsetMetaData()))
     Option(join.labelParts).foreach(_.labels.iterator().toScala.foreach(_.groupBy.unsetMetaData()))
     join.unsetOnlineExternalParts()
   }
 
-  implicit class LabelJoinNodeIsPlanNode(node: LabelJoinNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      unsetNestedMetadata(result.join)
-      result
-    })
-  }
-
-  implicit class JoinDerivationNodeIsPlanNode(node: JoinDerivationNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      unsetNestedMetadata(result.join)
-      result.join.unsetLabelParts()
-      result
-    })
-  }
-
-  implicit class JoinMergeNodeIsPlanNode(node: JoinMergeNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      unsetNestedMetadata(result.join)
-      result.join.unsetDerivations()
-      result.join.unsetLabelParts()
-      result
-    })
-  }
-
-  implicit class JoinPartNodeIsPlanNode(node: JoinPartNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      result.joinPart.groupBy.unsetMetaData()
-      result
-    })
-  }
-
-  implicit class JoinBootstrapNodeIsPlanNode(node: JoinBootstrapNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      unsetNestedMetadata(result.join)
-      result
-    })
-  }
-
-  implicit class SourceWithFilterNodeIsPlanNode(node: SourceWithFilterNode) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      result.unsetMetaData()
-      result
-    })
-  }
-
-  implicit class JoinIsPlanNode(node: Join) extends PlanNode {
-    override def metaData: MetaData = node.metaData
-    override def contents: Any = node
-    override def semanticHash: String = ThriftJsonCodec.hexDigest({
-      val result = node.deepCopy()
-      unsetNestedMetadata(node)
-      result
-    })
-  }
-
 }