Fix NPE when fetching a non-exist key from a GroupBy w/ derivation

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

@@ -63,7 +63,7 @@ class FetcherBase(kvStore: KVStore,
                                        context: Metrics.Context,
                                        totalResponseValueBytes: Int,
                                        keys: Map[String, Any] // The keys are used only for caching
-  ): Map[String, AnyRef] = {
+  ): Option[Map[String, AnyRef]] = {
     val (servingInfo, batchResponseMaxTs) = getServingInfo(oldServingInfo, batchResponses)
 
     // Batch metrics
@@ -114,7 +114,7 @@ class FetcherBase(kvStore: KVStore,
       ) {
         if (debug) logger.info("Both batch and streaming data are null")
         context.distribution("group_by.latency.millis", System.currentTimeMillis() - startTimeMs)
-        return null
+        return None
       }
 
       // Streaming metrics
@@ -238,7 +238,7 @@ class FetcherBase(kvStore: KVStore,
     }
 
     context.distribution("group_by.latency.millis", System.currentTimeMillis() - startTimeMs)
-    responseMap
+    Some(responseMap)
   }
 
   def reportKvResponse(ctx: Metrics.Context,
@@ -487,7 +487,7 @@ class FetcherBase(kvStore: KVStore,
                                            multiGetMillis,
                                            context,
                                            totalResponseValueBytes,
-                                           request.keys)
+                                           request.keys).getOrElse(Map.empty)
                 } catch {
                   case ex: Exception =>
                     // not all exceptions are due to stale schema, so we want to control how often we hit kv store

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

@@ -187,7 +187,11 @@ class FetcherTest extends TestCase {
         )
       ),
       accuracy = Accuracy.TEMPORAL,
-      metaData = Builders.MetaData(name = "unit_test/fetcher_mutations_gb", namespace = namespace, team = "chronon")
+      metaData = Builders.MetaData(name = "unit_test/fetcher_mutations_gb", namespace = namespace, team = "chronon"),
+      derivations=Seq(
+        Builders.Derivation(name = "*", expression = "*"),
+        Builders.Derivation(name = "rating_average_1d_same", expression = "rating_average_1d")
+      )
     )
 
     val joinConf = Builders.Join(
@@ -689,6 +693,38 @@ class FetcherTest extends TestCase {
     assertEquals(joinConf.joinParts.size() + derivationExceptionTypes.size, responseMap.size)
     assertTrue(responseMap.keys.forall(_.endsWith("_exception")))
   }
+
+  def testTemporalFetchGroupByNonExistKey(): Unit = {
+    val namespace = "non_exist_key_group_by_fetch"
+    val joinConf = generateMutationData(namespace)
+    val endDs = "2021-04-10"
+    val spark: SparkSession = SparkSessionBuilder.build(sessionName + "_" + Random.alphanumeric.take(6).mkString, local = true)
+    val tableUtils = TableUtils(spark)
+    val kvStoreFunc = () => OnlineUtils.buildInMemoryKVStore("FetcherTest")
+    val inMemoryKvStore = kvStoreFunc()
+    val mockApi = new MockApi(kvStoreFunc, namespace)
+    @transient lazy val fetcher = mockApi.buildFetcher(debug=false)
+
+    joinConf.joinParts.toScala.foreach(jp =>
+      OnlineUtils.serve(tableUtils,
+        inMemoryKvStore,
+        kvStoreFunc,
+        namespace,
+        endDs,
+        jp.groupBy,
+        dropDsOnWrite = true))
+
+    // a random key that doesn't exist
+    val nonExistKey = 123L
+    val request = Request("unit_test/fetcher_mutations_gb",
+      Map("listing_id" -> nonExistKey.asInstanceOf[AnyRef]))
+    val response = fetcher.fetchGroupBys(Seq(request))
+    val result = Await.result(response, Duration(10, SECONDS))
+
+    // result should be "null" if the key is not found
+    val expected: Map[String, AnyRef] = Map("rating_average_1d_same" -> null)
+    assertEquals(expected, result.head.values.get)
+  }
 }
 
 object FetcherTestUtil {