[SPARK-57322][SDP] Implement SCD2 Batch Processor; Reconcile StartAt/EndAt

sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/ChangeArgs.scala

@@ -153,21 +153,28 @@ object ScdType {
 /**
  * Configuration for an AutoCDC flow.
  *
- * @param keys            The column(s) that uniquely identify a row in the source data.
- * @param sequencing      Expression ordering CDC events to correctly resolve out-of-order
- *                        arrivals. Must be a sortable type.
- * @param deleteCondition Expression that marks a source row as a DELETE. When None, all
- *                        rows are treated as upserts.
- * @param storedAsScdType The SCD strategy these args should be applied to.
- * @param columnSelection Which source columns to select in the target table. None means
- *                        all columns.
+ * @param keys                   The column(s) that uniquely identify a row in the source data.
+ * @param sequencing             Expression ordering CDC events to correctly resolve out-of-order
+ *                               arrivals. Must be a sortable type.
+ * @param deleteCondition        Expression that marks a source row as a DELETE. When None, all
+ *                               rows are treated as upserts.
+ * @param storedAsScdType        The SCD strategy these args should be applied to.
+ * @param columnSelection        Which source columns to select in the target table. None means
+ *                               all columns.
+ * @param trackHistorySelection  SCD2 only. Selects the user-data columns whose values define a
+ *                               run: two consecutive upsert events for the same key are
+ *                               coalesced into the same run iff they agree on every selected
+ *                               column. None means every eligible user column (i.e. every
+ *                               source column that is neither a key nor a framework column) is
+ *                               considered tracked. Ignored under SCD1.
  */
 case class ChangeArgs(
     keys: Seq[UnqualifiedColumnName],
     sequencing: Column,
     storedAsScdType: ScdType,
     deleteCondition: Option[Column] = None,
-    columnSelection: Option[ColumnSelection] = None
+    columnSelection: Option[ColumnSelection] = None,
+    trackHistorySelection: Option[ColumnSelection] = None
 ) {
   ChangeArgs.validateNonEmptyKeys(keys)
 }

sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala

@@ -439,9 +439,9 @@ case class Scd2BatchProcessor(
    *   a dataframe with the same schema as the input. Every closed non-tombstone row that
    *   was bisected has been replaced by its head + tail pair; every other row is carried
    *   through as-is. Each output row can be classified as one of: {decomposition head,
-   *   decomposition tail, instantaneous delete, open upsert, closed-and-unbisected row}. It's
-   *   possible that some of the returned decomposition tails are logically redundant, as
-   *   deletion markers that are immediately overtaken by a succeeding row.
+   *   decomposition tail, tombstone, open upsert, closed-and-unbisected row}. It's possible
+   *   that some of the returned decomposition tails are logically redundant, as deletion
+   *   markers that are immediately overtaken by a succeeding row.
    */
   private[autocdc] def decomposeOutOfOrderRows(rowsToDecomposePerKey: DataFrame): DataFrame = {
     val recordStartAtField =
@@ -532,8 +532,8 @@ case class Scd2BatchProcessor(
 
   /**
    * Asserts that every row in `decomposedRowsPerKey` conforms to one of the four canonical
-   * post-decomposition shapes - instantaneous delete, open upsert, closed upsert, or
-   * decomposition tail - and is otherwise a structural identity transform.
+   * post-decomposition shapes - tombstone, open upsert, closed upsert, or decomposition
+   * tail - and is otherwise a structural identity transform.
    *
    * @param decomposedRowsPerKey
    *   the output of [[decomposeOutOfOrderRows]]: a dataframe conforming to the canonical
@@ -556,7 +556,7 @@ case class Scd2BatchProcessor(
     val row = Scd2IntervalColumns(recordStartAtField, startAtCol, endAtCol)
     val isWellFormedRow =
       RowClassifier.isDecompositionTail(row) ||
-        RowClassifier.isDeleteRepresentingRow(row) ||
+        RowClassifier.isTombstone(row) ||
         RowClassifier.isUpsertRepresentingRow(row)
 
     def stringOrNullLit(c: Column): Column = F.coalesce(c.cast(StringType), F.lit("null"))
@@ -639,6 +639,193 @@ case class Scd2BatchProcessor(
       .filter(!isRedundantAtSameEffectiveSequence)
       .drop(Scd2BatchProcessor.nextEffectiveRecordStartAtColName)
   }
+
+  /**
+   * Recompute every row's [[startAtColName]] and [[endAtColName]] over the per-key chronological
+   * window so the dataframe reflects the canonical SCD2 timeline that the downstream aux- and
+   * target-table merges consume.
+   *
+   * Decomposition tails and tombstones round-trip unchanged. An open upsert may close at its
+   * successor's effective sequence (becoming closed); a closed upsert may have its endAt cleared
+   * when absorbed into a run as a no-op continuation (becoming open). [[recordStartAtFieldName]]
+   * is never modified.
+   *
+   * @param decomposedAndCleanedDf
+   *   the output of [[dropRedundantRowsPostDecomposition]]: a dataframe conforming to the
+   *   canonical SCD2 row schema `[user_cols..., [[startAtColName]], [[endAtColName]],
+   *   [[cdcMetadataColName]]]` where every row is in one of the four canonical post-
+   *   decomposition shapes (decomposition tail, tombstone, open upsert, closed upsert).
+   *   If a row is a closed upsert in the input, it is assumed to not be bisected by any other
+   *   row in the input.
+   * @return
+   *   a dataframe with the same schema and row count as the input, with each row's
+   *   [[startAtColName]] / [[endAtColName]] replaced by their reconciled values.
+   */
+  private[autocdc] def reconcileStartAndEndAt(
+      decomposedAndCleanedDf: DataFrame): DataFrame = {
+    val trackedHistoryColumns = computeTrackedHistoryColumns(decomposedAndCleanedDf)
+
+    val recordStartAtField =
+      Scd2BatchProcessor.recordStartAtOf(F.col(AutoCdcReservedNames.cdcMetadataColName))
+    val startAtCol = F.col(Scd2BatchProcessor.startAtColName)
+    val endAtCol = F.col(Scd2BatchProcessor.endAtColName)
+
+    // Decomposition tails carry no recordStartAt of their own, so they take the closing
+    // sequence (`endAt`) as their effective ordering position - the same convention used by
+    // [[orderChronologicallyPerKeyWindow]] and [[dropRedundantRowsPostDecomposition]].
+    val current = Scd2IntervalColumns(recordStartAtField, startAtCol, endAtCol)
+    val previous = current.lagBy(1, orderChronologicallyPerKeyWindow)
+    val next = current.leadBy(1, orderChronologicallyPerKeyWindow)
+
+    // A row is the last in its per-key window when `LEAD(1)` has no successor; a constant
+    // literal is sufficient since we only care whether one exists.
+    val isLastRowInKeyWindow =
+      F.lead(F.lit(true), 1).over(orderChronologicallyPerKeyWindow).isNull
+
+    // The current row's tracked-history equality is computed against both its predecessor and
+    // its successor so the same window scan can decide both run-head start (LAG-side) and no-op
+    // continuation closure (LEAD-side) without an extra pass. The comparison is null-safe
+    // (`<=>`), so two rows with matching null values in the same tracked column register as
+    // equal. An empty tracked-history column set collapses to a constant `true`, which makes
+    // every consecutive upsert pair a no-op continuation - the correct degenerate behavior when
+    // the user tracks nothing.
+    val areTrackedColumnsEqualInPreviousRow = trackedHistoryColumns
+      .map { c =>
+        val col = F.col(QuotingUtils.quoteIdentifier(c))
+        col <=> F.lag(col, 1).over(orderChronologicallyPerKeyWindow)
+      }
+      .reduceOption(_ && _)
+      .getOrElse(F.lit(true))
+
+    val areTrackedColumnsEqualInNextRow = trackedHistoryColumns
+      .map { c =>
+        val col = F.col(QuotingUtils.quoteIdentifier(c))
+        col <=> F.lead(col, 1).over(orderChronologicallyPerKeyWindow)
+      }
+      .reduceOption(_ && _)
+      .getOrElse(F.lit(true))
+
+    // Reconciliation of start/end at is dependent on the class of row being reconciled. Build
+    // row classification predicates.
+    val isDecompositionTail = RowClassifier.isDecompositionTail(current)
+    val isUpsertRepresentingRow = RowClassifier.isUpsertRepresentingRow(current)
+
+    // From the previous row's perspective, the current row is its successor.
+    val previousIsNoOpUpsertWithCurrent =
+      RowClassifier.isNoOpUpsertContinuation(
+        row = previous,
+        next = current,
+        areTrackedColumnsEqualInNextRow = areTrackedColumnsEqualInPreviousRow
+      )
+
+    // "Window-local run head" means the current row begins a new run within the affected
+    // window. The first row in the window is automatically considered local-run-head since
+    // there's no predecessor to coalesce with. A non-first row is a local run head iff its
+    // predecessor is not a no-op continuation that absorbs it.
+    val isWindowLocalUpsertRunHead =
+      isUpsertRepresentingRow && !previousIsNoOpUpsertWithCurrent
+    val isFirstRowInKeyWindow = previous.effectiveRecordStartAt.isNull
+    val runHeadStartAt =
+      F.when(
+        isWindowLocalUpsertRunHead,
+        // The first row in the window may be a window-local run head but not a global run
+        // head (e.g., an aux anchor row pulled in for left context). In that case, `startAt`
+        // may be strictly less than `recordStartAt`, encoding the true global run start, and
+        // we propagate it forward to later in-window continuations of the same run.
+        // For every later window-local upsert run head, `recordStartAt` is the run start.
+        F.when(isFirstRowInKeyWindow, startAtCol).otherwise(recordStartAtField)
+      )
+
+    // Propagate the run head's `startAt` forward to every row in the run via a running
+    // `last(...)` over `[unboundedPreceding, currentRow]`. `runHeadStartAt` is non-null
+    // only on run heads, and `ignoreNulls = true` makes intermediate rows inherit the most
+    // recent head's value.
+    val runStartAt =
+      F.last(runHeadStartAt, ignoreNulls = true).over(
+        orderChronologicallyPerKeyWindow.rowsBetween(
+          Window.unboundedPreceding,
+          Window.currentRow
+        )
+      )
+
+    val currentIsNoOpUpsertWithNext =
+      RowClassifier.isNoOpUpsertContinuation(
+        row = current,
+        next = next,
+        areTrackedColumnsEqualInNextRow = areTrackedColumnsEqualInNextRow
+      )
+
+    val finalStartAt =
+      F.when(isDecompositionTail, F.lit(null).cast(resolvedSequencingType))
+        .when(isUpsertRepresentingRow, runStartAt)
+        .otherwise(startAtCol)
+
+    val finalEndAt =
+      F.when(isDecompositionTail, endAtCol)
+        .when(isLastRowInKeyWindow, endAtCol)
+        // A no-op continuation collapses into its run head, so the row's visible interval
+        // disappears and `endAt` is reset to null to route the row to the aux table.
+        .when(currentIsNoOpUpsertWithNext, F.lit(null).cast(resolvedSequencingType))
+        // The row already closes strictly before the next event (e.g., a tombstone or a
+        // closed upsert that ended before the next event arrived), so there is nothing to
+        // re-close.
+        .when(
+          RowClassifier.rowClosesStrictlyBeforeNextRow(endAtCol, next.effectiveRecordStartAt),
+          endAtCol)
+        .otherwise(next.effectiveRecordStartAt)
+
+    // Stage the recomputed start/end values into temporary columns first, then project them
+    // back over the originals via `select`. Both staged values reference the original
+    // `startAt`/`endAt`, so we cannot replace the originals in a single `withColumn` step.
+    val staged = decomposedAndCleanedDf
+      .withColumn(Scd2BatchProcessor.finalStartAtColName, finalStartAt)
+      .withColumn(Scd2BatchProcessor.finalEndAtColName, finalEndAt)
+
+    // Determine columns for selection from staged dataframe using the original input dataframe's
+    // schema. The final startAt/endAt should be remapped, all other columns should pass through
+    // as-is.
+    val outputColumns = decomposedAndCleanedDf.columns.map {
+      case col if col == Scd2BatchProcessor.startAtColName =>
+        val startAtMetadata = decomposedAndCleanedDf.schema(col).metadata
+        F.col(Scd2BatchProcessor.finalStartAtColName)
+          .as(Scd2BatchProcessor.startAtColName, startAtMetadata)
+      case col if col == Scd2BatchProcessor.endAtColName =>
+        val endAtMetadata = decomposedAndCleanedDf.schema(col).metadata
+        F.col(Scd2BatchProcessor.finalEndAtColName)
+          .as(Scd2BatchProcessor.endAtColName, endAtMetadata)
+      case col =>
+        F.col(QuotingUtils.quoteIdentifier(col))
+    }
+    staged.select(outputColumns.toImmutableArraySeq: _*)
+  }
+
+  /**
+   * Return the schema field names of columns selected for history-tracking on `df`:
+   * the eligible user-data columns (those not in [[ChangeArgs.keys]] or the framework
+   * reserved set) filtered through [[ChangeArgs.trackHistorySelection]].
+   */
+  private def computeTrackedHistoryColumns(df: DataFrame): Seq[String] = {
+    val conf = df.sparkSession.sessionState.conf
+    val resolver = conf.resolver
+
+    val keyColNames = changeArgs.keys.map(_.name)
+    val reservedColNames = Scd2BatchProcessor.reservedFrameworkColNames
+
+    val eligibleSchema = StructType(df.schema.fields.filterNot { field =>
+      reservedColNames.exists(resolver(_, field.name)) ||
+        keyColNames.exists(resolver(_, field.name))
+    })
+
+    ColumnSelection
+      .applyToSchema(
+        schemaName = "trackHistorySelection",
+        schema = eligibleSchema,
+        columnSelection = changeArgs.trackHistorySelection,
+        caseSensitive = conf.caseSensitiveAnalysis
+      )
+      .fieldNames
+      .toImmutableArraySeq
+  }
 }
 
 /**
@@ -833,6 +1020,18 @@ object Scd2BatchProcessor {
   private[autocdc] val nextEffectiveRecordStartAtColName: String =
     s"${AutoCdcReservedNames.prefix}next_effective_record_start_at"
 
+  /**
+   * Names of temporary columns used by [[reconcileStartAndEndAt]] to stage the recomputed
+   * [[startAtColName]] / [[endAtColName]] values before projecting them back over the originals.
+   *
+   * Temporary in that the columns have no observable side effect or persistence across
+   * microbatches.
+   */
+  private val finalStartAtColName: String =
+    s"${AutoCdcReservedNames.prefix}final_start_at"
+  private val finalEndAtColName: String =
+    s"${AutoCdcReservedNames.prefix}final_end_at"
+
   /**
    * Name of the temporary column used to identify the sequence associated with the anchor
    * row found in the auxiliary table for the incoming microbatch. Since sequences must be unique
@@ -896,6 +1095,20 @@ private[autocdc] case class Scd2IntervalColumns(
    * [[Scd2BatchProcessor.orderChronologicallyPerKeyWindow]].
    */
   def effectiveRecordStartAt: Column = F.coalesce(recordStartAt, endAt)
+
+  /** The same interval columns read from the row `offset` positions earlier in `window`. */
+  def lagBy(offset: Int, window: WindowSpec): Scd2IntervalColumns =
+    Scd2IntervalColumns(
+      F.lag(recordStartAt, offset).over(window),
+      F.lag(startAt, offset).over(window),
+      F.lag(endAt, offset).over(window))
+
+  /** The same interval columns read from the row `offset` positions later in `window`. */
+  def leadBy(offset: Int, window: WindowSpec): Scd2IntervalColumns =
+    Scd2IntervalColumns(
+      F.lead(recordStartAt, offset).over(window),
+      F.lead(startAt, offset).over(window),
+      F.lead(endAt, offset).over(window))
 }
 
 object RowClassifier {
@@ -942,18 +1155,50 @@ object RowClassifier {
     isOpenUpsert(row) || isClosedUpsert(row)
 
   /**
-   * Delete-representing row, encoded as an instantaneous (zero-width) interval at `recordStartAt`
-   * (`startAt == endAt == recordStartAt`). Never materializes in the target table.
+   * Tombstone (delete-boundary) row, encoded as an instantaneous interval at
+   * `recordStartAt`. Never materializes in the target table, only in the aux table.
    *
-   * User-data column values on these rows are not part of the SCD2 contract: they may reflect
-   * the originating delete event, the values of the upsert whose closed-interval row was bisected
-   * (when promoted from a decomposition tail), or be null altogether. Reconciliation does not
-   * consume these values for any semantic decision.
+   * User-data column values on tombstones are not part of the SCD2 contract: they may
+   * reflect the originating delete event, the values of the upsert whose closed-interval
+   * row was bisected (when the tombstone was promoted from a decomposition tail), or be
+   * null altogether. Reconciliation does not consume these values for any semantic
+   * decision.
    */
-  private[autocdc] def isDeleteRepresentingRow(row: Scd2IntervalColumns): Column =
+  private[autocdc] def isTombstone(row: Scd2IntervalColumns): Column =
     row.recordStartAt.isNotNull &&
       row.startAt.isNotNull &&
       row.endAt.isNotNull &&
       row.startAt === row.recordStartAt &&
       row.endAt === row.recordStartAt
+
+  /**
+   * Whether a row closes (`endAt`) strictly before the next chronologically-ordered row for the
+   * same key begins (`nextEffectiveRecordStartAt`), leaving a gap in the visible timeline.
+   */
+  private[autocdc] def rowClosesStrictlyBeforeNextRow(
+      endAt: Column,
+      nextEffectiveRecordStartAt: Column
+  ): Column =
+    endAt.isNotNull && endAt < nextEffectiveRecordStartAt
+
+  /**
+   * Whether `row` carries no new information beyond its immediate successor `next` and so
+   * collapses into that successor's run instead of standing as its own visible interval. It is
+   * the caller's responsibility to pass `row` and `next` as successive rows in chronological
+   * order, and `areTrackedColumnsEqualInNextRow` as true iff the two rows hold equal values for
+   * every tracked-history column.
+   *
+   * Returns true iff `row` and `next` are both upsert-representing, `row`'s interval reaches
+   * `next` without leaving a gap, and the two are tracked-history equal. It is false whenever
+   * there is no successor (the last row in a key window) or either row is not upsert-representing.
+   */
+  private[autocdc] def isNoOpUpsertContinuation(
+      row: Scd2IntervalColumns,
+      next: Scd2IntervalColumns,
+      areTrackedColumnsEqualInNextRow: Column
+  ): Column =
+    isUpsertRepresentingRow(row) &&
+      isUpsertRepresentingRow(next) &&
+      !rowClosesStrictlyBeforeNextRow(row.endAt, next.effectiveRecordStartAt) &&
+      areTrackedColumnsEqualInNextRow
 }