Swap out AvroCodecOutput + KV store output types from case classes to PoJos (#294)

## Summary
We use case classes for many of our intermediate types in Flink and a
few of these are persisted in Flink state. In the tiled setup the
TimestampIR as an example. In untiled & tiled the Input and Output types
of the AsyncKVWriter are persisted to state. In the future if we do need
to update these intermediate types to include additional fields (like
we're thinking of doing to support tiling) it will not be possible as
Flink doesn't support [state schema evolution for case
classes](https://nightlies.apache.org/flink/flink-docs-release-1.17/docs/dev/datastream/fault-tolerance/serialization/schema_evolution/).
Due to that we'll need to do a time consuming migration where we spin up
parallel operators with the new types, dual write and then cut over in a
subsequent job. Instead we can try and set this up before we're in prod
to hopefully minimize running into this issue down the line.

This PR essentially swaps the case classes for Scala PoJo types for the
tiled aggregations, AvroCodecOutput (as that feeds into KV store writer)
and the KV store write response. In a subsequent PR we can update the
TimestampTile to include the startTs of the tile (and plumb the latestTs
through to the sink so we can track e2e lag).

(Choose the Scala PoJo route here as there's a bit of interplay with
some of our aggregator libs and other Scala code related to these
classes - e.g. the FlinkRowAggregatorFunction and there's a bunch of
casting needed to interop. The Flink 2.0 migration will likely be a
decent sized chunk given all our Flink code is in Scala and I think we
can bite it off then)

## Checklist
- [ ] Added Unit Tests
- [X] Covered by existing CI
- [X] Integration tested - Confirmed that this works by kicking off our
TestFlinkJob on the cluster
- [ ] Documentation update



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced new data types (`AvroCodecOutput`, `TimestampedTile`,
`TimestampedIR`, `WriteResponse`) for improved data handling in Flink
jobs
	- Enhanced schema evolution support for stateful data processing

- **Refactor**
- Replaced `PutRequest` with `AvroCodecOutput` across multiple Flink
processing components
- Updated method signatures and data type handling in various
Flink-related classes
	- Simplified timestamp retrieval and data structure interactions

- **Chores**
	- Reorganized type imports and package structure
	- Updated test cases to align with new data type implementations

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: piyush-zlai

flink/src/main/scala/ai/chronon/flink/AsyncKVStoreWriter.scala

@@ -1,5 +1,7 @@
 package ai.chronon.flink
 
+import ai.chronon.flink.types.AvroCodecOutput
+import ai.chronon.flink.types.WriteResponse
 import ai.chronon.online.Api
 import ai.chronon.online.KVStore
 import ai.chronon.online.KVStore.PutRequest
@@ -19,14 +21,12 @@ import scala.concurrent.Future
 import scala.util.Failure
 import scala.util.Success
 
-case class WriteResponse(putRequest: PutRequest, status: Boolean)
-
 object AsyncKVStoreWriter {
   private val kvStoreConcurrency = 10
   private val defaultTimeoutMillis = 1000L
 
-  def withUnorderedWaits(inputDS: DataStream[PutRequest],
-                         kvStoreWriterFn: RichAsyncFunction[PutRequest, WriteResponse],
+  def withUnorderedWaits(inputDS: DataStream[AvroCodecOutput],
+                         kvStoreWriterFn: RichAsyncFunction[AvroCodecOutput, WriteResponse],
                          featureGroupName: String,
                          timeoutMillis: Long = defaultTimeoutMillis,
                          capacity: Int = kvStoreConcurrency): DataStream[WriteResponse] = {
@@ -69,7 +69,7 @@ object AsyncKVStoreWriter {
   * @param featureGroupName Name of the FG we're writing to
   */
 class AsyncKVStoreWriter(onlineImpl: Api, featureGroupName: String)
-    extends RichAsyncFunction[PutRequest, WriteResponse] {
+    extends RichAsyncFunction[AvroCodecOutput, WriteResponse] {
   @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
   @transient private var kvStore: KVStore = _
@@ -96,14 +96,17 @@ class AsyncKVStoreWriter(onlineImpl: Api, featureGroupName: String)
     kvStore = getKVStore
   }
 
-  override def timeout(input: PutRequest, resultFuture: ResultFuture[WriteResponse]): Unit = {
+  override def timeout(input: AvroCodecOutput, resultFuture: ResultFuture[WriteResponse]): Unit = {
     logger.error(s"Timed out writing to KV Store for object: $input")
     errorCounter.inc()
-    resultFuture.complete(util.Arrays.asList[WriteResponse](WriteResponse(input, status = false)))
+    resultFuture.complete(
+      util.Arrays.asList[WriteResponse](
+        new WriteResponse(input.keyBytes, input.valueBytes, input.dataset, input.tsMillis, status = false)))
   }
 
-  override def asyncInvoke(input: PutRequest, resultFuture: ResultFuture[WriteResponse]): Unit = {
-    val resultFutureRequested: Future[Seq[Boolean]] = kvStore.multiPut(Seq(input))
+  override def asyncInvoke(input: AvroCodecOutput, resultFuture: ResultFuture[WriteResponse]): Unit = {
+    val putRequest = PutRequest(input.keyBytes, input.valueBytes, input.dataset, Some(input.tsMillis))
+    val resultFutureRequested: Future[Seq[Boolean]] = kvStore.multiPut(Seq(putRequest))
     resultFutureRequested.onComplete {
       case Success(l) =>
         val succeeded = l.forall(identity)
@@ -113,14 +116,18 @@ class AsyncKVStoreWriter(onlineImpl: Api, featureGroupName: String)
           errorCounter.inc()
           logger.error(s"Failed to write to KVStore for object: $input")
         }
-        resultFuture.complete(util.Arrays.asList[WriteResponse](WriteResponse(input, status = succeeded)))
+        resultFuture.complete(
+          util.Arrays.asList[WriteResponse](
+            new WriteResponse(input.keyBytes, input.valueBytes, input.dataset, input.tsMillis, status = succeeded)))
       case Failure(exception) =>
         // this should be rare and indicates we have an uncaught exception
         // in the KVStore - we log the exception and skip the object to
         // not fail the app
         errorCounter.inc()
         logger.error(s"Caught exception writing to KVStore for object: $input", exception)
-        resultFuture.complete(util.Arrays.asList[WriteResponse](WriteResponse(input, status = false)))
+        resultFuture.complete(
+          util.Arrays.asList[WriteResponse](
+            new WriteResponse(input.keyBytes, input.valueBytes, input.dataset, input.tsMillis, status = false)))
     }
   }
 }

flink/src/main/scala/ai/chronon/flink/AvroCodecFn.scala

@@ -5,10 +5,10 @@ import ai.chronon.api.DataModel
 import ai.chronon.api.Extensions.GroupByOps
 import ai.chronon.api.Query
 import ai.chronon.api.{StructType => ChrononStructType}
-import ai.chronon.flink.window.TimestampedTile
+import ai.chronon.flink.types.AvroCodecOutput
+import ai.chronon.flink.types.TimestampedTile
 import ai.chronon.online.AvroConversions
 import ai.chronon.online.GroupByServingInfoParsed
-import ai.chronon.online.KVStore.PutRequest
 import org.apache.flink.api.common.functions.RichFlatMapFunction
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.metrics.Counter
@@ -89,7 +89,7 @@ sealed abstract class BaseAvroCodecFn[IN, OUT] extends RichFlatMapFunction[IN, O
   * @tparam T The input data type
   */
 case class AvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParsed)
-    extends BaseAvroCodecFn[Map[String, Any], PutRequest] {
+    extends BaseAvroCodecFn[Map[String, Any], AvroCodecOutput] {
 
   override def open(configuration: Configuration): Unit = {
     super.open(configuration)
@@ -101,7 +101,7 @@ case class AvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParsed)
 
   override def close(): Unit = super.close()
 
-  override def flatMap(value: Map[String, Any], out: Collector[PutRequest]): Unit =
+  override def flatMap(value: Map[String, Any], out: Collector[AvroCodecOutput]): Unit =
     try {
       out.collect(avroConvertMapToPutRequest(value))
     } catch {
@@ -113,11 +113,11 @@ case class AvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParsed)
         avroConversionErrorCounter.inc()
     }
 
-  def avroConvertMapToPutRequest(in: Map[String, Any]): PutRequest = {
+  def avroConvertMapToPutRequest(in: Map[String, Any]): AvroCodecOutput = {
     val tsMills = in(timeColumnAlias).asInstanceOf[Long]
     val keyBytes = keyToBytes(keyColumns.map(in(_)))
     val valueBytes = valueToBytes(valueColumns.map(in(_)))
-    PutRequest(keyBytes, valueBytes, streamingDataset, Some(tsMills))
+    new AvroCodecOutput(keyBytes, valueBytes, streamingDataset, tsMills)
   }
 }
 
@@ -129,7 +129,7 @@ case class AvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParsed)
   * @tparam T The input data type
   */
 case class TiledAvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParsed)
-    extends BaseAvroCodecFn[TimestampedTile, PutRequest] {
+    extends BaseAvroCodecFn[TimestampedTile, AvroCodecOutput] {
   override def open(configuration: Configuration): Unit = {
     super.open(configuration)
     val metricsGroup = getRuntimeContext.getMetricGroup
@@ -140,7 +140,7 @@ case class TiledAvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParse
   }
   override def close(): Unit = super.close()
 
-  override def flatMap(value: TimestampedTile, out: Collector[PutRequest]): Unit =
+  override def flatMap(value: TimestampedTile, out: Collector[AvroCodecOutput]): Unit =
     try {
       out.collect(avroConvertTileToPutRequest(value))
     } catch {
@@ -152,7 +152,7 @@ case class TiledAvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParse
         avroConversionErrorCounter.inc()
     }
 
-  def avroConvertTileToPutRequest(in: TimestampedTile): PutRequest = {
+  def avroConvertTileToPutRequest(in: TimestampedTile): AvroCodecOutput = {
     val tsMills = in.latestTsMillis
 
     // 'keys' is a map of (key name in schema -> key value), e.g. Map("card_number" -> "4242-4242-4242-4242")
@@ -170,6 +170,6 @@ case class TiledAvroCodecFn[T](groupByServingInfoParsed: GroupByServingInfoParse
         |streamingDataset=$streamingDataset""".stripMargin
     )
 
-    PutRequest(keyBytes, valueBytes, streamingDataset, Some(tsMills))
+    new AvroCodecOutput(keyBytes, valueBytes, streamingDataset, tsMills)
   }
 }

flink/src/main/scala/ai/chronon/flink/FlinkJob.scala

@@ -7,14 +7,15 @@ import ai.chronon.api.DataType
 import ai.chronon.api.Extensions.GroupByOps
 import ai.chronon.api.Extensions.SourceOps
 import ai.chronon.flink.FlinkJob.watermarkStrategy
+import ai.chronon.flink.types.AvroCodecOutput
+import ai.chronon.flink.types.TimestampedTile
+import ai.chronon.flink.types.WriteResponse
 import ai.chronon.flink.window.AlwaysFireOnElementTrigger
 import ai.chronon.flink.window.FlinkRowAggProcessFunction
 import ai.chronon.flink.window.FlinkRowAggregationFunction
 import ai.chronon.flink.window.KeySelector
-import ai.chronon.flink.window.TimestampedTile
 import ai.chronon.online.Api
 import ai.chronon.online.GroupByServingInfoParsed
-import ai.chronon.online.KVStore.PutRequest
 import ai.chronon.online.MetadataStore
 import ai.chronon.online.SparkConversions
 import ai.chronon.online.TopicInfo
@@ -58,7 +59,7 @@ import scala.concurrent.duration.FiniteDuration
   * @tparam T - The input data type
   */
 class FlinkJob[T](eventSrc: FlinkSource[T],
-                  sinkFn: RichAsyncFunction[PutRequest, WriteResponse],
+                  sinkFn: RichAsyncFunction[AvroCodecOutput, WriteResponse],
                   groupByServingInfoParsed: GroupByServingInfoParsed,
                   encoder: Encoder[T],
                   parallelism: Int) {
@@ -117,7 +118,7 @@ class FlinkJob[T](eventSrc: FlinkSource[T],
       .name(s"Spark expression eval with timestamps for $groupByName")
       .setParallelism(sourceStream.parallelism)
 
-    val putRecordDS: DataStream[PutRequest] = sparkExprEvalDSWithWatermarks
+    val putRecordDS: DataStream[AvroCodecOutput] = sparkExprEvalDSWithWatermarks
       .flatMap(AvroCodecFn[T](groupByServingInfoParsed))
       .uid(s"avro-conversion-$groupByName")
       .name(s"Avro conversion for $groupByName")
@@ -221,7 +222,7 @@ class FlinkJob[T](eventSrc: FlinkSource[T],
       .name(s"Tiling Side Output Late Data for $groupByName")
       .setParallelism(sourceStream.parallelism)
 
-    val putRecordDS: DataStream[PutRequest] = tilingDS
+    val putRecordDS: DataStream[AvroCodecOutput] = tilingDS
       .flatMap(new TiledAvroCodecFn[T](groupByServingInfoParsed))
       .uid(s"avro-conversion-01-$groupByName")
       .name(s"Avro conversion for $groupByName")

flink/src/main/scala/ai/chronon/flink/MetricsSink.scala

@@ -1,4 +1,5 @@
 package ai.chronon.flink
+import ai.chronon.flink.types.WriteResponse
 import com.codahale.metrics.ExponentiallyDecayingReservoir
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.dropwizard.metrics.DropwizardHistogramWrapper
@@ -30,7 +31,7 @@ class MetricsSink(groupByName: String) extends RichSinkFunction[WriteResponse] {
   }
 
   override def invoke(value: WriteResponse, context: SinkFunction.Context): Unit = {
-    val eventCreatedToSinkTime = System.currentTimeMillis() - value.putRequest.tsMillis.get
+    val eventCreatedToSinkTime = System.currentTimeMillis() - value.tsMillis
     eventCreatedToSinkTimeHistogram.update(eventCreatedToSinkTime)
   }
 }

flink/src/main/scala/ai/chronon/flink/TestFlinkJob.scala

@@ -16,16 +16,22 @@ import ai.chronon.api.StringType
 import ai.chronon.api.TimeUnit
 import ai.chronon.api.Window
 import ai.chronon.api.{StructType => ApiStructType}
+import ai.chronon.flink.types.WriteResponse
 import ai.chronon.online.Api
 import ai.chronon.online.AvroCodec
 import ai.chronon.online.AvroConversions
 import ai.chronon.online.Extensions.StructTypeOps
 import ai.chronon.online.GroupByServingInfoParsed
+import org.apache.flink.api.common.serialization.DeserializationSchema
+import org.apache.flink.api.common.serialization.SerializationSchema
 import org.apache.flink.api.scala.createTypeInformation
+import org.apache.flink.metrics.groups.UnregisteredMetricsGroup
 import org.apache.flink.streaming.api.functions.sink.SinkFunction
 import org.apache.flink.streaming.api.functions.source.SourceFunction
 import org.apache.flink.streaming.api.scala.DataStream
 import org.apache.flink.streaming.api.scala.StreamExecutionEnvironment
+import org.apache.flink.util.SimpleUserCodeClassLoader
+import org.apache.flink.util.UserCodeClassLoader
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.avro.AvroDeserializationSupport
 import org.apache.spark.sql.types.StructType
@@ -71,11 +77,20 @@ class PrintSink extends SinkFunction[WriteResponse] {
   @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
   override def invoke(value: WriteResponse, context: SinkFunction.Context): Unit = {
-    val elapsedTime = System.currentTimeMillis() - value.putRequest.tsMillis.get
+    val elapsedTime = System.currentTimeMillis() - value.tsMillis
     logger.info(s"Received write response with status ${value.status}; elapsedTime = $elapsedTime ms")
   }
 }
 
+class DummyInitializationContext
+    extends SerializationSchema.InitializationContext
+    with DeserializationSchema.InitializationContext {
+  override def getMetricGroup = new UnregisteredMetricsGroup
+
+  override def getUserCodeClassLoader: UserCodeClassLoader =
+    SimpleUserCodeClassLoader.create(classOf[DummyInitializationContext].getClassLoader)
+}
+
 object TestFlinkJob {
   val fields: Array[(String, DataType)] = Array(
     "id" -> StringType,
@@ -92,6 +107,7 @@ object TestFlinkJob {
 
   def makeSource(mockPartitionCount: Int): FlinkSource[Row] = {
     val avroCodec = AvroCodec.of(e2eTestEventAvroSchema)
+    avroDeserializationSchema.open(new DummyInitializationContext)
     val startTs = System.currentTimeMillis()
     val elements: Seq[Map[String, AnyRef]] = (0 until 10).map(i =>
       Map("id" -> s"test$i",

flink/src/main/scala/ai/chronon/flink/types/FlinkTypes.scala

@@ -0,0 +1,120 @@
+package ai.chronon.flink.types
+
+import java.util
+import java.util.Objects
+
+// This file contains PoJo classes that are persisted while taking checkpoints in Chronon's Flink jobs. This falls primarily
+// in two buckets - tiled state and KV store incoming / outgoing records. The classes used in these cases need to allow for state
+// schema evolution (https://nightlies.apache.org/flink/flink-docs-release-1.17/docs/dev/datastream/fault-tolerance/serialization/schema_evolution/)
+// This allows us to add / remove fields without requiring us to migrate the state using dual write / read patterns.
+
+/**
+  * Combines the IR (intermediate result) with the timestamp of the event being processed.
+  * We need the timestamp of the event processed so we can calculate processing lag down the line.
+  *
+  * Example: for a GroupBy with 2 windows, we'd have TimestampedTile( [IR for window 1, IR for window 2], timestamp ).
+  *
+  * @param ir the array of partial aggregates
+  * @param latestTsMillis timestamp of the current event being processed
+  */
+class TimestampedIR(var ir: Array[Any], var latestTsMillis: Option[Long]) {
+  def this() = this(Array(), None)
+
+  override def toString: String =
+    s"TimestampedIR(ir=${ir.mkString(", ")}, latestTsMillis=$latestTsMillis)"
+
+  override def hashCode(): Int =
+    Objects.hash(ir.deep, latestTsMillis)
+
+  override def equals(other: Any): Boolean =
+    other match {
+      case e: TimestampedIR =>
+        util.Arrays.deepEquals(ir.asInstanceOf[Array[AnyRef]],
+                               e.ir.asInstanceOf[Array[AnyRef]]) && latestTsMillis == e.latestTsMillis
+      case _ => false
+    }
+}
+
+/**
+  * Combines the entity keys, the encoded IR (intermediate result), and the timestamp of the event being processed.
+  *
+  * We need the timestamp of the event processed so we can calculate processing lag down the line.
+  *
+  * @param keys the GroupBy entity keys
+  * @param tileBytes encoded tile IR
+  * @param latestTsMillis timestamp of the current event being processed
+  */
+class TimestampedTile(var keys: List[Any], var tileBytes: Array[Byte], var latestTsMillis: Long) {
+  def this() = this(List(), Array(), 0L)
+
+  override def toString: String =
+    s"TimestampedTile(keys=${keys.mkString(", ")}, tileBytes=${java.util.Base64.getEncoder
+      .encodeToString(tileBytes)}, latestTsMillis=$latestTsMillis)"
+
+  override def hashCode(): Int =
+    Objects.hash(keys.toArray.deep, tileBytes, latestTsMillis.asInstanceOf[java.lang.Long])
+
+  override def equals(other: Any): Boolean =
+    other match {
+      case e: TimestampedTile =>
+        util.Arrays.deepEquals(keys.toArray.asInstanceOf[Array[AnyRef]], e.keys.toArray.asInstanceOf[Array[AnyRef]]) &&
+          util.Arrays.equals(tileBytes, e.tileBytes) &&
+          latestTsMillis == e.latestTsMillis
+      case _ => false
+    }
+}
+
+/**
+  * Output emitted by the AvroCodecFn operator. This is fed into the Async KV store writer and objects of this type are persisted
+  * while taking checkpoints.
+  */
+class AvroCodecOutput(var keyBytes: Array[Byte], var valueBytes: Array[Byte], var dataset: String, var tsMillis: Long) {
+  def this() = this(Array(), Array(), "", 0L)
+
+  override def hashCode(): Int =
+    Objects.hash(
+      keyBytes,
+      valueBytes,
+      dataset,
+      tsMillis.asInstanceOf[java.lang.Long]
+    )
+
+  override def equals(other: Any): Boolean =
+    other match {
+      case o: AvroCodecOutput =>
+        util.Arrays.equals(keyBytes, o.keyBytes) &&
+          util.Arrays.equals(valueBytes, o.valueBytes) &&
+          dataset == o.dataset &&
+          tsMillis == o.tsMillis
+      case _ => false
+    }
+}
+
+/**
+  * Output records emitted by the AsyncKVStoreWriter. Objects of this type are persisted while taking checkpoints.
+  */
+class WriteResponse(var keyBytes: Array[Byte],
+                    var valueBytes: Array[Byte],
+                    var dataset: String,
+                    var tsMillis: Long,
+                    var status: Boolean) {
+  def this() = this(Array(), Array(), "", 0L, false)
+
+  override def hashCode(): Int =
+    Objects.hash(keyBytes,
+                 valueBytes,
+                 dataset,
+                 tsMillis.asInstanceOf[java.lang.Long],
+                 status.asInstanceOf[java.lang.Boolean])
+
+  override def equals(other: Any): Boolean =
+    other match {
+      case o: WriteResponse =>
+        util.Arrays.equals(keyBytes, o.keyBytes) &&
+          util.Arrays.equals(valueBytes, o.valueBytes) &&
+          dataset == o.dataset &&
+          tsMillis == o.tsMillis &&
+          status == o.status
+      case _ => false
+    }
+}