cherry pick Don't request null keys from KVStore #774 (#297)

## Summary

Cherry pick: https://github.com/airbnb/chronon/pull/774/files

## 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**
	- Enhanced group-by request processing with improved key validation
- Added a new method for parsing group-by responses with more robust
error handling

- **Tests**
	- Added comprehensive test cases for new response parsing method
	- Verified handling of scenarios with null keys and missing keys

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

Co-authored-by: ezvz <[email protected]>

Author: varant-zlai

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

@@ -323,57 +323,63 @@ class FetcherBase(kvStore: KVStore,
   // 4. Finally converted to outputSchema
   def fetchGroupBys(requests: scala.collection.Seq[Request]): Future[scala.collection.Seq[Response]] = {
     // split a groupBy level request into its kvStore level requests
-    val groupByRequestToKvRequest: Seq[(Request, Try[GroupByRequestMeta])] = requests.iterator.map { request =>
-      val groupByRequestMetaTry: Try[GroupByRequestMeta] = getGroupByServingInfo(request.name)
-        .map { groupByServingInfo =>
-          val context =
-            request.context.getOrElse(Metrics.Context(Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
-          context.increment("group_by_request.count")
-          var batchKeyBytes: Array[Byte] = null
-          var streamingKeyBytes: Array[Byte] = null
-          try {
-            // The formats of key bytes for batch requests and key bytes for streaming requests may differ based
-            // on the KVStore implementation, so we encode each distinctly.
-            batchKeyBytes =
-              kvStore.createKeyBytes(request.keys, groupByServingInfo, groupByServingInfo.groupByOps.batchDataset)
-            streamingKeyBytes =
-              kvStore.createKeyBytes(request.keys, groupByServingInfo, groupByServingInfo.groupByOps.streamingDataset)
-          } catch {
-            // TODO: only gets hit in cli path - make this code path just use avro schema to decode keys directly in cli
-            // TODO: Remove this code block
-            case ex: Exception =>
-              val castedKeys = groupByServingInfo.keyChrononSchema.fields.map {
-                case StructField(name, typ) => name -> ColumnAggregator.castTo(request.keys.getOrElse(name, null), typ)
-              }.toMap
-              try {
-                batchKeyBytes =
-                  kvStore.createKeyBytes(castedKeys, groupByServingInfo, groupByServingInfo.groupByOps.batchDataset)
-                streamingKeyBytes =
-                  kvStore.createKeyBytes(castedKeys, groupByServingInfo, groupByServingInfo.groupByOps.streamingDataset)
-              } catch {
-                case exInner: Exception =>
-                  exInner.addSuppressed(ex)
-                  throw new RuntimeException("Couldn't encode request keys or casted keys", exInner)
-              }
-          }
-          val batchRequest = GetRequest(batchKeyBytes, groupByServingInfo.groupByOps.batchDataset)
-          val streamingRequestOpt = groupByServingInfo.groupByOps.inferredAccuracy match {
-            // fetch batch(ir) and streaming(input) and aggregate
-            case Accuracy.TEMPORAL =>
-              Some(
-                GetRequest(streamingKeyBytes,
-                           groupByServingInfo.groupByOps.streamingDataset,
-                           Some(groupByServingInfo.batchEndTsMillis)))
-            // no further aggregation is required - the value in KvStore is good as is
-            case Accuracy.SNAPSHOT => None
+    val groupByRequestToKvRequest: Seq[(Request, Try[GroupByRequestMeta])] = requests.iterator
+      .filter(r => r.keys == null || r.keys.values == null || r.keys.values.exists(_ != null))
+      .map { request =>
+        val groupByRequestMetaTry: Try[GroupByRequestMeta] = getGroupByServingInfo(request.name)
+          .map { groupByServingInfo =>
+            val context =
+              request.context.getOrElse(
+                Metrics.Context(Metrics.Environment.GroupByFetching, groupByServingInfo.groupBy))
+            context.increment("group_by_request.count")
+            var batchKeyBytes: Array[Byte] = null
+            var streamingKeyBytes: Array[Byte] = null
+            try {
+              // The formats of key bytes for batch requests and key bytes for streaming requests may differ based
+              // on the KVStore implementation, so we encode each distinctly.
+              batchKeyBytes =
+                kvStore.createKeyBytes(request.keys, groupByServingInfo, groupByServingInfo.groupByOps.batchDataset)
+              streamingKeyBytes =
+                kvStore.createKeyBytes(request.keys, groupByServingInfo, groupByServingInfo.groupByOps.streamingDataset)
+            } catch {
+              // TODO: only gets hit in cli path - make this code path just use avro schema to decode keys directly in cli
+              // TODO: Remove this code block
+              case ex: Exception =>
+                val castedKeys = groupByServingInfo.keyChrononSchema.fields.map {
+                  case StructField(name, typ) =>
+                    name -> ColumnAggregator.castTo(request.keys.getOrElse(name, null), typ)
+                }.toMap
+                try {
+                  batchKeyBytes =
+                    kvStore.createKeyBytes(castedKeys, groupByServingInfo, groupByServingInfo.groupByOps.batchDataset)
+                  streamingKeyBytes = kvStore.createKeyBytes(castedKeys,
+                                                             groupByServingInfo,
+                                                             groupByServingInfo.groupByOps.streamingDataset)
+                } catch {
+                  case exInner: Exception =>
+                    exInner.addSuppressed(ex)
+                    throw new RuntimeException("Couldn't encode request keys or casted keys", exInner)
+                }
+            }
+            val batchRequest = GetRequest(batchKeyBytes, groupByServingInfo.groupByOps.batchDataset)
+            val streamingRequestOpt = groupByServingInfo.groupByOps.inferredAccuracy match {
+              // fetch batch(ir) and streaming(input) and aggregate
+              case Accuracy.TEMPORAL =>
+                Some(
+                  GetRequest(streamingKeyBytes,
+                             groupByServingInfo.groupByOps.streamingDataset,
+                             Some(groupByServingInfo.batchEndTsMillis)))
+              // no further aggregation is required - the value in KvStore is good as is
+              case Accuracy.SNAPSHOT => None
+            }
+            GroupByRequestMeta(groupByServingInfo, batchRequest, streamingRequestOpt, request.atMillis, context)
           }
-          GroupByRequestMeta(groupByServingInfo, batchRequest, streamingRequestOpt, request.atMillis, context)
+        if (groupByRequestMetaTry.isFailure) {
+          request.context.foreach(_.increment("group_by_serving_info_failure.count"))
         }
-      if (groupByRequestMetaTry.isFailure) {
-        request.context.foreach(_.increment("group_by_serving_info_failure.count"))
+        request -> groupByRequestMetaTry
       }
-      request -> groupByRequestMetaTry
-    }.toSeq
+      .toSeq
 
     // If caching is enabled, we check if any of the GetRequests are already cached. If so, we store them in a Map
     // and avoid the work of re-fetching them. It is mainly for batch data requests.
@@ -583,28 +589,8 @@ class FetcherBase(kvStore: KVStore,
                 case Right(keyMissingException) => {
                   Map(keyMissingException.requestName + "_exception" -> keyMissingException.getMessage)
                 }
-                case Left(PrefixedRequest(prefix, groupByRequest)) => {
-                  responseMap
-                    .getOrElse(groupByRequest,
-                               Failure(new IllegalStateException(
-                                 s"Couldn't find a groupBy response for $groupByRequest in response map")))
-                    .map { valueMap =>
-                      if (valueMap != null) {
-                        valueMap.map { case (aggName, aggValue) => prefix + "_" + aggName -> aggValue }
-                      } else {
-                        Map.empty[String, AnyRef]
-                      }
-                    }
-                    // prefix feature names
-                    .recover { // capture exception as a key
-                      case ex: Throwable =>
-                        if (debug || Math.random() < 0.001) {
-                          logger.error(s"Failed to fetch $groupByRequest", ex)
-                        }
-                        Map(groupByRequest.name + "_exception" -> ex.traceString)
-                    }
-                    .get
-                }
+                case Left(PrefixedRequest(prefix, groupByRequest)) =>
+                  parseGroupByResponse(prefix, groupByRequest, responseMap)
               }.toMap
             }
             joinValuesTry match {
@@ -624,6 +610,39 @@ class FetcherBase(kvStore: KVStore,
       }
   }
 
+  def parseGroupByResponse(prefix: String,
+                           groupByRequest: Request,
+                           responseMap: Map[Request, Try[Map[String, AnyRef]]]): Map[String, AnyRef] = {
+    // Group bys with all null keys won't be requested from the KV store and we don't expect a response.
+    val isRequiredRequest = groupByRequest.keys.values.exists(_ != null) || groupByRequest.keys.isEmpty
+
+    val response: Try[Map[String, AnyRef]] = responseMap.get(groupByRequest) match {
+      case Some(value) => value
+      case None =>
+        if (isRequiredRequest)
+          Failure(new IllegalStateException(s"Couldn't find a groupBy response for $groupByRequest in response map"))
+        else Success(null)
+    }
+
+    response
+      .map { valueMap =>
+        if (valueMap != null) {
+          valueMap.map { case (aggName, aggValue) => prefix + "_" + aggName -> aggValue }
+        } else {
+          Map.empty[String, AnyRef]
+        }
+      }
+      // prefix feature names
+      .recover { // capture exception as a key
+        case ex: Throwable =>
+          if (debug || Math.random() < 0.001) {
+            println(s"Failed to fetch $groupByRequest with \n${ex.traceString}")
+          }
+          Map(groupByRequest.name + "_exception" -> ex.traceString)
+      }
+      .get
+  }
+
   /**
     * Fetch method to simulate a random access interface for Chronon
     * by distributing requests to relevant GroupBys. This is a batch

online/src/test/scala/ai/chronon/online/test/FetcherBaseTest.scala

@@ -27,6 +27,7 @@ import ai.chronon.online.Fetcher.Response
 import ai.chronon.online.FetcherCache.BatchResponses
 import ai.chronon.online.KVStore.TimedValue
 import ai.chronon.online._
+import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
 import org.mockito.Answers
@@ -224,4 +225,47 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
     assertFalse(fetcherBaseWithFlagStore.isCachingEnabled(buildGroupByWithCustomJson("test_groupby_2")))
     assertTrue(fetcherBaseWithFlagStore.isCachingEnabled(buildGroupByWithCustomJson("test_groupby_3")))
   }
+
+  it should "fetch in the happy case" in {
+    val baseFetcher = new FetcherBase(mock[KVStore])
+    val request = Request(name = "name", keys = Map("email" -> "email"), atMillis = None, context = None)
+    val response: Map[Request, Try[Map[String, AnyRef]]] = Map(
+      request -> Success(Map(
+        "key" -> "value"
+      ))
+    )
+
+    val result = baseFetcher.parseGroupByResponse("prefix", request, response)
+    assertEquals(result, Map("prefix_key" -> "value"))
+  }
+
+  it should "Not fetch with null keys" in {
+    val baseFetcher = new FetcherBase(mock[KVStore])
+    val request = Request(name = "name", keys = Map("email" -> null), atMillis = None, context = None)
+    val request2 = Request(name = "name2", keys = Map("email" -> null), atMillis = None, context = None)
+
+    val response: Map[Request, Try[Map[String, AnyRef]]] = Map(
+      request2 -> Success(Map(
+        "key" -> "value"
+      ))
+    )
+
+    val result = baseFetcher.parseGroupByResponse("prefix", request, response)
+    result shouldBe Map()
+  }
+
+  it should "parse with missing keys" in {
+    val baseFetcher = new FetcherBase(mock[KVStore])
+    val request = Request(name = "name", keys = Map("email" -> "email"), atMillis = None, context = None)
+    val request2 = Request(name = "name2", keys = Map("email" -> "email"), atMillis = None, context = None)
+
+    val response: Map[Request, Try[Map[String, AnyRef]]] = Map(
+      request2 -> Success(Map(
+        "key" -> "value"
+      ))
+    )
+
+    val result = baseFetcher.parseGroupByResponse("prefix", request, response)
+    result.keySet shouldBe Set("name_exception")
+  }
 }