chore: refactor additional stat tracking (#69222)

Author: Jonathan Pearlin

airbyte-cdk/bulk/changelog.md

@@ -1,3 +1,9 @@
+## Version 0.1.68
+
+**Load CDK**
+
+* Refactor additional state statistic tracking.
+
 ## Version 0.1.67
 
 **Load CDK**

airbyte-cdk/bulk/core/load/src/integrationTest/kotlin/io/airbyte/cdk/load/mock_integration_test/MockBasicFunctionalityIntegrationTest.kt

@@ -80,7 +80,6 @@ abstract class BaseMockBasicFunctionalityIntegrationTest(
             } else {
                 UnknownTypesBehavior.PASS_THROUGH
             },
-        includesAdditionalStats = false,
     ) {
     @Test
     override fun testBasicWrite() {

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/aggregate/AggregateStore.kt

@@ -29,8 +29,8 @@ class AggregateStore(
         val (_, agg, counts, bytes, timeTrigger, countTrigger, bytesTrigger) = getOrCreate(key)
 
         agg.accept(record)
-        counts.increment(record.partitionKey, 1)
-        bytes.increment(record.partitionKey, record.sizeBytes)
+        counts.increment(record.partitionKey, 1.0)
+        bytes.increment(record.partitionKey, record.sizeBytes.toDouble())
         countTrigger.increment(1)
         bytesTrigger.increment(record.sizeBytes)
         timeTrigger.update(record.emittedAtMs)

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/stages/ParseStage.kt

@@ -4,22 +4,28 @@
 
 package io.airbyte.cdk.load.dataflow.stages
 
+import io.airbyte.cdk.load.data.AirbyteValue
 import io.airbyte.cdk.load.dataflow.pipeline.DataFlowStage
 import io.airbyte.cdk.load.dataflow.pipeline.DataFlowStageIO
+import io.airbyte.cdk.load.dataflow.state.PartitionKey
 import io.airbyte.cdk.load.dataflow.transform.RecordDTO
-import io.airbyte.cdk.load.dataflow.transform.RecordMunger
+import io.airbyte.cdk.load.dataflow.transform.medium.ConversionInput
+import io.airbyte.cdk.load.dataflow.transform.medium.JsonConverter
+import io.airbyte.cdk.load.dataflow.transform.medium.ProtobufConverter
+import io.airbyte.cdk.load.message.DestinationRecordProtobufSource
+import io.airbyte.cdk.load.message.DestinationRecordRaw
 import jakarta.inject.Named
 import jakarta.inject.Singleton
 
 @Named("parse")
 @Singleton
 class ParseStage(
-    private val munger: RecordMunger,
+    private val jsonConverter: JsonConverter,
+    private val protobufConverter: ProtobufConverter
 ) : DataFlowStage {
     override suspend fun apply(input: DataFlowStageIO): DataFlowStageIO {
         val raw = input.raw!!
-        val fields = munger.transformForDest(raw)
-
+        val fields = transform(raw, input.partitionKey!!)
         return input.apply {
             munged =
                 RecordDTO(
@@ -30,4 +36,15 @@ class ParseStage(
                 )
         }
     }
+
+    private fun transform(
+        msg: DestinationRecordRaw,
+        partitionKey: PartitionKey
+    ): Map<String, AirbyteValue> {
+        val input = ConversionInput(msg = msg, partitionKey = partitionKey)
+        return when (msg.rawData) {
+            is DestinationRecordProtobufSource -> protobufConverter.convert(input)
+            else -> jsonConverter.convert(input)
+        }
+    }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/StateHistogram.kt

@@ -22,18 +22,22 @@ data class PartitionKey(
     val id: String,
 )
 
-open class Histogram<T>(private val map: ConcurrentMap<T, Long> = ConcurrentHashMap()) {
-    fun increment(key: T, quantity: Long): Histogram<T> =
-        this.apply { map.merge(key, quantity, Long::plus) }
+open class Histogram<T>(private val map: ConcurrentMap<T, Double> = ConcurrentHashMap()) {
+    fun increment(key: T, quantity: Double): Histogram<T> =
+        this.apply { map.merge(key, quantity, Double::plus) }
 
     fun merge(other: Histogram<T>): Histogram<T> =
-        this.apply { other.map.forEach { map.merge(it.key, it.value, Long::plus) } }
+        this.apply { other.map.forEach { map.merge(it.key, it.value, Double::plus) } }
 
-    fun get(key: T): Long? = map[key]
+    fun get(key: T): Double? = map[key]
 
-    fun remove(key: T): Long? = map.remove(key)
+    fun remove(key: T): Double? = map.remove(key)
+
+    fun toMap(): Map<T, Double> = map.toMap()
 }
 
+typealias AdditionalStatsHistogram = Histogram<String>
+
 typealias StateHistogram = Histogram<StateKey>
 
 typealias PartitionHistogram = Histogram<PartitionKey>

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/StateHistogramStore.kt

@@ -19,29 +19,29 @@ class StateHistogramStore {
     }
 
     fun acceptExpectedCounts(key: StateKey, count: Long): StateHistogram {
-        val inner = ConcurrentHashMap<StateKey, Long>()
-        inner[key] = count
+        val inner = ConcurrentHashMap<StateKey, Double>()
+        inner[key] = count.toDouble()
 
         return expected.merge(StateHistogram(inner))
     }
 
     fun isComplete(key: StateKey): Boolean {
         val expectedCount = expected.get(key)
-        val flushedCount = key.partitionKeys.sumOf { flushed.get(it) ?: 0 }
+        val flushedCount = key.partitionKeys.sumOf { flushed.get(it) ?: 0.0 }
 
         return expectedCount == flushedCount
     }
 
     // mirrors isComplete. Purely for debugging purposes.
     fun whyIsStateIncomplete(key: StateKey): String {
         val expectedCount = expected.get(key)
-        val partitionFlushCounts = key.partitionKeys.map { flushed.get(it) ?: 0 }
+        val partitionFlushCounts = key.partitionKeys.map { flushed.get(it) ?: 0.0 }
         val flushedCount = partitionFlushCounts.sum()
         return "expectedCount $expectedCount does not equal flushedCount $flushedCount (by partition: $partitionFlushCounts)"
     }
 
     fun remove(key: StateKey): Long? {
         key.partitionKeys.forEach { flushed.remove(it) }
-        return expected.remove(key)
+        return expected.remove(key)?.toLong()
     }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/stats/CommitedStatsStore.kt

@@ -51,8 +51,8 @@ class CommittedStatsStore {
     internal fun removeLiveStats(s: DestinationStream.Descriptor, p: PartitionKey): EmissionStats? =
         liveStats[s]?.let {
             EmissionStats(
-                count = it.counts.remove(p) ?: 0,
-                bytes = it.bytes.remove(p) ?: 0,
+                count = it.counts.remove(p)?.toLong() ?: 0,
+                bytes = it.bytes.remove(p)?.toLong() ?: 0,
             )
         }
 }

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/stats/StateAdditionalStatsStore.kt

@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.dataflow.state.stats
+
+import io.airbyte.cdk.load.command.DestinationStream
+import io.airbyte.cdk.load.dataflow.state.AdditionalStatsHistogram
+import io.airbyte.cdk.load.dataflow.state.Histogram
+import io.airbyte.cdk.load.dataflow.state.PartitionKey
+import jakarta.inject.Singleton
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * A singleton class responsible for managing additional state statistics, specifically histograms
+ * of metrics associated with partitions and destination streams. This class provides methods to add
+ * metrics values, retrieve combined statistics, and maintain metrics completeness with default
+ * values.
+ */
+@Singleton
+class StateAdditionalStatsStore {
+
+    private val store =
+        ConcurrentHashMap<
+            PartitionKey, ConcurrentHashMap<DestinationStream.Descriptor, AdditionalStatsHistogram>
+        >()
+
+    /**
+     * Adds a value to the histogram associated with the given partition and stream descriptor for
+     * the specified metric.
+     *
+     * @param partitionKey The key representing the partition for which the histogram is maintained.
+     * @param streamDescriptor The descriptor identifying the specific destination stream.
+     * @param metric The metric whose value is to be incremented in the histogram.
+     * @param value The value to increment for the specified metric.
+     */
+    fun add(
+        partitionKey: PartitionKey,
+        streamDescriptor: DestinationStream.Descriptor,
+        metric: ObservabilityMetrics,
+        value: Double
+    ) {
+        store.putIfAbsent(partitionKey, ConcurrentHashMap())
+        store[partitionKey]?.let { partitionStats ->
+            partitionStats.putIfAbsent(streamDescriptor, Histogram())
+            partitionStats[streamDescriptor]?.increment(metric.metricName, value)
+        }
+    }
+
+    /**
+     * Drains the accumulated statistics for the specified destination stream across the provided
+     * partition keys, returning a combined histogram of all relevant metrics. For any metric values
+     * not present, default values are added. This is to ensure that the metrics are always
+     * published as part of the state message with a zero value if no value has been recorded since
+     * the last state message publish event.
+     *
+     * @param partitionKeys The list of partition keys to process, each representing a distinct
+     * partition.
+     * @param stream The descriptor of the destination stream whose statistics are to be drained.
+     * @return A `Histogram` containing the combined metrics for the specified destination stream.
+     */
+    fun drain(
+        partitionKeys: List<PartitionKey>,
+    ): Map<DestinationStream.Descriptor, AdditionalStatsHistogram> {
+        val accumulator =
+            mutableMapOf<DestinationStream.Descriptor, AdditionalStatsHistogram>().withDefault {
+                populateWithDefaultValues(AdditionalStatsHistogram())
+            }
+
+        return partitionKeys
+            .mapNotNull { store.remove(it) }
+            .fold(accumulator) { acc, perStreamStats ->
+                perStreamStats.forEach { (stream, histogram) ->
+                    acc.merge(stream, histogram, AdditionalStatsHistogram::merge)
+                }
+                addDefaultValues(acc)
+            }
+    }
+
+    /**
+     * Adds default values for metrics that are not present in the given histogram. This ensures
+     * that all metrics defined in `ObservabilityMetrics` are represented in the histogram with a
+     * default value of 0.0 if they are missing.
+     *
+     * @param histogram The histogram to which default values should be added,
+     * ```
+     *                  representing metric counts by their names.
+     * @return
+     * ```
+     * The updated histogram with default values for missing metrics.
+     */
+    private fun addDefaultValues(
+        stats: MutableMap<DestinationStream.Descriptor, AdditionalStatsHistogram>
+    ): MutableMap<DestinationStream.Descriptor, AdditionalStatsHistogram> {
+        stats.forEach { (_, histogram) -> populateWithDefaultValues(histogram) }
+        return stats
+    }
+
+    private fun populateWithDefaultValues(
+        histogram: AdditionalStatsHistogram
+    ): AdditionalStatsHistogram {
+        println("Populating...")
+        ObservabilityMetrics.entries.forEach {
+            if (histogram.get(it.metricName) == null) {
+                histogram.increment(it.metricName, 0.0)
+            }
+        }
+        return histogram
+    }
+
+    /**
+     * Enum representing the available observability metrics. Each metric is associated with a
+     * specific metric name used for tracking system behavior and performance.
+     *
+     * @property metricName The name of the metric used in tracking.
+     */
+    enum class ObservabilityMetrics(val metricName: String) {
+        NULLED_VALUE_COUNT("nulledValueCount"),
+        TRUNCATED_VALUE_COUNT("truncatedValueCount")
+    }
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/stats/StateStatsEnricher.kt

@@ -7,8 +7,8 @@ package io.airbyte.cdk.load.dataflow.state.stats
 import com.google.common.annotations.VisibleForTesting
 import io.airbyte.cdk.load.command.DestinationStream
 import io.airbyte.cdk.load.command.NamespaceMapper
+import io.airbyte.cdk.load.dataflow.state.PartitionKey
 import io.airbyte.cdk.load.dataflow.state.StateKey
-import io.airbyte.cdk.load.dataflow.stats.MetricTracker
 import io.airbyte.cdk.load.message.CheckpointMessage
 import io.airbyte.cdk.load.message.GlobalCheckpoint
 import io.airbyte.cdk.load.message.GlobalSnapshotCheckpoint
@@ -20,7 +20,7 @@ import jakarta.inject.Singleton
 class StateStatsEnricher(
     private val statsStore: CommittedStatsStore,
     private val namespaceMapper: NamespaceMapper,
-    private val metricTracker: MetricTracker,
+    private val stateAdditionalStatsStore: StateAdditionalStatsStore,
 ) {
     // Enriches provided state message with stats associated with the given state key.
     fun enrich(msg: CheckpointMessage, key: StateKey): CheckpointMessage {
@@ -36,13 +36,15 @@ class StateStatsEnricher(
     fun enrichTopLevelDestinationStats(
         msg: CheckpointMessage,
         desc: DestinationStream.Descriptor,
+        partitionKeys: List<PartitionKey>,
         count: Long
     ): CheckpointMessage {
         // TODO: set this using the count above once we get to total rejected
         // records.
         msg.updateStats(
             destinationStats = msg.sourceStats,
-            additionalStats = metricTracker.drain(desc)
+            // Use getValue to ensure the default histogram is returned if the stream is no present
+            additionalStats = stateAdditionalStatsStore.drain(partitionKeys).getValue(desc).toMap(),
         )
 
         return msg
@@ -83,7 +85,7 @@ class StateStatsEnricher(
             )
         val (committed, cumulative) = statsStore.commitStats(desc, key)
 
-        enrichTopLevelDestinationStats(msg, desc, committed.count)
+        enrichTopLevelDestinationStats(msg, desc, key.partitionKeys, committed.count)
         enrichTopLevelStats(msg, cumulative)
 
         return msg

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/dataflow/state/stats/socket/EmittedStatsStoreImpl.kt

@@ -32,8 +32,8 @@ class EmittedStatsStoreImpl(
         count: Long,
         bytes: Long,
     ) {
-        readCounts.increment(s, count)
-        readBytes.increment(s, bytes)
+        readCounts.increment(s, count.toDouble())
+        readBytes.increment(s, bytes.toDouble())
     }
 
     override fun getStats(): List<AirbyteMessage> =
@@ -45,8 +45,8 @@ class EmittedStatsStoreImpl(
     @VisibleForTesting
     internal fun get(s: DestinationStream.Descriptor) =
         EmissionStats(
-            count = readCounts.get(s) ?: 0,
-            bytes = readBytes.get(s) ?: 0,
+            count = readCounts.get(s)?.toLong() ?: 0,
+            bytes = readBytes.get(s)?.toLong() ?: 0,
         )
 
     @VisibleForTesting