Load CDK: Improved new interface bookkeeping by johnny-schmidt · Pull Request #56346 · airbytehq/airbyte

Johnny Schmidt (johnny-schmidt) · 2025-03-22T23:54:59Z

What

This is all quality-of-life. Before we were just counting COMPLETE and PERSISTED (ackable) records. Now we keep detailed information per stage/partition about exactly what we've counted. Additionally I cleaned up Batch.State (now BatchState) and deprecated all the old stuff.

It also adds some detailed opt-in debug logging that can be turned on when needed. Similar to the detailed logging in the old interface, but more user-friendly. Ie:

Stream test20250322ozzl:testInsertRecords: Records Read: 1000000 (done: true)
ObjectLoaderPartFormatter[0](PROCESSED): 500000 records (done: true, saw eos at input 500000)
ObjectLoaderPartFormatter[1](PROCESSED): 500000 records (done: true, saw eos at input 500000)
ObjectLoaderPartFormatter[TOTAL](PROCESSED): 1000000 records (done: true, saw eos at input 1000000)
ObjectLoaderPartLoader[0](STAGED): 148798 records  (inputs: 6, done: false)
ObjectLoaderPartLoader[1](STAGED): 198378 records  (inputs: 8, done: true, saw eos at input 8)
ObjectLoaderPartLoader[2](STAGED): 202489 records  (inputs: 9, done: true, saw eos at input 9)
ObjectLoaderPartLoader[3](STAGED): 148740 records  (inputs: 6, done: false)
ObjectLoaderPartLoader[4](STAGED): 198378 records  (inputs: 8, done: true, saw eos at input 8)
ObjectLoaderPartLoader[TOTAL](STAGED): 896783 records  (inputs: 37, done: true, saw eos at input 25)
ObjectLoaderUploadCompleter[0](COMPLETE): 4084 records  (inputs: 1, done: false)
ObjectLoaderUploadCompleter[TOTAL](COMPLETE): 4084 records  (inputs: 1, done: false)

I found this critical for debugging an issue exposed by the bulk loader work. It also enables asserting against bad counts / out-of-order processing, which caused a specific issue during testing. (Bug-fix/guards for that impending. The issue doesn't effect anyone using DirectLoader, which is currently all our production connectors.)

Review guide

Mostly name changes. The functional stuff & new logging is in the StreamManager. Also LoadPipelineStepTask and UpdateBatchStateTask are tweaked to collect/pass the new information, and I added BatchState to all of the steps in ObjectLoader.

StreamManager already had rigorous tests, which still pass. I also re-tested S3, S3DataLake, and the (unmerged) MSSQL changes with this, both ITs and Perf Tests.

(NOTE: This is on top of the S3 Data Lake setup fix because the impending bulk load PR depends on both.)

vercel · 2025-03-22T23:55:03Z

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
airbyte-docs	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Mar 27, 2025 0:07am

Francis Genet (frifriSF59) · 2025-03-24T19:54:11Z

Haha, that was added automatically by the formatter, which I guess doesn't know what year it is?

Yeah, I have seen that before

Francis Genet (frifriSF59) · 2025-03-24T20:59:16Z

+            if (taskCompletionCounts.containsKey(taskKey)) {
+                throw IllegalStateException(
+                    """"$taskKey received input after seeing end-of-stream
+                        (checkpointCounts=$checkpointCounts, inputCount=$inputCount, sawEosAt=${taskCompletionCounts[taskKey]})                         if (stateStore.streamsEnded.contains(input.key.stream)) {\n" +


Is this formatting intended?

Johnny Schmidt (johnny-schmidt) · 2025-03-24T21:00:36Z

+    val records: Long,
+// TODO: File bytes moved
+) {
+    operator fun plus(other: CheckpointValue): CheckpointValue {


It overloads +

Francis Genet (frifriSF59) · 2025-03-24T21:04:29Z

-                    acc.third + checkpointCount.recordsCompleted.get(),
-                )
-            }
+        val readCount = counts.map { it.value.records }.sum()


anyway we can do something like this intsead?

fun countRecordsInState(state: BatchState): Long { return namedCheckpointCounts .filter { (key, _) -> key.second == state } .values .fold(0L) { accumulator, checkpointMap -> accumulator + checkpointMap .filterKeys { it.id <= checkpointId.id } .values .sumOf { it.records } } }

and then:

val persistedCount = countRecordsInState(BatchState.PERSISTED) if (persistedCount == totalRecordsRead) { return true } val completedCount = countRecordsInState(BatchState.COMPLETE) return completedCount == totalRecordsRead

I'll try it

Francis Genet (frifriSF59) · 2025-03-24T21:12:25Z

                        } else {
-                            stateStore.remove(input.key)?.close()
+                            stateStore.stateWithCounts.remove(input.key)?.let {
+                                if (inputCount > 0 || it.checkpointCounts.isNotEmpty()) {


Francis Genet (frifriSF59) · 2025-03-24T21:13:19Z

    val batchSizeOverride: Long? = null,
 ) {
-    // A map of (TaskIndex, Stream) -> Count_of_closed streams to ensure eos is not forwared from
+    // A map of (TaskIndex, Stream) ->  streams to ensure eos is not forwared from


forwarded

Francis Genet (frifriSF59) · 2025-03-24T21:15:10Z

-                            )
-                        }
+                    manager.incrementCheckpointCounts(
+                        message.taskName,


Would it make more sense to pass the message instead of individual attributes?

yeah, at this point it's coupled enough i might as well

Francis Genet (frifriSF59) · 2025-03-24T21:16:51Z


-// @Singleton
-class S3V2RoundRobinInputPartitioner : RoundRobinInputPartitioner()
+// @Singleton class S3V2RoundRobinInputPartitioner : RoundRobinInputPartitioner()


Should this just be removed?

Sorry I meant to add an explanation lol. This is here so that S3 is ready to have the object loader just switched on when I have bandwidth to release it.

Francis Genet (frifriSF59) · 2025-03-24T21:17:25Z

          fileName: glue_assume_role.json
          secretStore:
-            type: GSM
+            type: GSMs


Is this change intended?

definitely not

Edward Gao (edgao)

logic seems right, though I only skimmed the debug log piece. Some assorted comments on readability.

Edward Gao (edgao) · 2025-03-24T21:38:42Z

-                                val output = batchAccumulator.finish(stored.accumulatorState).output
-                                handleOutput(key, stored.checkpointCounts, output)
-                                stored.close()
+                            log.warn { "Finishing state for $key remaining after end-of-stream" }


checking: this is the case where we hit end of input, but there's still some not-yet-persisted data sitting around?

(if yes, nit: let's use INFO, since this is likely to happen on every stream)

Yes, it's actually normal, it shouldn't be a warning.

also s/after/at, so it sounds less like an error/edge case

Edward Gao (edgao) · 2025-03-24T21:40:34Z

+     * end (it is a critical error to receive input for a stream that has ended, as it means that
+     * something is likely wrong with our bookkeeping.)
+     */
+    data class StateStore<K1, S : AutoCloseable>(


nonblocking: does this bookkeeping have noticeable performance impact? since we're doing some map lookups per record

No, map lookups in non-concurrent containers in java are pretty heavily optimized. What matters is that we're avoiding concurrent lookups (by passing this data structure through the fold).

Edward Gao (edgao) · 2025-03-24T21:42:34Z

+     * something is likely wrong with our bookkeeping.)
+     */
+    data class StateStore<K1, S : AutoCloseable>(
+        val stateWithCounts: MutableMap<K1, StateWithCounts<S>> = mutableMapOf(),


do these need to be concurrent maps? (I think not? afaict StateStore is only ever used within a single thread)

no. because they're passed as an argument to the fold operation, we shouldn't need to do anything fancy

Edward Gao (edgao) · 2025-03-24T21:46:55Z

+    sealed interface PartResult : WithBatchState {
        val objectKey: String
+        override val state: BatchState
+            get() = BatchState.STAGED


nit: this can probably just be override val state = STAGED? (i.e. it doesn't need to be recomputed on every access)

(the compiler probably optimizes it out anyway 🤷 but might as well make it more explicit + less verbose)

Yeah probably better to move it to the class level

Edward Gao (edgao) · 2025-03-24T21:47:12Z

          fileName: glue_assume_role.json
          secretStore:
-            type: GSM
+            type: GSMs


this is probably a typo?

Edward Gao (edgao) · 2025-03-24T21:49:39Z

+    val records: Long,
+// TODO: File bytes moved
+) {
+    operator fun plus(other: CheckpointValue): CheckpointValue {


I'm pretty impressed by intellij's ability to cmd-click on a + and have it Just Work (i.e. correctly jump to here)!

Edward Gao (edgao) · 2025-03-24T21:54:16Z

+    ) {
+        checkpointCounts.forEach { (checkpointId, recordCount) ->
+            val taskKey = TaskKey(taskName, part)
+            if (taskCompletionCounts.containsKey(taskKey)) {


why is this happening inside the .forEach loop?

Oof, good question.

Edward Gao (edgao) · 2025-03-24T21:57:29Z

+                        This indicates data was processed out of order and future bookkeeping might be corrupt. Failing hard."""
+                )
+            }
+            namedCheckpointCounts


nit: I think this would be easier to read if you pulled the getOrPut out of the loop i.e.

val x = namedCheckpointCounts.getOrPut() {} checkpointCounts.forEach { x.merge() }

(at least I think that's equivalent?)

yeah i agree that's better and probably more efficient also

Johnny Schmidt (johnny-schmidt) requested a review from a team as a code owner March 22, 2025 23:55

Octavia Squidington III (octavia-squidington-iii) added connectors/destination/s3 connectors/destination/mssql connectors/destination/s3-data-lake labels Mar 22, 2025

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch 2 times, most recently from e06a1c7 to 5c6789f Compare March 23, 2025 00:27

Johnny Schmidt (johnny-schmidt) changed the base branch from master to jschmidt/dest-s3-datalake/bugfix-setup-doesn-not-await-start March 23, 2025 00:27

vercel Bot deployed to Preview March 23, 2025 06:48 View deployment

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/dest-s3-datalake/bugfix-setup-doesn-not-await-start branch from 92acee8 to 3677a7e Compare March 24, 2025 19:19

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 5c6789f to 2bd4f8c Compare March 24, 2025 19:23

vercel Bot deployed to Preview March 24, 2025 19:30 View deployment

Base automatically changed from jschmidt/dest-s3-datalake/bugfix-setup-doesn-not-await-start to master March 24, 2025 19:34

Francis Genet (frifriSF59) reviewed Mar 24, 2025

View reviewed changes

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 2bd4f8c to 5c1a49c Compare March 24, 2025 20:44

vercel Bot deployed to Preview March 24, 2025 20:51 View deployment

Francis Genet (frifriSF59) reviewed Mar 24, 2025

View reviewed changes

Johnny Schmidt (johnny-schmidt) commented Mar 24, 2025

View reviewed changes

Francis Genet (frifriSF59) reviewed Mar 24, 2025

View reviewed changes

Edward Gao (edgao) approved these changes Mar 24, 2025

View reviewed changes

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 5c1a49c to 16669f7 Compare March 25, 2025 18:03

Johnny Schmidt (johnny-schmidt) enabled auto-merge (squash) March 25, 2025 18:03

vercel Bot deployed to Preview March 25, 2025 18:09 View deployment

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 16669f7 to fce0ca7 Compare March 25, 2025 18:38

vercel Bot deployed to Preview March 25, 2025 18:45 View deployment

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from fce0ca7 to 3c2ac84 Compare March 25, 2025 22:06

vercel Bot deployed to Preview March 25, 2025 22:12 View deployment

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 3c2ac84 to 1827227 Compare March 26, 2025 20:28

vercel Bot deployed to Preview March 26, 2025 20:34 View deployment

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 1827227 to 93e8c2f Compare March 26, 2025 22:01

vercel Bot deployed to Preview March 26, 2025 22:07 View deployment

Load CDK: Improved new interface bookkeeping

00d7882

Johnny Schmidt (johnny-schmidt) force-pushed the jschmidt/load-cdk/improved-new-interface-bookkeeping branch from 93e8c2f to 00d7882 Compare March 27, 2025 00:01

vercel Bot deployed to Preview March 27, 2025 00:07 View deployment

Johnny Schmidt (johnny-schmidt) merged commit 02879f6 into master Mar 27, 2025

Johnny Schmidt (johnny-schmidt) deleted the jschmidt/load-cdk/improved-new-interface-bookkeeping branch March 27, 2025 00:15

Conversation

Johnny Schmidt (johnny-schmidt) commented Mar 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Review guide

Uh oh!

vercel Bot commented Mar 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Francis Genet (frifriSF59) Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Edward Gao (edgao) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch was successfully deployed

Reviewers

Johnny Schmidt (johnny-schmidt) commented Mar 22, 2025 •

edited

Loading

vercel Bot commented Mar 22, 2025 •

edited

Loading

Francis Genet (frifriSF59) Mar 24, 2025 •

edited

Loading