Summary upload (#50)

## Summary
Creates a Summary Uploader which uploads summary data to a KVStore.

## Checklist
- [ x ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update



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

- **New Features**
- Introduced `SummaryUploader` for uploading summary statistics to a
key-value store.
	- Added `MockKVStore` for testing key-value store operations.
- Implemented `KVStoreSemaphore` for managing concurrent access to
resources.

- **Enhancements**
	- Increased data volume in tests to improve testing scenarios.
- Integrated `SummaryUploader` in `DriftTest` for uploading summary data
during tests.
- Enhanced control over concurrent reads and writes to DynamoDB with
updated `DynamoDBKVStoreImpl`.
- Refined error handling and flow in `multiPut` operations for better
robustness.
- Updated Spark dependency from `3.5.0` to `3.5.1` for improved
stability.
	- Added a new constant `DriftStatsTable` for drift statistics.

- **Bug Fixes**
	- Improved error handling for upload failures in `SummaryUploader`.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: nikhil-zlai <[email protected]>
Co-authored-by: Chewy Shaw <[email protected]>
Co-authored-by: Piyush Narang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: ken-zlai <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: ezvz <[email protected]>

Author: 8 people

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

@@ -62,4 +62,5 @@ object Constants {
   val LabelViewPropertyFeatureTable: String = "feature_table"
   val LabelViewPropertyKeyLabelTable: String = "label_table"
   val ChrononRunDs: String = "CHRONON_RUN_DS"
+  val DriftStatsTable: String = "drift_statistics"
 }

build.sbt

@@ -157,8 +157,9 @@ lazy val spark = project
     crossScalaVersions := supportedVersions,
     libraryDependencies ++= spark_all_provided,
     libraryDependencies ++= spark_all.map(_ % "test"),
-    libraryDependencies += "jakarta.servlet" % "jakarta.servlet-api" % "4.0.3"
-  )
+    libraryDependencies += "jakarta.servlet" % "jakarta.servlet-api" % "4.0.3",
+    libraryDependencies += "com.google.guava" % "guava" % "33.3.1-jre"
+)
 
 lazy val flink = project
   .dependsOn(aggregator.%("compile->compile;test->test"), online)
@@ -189,6 +190,7 @@ lazy val cloud_aws = project
     libraryDependencies += "io.circe" %% "circe-core" % circeVersion % "test",
     libraryDependencies += "io.circe" %% "circe-generic" % circeVersion % "test",
     libraryDependencies += "io.circe" %% "circe-parser" % circeVersion % "test",
+    libraryDependencies += "com.google.guava" % "guava" % "33.3.1-jre",
     libraryDependencies ++= spark_all
   )
 

cloud_aws/src/main/scala/ai/chronon/integrations/aws/DynamoDBKVStoreImpl.scala

@@ -9,6 +9,7 @@ import ai.chronon.online.KVStore.ListValue
 import ai.chronon.online.KVStore.TimedValue
 import ai.chronon.online.Metrics
 import ai.chronon.online.Metrics.Context
+import com.google.common.util.concurrent.RateLimiter
 import software.amazon.awssdk.core.SdkBytes
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition
@@ -30,6 +31,7 @@ import software.amazon.awssdk.services.dynamodb.model.ScanResponse
 
 import java.time.Instant
 import java.util
+import java.util.concurrent.ConcurrentHashMap
 import scala.concurrent.Future
 import scala.jdk.CollectionConverters._
 import scala.util.Success
@@ -64,6 +66,8 @@ object DynamoDBKVStoreConstants {
 
 class DynamoDBKVStoreImpl(dynamoDbClient: DynamoDbClient) extends KVStore {
   import DynamoDBKVStoreConstants._
+  private val readRateLimiters = new ConcurrentHashMap[String, RateLimiter]()
+  private val writeRateLimiters = new ConcurrentHashMap[String, RateLimiter]()
 
   protected val metricsContext: Metrics.Context = Metrics.Context(Metrics.Environment.KVStore).withSuffix("dynamodb")
 
@@ -88,6 +92,9 @@ class DynamoDBKVStoreImpl(dynamoDbClient: DynamoDbClient) extends KVStore {
     val rcu = getCapacityUnits(props, readCapacityUnits, defaultReadCapacityUnits)
     val wcu = getCapacityUnits(props, writeCapacityUnits, defaultWriteCapacityUnits)
 
+    readRateLimiters.put(dataset, RateLimiter.create(rcu))
+    writeRateLimiters.put(dataset, RateLimiter.create(wcu))
+
     val request =
       CreateTableRequest.builder
         .attributeDefinitions(keyAttributes.toList.asJava)
@@ -137,6 +144,7 @@ class DynamoDBKVStoreImpl(dynamoDbClient: DynamoDbClient) extends KVStore {
     val getItemResults = getItemRequestPairs.map {
       case (req, getItemReq) =>
         Future {
+          readRateLimiters.computeIfAbsent(req.dataset, _ => RateLimiter.create(defaultReadCapacityUnits)).acquire()
           val item: Try[util.Map[String, AttributeValue]] =
             handleDynamoDbOperation(metricsContext.withSuffix("multiget"), req.dataset) {
               dynamoDbClient.getItem(getItemReq).item()
@@ -151,6 +159,7 @@ class DynamoDBKVStoreImpl(dynamoDbClient: DynamoDbClient) extends KVStore {
     val queryResults = queryRequestPairs.map {
       case (req, queryRequest) =>
         Future {
+          readRateLimiters.computeIfAbsent(req.dataset, _ => RateLimiter.create(defaultReadCapacityUnits)).acquire()
           val responses = handleDynamoDbOperation(metricsContext.withSuffix("query"), req.dataset) {
             dynamoDbClient.query(queryRequest).items()
           }
@@ -218,12 +227,10 @@ class DynamoDBKVStoreImpl(dynamoDbClient: DynamoDbClient) extends KVStore {
     val futureResponses = datasetToWriteRequests.map {
       case (dataset, putItemRequest) =>
         Future {
+          writeRateLimiters.computeIfAbsent(dataset, _ => RateLimiter.create(defaultWriteCapacityUnits)).acquire()
           handleDynamoDbOperation(metricsContext.withSuffix("multiput"), dataset) {
             dynamoDbClient.putItem(putItemRequest)
-          }.recover {
-            case _: Exception => false
-          }
-          true
+          }.isSuccess
         }
     }
     Future.sequence(futureResponses)

online/src/main/scala/ai/chronon/online/Api.scala

@@ -64,7 +64,6 @@ object KVStore {
 trait KVStore {
   @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
   implicit val executionContext: ExecutionContext = FlexibleExecutionContext.buildExecutionContext
-
   def create(dataset: String): Unit
 
   def create(dataset: String, props: Map[String, Any]): Unit = create(dataset)

spark/src/main/scala/ai/chronon/spark/stats/drift/SummaryUploader.scala

@@ -0,0 +1,60 @@
+package ai.chronon.spark.stats.drift
+import ai.chronon.api.Constants
+import ai.chronon.online.KVStore
+import ai.chronon.online.KVStore.PutRequest
+import ai.chronon.spark.TableUtils
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.types
+
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
+
+class SummaryUploader(summaryDF: DataFrame, kvStoreFunc: () => KVStore, putsPerRequest: Int = 100)(implicit
+    tu: TableUtils)
+    extends Serializable {
+  val completed_schema: types.StructType = types.StructType(
+    Seq(
+      types.StructField(tu.partitionColumn, types.StringType, nullable = false)
+    )
+  )
+  private val statsTableName = Constants.DriftStatsTable
+
+  def run(): Unit = {
+    // Validate schema
+    val requiredColumns = Seq("keyBytes", "valueBytes", "timestamp")
+    val missingColumns = requiredColumns.filterNot(summaryDF.columns.contains)
+    require(missingColumns.isEmpty, s"Missing required columns: ${missingColumns.mkString(", ")}")
+
+    summaryDF.rdd.foreachPartition(rows => {
+      val kvStore: KVStore = kvStoreFunc()
+      val putRequests = new scala.collection.mutable.ArrayBuffer[PutRequest]
+      for (row <- rows) {
+        putRequests += PutRequest(
+          Option(row.getAs[Array[Byte]]("keyBytes")).getOrElse(Array.empty[Byte]),
+          Option(row.getAs[Array[Byte]]("valueBytes")).getOrElse(Array.empty[Byte]),
+          statsTableName,
+          Option(row.getAs[Long]("timestamp"))
+        )
+      }
+
+      val futureResults = putRequests.grouped(putsPerRequest).map { batch =>
+        kvStore.multiPut(batch.toList).map { result =>
+          if (!result.forall(identity)) {
+            throw new RuntimeException(s"Failed to put ${result.count(!_)} records")
+          }
+        }
+      }
+
+      val aggregatedFuture = Future.sequence(futureResults.toSeq)
+      aggregatedFuture.onComplete {
+        case scala.util.Success(_) => // All operations completed successfully
+        case scala.util.Failure(e: IllegalArgumentException) =>
+          throw new IllegalArgumentException(s"Invalid request data: ${e.getMessage}", e)
+        case scala.util.Failure(e: java.io.IOException) =>
+          throw new RuntimeException(s"KVStore I/O error: ${e.getMessage}", e)
+        case scala.util.Failure(e) =>
+          throw new RuntimeException(s"Failed to upload summary statistics: ${e.getMessage}", e)
+      }
+    })
+  }
+}

spark/src/test/scala/ai/chronon/spark/test/MockKVStore.scala

@@ -0,0 +1,34 @@
+package ai.chronon.spark.test
+
+import ai.chronon.online.KVStore
+
+import scala.collection.mutable
+import scala.concurrent.Future
+
+class MockKVStore() extends KVStore with Serializable {
+    val num_puts: mutable.Map[String,Int] = collection.mutable.Map[String, Int]()
+
+    def bulkPut(sourceOfflineTable: String,destinationOnlineDataSet: String,partition: String): Unit = 
+      throw new UnsupportedOperationException("Not implemented in mock")
+  def create(dataset: String): Unit = 
+    {
+      num_puts(dataset) = 0
+    }
+  def multiGet(requests: Seq[ai.chronon.online.KVStore.GetRequest]): scala.concurrent.Future[Seq[ai.chronon.online.KVStore.GetResponse]] =
+    throw new UnsupportedOperationException("Not implemented in mock")
+  def multiPut(keyValueDatasets: Seq[ai.chronon.online.KVStore.PutRequest]): scala.concurrent.Future[Seq[Boolean]] = {
+    logger.info(s"Triggering multiput for ${keyValueDatasets.size}: rows")
+    for (req <- keyValueDatasets if (!req.keyBytes.isEmpty && !req.valueBytes.isEmpty)) num_puts(req.dataset) += 1
+
+    val futureResponses = keyValueDatasets.map { req =>
+        if (!req.keyBytes.isEmpty && !req.valueBytes.isEmpty) Future{true}
+        else Future{false}
+    }
+    Future.sequence(futureResponses)
+  }
+
+  def show(): Unit = {
+    num_puts.foreach(x => logger.info(s"Ran ${x._2} non-empty put actions for dataset ${x._1}"))
+    
+  }
+}

spark/src/test/scala/ai/chronon/spark/test/stats/drift/DriftTest.scala

@@ -3,13 +3,17 @@ package ai.chronon.spark.test.stats.drift
 import ai.chronon.aggregator.test.Column
 import ai.chronon.api
 import ai.chronon.api.ColorPrinter.ColorString
+import ai.chronon.api.Constants
 import ai.chronon.api.Extensions.MetadataOps
+import ai.chronon.online.KVStore
 import ai.chronon.spark.Extensions._
 import ai.chronon.spark.SparkSessionBuilder
 import ai.chronon.spark.TableUtils
 import ai.chronon.spark.stats.drift.Summarizer
 import ai.chronon.spark.stats.drift.SummaryPacker
+import ai.chronon.spark.stats.drift.SummaryUploader
 import ai.chronon.spark.test.DataFrameGen
+import ai.chronon.spark.test.MockKVStore
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.SparkSession
 import org.scalatest.flatspec.AnyFlatSpec
@@ -36,6 +40,18 @@ class DriftTest extends AnyFlatSpec with Matchers {
       val packer = new SummaryPacker("drift_test_basic", summaryExprs, summarizer.tileSize, summarizer.sliceColumns)
       val (packed, _) = packer.packSummaryDf(result)
       packed.show()
+
+      val props = Map("is-time-sorted" -> "true")
+
+      val kvStore: () => KVStore = () => {
+        val result = new MockKVStore()
+        result.create(Constants.DriftStatsTable, props)
+        result
+      }
+      
+      val uploader = new SummaryUploader(packed,kvStore)
+      uploader.run()
+      //kvStore.show()
     }
   }