fix the element wise aggregation

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

@@ -174,6 +174,50 @@ class SawtoothAggregatorTest extends TestCase {
     timer.publish("comparison")
   }
 
+  def testElementWiseAverageWithMerge(): Unit = {
+    val columns = Seq(
+      Column("ts", LongType, 1),
+      Column("embeddings", ListType(DoubleType), 1000, chunkSize = 3, nullRate = -1, fixedListSize = true)
+    )
+    val events = CStream.gen(columns, 100).rows
+    val schema = columns.map(_.schema)
+
+    val queries = Array(System.currentTimeMillis())
+    val aggregations = Seq(
+      Builders.Aggregation(
+        Operation.AVERAGE,
+        "embeddings",
+        Seq(new Window(1, TimeUnit.DAYS)),
+        elementWise = Some(true)
+      )
+    )
+
+    val sawtoothIrs = sawtoothAggregate(events, queries, aggregations, schema)
+    assertNotNull(sawtoothIrs)
+    assertEquals(1, sawtoothIrs.length)
+
+    // Compare with naive aggregator
+    val unpacked = aggregations.flatMap(_.unpack.map(_.window)).toArray
+    val tailHops = unpacked.map(FiveMinuteResolution.calculateTailHop)
+    val rowAgg = new RowAggregator(schema, aggregations.flatMap(_.unpack))
+    val naiveAggregator = new NaiveAggregator(
+      rowAgg,
+      unpacked,
+      tailHops
+    )
+    val naiveIrs = naiveAggregator.aggregate(events, queries)
+
+    assertEquals(naiveIrs.length, queries.length)
+    assertEquals(sawtoothIrs.length, queries.length)
+    for (i <- queries.indices) {
+      val naiveFinal = rowAgg.finalize(naiveIrs(i))
+      val sawtoothFinal = rowAgg.finalize(sawtoothIrs(i))
+      naiveFinal.zip(sawtoothFinal).foreach { case (expected, result) =>
+        assertEquals(expected.asInstanceOf[Double], result.asInstanceOf[Double], 1e-9)
+      }
+    }
+  }
+
 }
 
 object SawtoothAggregatorTest {

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

@@ -259,7 +259,8 @@ object Extensions {
                 _.toScala.toMap
               )
               .orNull,
-            bucket
+            bucket,
+            Option(agg.elementWise).getOrElse(false)
           )
           for (window <- windows) {
             perWindow += WindowMapping(
@@ -271,7 +272,8 @@ object Extensions {
                                            _.toScala.toMap
                                          )
                                          .orNull,
-                                       bucket),
+                                       bucket,
+                                       Option(agg.elementWise).getOrElse(false)),
               counter
             )
           }