anew commented on code in PR #57495: URL: https://github.com/apache/spark/pull/57495#discussion_r3660841092
########## sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandlerSuite.scala: ########## @@ -0,0 +1,1087 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.pipelines.autocdc + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.{functions => F, AnalysisException, QueryTest, Row} +import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ + +/** + * End-to-end unit tests for [[Scd2ForeachBatchHandler]]. Unlike the focused unit suites that + * exercise individual [[Scd2BatchProcessor]] stages in isolation, these drive the entire + * microbatch reconciliation pipeline - validation, preprocessing, affected-row pull-in from both + * side tables, decomposition, start/end reconciliation, aux routing, and the two `MERGE INTO` + * writes - through the public `execute` entrypoint against an in-memory v2 catalog. + * + * They are the first place the cross-microbatch stateful behaviors (out-of-order arrivals, no-op + * run coalescing across batches, tombstone-driven suppression, and auxiliary-table garbage + * collection) can be observed against materialized target and auxiliary tables, so the idempotency + * / GC / cross-batch scenarios are emphasized here. + * + * The default flow tracks every persisted user column (`value`) under key `id`, sequences by + * `seq`, and treats `is_delete = true` rows as deletes. + */ +class Scd2ForeachBatchHandlerSuite + extends QueryTest + with SharedSparkSession + with BeforeAndAfter + with AutoCdcCatalogExecutionTestBase { + + private val sourceSchema = new StructType() + .add("id", IntegerType) + .add("value", StringType) + .add("seq", LongType) + .add("is_delete", BooleanType) + + /** The SCD2 cdc-metadata struct carries a single `recordStartAt` field (unlike SCD1's two). */ + private val scd2MetadataSchema: StructType = Scd2BatchProcessor.cdcMetadataColSchema(LongType) + + /** Canonical SCD2 row schema: persisted user columns + framework start/end + cdc metadata. */ + private val canonicalSchema = new StructType() + .add("id", IntegerType) + .add("value", StringType) + .add(Scd2BatchProcessor.startAtColName, LongType, nullable = true) + .add(Scd2BatchProcessor.endAtColName, LongType, nullable = true) + .add(AutoCdcReservedNames.cdcMetadataColName, scd2MetadataSchema, nullable = false) + + /** Auxiliary table schema: canonical schema plus the aux-only logical-delete marker column. */ + private val auxSchema = canonicalSchema + .add(Scd2BatchProcessor.deletedByBatchIdColName, LongType, nullable = true) + + /** Target table schema is exactly the canonical schema. */ + private val targetSchema = canonicalSchema + + private val processor = Scd2BatchProcessor( + changeArgs = ChangeArgs( + keys = Seq(UnqualifiedColumnName("id")), + sequencing = F.col("seq"), + storedAsScdType = ScdType.Type2, + deleteCondition = Some(F.col("is_delete")), + // Persist only id + value; seq / is_delete are control columns that need not be included. + columnSelection = Some( + ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName("seq"), UnqualifiedColumnName("is_delete")) + ) + ) + ), + resolvedSequencingType = LongType + ) + + private def createAuxTable(seedRows: Row*): Unit = + createTable(defaultAuxIdent, defaultAuxTableIdentifier, auxSchema, seedRows: _*) + + private def createTargetTable(seedRows: Row*): Unit = + createTable(defaultTargetIdent, defaultTargetTableIdentifier, targetSchema, seedRows: _*) + + private def auxTable: DataFrame = spark.read.table(defaultAuxTableIdentifier.quotedString) + + private def targetTable: DataFrame = spark.read.table(defaultTargetTableIdentifier.quotedString) + + private def execWith(p: Scd2BatchProcessor): Scd2ForeachBatchHandler = Scd2ForeachBatchHandler( + batchProcessor = p, + auxiliaryTableIdentifier = defaultAuxTableIdentifier, + targetTableIdentifier = defaultTargetTableIdentifier + ) + + private def exec: Scd2ForeachBatchHandler = execWith(processor) + + /** A source UPSERT event: `(id, value, seq, is_delete = false)`. */ + private def upsert(id: Int, value: String, seq: Long): Row = Row(id, value, seq, false) + + /** A source DELETE event: `(id, null, seq, is_delete = true)`. */ + private def del(id: Int, seq: Long): Row = Row(id, null, seq, true) + + /** The cdc-metadata struct value for a given `recordStartAt`. */ + private def meta(recordStartAt: Long): Row = Row(recordStartAt) + + /** A canonical target row `(id, value, startAt, endAt, meta(recordStartAt))`. */ + private def targetRow( + id: Int, + value: String, + startAt: java.lang.Long, + endAt: java.lang.Long, + recordStartAt: Long): Row = + Row(id, value, startAt, endAt, meta(recordStartAt)) + + /** A canonical aux row `(id, value, startAt, endAt, meta(recordStartAt), deletedByBatchId)`. */ + private def auxRow( + id: Int, + value: String, + startAt: java.lang.Long, + endAt: java.lang.Long, + recordStartAt: Long, + deletedByBatchId: java.lang.Long): Row = + Row(id, value, startAt, endAt, meta(recordStartAt), deletedByBatchId) + + /** Run a microbatch of source rows through the default handler. */ + private def runBatch(batchId: Long)(rows: Row*): Unit = + exec.execute(microbatchOf(sourceSchema)(rows: _*), batchId) + + /** + * Run `rows` as batch `batchId`, capture both tables, then replay the identical batch under the + * same `batchId` and assert both tables are byte-for-byte unchanged. Models a crash/redelivery + * where a committed microbatch is reprocessed. + */ + private def assertReplayStable(batchId: Long)(rows: Row*): Unit = { + runBatch(batchId)(rows: _*) + val targetAfterFirst = targetTable.collect().toSeq + val auxAfterFirst = auxTable.collect().toSeq + + runBatch(batchId)(rows: _*) + checkAnswer(targetTable, targetAfterFirst) + checkAnswer(auxTable, auxAfterFirst) + } + + /** + * Replicate [[Scd2ForeachBatchHandler.execute]] but stop after the auxiliary-table merge, + * skipping the target-table merge. Models a crash *between* the two merges: the aux table has + * committed this `batchId`'s logical deletes / inserts, but the target table has not yet been + * updated. On recovery Structured Streaming reruns the same `batchId`, which + * [[Scd2BatchProcessor.deletedByBatchIdColName]] is designed to make idempotent. + */ + private def runBatchAuxMergeOnly(batchId: Long)(rows: Row*): Unit = { Review Comment: You are right, it's much better to share that, refactored. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
