Restoring aggregator files to main

Author: chewy-zlai

aggregator/src/test/scala/ai/chronon/aggregator/test/DataGen.scala

@@ -25,7 +25,7 @@ import scala.reflect.ClassTag
 import scala.util.Random
 
 // utility classes to generate random data
-abstract class CStream[T: ClassTag] {
+abstract class CStream[+T: ClassTag] {
   def next(): T
 
   // roll a dice that gives max to min uniformly, with nulls interspersed as per null rate
@@ -105,12 +105,12 @@ object CStream {
       .toArray
   }
 
-  class PartitionStream(count: Int, partitionSpec: PartitionSpec) extends CStream[Any] {
+  class PartitionStream(count: Int, partitionSpec: PartitionSpec) extends CStream[String] {
     val keys: Array[String] = genPartitions(count, partitionSpec)
     override def next(): String = Option(roll(keys.length, nullRate = 0)).map(dice => keys(dice.toInt)).get
   }
 
-  class StringStream(count: Int, prefix: String, nullRate: Double = 0.2) extends CStream[Any] {
+  class StringStream(count: Int, prefix: String, nullRate: Double = 0.2) extends CStream[String] {
     val keyCount: Int = (count * (1 - nullRate)).ceil.toInt
     val keys: Array[String] = {
       val fullKeySet = (1 until (count + 1)).map(i => s"$prefix$i")
@@ -121,7 +121,7 @@ object CStream {
   }
 
   class TimeStream(window: Window, roundMillis: Long = 1, maxTs: Long = System.currentTimeMillis())
-      extends CStream[Any] {
+      extends CStream[Long] {
     private val minTs = maxTs - window.millis
 
     def round(v: Long): Long = (v / roundMillis) * roundMillis
@@ -131,17 +131,17 @@ object CStream {
     }
   }
 
-  class IntStream(max: Int = 10000, nullRate: Double = 0.1) extends CStream[Any] {
+  class IntStream(max: Int = 10000, nullRate: Double = 0.1) extends CStream[Integer] {
     override def next(): Integer =
       Option(roll(max, 1, nullRate = nullRate)).map(dice => Integer.valueOf(dice.toInt)).orNull
   }
 
-  class LongStream(max: Int = 10000, nullRate: Double = 0.1) extends CStream[Any] {
+  class LongStream(max: Int = 10000, nullRate: Double = 0.1) extends CStream[JLong] {
     override def next(): JLong =
       Option(roll(max, 1, nullRate = nullRate)).map(java.lang.Long.valueOf(_)).orNull
   }
 
-  class DoubleStream(max: Double = 10000, nullRate: Double = 0.1) extends CStream[Any] {
+  class DoubleStream(max: Double = 10000, nullRate: Double = 0.1) extends CStream[JDouble] {
     override def next(): JDouble =
       Option(rollDouble(max, 1, nullRate = nullRate)).map(java.lang.Double.valueOf(_)).orNull
   }

aggregator/src/test/scala/ai/chronon/aggregator/test/SawtoothAggregatorTest.scala

@@ -50,7 +50,7 @@ class SawtoothAggregatorTest extends TestCase {
 
   def testTailAccuracy(): Unit = {
     val timer = new Timer
-    val queries = CStream.genTimestamps(new Window(30, TimeUnit.DAYS), 10000, 5 * 60 * 1000).map(_.asInstanceOf[Long])
+    val queries = CStream.genTimestamps(new Window(30, TimeUnit.DAYS), 10000, 5 * 60 * 1000)
 
     val columns = Seq(Column("ts", LongType, 180), Column("num", LongType, 1000))
     val events = CStream.gen(columns, 10000).rows
@@ -121,7 +121,7 @@ class SawtoothAggregatorTest extends TestCase {
 
   def testRealTimeAccuracy(): Unit = {
     val timer = new Timer
-    val queries = CStream.genTimestamps(new Window(1, TimeUnit.DAYS), 1000).map(_.asInstanceOf[Long])
+    val queries = CStream.genTimestamps(new Window(1, TimeUnit.DAYS), 1000)
     val columns = Seq(Column("ts", LongType, 180),
                       Column("num", LongType, 1000),
                       Column("age", LongType, 100),

aggregator/src/test/scala/ai/chronon/aggregator/test/SawtoothOnlineAggregatorTest.scala

@@ -37,7 +37,7 @@ class SawtoothOnlineAggregatorTest extends TestCase {
   def testConsistency(): Unit = {
     val queryEndTs = TsUtils.round(System.currentTimeMillis(), WindowUtils.Day.millis)
     val batchEndTs = queryEndTs - WindowUtils.Day.millis
-    val queries = CStream.genTimestamps(new Window(1, TimeUnit.DAYS), 1000).map(_.asInstanceOf[Long])
+    val queries = CStream.genTimestamps(new Window(1, TimeUnit.DAYS), 1000)
     val eventCount = 10000
 
     val columns = Seq(Column("ts", LongType, 60),

aggregator/src/test/scala/ai/chronon/aggregator/test/TwoStackLiteAggregatorTest.scala

@@ -65,7 +65,7 @@ class TwoStackLiteAggregatorTest extends TestCase{
 
   def testAgainstSawtooth(): Unit = {
     val timer = new Timer
-    val queries = CStream.genTimestamps(new Window(30, TimeUnit.DAYS), 100000, 5 * 60 * 1000).map(_.asInstanceOf[Long])
+    val queries = CStream.genTimestamps(new Window(30, TimeUnit.DAYS), 100000, 5 * 60 * 1000)
 
     val columns = Seq(Column("ts", LongType, 180), Column("num", LongType, 1000))
     val events = CStream.gen(columns, 10000).rows