observability script for demo

docker-init/demo/Dockerfile

@@ -0,0 +1,35 @@
+FROM apache/spark:3.5.3-scala2.12-java17-ubuntu
+
+# Switch to root to install Java 17
+USER root
+
+# Install Amazon Corretto 17
+RUN apt-get update && \
+    apt-get install -y wget software-properties-common gnupg2 && \
+    wget -O- https://apt.corretto.aws/corretto.key --https-only | apt-key add - && \
+    add-apt-repository 'deb https://apt.corretto.aws stable main' && \
+    apt-get update && \
+    apt-get install -y java-17-amazon-corretto-jdk && \
+    update-alternatives --set java /usr/lib/jvm/java-17-amazon-corretto/bin/java && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Create directory and set appropriate permissions
+RUN mkdir -p /opt/chronon/jars && \
+    chown -R 185:185 /opt/chronon && \
+    chmod 755 /opt/chronon/jars
+
+# Set JAVA_HOME
+ENV JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto
+ENV PATH=$PATH:$JAVA_HOME/bin
+
+# Switch back to spark user
+USER 185
+
+# Set environment variables for Spark classpath
+ENV SPARK_CLASSPATH="/opt/spark/jars/*"
+ENV SPARK_DIST_CLASSPATH="/opt/spark/jars/*"
+ENV SPARK_EXTRA_CLASSPATH="/opt/spark/jars/*:/opt/chronon/jars/*"
+ENV HADOOP_CLASSPATH="/opt/spark/jars/*"
+
+CMD ["tail", "-f", "/dev/null"]

docker-init/demo/README.md

@@ -0,0 +1,2 @@
+run build.sh once, and you can repeatedly exec 
+sbt spark/assembly + run.sh on iterations to the chronon code.

docker-init/demo/build.sh

@@ -0,0 +1 @@
+docker build -t obs .

docker-init/demo/run.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# Stop and remove existing container
+if docker ps -a | grep -q spark-app; then
+  docker stop spark-app || echo "Failed to stop container"
+  docker rm spark-app || echo "Failed to remove container"
+fi
+
+CHRONON_JAR_PATH="${CHRONON_JAR_PATH:-$HOME/repos/chronon/spark/target/scala-2.12}"
+
+if [ ! -d "$CHRONON_JAR_PATH" ]; then
+  echo "Error: JAR directory not found: $CHRONON_JAR_PATH"
+  exit 1
+fi
+
+# Run new container
+docker run -d \
+  --name spark-app \
+  -v "$CHRONON_JAR_PATH":/opt/chronon/jars \
+  obs
+
+# Submit with increased memory
+docker exec spark-app \
+  /opt/spark/bin/spark-submit \
+  --master "local[*]" \
+  --driver-memory 8g \
+  --conf "spark.driver.maxResultSize=6g" \
+  --conf "spark.driver.memory=8g" \
+  --driver-class-path "/opt/spark/jars/*:/opt/chronon/jars/*" \
+  --conf "spark.driver.host=localhost" \
+  --conf "spark.driver.bindAddress=0.0.0.0" \
+  --class ai.chronon.spark.scripts.ObservabilityDemo \
+  /opt/chronon/jars/spark-assembly-0.1.0-SNAPSHOT.jar

docker-init/generate_anomalous_data.py

@@ -5,7 +5,7 @@
 from pyspark.sql.types import StructType, StructField, DoubleType, IntegerType, StringType, TimestampType, BooleanType
 
 # Initialize Spark session
-spark = SparkSession.builder.appName("FraudClassificationSchema").getOrCreate()
+spark = SparkSession.builder.appName("FraudClassificationSchema").config("spark.log.level", "WARN").getOrCreate()
 
 def time_to_value(t, base_value, amplitude, noise_level, scale=1):
     if scale is None:

docker-init/start.sh

@@ -1,7 +1,13 @@
 #!/bin/bash
+
+start_time=$(date +%s)
 if ! python3.8 generate_anomalous_data.py; then
     echo "Error: Failed to generate anomalous data" >&2
     exit 1
+else
+    end_time=$(date +%s)
+    elapsed_time=$((end_time - start_time))
+    echo "Anomalous data generated successfully! Took $elapsed_time seconds."
 fi
 
 
@@ -11,7 +17,6 @@ if [[ ! -f $SPARK_JAR ]] || [[ ! -f $CLOUD_AWS_JAR ]]; then
     exit 1
 fi
 
-
 # Load up metadata into DynamoDB
 echo "Loading metadata.."
 if ! java -cp $SPARK_JAR:$CLASSPATH ai.chronon.spark.Driver metadata-upload --conf-path=/chronon_sample/production/ --online-jar=$CLOUD_AWS_JAR --online-class=$ONLINE_CLASS; then
@@ -32,9 +37,11 @@ fi
 echo "DynamoDB Table created successfully!"
 
 
-# Load up summary data into DynamoDB
-echo "Loading Summary.."
-if ! java --add-opens=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/sun.security.action=ALL-UNNAMED \
+start_time=$(date +%s)
+
+if ! java \
+  --add-opens=java.base/sun.nio.ch=ALL-UNNAMED \
+  --add-opens=java.base/sun.security.action=ALL-UNNAMED \
   -cp $SPARK_JAR:$CLASSPATH ai.chronon.spark.Driver summarize-and-upload \
   --online-jar=$CLOUD_AWS_JAR \
   --online-class=$ONLINE_CLASS \
@@ -43,8 +50,11 @@ if ! java --add-opens=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/sun
   --time-column=transaction_time; then
   echo "Error: Failed to load summary data into DynamoDB" >&2
   exit 1
+else
+  end_time=$(date +%s)
+  elapsed_time=$((end_time - start_time))
+  echo "Summary load completed successfully! Took $elapsed_time seconds."
 fi
-echo "Summary load completed successfully!"
 
 # Add these java options as without them we hit the below error:
 # throws java.lang.ClassFormatError accessible: module java.base does not "opens java.lang" to unnamed module @36328710

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

@@ -74,7 +74,7 @@ class DriftStore(kvStore: KVStore,
   def getSummaries(joinConf: api.Join,
                    startMs: Option[Long],
                    endMs: Option[Long],
-                   columnPrefix: Option[String] = None): Future[Seq[TileSummaryInfo]] = {
+                   columnPrefix: Option[String]): Future[Seq[TileSummaryInfo]] = {
 
     val serializer: TSerializer = compactSerializer
     val tileKeyMap = tileKeysForJoin(joinConf, columnPrefix)

spark/src/main/resources/logback.xml

@@ -0,0 +1,12 @@
+<configuration>
+    <appender name="CONSOLE" class="ch.qos.logback.core.ConsoleAppender">
+        <encoder>
+            <pattern>[%date] {%logger{0}} %level - %message%n</pattern>
+        </encoder>
+    </appender>
+    <logger name="ai.chronon" level="INFO" />
+
+    <root level="INFO">
+        <appender-ref ref="CONSOLE"/>
+    </root>
+</configuration>

spark/src/main/scala/ai/chronon/spark/Driver.scala

@@ -147,7 +147,9 @@ object Driver {
 
     protected def buildSparkSession(): SparkSession = {
       if (localTableMapping.nonEmpty) {
-        val localSession = SparkSessionBuilder.build(subcommandName(), local = true, localWarehouseLocation.toOption)
+        val localSession = SparkSessionBuilder.build(subcommandName(),
+                                                     local = true,
+                                                     localWarehouseLocation = localWarehouseLocation.toOption)
         localTableMapping.foreach {
           case (table, filePath) =>
             val file = new File(filePath)

spark/src/main/scala/ai/chronon/spark/SparkSessionBuilder.scala

@@ -34,6 +34,7 @@ object SparkSessionBuilder {
   // we would want to share locally generated warehouse during CI testing
   def build(name: String,
             local: Boolean = false,
+            hiveSupport: Boolean = true,
             localWarehouseLocation: Option[String] = None,
             additionalConfig: Option[Map[String, String]] = None,
             enforceKryoSerializer: Boolean = true): SparkSession = {
@@ -44,7 +45,10 @@ object SparkSessionBuilder {
     var baseBuilder = SparkSession
       .builder()
       .appName(name)
-      .enableHiveSupport()
+
+    if (hiveSupport) baseBuilder = baseBuilder.enableHiveSupport()
+
+    baseBuilder = baseBuilder
       .config("spark.sql.session.timeZone", "UTC")
       //otherwise overwrite will delete ALL partitions, not just the ones it touches
       .config("spark.sql.sources.partitionOverwriteMode", "dynamic")

spark/src/main/scala/ai/chronon/spark/TableUtils.scala

@@ -76,7 +76,7 @@ case class TableUtils(sparkSession: SparkSession) {
     sparkSession.conf.get("spark.chronon.backfill.small_mode_cutoff", "5000").toInt
   val backfillValidationEnforced: Boolean =
     sparkSession.conf.get("spark.chronon.backfill.validation.enabled", "true").toBoolean
-  // Threshold to control whether or not to use bloomfilter on join backfill. If the backfill row approximate count is under this threshold, we will use bloomfilter.
+  // Threshold to control whether to use bloomfilter on join backfill. If the backfill row approximate count is under this threshold, we will use bloomfilter.
   // default threshold is 100K rows
   val bloomFilterThreshold: Long =
     sparkSession.conf.get("spark.chronon.backfill.bloomfilter.threshold", "1000000").toLong

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

@@ -0,0 +1,192 @@
+package ai.chronon.spark.scripts
+
+import ai.chronon
+import ai.chronon.api.ColorPrinter.ColorString
+import ai.chronon.api.Constants
+import ai.chronon.api.DriftMetric
+import ai.chronon.api.Extensions.MetadataOps
+import ai.chronon.api.PartitionSpec
+import ai.chronon.api.TileDriftSeries
+import ai.chronon.api.TileSummarySeries
+import ai.chronon.api.Window
+import ai.chronon.online.KVStore
+import ai.chronon.online.stats.DriftStore
+import ai.chronon.spark.SparkSessionBuilder
+import ai.chronon.spark.TableUtils
+import ai.chronon.spark.stats.drift.Summarizer
+import ai.chronon.spark.stats.drift.SummaryUploader
+import ai.chronon.spark.stats.drift.scripts.PrepareData
+import ai.chronon.spark.utils.InMemoryKvStore
+import ai.chronon.spark.utils.MockApi
+import org.rogach.scallop.ScallopConf
+import org.rogach.scallop.ScallopOption
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+
+import java.util.concurrent.TimeUnit
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
+import scala.util.ScalaJavaConversions.IteratorOps
+
+object ObservabilityDemo {
+  @transient lazy val logger: Logger = LoggerFactory.getLogger(getClass)
+
+  def time(message: String)(block: => Unit): Unit = {
+    logger.info(s"$message..".yellow)
+    val start = System.currentTimeMillis()
+    block
+    val end = System.currentTimeMillis()
+    logger.info(s"$message took ${end - start} ms".green)
+  }
+
+  class Conf(arguments: Seq[String]) extends ScallopConf(arguments) {
+    val startDs: ScallopOption[String] = opt[String](
+      name = "start-ds",
+      default = Some("2023-01-01"),
+      descr = "Start date in YYYY-MM-DD format"
+    )
+
+    val endDs: ScallopOption[String] = opt[String](
+      name = "end-ds",
+      default = Some("2023-02-30"),
+      descr = "End date in YYYY-MM-DD format"
+    )
+
+    val rowCount: ScallopOption[Int] = opt[Int](
+      name = "row-count",
+      default = Some(700000),
+      descr = "Number of rows to generate"
+    )
+
+    val namespace: ScallopOption[String] = opt[String](
+      name = "namespace",
+      default = Some("observability_demo"),
+      descr = "Namespace for the demo"
+    )
+
+    verify()
+  }
+
+  def main(args: Array[String]): Unit = {
+
+    val config = new Conf(args)
+    val startDs = config.startDs()
+    val endDs = config.endDs()
+    val rowCount = config.rowCount()
+    val namespace = config.namespace()
+
+    val spark = SparkSessionBuilder.build(namespace, local = true)
+    implicit val tableUtils: TableUtils = TableUtils(spark)
+    tableUtils.createDatabase(namespace)
+
+    // generate anomalous data (join output)
+    val prepareData = PrepareData(namespace)
+    val join = prepareData.generateAnomalousFraudJoin
+
+    time("Preparing data") {
+      val df = prepareData.generateFraudSampleData(rowCount, startDs, endDs, join.metaData.loggedTable)
+      df.show(10, truncate = false)
+    }
+
+    time("Summarizing data") {
+      // compute summary table and packed table (for uploading)
+      Summarizer.compute(join.metaData, ds = endDs, useLogs = true)
+    }
+
+    val packedTable = join.metaData.packedSummaryTable
+    // mock api impl for online fetching and uploading
+    val kvStoreFunc: () => KVStore = () => {
+      // cannot reuse the variable - or serialization error
+      val result = InMemoryKvStore.build(namespace, () => null)
+      result
+    }
+    val api = new MockApi(kvStoreFunc, namespace)
+
+    // create necessary tables in kvstore
+    val kvStore = api.genKvStore
+    kvStore.create(Constants.MetadataDataset)
+    kvStore.create(Constants.TiledSummaryDataset)
+
+    // upload join conf
+    api.buildFetcher().putJoinConf(join)
+
+    time("Uploading summaries") {
+      val uploader = new SummaryUploader(tableUtils.loadTable(packedTable), api)
+      uploader.run()
+    }
+
+    // test drift store methods
+    val driftStore = new DriftStore(api.genKvStore)
+
+    // TODO: Wire up drift store into hub and create an endpoint
+
+    // fetch keys
+    val tileKeys = driftStore.tileKeysForJoin(join)
+    val tileKeysSimple = tileKeys.mapValues(_.map(_.column).toSeq)
+    tileKeysSimple.foreach { case (k, v) => logger.info(s"$k -> [${v.mkString(", ")}]") }
+
+    // fetch summaries
+    val startMs = PartitionSpec.daily.epochMillis(startDs)
+    val endMs = PartitionSpec.daily.epochMillis(endDs)
+    val summariesFuture = driftStore.getSummaries(join, Some(startMs), Some(endMs), None)
+    val summaries = Await.result(summariesFuture, Duration.create(10, TimeUnit.SECONDS))
+    logger.info(summaries.toString())
+
+    var driftSeries: Seq[TileDriftSeries] = null
+    // fetch drift series
+    time("Fetching drift series") {
+      val driftSeriesFuture = driftStore.getDriftSeries(
+        join.metaData.nameToFilePath,
+        DriftMetric.JENSEN_SHANNON,
+        lookBack = new Window(7, chronon.api.TimeUnit.DAYS),
+        startMs,
+        endMs
+      )
+      driftSeries = Await.result(driftSeriesFuture.get, Duration.create(10, TimeUnit.SECONDS))
+    }
+
+    val (nulls, totals) = driftSeries.iterator.foldLeft(0 -> 0) {
+      case ((nulls, total), s) =>
+        val currentNulls = s.getPercentileDriftSeries.iterator().toScala.count(_ == null)
+        val currentCount = s.getPercentileDriftSeries.size()
+        (nulls + currentNulls, total + currentCount)
+    }
+
+    logger.info(s"""drift totals: $totals
+           |drift nulls: $nulls
+           |""".stripMargin.red)
+
+    logger.info("Drift series fetched successfully".green)
+
+    var summarySeries: Seq[TileSummarySeries] = null
+
+    time("Fetching summary series") {
+      val summarySeriesFuture = driftStore.getSummarySeries(
+        join.metaData.nameToFilePath,
+        startMs,
+        endMs
+      )
+      summarySeries = Await.result(summarySeriesFuture.get, Duration.create(10, TimeUnit.SECONDS))
+    }
+
+    val (summaryNulls, summaryTotals) = summarySeries.iterator.foldLeft(0 -> 0) {
+      case ((nulls, total), s) =>
+        if (s.getPercentiles == null) {
+          (nulls + 1) -> (total + 1)
+        } else {
+          val currentNulls = s.getPercentiles.iterator().toScala.count(_ == null)
+          val currentCount = s.getPercentiles.size()
+          (nulls + currentNulls, total + currentCount)
+        }
+    }
+
+    println(s"""summary ptile totals: $summaryTotals
+           |summary ptile nulls: $summaryNulls
+           |""".stripMargin)
+
+    logger.info("Summary series fetched successfully".green)
+
+    spark.stop()
+    System.exit(0)
+  }
+}