Write daily IRs for incremental aggregation

aggregator/src/main/scala/ai/chronon/aggregator/base/SimpleAggregators.scala

@@ -117,6 +117,57 @@ class UniqueCount[T](inputType: DataType) extends SimpleAggregator[T, util.HashS
   }
 }
 
+class AverageIR extends SimpleAggregator[Array[Any], Array[Any], Double] {
+  override def outputType: DataType = DoubleType
+
+  override def irType: DataType =
+    StructType(
+      "AvgIr",
+      Array(StructField("sum", DoubleType), StructField("count", IntType))
+    )
+
+  override def prepare(input: Array[Any]): Array[Any] = {
+    Array(input(0).asInstanceOf[Double], input(1).asInstanceOf[Int])
+  }
+
+  // mutating
+  override def update(ir: Array[Any], input: Array[Any]): Array[Any] = {
+    val inputSum = input(0).asInstanceOf[Double]
+    val inputCount = input(1).asInstanceOf[Int]
+    ir.update(0, ir(0).asInstanceOf[Double] + inputSum)
+    ir.update(1, ir(1).asInstanceOf[Int] + inputCount)
+    ir
+  }
+
+  // mutating
+  override def merge(ir1: Array[Any], ir2: Array[Any]): Array[Any] = {
+    ir1.update(0, ir1(0).asInstanceOf[Double] + ir2(0).asInstanceOf[Double])
+    ir1.update(1, ir1(1).asInstanceOf[Int] + ir2(1).asInstanceOf[Int])
+    ir1
+  }
+
+  override def finalize(ir: Array[Any]): Double =
+    ir(0).asInstanceOf[Double] / ir(1).asInstanceOf[Int].toDouble
+
+  override def delete(ir: Array[Any], input: Array[Any]): Array[Any] = {
+    val inputSum = input(0).asInstanceOf[Double]
+    val inputCount = input(1).asInstanceOf[Int]
+    ir.update(0, ir(0).asInstanceOf[Double] - inputSum)
+    ir.update(1, ir(1).asInstanceOf[Int] - inputCount)
+    ir
+  }
+
+  override def clone(ir: Array[Any]): Array[Any] = {
+    val arr = new Array[Any](ir.length)
+    ir.copyToArray(arr)
+    arr
+  }
+
+  override def isDeletable: Boolean = true
+}
+
+
+
 class Average extends SimpleAggregator[Double, Array[Any], Double] {
   override def outputType: DataType = DoubleType
 

aggregator/src/main/scala/ai/chronon/aggregator/row/ColumnAggregator.scala

@@ -216,6 +216,13 @@ object ColumnAggregator {
   private def toJavaDouble[A: Numeric](inp: Any) =
     implicitly[Numeric[A]].toDouble(inp.asInstanceOf[A]).asInstanceOf[java.lang.Double]
 
+
+  private def toStructArray(inp: Any): Array[Any] = inp match {
+    case r: org.apache.spark.sql.Row => r.toSeq.toArray
+    case null => null
+    case other => throw new IllegalArgumentException(s"Expected Row, got: $other")
+  }
+
   def construct(baseInputType: DataType,
                 aggregationPart: AggregationPart,
                 columnIndices: ColumnIndices,
@@ -341,6 +348,7 @@ object ColumnAggregator {
           case ShortType  => simple(new Average, toDouble[Short])
           case DoubleType => simple(new Average)
           case FloatType  => simple(new Average, toDouble[Float])
+          case StructType(name, fields) => simple(new AverageIR, toStructArray)
           case _          => mismatchException
         }
 

aggregator/src/main/scala/ai/chronon/aggregator/row/RowAggregator.scala

@@ -24,7 +24,10 @@ import scala.collection.Seq
 
 // The primary API of the aggregator package.
 // the semantics are to mutate values in place for performance reasons
-class RowAggregator(val inputSchema: Seq[(String, DataType)], val aggregationParts: Seq[AggregationPart])
+// userAggregationParts is used when incrementalMode = True.
+class RowAggregator(val inputSchema: Seq[(String, DataType)],
+                    val aggregationParts: Seq[AggregationPart],
+                    val userInputAggregationParts: Option[Seq[AggregationPart]] = None )
     extends Serializable
     with SimpleAggregator[Row, Array[Any], Array[Any]] {
 
@@ -70,11 +73,25 @@ class RowAggregator(val inputSchema: Seq[(String, DataType)], val aggregationPar
     .toArray
     .zip(columnAggregators.map(_.irType))
 
-  val outputSchema: Array[(String, DataType)] = aggregationParts
+  val incrementalOutputSchema: Array[(String, DataType)] = aggregationParts
+    .map(_.incrementalOutputColumnName)
+    .toArray
+    .zip(columnAggregators.map(_.irType))
+
+  val aggregationPartsOutputSchema: Array[(String, DataType)] = aggregationParts
     .map(_.outputColumnName)
     .toArray
     .zip(columnAggregators.map(_.outputType))
 
+  val outputSchema: Array[(String, DataType)] = userInputAggregationParts
+    .map{ parts =>
+      parts
+        .map(_.outputColumnName)
+        .toArray
+        .zip(columnAggregators.map(_.outputType))
+    }.getOrElse(aggregationPartsOutputSchema)
+
+
   val isNotDeletable: Boolean = columnAggregators.forall(!_.isDeletable)
 
   // this will mutate in place

api/py/ai/chronon/group_by.py

@@ -363,6 +363,7 @@ def GroupBy(
     derivations: Optional[List[ttypes.Derivation]] = None,
     deprecation_date: Optional[str] = None,
     description: Optional[str] = None,
+    is_incremental: Optional[bool] = False,
     **kwargs,
 ) -> ttypes.GroupBy:
     """
@@ -570,6 +571,7 @@ def _normalize_source(source):
         backfillStartDate=backfill_start_date,
         accuracy=accuracy,
         derivations=derivations,
+        isIncremental=is_incremental,
     )
     validate_group_by(group_by)
     return group_by

api/py/test/sample/scripts/spark_submit.sh

@@ -28,13 +28,14 @@
 
 set -euxo pipefail
 CHRONON_WORKING_DIR=${CHRONON_TMPDIR:-/tmp}/${USER}
+echo $CHRONON_WORKING_DIR
 mkdir -p ${CHRONON_WORKING_DIR}
 export TEST_NAME="${APP_NAME}_${USER}_test"
 unset PYSPARK_DRIVER_PYTHON
 unset PYSPARK_PYTHON
 unset SPARK_HOME
 unset SPARK_CONF_DIR
-export LOG4J_FILE="${CHRONON_WORKING_DIR}/log4j_file"
+export LOG4J_FILE="${CHRONON_WORKING_DIR}/log4j.properties"
 cat > ${LOG4J_FILE} << EOF
 log4j.rootLogger=INFO, stdout
 log4j.appender.stdout=org.apache.log4j.ConsoleAppender
@@ -47,6 +48,9 @@ EOF
 $SPARK_SUBMIT_PATH \
 --driver-java-options " -Dlog4j.configuration=file:${LOG4J_FILE}" \
 --conf "spark.executor.extraJavaOptions= -XX:ParallelGCThreads=4 -XX:+UseParallelGC -XX:+UseCompressedOops" \
+--conf "spark.driver.extraJavaOptions=-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005 -Dlog4j.configuration=file:${LOG4J_FILE}" \
+--conf "spark.sql.warehouse.dir=/home/chaitu/projects/chronon/spark-warehouse" \
+--conf "javax.jdo.option.ConnectionURL=jdbc:derby:;databaseName=/home/chaitu/projects/chronon/hive-metastore/metastore_db;create=true" \
 --conf spark.sql.shuffle.partitions=${PARALLELISM:-4000} \
 --conf spark.dynamicAllocation.maxExecutors=${MAX_EXECUTORS:-1000} \
 --conf spark.default.parallelism=${PARALLELISM:-4000} \
@@ -77,3 +81,6 @@ tee ${CHRONON_WORKING_DIR}/${APP_NAME}_spark.log
 
 
 
+
+#--conf "spark.driver.extraJavaOptions=-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=5005 -Dlog4j.rootLogger=INFO,console" \
+

api/py/test/sample/teams.json

@@ -5,7 +5,7 @@
         },
         "common_env": {
             "VERSION": "latest",
-            "SPARK_SUBMIT_PATH": "[TODO]/path/to/spark-submit",
+            "SPARK_SUBMIT_PATH": "spark-submit",
             "JOB_MODE": "local[*]",
             "HADOOP_DIR": "[STREAMING-TODO]/path/to/folder/containing",
             "CHRONON_ONLINE_CLASS": "[ONLINE-TODO]your.online.class",

api/src/main/scala/ai/chronon/api/Extensions.scala

@@ -97,7 +97,7 @@ object Extensions {
     def cleanName: String = metaData.name.sanitize
 
     def outputTable = s"${metaData.outputNamespace}.${metaData.cleanName}"
-
+    def incrementalOutputTable = s"${metaData.outputNamespace}.${metaData.cleanName}_inc"
     def preModelTransformsTable = s"${metaData.outputNamespace}.${metaData.cleanName}_pre_mt"
     def outputLabelTable = s"${metaData.outputNamespace}.${metaData.cleanName}_labels"
     def outputFinalView = s"${metaData.outputNamespace}.${metaData.cleanName}_labeled"
@@ -178,8 +178,13 @@ object Extensions {
 
     def outputColumnName =
       s"${aggregationPart.inputColumn}_$opSuffix${aggregationPart.window.suffix}${bucketSuffix}"
+
+    def incrementalOutputColumnName =
+      s"${aggregationPart.inputColumn}_$opSuffix${bucketSuffix}"
+
   }
 
+
   implicit class AggregationOps(aggregation: Aggregation) {
 
     // one agg part per bucket per window

api/thrift/api.thrift

@@ -301,6 +301,7 @@ struct GroupBy {
     6: optional string backfillStartDate
     // Optional derivation list
     7: optional list<Derivation> derivations
+    8: optional bool isIncremental
 }
 
 struct JoinPart {

spark/src/main/scala/ai/chronon/spark/DataRange.scala

@@ -54,6 +54,11 @@ case class PartitionRange(start: String, end: String)(implicit tableUtils: Table
     }
   }
 
+  def daysBetween: Int = {
+    if (start == null || end == null) 0
+    else Stream.iterate(start)(tableUtils.partitionSpec.after).takeWhile(_ <= end).size
+  }
+
   def isSingleDay: Boolean = {
     start == end
   }

spark/src/main/scala/ai/chronon/spark/Driver.scala

@@ -467,7 +467,8 @@ object Driver {
         tableUtils,
         args.stepDays.toOption,
         args.startPartitionOverride.toOption,
-        !args.runFirstHole()
+        !args.runFirstHole(),
+        Option(args.groupByConf.isIncremental).getOrElse(false)
       )
 
       if (args.shouldExport()) {