observability script for demo

docker-init/demo/Dockerfile

@@ -0,0 +1,33 @@
+FROM apache/spark:latest
+
+# 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 | 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
+
+# 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,22 @@
+# Stop and remove existing container
+docker stop spark-app
+docker rm spark-app
+
+# Run new container
+docker run -d \
+  --name spark-app \
+  -v $HOME/repos/chronon/spark/target/scala-2.12:/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,18 +17,27 @@ if [[ ! -f $SPARK_JAR ]] || [[ ! -f $CLOUD_AWS_JAR ]]; then
     exit 1
 fi
 
+# Load up summary data into DynamoDB
+echo "Loading Summary.."
+#check if log file exists
+if [[ ! -f log4j.properties ]]; then
+  echo "Error: log4j.properties file not found" >&2
+  exit 1
+else
+  echo "log4j.properties file found!"
+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
+if ! java -Dlog4j.configurationFile=log4j.properties -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
   echo "Error: Failed to load metadata into DynamoDB" >&2
   exit 1
 fi
 echo "Metadata load completed successfully!"
 
 # Initialize DynamoDB
 echo "Initializing DynamoDB Table .."
-if ! output=$(java -cp $SPARK_JAR:$CLASSPATH ai.chronon.spark.Driver create-summary-dataset \
+if ! output=$(java -Dlog4j.configurationFile=log4j.properties -cp $SPARK_JAR:$CLASSPATH ai.chronon.spark.Driver create-summary-dataset \
   --online-jar=$CLOUD_AWS_JAR \
   --online-class=$ONLINE_CLASS 2>&1); then
   echo "Error: Failed to bring up DynamoDB table" >&2
@@ -32,9 +47,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 -Dlog4j.configurationFile=log4j.properties \
+  --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 +60,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/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
@@ -327,9 +327,9 @@ case class TableUtils(sparkSession: SparkSession) {
         sql(creationSql)
       } catch {
         case _: TableAlreadyExistsException =>
-          logger.info(s"Table $tableName already exists, skipping creation")
+          println(s"Table $tableName already exists, skipping creation")
         case e: Exception =>
-          logger.error(s"Failed to create table $tableName", e)
+          println(s"Failed to create table $tableName", e)
           throw e
       }
     }
@@ -357,6 +357,7 @@ case class TableUtils(sparkSession: SparkSession) {
       // so that an exception will be thrown below
       dfRearranged
     }
+    println(s"Repartitioning and writing into table $tableName".yellow)
     repartitionAndWrite(finalizedDf, tableName, saveMode, stats, sortByCols)
   }
 

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

@@ -0,0 +1,154 @@
+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 java.util.concurrent.TimeUnit
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
+import scala.util.ScalaJavaConversions.IteratorOps
+
+object ObservabilityDemo {
+
+  def Time(message: String)(block: => Unit): Unit = {
+    println(s"$message..".yellow)
+    val start = System.currentTimeMillis()
+    block
+    val end = System.currentTimeMillis()
+    println(s"$message took ${end - start} ms".green)
+  }
+
+  def main(args: Array[String]): Unit = {
+
+    val startDs = "2023-01-01"
+    val endDs = "2023-02-30"
+    val rowCount = 700000
+    val namespace = "observability_demo"
+    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) => println(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))
+    println(summaries)
+
+    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)
+    }
+
+    println(s"""drift totals: $totals
+           |drift nulls: $nulls
+           |""".stripMargin.red)
+
+    println("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)
+
+    println("Summary series fetched successfully".green)
+  }
+
+}