test fixes

Author: nikhil-zlai

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

@@ -4,16 +4,6 @@ import ai.chronon.api.thrift.protocol.{TBinaryProtocol, TCompactProtocol}
 import ai.chronon.api.thrift.{TDeserializer, TSerializer}
 
 object SerdeUtils {
-  @transient
-  lazy val binarySerializer: ThreadLocal[TSerializer] = new ThreadLocal[TSerializer] {
-    override def initialValue(): TSerializer = new TSerializer(new TBinaryProtocol.Factory())
-  }
-
-  @transient
-  lazy val binaryDeserializer: ThreadLocal[TDeserializer] = new ThreadLocal[TDeserializer] {
-    override def initialValue(): TDeserializer = new TDeserializer(new TBinaryProtocol.Factory())
-  }
-
   @transient
   lazy val compactSerializer: ThreadLocal[TSerializer] = new ThreadLocal[TSerializer] {
     override def initialValue(): TSerializer = new TSerializer(new TCompactProtocol.Factory())

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

@@ -45,6 +45,7 @@ class TTLCache[I, O](f: I => O,
 
   case class Entry(value: O, updatedAtMillis: Long, var markedForUpdate: AtomicBoolean = new AtomicBoolean(false))
   @transient implicit lazy val logger: Logger = LoggerFactory.getLogger(getClass)
+
   private val updateWhenNull =
     new function.BiFunction[I, Entry, Entry] {
       override def apply(t: I, u: Entry): Entry = {

online/src/main/scala/ai/chronon/online/fetcher/GroupByResponseHandler.scala

@@ -186,7 +186,7 @@ class GroupByResponseHandler(fetchContext: FetchContext, metadataStore: Metadata
             val groupByFlag: Option[Boolean] = Option(fetchContext.flagStore)
               .map(_.isSet(
                 "disable_streaming_decoding_error_throws",
-                Map("groupby_streaming_dataset" -> servingInfo.groupByServingInfo.groupBy.getMetaData.getName).toJava))
+                Map("group_by_streaming_dataset" -> servingInfo.groupByServingInfo.groupBy.getMetaData.getName).toJava))
             if (groupByFlag.getOrElse(fetchContext.disableErrorThrows)) {
               Array.empty[TiledIr]
             } else {

online/src/main/scala/ai/chronon/online/fetcher/JoinPartFetcher.scala

@@ -31,7 +31,7 @@ class JoinPartFetcher(fetchContext: FetchContext, metadataStore: MetadataStore)
 
   @transient implicit lazy val logger: Logger = LoggerFactory.getLogger(getClass)
 
-  private val groupByFetcher = new GroupByFetcher(fetchContext, metadataStore)
+  private[online] val groupByFetcher = new GroupByFetcher(fetchContext, metadataStore)
   private implicit val executionContext: ExecutionContext = fetchContext.getOrCreateExecutionContext
 
   def fetchGroupBys(requests: Seq[Request]): Future[Seq[Response]] = {

online/src/main/scala/ai/chronon/online/stats/DriftStore.scala

@@ -8,6 +8,7 @@ import ai.chronon.observability._
 import ai.chronon.online.KVStore
 import ai.chronon.online.KVStore.GetRequest
 import ai.chronon.online.fetcher.{FetchContext, MetadataStore}
+import org.slf4j.LoggerFactory
 
 import scala.collection.Seq
 import scala.concurrent.{ExecutionContext, Future}
@@ -21,6 +22,8 @@ class DriftStore(kvStore: KVStore,
   private val metadataStore = new MetadataStore(fetchContext)
   implicit private val executionContext: ExecutionContext = fetchContext.getOrCreateExecutionContext
 
+  @transient private lazy val logger = LoggerFactory.getLogger(this.getClass)
+
   def tileKeysForJoin(join: api.Join,
                       slice: Option[String] = None,
                       columnNamePrefix: Option[String] = None): Map[String, Array[TileKey]] = {
@@ -138,7 +141,9 @@ class DriftStore(kvStore: KVStore,
         _ match {
           case Success(responseContext) => Some(responseContext)
           // TODO instrument failures
-          case Failure(exception) => exception.printStackTrace(); None
+          case Failure(exception) =>
+            logger.error("Failed to fetch summary response", exception)
+            None
         }
       }
 

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

@@ -26,7 +26,7 @@ import ai.chronon.online.fetcher.Fetcher.Request
 import ai.chronon.online.fetcher.Fetcher.Response
 import ai.chronon.online.fetcher.FetcherCache.BatchResponses
 import ai.chronon.online.KVStore.TimedValue
-import ai.chronon.online.fetcher.{FetchContext, MetadataStore}
+import ai.chronon.online.fetcher.{FetchContext, GroupByFetcher, MetadataStore}
 import ai.chronon.online.{fetcher, _}
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
@@ -57,7 +57,8 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
   val HostKey = "host"
   val GuestId: AnyRef = 123.asInstanceOf[AnyRef]
   val HostId = "456"
-  var fetcherBase: fetcher.JoinPartFetcher = _
+  var joinPartFetcher: fetcher.JoinPartFetcher = _
+  var groupByFetcher: fetcher.GroupByFetcher = _
   var kvStore: KVStore = _
   var fetchContext: FetchContext = _
   var metadataStore: MetadataStore = _
@@ -69,32 +70,32 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
     // the mock to prevent hanging.
     when(kvStore.executionContext).thenReturn(ExecutionContext.global)
     fetchContext = FetchContext(kvStore)
-    metadataStore = new MetadataStore(fetchContext)
-    fetcherBase = spy[fetcher.JoinPartFetcher](new fetcher.JoinPartFetcher(fetchContext, metadataStore))
+    metadataStore = spy[fetcher.MetadataStore](new MetadataStore(fetchContext))
+    joinPartFetcher = spy[fetcher.JoinPartFetcher](new fetcher.JoinPartFetcher(fetchContext, metadataStore))
+    groupByFetcher = spy[fetcher.GroupByFetcher](new GroupByFetcher(fetchContext, metadataStore))
   }
 
   it should "fetch columns single query" in {
     // Fetch a single query
     val keyMap = Map(GuestKey -> GuestId)
     val query = ColumnSpec(GroupBy, Column, None, Some(keyMap))
-
     doAnswer(new Answer[Future[Seq[fetcher.Fetcher.Response]]] {
       def answer(invocation: InvocationOnMock): Future[Seq[Response]] = {
         val requests = invocation.getArgument(0).asInstanceOf[Seq[Request]]
         val request = requests.head
         val response = Response(request, Success(Map(request.name -> "100")))
         Future.successful(Seq(response))
       }
-    }).when(fetcherBase).fetchGroupBys(any())
+    }).when(groupByFetcher).fetchGroupBys(any())
 
     // Map should contain query with valid response
-    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(query)), 1.second)
+    val queryResults = Await.result(groupByFetcher.fetchColumns(Seq(query)), 1.second)
     queryResults.contains(query) shouldBe true
     queryResults.get(query).map(_.values) shouldBe Some(Success(Map(s"$GroupBy.$Column" -> "100")))
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(groupByFetcher, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequest = requestsCaptor.getValue.asInstanceOf[Seq[Request]].headOption
     actualRequest shouldNot be(None)
     actualRequest.get.name shouldBe s"${query.groupByName}.${query.columnName}"
@@ -114,18 +115,18 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
         val responses = requests.map(r => Response(r, Success(Map(r.name -> "100"))))
         Future.successful(responses)
       }
-    }).when(fetcherBase).fetchGroupBys(any())
+    }).when(groupByFetcher).fetchGroupBys(any())
 
     // Map should contain query with valid response
-    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(guestQuery, hostQuery)), 1.second)
+    val queryResults = Await.result(groupByFetcher.fetchColumns(Seq(guestQuery, hostQuery)), 1.second)
     queryResults.contains(guestQuery) shouldBe true
     queryResults.get(guestQuery).map(_.values) shouldBe Some(Success(Map(s"${GuestKey}_$GroupBy.$Column" -> "100")))
     queryResults.contains(hostQuery) shouldBe true
     queryResults.get(hostQuery).map(_.values) shouldBe Some(Success(Map(s"${HostKey}_$GroupBy.$Column" -> "100")))
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(groupByFetcher, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequests = requestsCaptor.getValue.asInstanceOf[Seq[Request]]
     actualRequests.length shouldBe 2
     actualRequests.head.name shouldBe s"${guestQuery.groupByName}.${guestQuery.columnName}"
@@ -143,10 +144,10 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
       def answer(invocation: InvocationOnMock): Future[Seq[Response]] = {
         Future.successful(Seq())
       }
-    }).when(fetcherBase).fetchGroupBys(any())
+    }).when(groupByFetcher).fetchGroupBys(any())
 
     // Map should contain query with Failure response
-    val queryResults = Await.result(fetcherBase.fetchColumns(Seq(query)), 1.second)
+    val queryResults = Await.result(groupByFetcher.fetchColumns(Seq(query)), 1.second)
     queryResults.contains(query) shouldBe true
     queryResults.get(query).map(_.values) match {
       case Some(Failure(_: IllegalStateException)) => succeed
@@ -155,7 +156,7 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
 
     // GroupBy request sent to KV store for the query
     val requestsCaptor = ArgumentCaptor.forClass(classOf[Seq[_]])
-    verify(fetcherBase, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
+    verify(groupByFetcher, times(1)).fetchGroupBys(requestsCaptor.capture().asInstanceOf[Seq[Request]])
     val actualRequest = requestsCaptor.getValue.asInstanceOf[Seq[Request]].headOption
     actualRequest shouldNot be(None)
     actualRequest.get.name shouldBe query.groupByName + "." + query.columnName
@@ -166,16 +167,16 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
   it should "get serving info should call update serving info if batch response is from kv store" in {
     val oldServingInfo = mock[GroupByServingInfoParsed]
     val updatedServingInfo = mock[GroupByServingInfoParsed]
-    doReturn(updatedServingInfo).when(fetcherBase).getServingInfo(any(), any())
+    doReturn(updatedServingInfo).when(joinPartFetcher).getServingInfo(any(), any())
 
     val batchTimedValuesSuccess = Success(Seq(TimedValue(Array(1.toByte), 2000L)))
     val kvStoreBatchResponses = BatchResponses(batchTimedValuesSuccess)
 
-    val result = fetcherBase.getServingInfo(oldServingInfo, kvStoreBatchResponses)
+    val result = joinPartFetcher.getServingInfo(oldServingInfo, kvStoreBatchResponses)
 
     // updateServingInfo is called
     result shouldEqual updatedServingInfo
-    verify(fetcherBase).getServingInfo(any(), any())
+    verify(joinPartFetcher).getServingInfo(any(), any())
   }
 
   // If a batch response is cached, the serving info should be refreshed. This is needed to prevent
@@ -194,23 +195,23 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
     doReturn(groupByOpsMock).when(oldServingInfo).groupByOps
 
     val cachedBatchResponses = BatchResponses(mock[FinalBatchIr])
-    val result = fetcherBase.getServingInfo(oldServingInfo, cachedBatchResponses)
+    val result = groupByFetcher.getServingInfo(oldServingInfo, cachedBatchResponses)
 
     // FetcherBase.updateServingInfo is not called, but getGroupByServingInfo.refresh() is.
     result shouldEqual oldServingInfo
     verify(ttlCache).refresh(any())
-    verify(fetcherBase, never()).getServingInfo(any(), any())
+    verify(ttlCache, never()).apply(any())
   }
 
-  it should "is caching enabled correctly determine if cache is enabled" in {
+  it should "determine if caching is enabled correctly" in {
     val flagStore: FlagStore = (flagName: String, attributes: java.util.Map[String, String]) => {
       flagName match {
         case "enable_fetcher_batch_ir_cache" =>
-          attributes.get("groupby_streaming_dataset") match {
+          attributes.get("group_by_streaming_dataset") match {
             case "test_groupby_2" => false
             case "test_groupby_3" => true
             case other @ _ =>
-              fail(s"Unexpected groupby_streaming_dataset: $other")
+              fail(s"Unexpected group_by_streaming_dataset: $other")
               false
           }
         case _ => false
@@ -226,10 +227,7 @@ class FetcherBaseTest extends AnyFlatSpec with MockitoSugar with Matchers with M
       spy[fetcher.JoinPartFetcher](new fetcher.JoinPartFetcher(fetchContext, new MetadataStore(fetchContext)))
     when(fetcherBaseWithFlagStore.isCacheSizeConfigured).thenReturn(true)
 
-    def buildGroupByWithCustomJson(name: String): GroupBy = Builders.GroupBy(metaData = Builders.MetaData(name = name))
-
     // no name set
-    assertFalse(fetchContext.isCachingEnabled(null))
     assertFalse(fetchContext.isCachingEnabled("test_groupby_2"))
     assertTrue(fetchContext.isCachingEnabled("test_groupby_3"))
   }

spark/src/main/scala/ai/chronon/spark/scripts/DataServer.scala

@@ -37,7 +37,7 @@ class DataServer(driftSeries: Seq[TileDriftSeries], summarySeries: Seq[TileSumma
     private def convertToBytesMap[T <: TBase[_, _]: Manifest: ClassTag](
         series: T,
         keyF: T => TileSeriesKey): Map[String, String] = {
-      val serializerInstance = SerdeUtils.binarySerializer.get()
+      val serializerInstance = SerdeUtils.compactSerializer.get()
       val encoder = Base64.getEncoder
       val keyBytes = serializerInstance.serialize(keyF(series))
       val valueBytes = serializerInstance.serialize(series)

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

@@ -3,7 +3,7 @@ package ai.chronon.spark.stats.drift
 import ai.chronon.api.ColorPrinter.ColorString
 import ai.chronon.api.Extensions._
 import ai.chronon.api.ScalaJavaConversions._
-import ai.chronon.api.SerdeUtils.binarySerializer
+import ai.chronon.api.SerdeUtils.compactSerializer
 import ai.chronon.api._
 import ai.chronon.observability.Cardinality
 import ai.chronon.observability.TileKey
@@ -325,7 +325,7 @@ class SummaryPacker(confPath: String,
 
     val packedRdd: RDD[sql.Row] = df.rdd.flatMap(func).map { tileRow =>
       // pack into bytes
-      val serializer = binarySerializer.get()
+      val serializer = compactSerializer.get()
 
       val partition = tileRow.partition
       val timestamp = tileRow.tileTs

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

@@ -48,7 +48,7 @@ class DriftTest extends AnyFlatSpec with Matchers {
       logger.info(s"    ${pad(f.name)} : ${f.dataType.typeName}".yellow)
     }
 
-    df.show(10, truncate = false)
+    df.show(10)
   }
 
   "end_to_end" should "fetch prepare anomalous data, summarize, upload and fetch without failures" in {
@@ -57,7 +57,7 @@ class DriftTest extends AnyFlatSpec with Matchers {
     val prepareData = PrepareData(namespace)
     val join = prepareData.generateAnomalousFraudJoin
     val df = prepareData.generateFraudSampleData(600000, "2023-01-01", "2023-02-30", join.metaData.loggedTable)
-    df.show(10, truncate = false)
+    df.show(10)
 
     // mock api impl for online fetching and uploading
     val kvStoreFunc: () => KVStore = () => {