Don't request from KVStore if ANY of key is null (#968)

* Don't request null keys from KVStore

* add another val to avoid indentation

* run scalafmt

* change from ALL to ANY

* run scalafmt

* merge recent changes to fix unit tests

* fix unit tests

* rebase the latest changes

Author: xiaohui-sun

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

@@ -350,7 +350,10 @@ 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 =>
+    // don't send to KV store if ANY of key's value is null
+    val validRequests =
+      requests.filter(r => r.keys == null || r.keys.values == null || !r.keys.values.exists(_ == null))
+    val groupByRequestToKvRequest: Seq[(Request, Try[GroupByRequestMeta])] = validRequests.iterator.map { request =>
       val groupByServingInfoTry = getGroupByServingInfo(request.name)
         .recover {
           case ex: Throwable =>
@@ -665,28 +668,8 @@ class FetcherBase(kvStore: KVStore,
                   }
                   Map(fetchException.getRequestName + "_exception" -> fetchException.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(prefix + "_exception" -> ex.traceString)
-                    }
-                    .get
-                }
+                case Left(PrefixedRequest(prefix, groupByRequest)) =>
+                  parseGroupByResponse(prefix, groupByRequest, responseMap)
               }.toMap
             }
             joinValuesTry match {
@@ -706,6 +689,41 @@ class FetcherBase(kvStore: KVStore,
       }
   }
 
+  def parseGroupByResponse(prefix: String,
+                           groupByRequest: Request,
+                           responseMap: Map[Request, Try[Map[String, AnyRef]]]) = {
+
+    // Group bys with ANY null keys won't be requested from the KV store and we don't expect a response.
+    val isRequiredRequest =
+      (groupByRequest.keys.nonEmpty && !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) {
+            logger.error(s"Failed to fetch $groupByRequest", ex)
+          }
+          Map(prefix + "_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/FetcherBaseTest.scala

@@ -269,4 +269,50 @@ class FetcherBaseTest extends MockitoSugar with Matchers with MockitoHelper {
     val result2 = fetcherBase.checkLateBatchData(1710896400000L, "myGroupBy", 1710633600000L, tailHops2d, shortWindows)
     assertSame(result2, 0L)
   }
+
+  @Test
+  def testParsingGroupByResponse_HappyCase(): Unit = {
+    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)
+    result shouldBe Map("prefix_key" -> "value")
+  }
+
+  @Test
+  def testParsingGroupByResponse_NullKey(): Unit = {
+    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()
+  }
+
+  @Test
+  def testParsingGroupByResponse_MissingKey(): Unit = {
+    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("prefix_exception")
+  }
 }

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

@@ -401,7 +401,10 @@ class FetcherTest extends TestCase {
 
     val listingEventData = Seq(
       Row(1L, toTs("2021-04-10 03:10:00"), "2021-04-10"),
-      Row(2L, toTs("2021-04-10 03:10:00"), "2021-04-10")
+      Row(2L, toTs("2021-04-10 03:10:00"), "2021-04-10"),
+      Row(2L, toTs("2021-04-10 03:10:00"), "2021-04-10"),
+      Row(2L, toTs("2021-04-10 03:10:00"), "2021-04-10"),
+      Row(null, toTs("2021-04-10 03:10:00"), "2021-04-10")
     )
     val ratingEventData = Seq(
       // 1L listing id event data
@@ -422,7 +425,9 @@ class FetcherTest extends TestCase {
       Row(2L, toTs("2021-04-10 02:30:00"), 5, "2021-04-10"),
       Row(2L, toTs("2021-04-10 02:30:00"), 8, "2021-04-10"),
       Row(2L, toTs("2021-04-10 02:30:00"), 8, "2021-04-10"),
-      Row(2L, toTs("2021-04-07 00:30:00"), 10, "2021-04-10") // dated 4/10 but excluded from avg agg based on ts
+      Row(2L, toTs("2021-04-07 00:30:00"), 10, "2021-04-10"), // dated 4/10 but excluded from avg agg based on ts
+      Row(2L, toTs("2021-04-07 00:30:00"), 10, "2021-04-10"), // dated 4/10 but excluded from avg agg based on ts
+      Row(null, toTs("2021-04-10 02:30:00"), 8, "2021-04-10")
     )
     // Schemas
     // {..., event (generic event column), ...}