jose-torres commented on code in PR #57444:
URL: https://github.com/apache/spark/pull/57444#discussion_r3643035269
##########
sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala:
##########
@@ -942,6 +944,315 @@ case class Scd2BatchProcessor(
.fieldNames
.toImmutableArraySeq
}
+
+ /**
+ * Tag each post-reconciliation row with
[[Scd2BatchProcessor.shouldRouteToAuxTableColName]]:
+ * `true` for rows that belong in the auxiliary table (tombstones and hidden
no-op upserts),
+ * `false` for rows that do not. The flag's contract is solely about
aux-table membership; it
+ * makes no claim about whether a `false` row belongs in the target table.
+ *
+ * @param reconciledDf the canonical post-reconciliation rows for the
affected keys.
+ * @return `reconciledDf` with an added boolean
+ * [[Scd2BatchProcessor.shouldRouteToAuxTableColName]] column; no
rows are added or
+ * removed.
+ */
+ private[autocdc] def identifyAndTagAuxRows(reconciledDf: DataFrame):
DataFrame = {
+ val recordStartAt =
+
Scd2BatchProcessor.recordStartAtOf(F.col(AutoCdcReservedNames.cdcMetadataColName))
+ val startAt = F.col(Scd2BatchProcessor.startAtColName)
+ val endAt = F.col(Scd2BatchProcessor.endAtColName)
+ val current = Scd2IntervalColumns(recordStartAt, startAt, endAt)
+ val next = current.leadBy(1, orderChronologicallyPerKeyWindow)
+
+ val trackedHistoryColumns = computeTrackedHistoryColumns(reconciledDf)
+ val areTrackedColumnsEqualInNextRow = trackedHistoryColumns
+ .map { c =>
+ val col = F.col(QuotingUtils.quoteIdentifier(c))
+ col <=> F.lead(col, 1).over(orderChronologicallyPerKeyWindow)
+ }
+ .reduceOption(_ && _)
+ .getOrElse(F.lit(true))
+
+ val isTombstone = RowClassifier.isTombstone(current)
+ val isHiddenNoOpUpsert = RowClassifier.isNoOpUpsertContinuation(
+ row = current,
+ next = next,
+ areTrackedColumnsEqual = areTrackedColumnsEqualInNextRow
+ )
+
+ reconciledDf.withColumn(
+ Scd2BatchProcessor.shouldRouteToAuxTableColName,
+ isTombstone || isHiddenNoOpUpsert
+ )
+ }
+
+ /**
+ * Merge the reconciled rows that belong in the auxiliary table (tombstones
and hidden no-op
+ * upserts, as tagged by [[identifyAndTagAuxRows]]) onto the auxiliary
table, and logically
+ * delete any previously-affected aux row that did not survive
reconciliation.
+ *
+ * Idempotency across `foreachBatch` retries: rows this batch logically
deletes are not
+ * physically removed but stamped with
[[Scd2BatchProcessor.deletedByBatchIdColName]] `=
+ * batchId`, so a retry of the same `batchId` still observes them via
+ * [[findAffectedRowsFromAuxiliaryTable]] and re-derives the same
reconciliation output. Aux
+ * rows logically deleted by an older, already-committed batch
+ * ([[Scd2BatchProcessor.deletedByBatchIdColName]] `!= batchId`) are
physically
+ * garbage-collected as part of this same merge.
+ *
+ * @param reconciledDfWithAuxRowsTagged reconciled rows tagged by
[[identifyAndTagAuxRows]].
+ * @param originalAffectedRowsFromAuxiliaryTable the affected aux rows
pulled in for this
+ * microbatch, in canonical SCD2 row schema (i.e. as returned by
+ * [[findAffectedRowsFromAuxiliaryTable]], with the aux-only
+ * [[Scd2BatchProcessor.deletedByBatchIdColName]] already dropped).
+ * @param auxiliaryTableIdentifier the identifier of the auxiliary table to
merge into.
+ * @param batchId the underlying Spark streaming query's batchId, used to
stamp logical deletes
+ * and scope garbage collection.
+ */
+ private[autocdc] def mergeRowsIntoAuxiliaryTable(
+ reconciledDfWithAuxRowsTagged: DataFrame,
+ originalAffectedRowsFromAuxiliaryTable: DataFrame,
+ auxiliaryTableIdentifier: TableIdentifier,
+ batchId: Long): Unit = {
+ val resolver =
reconciledDfWithAuxRowsTagged.sparkSession.sessionState.conf.resolver
+ val deletedByBatchIdCol = Scd2BatchProcessor.deletedByBatchIdColName
+
+ val reconciledAuxRows = reconciledDfWithAuxRowsTagged
+ .filter(F.col(Scd2BatchProcessor.shouldRouteToAuxTableColName))
+ .drop(Scd2BatchProcessor.shouldRouteToAuxTableColName)
+
+ // All of the reconciledAuxRows will be landing in the aux table as alive
(non-deleted) rows,
+ // so tag them with a null deleted-by batch id.
+ val auxRowsToUpsert = reconciledAuxRows
+ .withColumn(deletedByBatchIdCol, F.lit(null).cast(LongType))
+
+ // Any aux row pulled in for reconciliation but absent from the
post-reconciliation aux rows
+ // was either dropped as redundant or promoted to the (now-visible) target
table. Either way
+ // it must leave the aux table; stamp it with this batch's id for
deleted-by batch id. On this
+ // batch's MERGE, the existing aux row will be considered logically
deleted. In a future
+ // microbatch's merge, it will be physically deleted.
+ val auxRowsToDelete = antiJoinRowsByRecordStartAtPerKey(
+ leftRows = originalAffectedRowsFromAuxiliaryTable,
+ rightRows = reconciledAuxRows
+ )
+ .withColumn(deletedByBatchIdCol, F.lit(batchId))
+
+ val mergeSource = auxRowsToUpsert
+ .unionByName(auxRowsToDelete)
+ .as("source")
+
+ // At this point [[mergeSource]] must have the same columns/shape as the
persisted aux table.
+ val auxTableColumns = mergeSource.columns.toImmutableArraySeq
+
+ // Build predicate for whether rows represent the same event (match). In
SCD2 each key can have
+ // multiple records, so a row's identity is defined by (keys, sequencing).
+ val auxIdentQuoted = auxiliaryTableIdentifier.quotedString
+ val meta = AutoCdcReservedNames.cdcMetadataColName
+
+ val mergeSourceRecordStartAt =
+ Scd2BatchProcessor.recordStartAtOf(F.col(s"source.`$meta`"))
+ val auxTableRecordStartAt =
+ Scd2BatchProcessor.recordStartAtOf(F.col(s"$auxIdentQuoted.`$meta`"))
+
+ val doKeysMatch = changeArgs.keys
+ .map(k => F.col(s"source.${k.quoted}") ===
F.col(s"$auxIdentQuoted.${k.quoted}"))
+ .reduce(_ && _)
+ val doRowsMatch = doKeysMatch && (mergeSourceRecordStartAt <=>
auxTableRecordStartAt)
+
+ // On updates, MERGE requires only non-key columns are updated (remapped).
For inserts, all of
+ // the row's columns must explicitly be mapped.
+ val keyNames = changeArgs.keys.map(_.name)
+ def upsertAssignments(columnName: String): (String, Column) = {
+ val quoted = QuotingUtils.quoteIdentifier(columnName)
+ s"$auxIdentQuoted.$quoted" -> F.col(s"source.$quoted")
+ }
+ val nonKeyUpdateAssignments = auxTableColumns
+ .filterNot(c => keyNames.exists(resolver(_, c)))
+ .map(upsertAssignments)
+ .toMap
+ val insertAssignments = auxTableColumns.map(upsertAssignments).toMap
+
+ // Physically garbage-collect aux rows logically deleted by some other,
already-committed
+ // batch. Such rows are always excluded when pulling in the current
microbatch's affected set,
+ // so they can never match against a row being merged in.
+ val auxTableDeletedByBatchId =
F.col(s"$auxIdentQuoted.`$deletedByBatchIdCol`")
+ val isGarbageCollectableAuxRow =
+ auxTableDeletedByBatchId.isNotNull && auxTableDeletedByBatchId =!=
F.lit(batchId)
+
+ // Whether a row in the MERGE source represents a row being [logically]
deleted, as opposed to
+ // being upserted.
+ val shouldLogicallyDeleteAuxRow =
F.col(s"source.`$deletedByBatchIdCol`").isNotNull
+
+ mergeSource
+ .mergeInto(auxIdentQuoted, doRowsMatch)
+ // Keys in source row match against existing row in aux table, and
declare intent to delete
+ // the corresponding row in the aux; mark the row as logically deleted
in the aux table.
+ .whenMatched(shouldLogicallyDeleteAuxRow)
+ .update(
+ Map(
+ s"$auxIdentQuoted.`$deletedByBatchIdCol`" ->
+ F.col(s"source.`$deletedByBatchIdCol`")
+ )
+ )
+ // Keys in source row match against existing row in aux table and does
not represent a row
+ // being deleted; update the data/operational columns
+ .whenMatched(!shouldLogicallyDeleteAuxRow)
+ .update(nonKeyUpdateAssignments)
+ // Keys in source do not match against an existing row in aux table and
does not represent a
+ // row being deleted; insert the new key's row.
+ .whenNotMatched(!shouldLogicallyDeleteAuxRow)
+ .insert(insertAssignments)
+ // If this is a row not affected by the current microbatch but is
eligible for garbage
+ // collection now, proactively hard-delete it.
+ // TODO: This GC triggers a full scan of the aux table; revisit whether
to GC only
+ // periodically rather than on every microbatch.
+ .whenNotMatchedBySource(isGarbageCollectableAuxRow)
Review Comment:
Noting for the record that this notMatchedBySource _is_ safe, even though
the incoming batch built from originalAffectedRowsFromAuxiliaryTable may not
have included all target rows in the first place, because
isGarbageCollectableAuxRow will exclude anything that's not already marked as
ready for deletion.
--
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]
