Allow using key columns as input to derivations

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

@@ -1085,17 +1085,18 @@ object Extensions {
     }
 
     // Used only during online fetching to reduce latency
-    def applyRenameOnlyDerivation(baseColumns: Map[String, Any]): Map[String, Any] = {
+    def applyRenameOnlyDerivation(keys: Map[String, Any], values: Map[String, Any]): Map[String, Any] = {
       assert(
         areDerivationsRenameOnly,
         s"Derivations contain more complex expressions than simple renames: ${derivations.map(d => (d.name, d.expression))}")
       val wildcardDerivations = if (derivationsContainStar) {
-        baseColumns.filterNot(derivationExpressionSet contains _._1)
+        values.filterNot(derivationExpressionSet contains _._1)
       } else {
         Map.empty[String, Any]
       }
-
-      wildcardDerivations ++ derivationsWithoutStar.map(d => d.name -> baseColumns.getOrElse(d.expression, null)).toMap
+      wildcardDerivations ++ derivationsWithoutStar
+        .map(d => d.name -> (keys ++ values).getOrElse(d.expression, null))
+        .toMap
     }
   }
 }

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

@@ -229,9 +229,16 @@ class Fetcher(val kvStore: KVStore,
             joinCodecTry match {
               case Success(joinCodec) =>
                 ctx.distribution("derivation_codec.latency.millis", System.currentTimeMillis() - derivationStartTs)
+                // try to fix request mistype
+                val castedKeys : Map[String, AnyRef]  = joinCodec.keySchema.fields.map {
+                  field => field.name -> ColumnAggregator.castTo(internalResponse.request.keys.get(field.name).orNull, field.fieldType)
+                }.toMap
+                val request = Request(name = internalResponse.request.name, keys = castedKeys,
+                  atMillis = internalResponse.request.atMillis, context = internalResponse.request.context)
+
                 val baseMap = internalMap ++ externalMap
                 val derivedMapTry: Try[Map[String, AnyRef]] = Try {
-                  applyDeriveFunc(joinCodec.deriveFunc, internalResponse.request, baseMap)
+                  applyDeriveFunc(joinCodec.deriveFunc, request, baseMap)
                 }
                 val derivedMap: Map[String, AnyRef] = derivedMapTry match {
                   case Success(derivedMap) => derivedMap
@@ -260,7 +267,7 @@ class Fetcher(val kvStore: KVStore,
                 val requestEndTs = System.currentTimeMillis()
                 ctx.distribution("derivation.latency.millis", requestEndTs - derivationStartTs)
                 ctx.distribution("overall.latency.millis", requestEndTs - ts)
-                ResponseWithContext(internalResponse.request, finalizedDerivedMap, baseMap)
+                ResponseWithContext(request, finalizedDerivedMap, baseMap)
               case Failure(exception) =>
                 // more validation logic will be covered in compile.py to avoid this case
                 ctx.incrementException(exception)

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

@@ -30,8 +30,8 @@ object OnlineDerivationUtil {
 
   def buildRenameOnlyDerivationFunction(derivationsScala: List[Derivation]): DerivationFunc = {
     {
-      case (_: Map[String, Any], values: Map[String, Any]) =>
-        reintroduceExceptions(derivationsScala.applyRenameOnlyDerivation(values), values)
+      case (keys: Map[String, Any], values: Map[String, Any]) =>
+        reintroduceExceptions(derivationsScala.applyRenameOnlyDerivation(keys, values), values)
     }
   }
 
@@ -104,12 +104,13 @@ object OnlineDerivationUtil {
       }
       val expressions: Seq[StructField] = baseExpressions ++ derivationsScala.derivationsWithoutStar.map { d =>
         {
-          if (baseValueSchema.typeOf(d.expression).isEmpty) {
+          val allSchema = StructType("all", (keySchema ++ baseValueSchema).toArray)
+          if (allSchema.typeOf(d.expression).isEmpty) {
             throw new IllegalArgumentException(
               s"Failed to run expression ${d.expression} for ${d.name}. Please ensure the derivation is " +
                 s"correct.")
           } else {
-            StructField(d.name, baseValueSchema.typeOf(d.expression).get)
+            StructField(d.name, allSchema.typeOf(d.expression).get)
           }
         }
       }

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

@@ -648,6 +648,30 @@ class FetcherTest extends TestCase {
     compareTemporalFetch(joinConf, "2021-04-10", namespace, consistencyCheck = false, dropDsOnWrite = true)
   }
 
+  def testTemporalFetchJoinDerivation(): Unit = {
+    val namespace = "derivation_fetch"
+    val joinConf = generateMutationData(namespace)
+    val derivations = Seq(Builders.Derivation(name = "*", expression = "*"),
+      Builders.Derivation(name = "unit_test_fetcher_mutations_gb_rating_sum_plus", expression = "unit_test_fetcher_mutations_gb_rating_sum + 1"),
+      Builders.Derivation(name = "listing_id_renamed", expression = "listing_id")
+    )
+    joinConf.setDerivations(derivations.toJava)
+
+    compareTemporalFetch(joinConf, "2021-04-10", namespace, consistencyCheck = false, dropDsOnWrite = true)
+  }
+
+  def testTemporalFetchJoinDerivationRenameOnly(): Unit = {
+    val namespace = "derivation_fetch_rename_only"
+    val joinConf = generateMutationData(namespace)
+    val derivations = Seq(Builders.Derivation(name = "*", expression = "*"),
+      Builders.Derivation(name = "listing_id_renamed", expression = "listing_id")
+    )
+    joinConf.setDerivations(derivations.toJava)
+
+    compareTemporalFetch(joinConf, "2021-04-10", namespace, consistencyCheck = false, dropDsOnWrite = true)
+  }
+
+
   def testTemporalFetchJoinGenerated(): Unit = {
     val namespace = "generated_fetch"
     val joinConf = generateRandomData(namespace)