fix(utilities): Re-inject _corrupt_record in ErrorTableAwareChainedTransformer after each transformer

hudi-utilities/src/main/java/org/apache/hudi/utilities/streamer/ErrorTableUtils.java

@@ -27,13 +27,18 @@
 import org.apache.hudi.common.util.ValidationUtils;
 import org.apache.hudi.config.HoodieErrorTableConfig;
 import org.apache.hudi.exception.HoodieException;
-import org.apache.hudi.exception.HoodieValidationException;
 import org.apache.hudi.utilities.ingestion.HoodieIngestionMetrics;
 
 import org.apache.hadoop.fs.FileSystem;
+import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Arrays;
 
@@ -43,6 +48,7 @@
 import static org.apache.spark.sql.functions.lit;
 
 public final class ErrorTableUtils {
+  private static final Logger LOG = LoggerFactory.getLogger(ErrorTableUtils.class);
   public static Option<BaseErrorTableWriter> getErrorTableWriter(HoodieStreamer.Config cfg,
                                                                  SparkSession sparkSession,
                                                                  TypedProperties props,
@@ -79,25 +85,61 @@ public static HoodieErrorTableConfig.ErrorWriteFailureStrategy getErrorWriteFail
     return HoodieErrorTableConfig.ErrorWriteFailureStrategy.valueOf(writeFailureStrategy);
   }
 
-  /**
-   * validates for constraints on ErrorRecordColumn when ErrorTable enabled configs are set.
-   * @param dataset
-   */
-  public static void validate(Dataset<Row> dataset) {
-    if (!isErrorTableCorruptRecordColumnPresent(dataset)) {
-      throw new HoodieValidationException(String.format("Invalid condition, columnName=%s "
-              + "is not present in transformer " + "output schema", ERROR_TABLE_CURRUPT_RECORD_COL_NAME));
-    }
-  }
-
   public static Dataset<Row> addNullValueErrorTableCorruptRecordColumn(Dataset<Row> dataset) {
     if (!isErrorTableCorruptRecordColumnPresent(dataset)) {
       dataset = dataset.withColumn(ERROR_TABLE_CURRUPT_RECORD_COL_NAME, lit(null));
     }
     return dataset;
   }
 
-  private static boolean isErrorTableCorruptRecordColumnPresent(Dataset<Row> dataset) {
+  public static boolean isErrorTableCorruptRecordColumnPresent(Dataset<Row> dataset) {
     return Arrays.stream(dataset.columns()).anyMatch(col -> col.equals(ERROR_TABLE_CURRUPT_RECORD_COL_NAME));
   }
+
+  /**
+   * Restores stashed {@code _corrupt_record} values onto a transformed dataset via positional
+   * RDD zip. Works when the transformer only projects columns (row count and partition layout
+   * unchanged). Falls back to {@link #addNullValueErrorTableCorruptRecordColumn} with a WARN
+   * if the zip fails (e.g. because the transformer filtered or repartitioned rows).
+   *
+   * <p><b>Limitation:</b> a transformer that reorders rows (e.g. {@code orderBy}) without
+   * changing the count will produce a successful zip with misaligned values. This is no worse
+   * than the null-re-injection alternative, but callers should be aware.
+   *
+   * @param sparkSession the active Spark session
+   * @param transformed  the dataset after the transformer ran (missing {@code _corrupt_record})
+   * @param stash        single-column dataset of pre-transform {@code _corrupt_record} values
+   * @return {@code transformed} with the {@code _corrupt_record} column restored
+   */
+  public static Dataset<Row> restoreCorruptRecordColumn(
+      SparkSession sparkSession, Dataset<Row> transformed, Dataset<Row> stash) {
+    try {
+      JavaRDD<Row> zipped = transformed.javaRDD()
+          .zip(stash.javaRDD())
+          .map(pair -> {
+            Row dataRow = pair._1();
+            Row stashRow = pair._2();
+            int size = dataRow.size();
+            Object[] fields = new Object[size + 1];
+            for (int i = 0; i < size; i++) {
+              fields[i] = dataRow.get(i);
+            }
+            fields[size] = stashRow.isNullAt(0) ? null : stashRow.get(0);
+            return RowFactory.create(fields);
+          });
+      StructType schema = transformed.schema()
+          .add(ERROR_TABLE_CURRUPT_RECORD_COL_NAME, DataTypes.StringType, true);
+      Dataset<Row> result = sparkSession.createDataFrame(zipped, schema);
+      // Force eager evaluation so zip failures (partition/element count mismatches)
+      // are caught here instead of surfacing downstream. Without this, the entire
+      // zip/map/createDataFrame chain is lazy and the catch block never triggers.
+      result.count();
+      return result;
+    } catch (Exception e) {
+      LOG.warn("Failed to restore {} column after transformer dropped it "
+          + "(row count or partitioning changed): {}",
+          ERROR_TABLE_CURRUPT_RECORD_COL_NAME, e.getMessage());
+      return addNullValueErrorTableCorruptRecordColumn(transformed);
+    }
+  }
 }

hudi-utilities/src/main/java/org/apache/hudi/utilities/transform/ErrorTableAwareChainedTransformer.java

@@ -25,6 +25,7 @@
 import org.apache.hudi.utilities.streamer.ErrorTableUtils;
 
 import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.Column;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
@@ -33,10 +34,14 @@
 import java.util.List;
 import java.util.function.Supplier;
 
+import static org.apache.hudi.utilities.streamer.BaseErrorTableWriter.ERROR_TABLE_CURRUPT_RECORD_COL_NAME;
+
 /**
  * A {@link Transformer} to chain other {@link Transformer}s and apply sequentially.
- * Adds errorTableCorruptRecordColumn at the beginning of transformations and validates
- * if that column is not dropped in any of the transformations.
+ * Adds errorTableCorruptRecordColumn at the beginning of transformations and preserves
+ * its values across transformers that drop it (e.g. custom column-projecting transformers).
+ * Values are stashed before each transformer and restored via positional RDD zip if the
+ * transformer dropped the column.
  */
 public class ErrorTableAwareChainedTransformer extends ChainedTransformer {
   public ErrorTableAwareChainedTransformer(List<String> configuredTransformers, Supplier<Option<HoodieSchema>> sourceSchemaSupplier) {
@@ -54,9 +59,29 @@ public Dataset<Row> apply(JavaSparkContext jsc, SparkSession sparkSession, Datas
     dataset = ErrorTableUtils.addNullValueErrorTableCorruptRecordColumn(dataset);
     for (TransformerInfo transformerInfo : transformers) {
       Transformer transformer = transformerInfo.getTransformer();
+
+      // Stash _corrupt_record values before the transformer can drop them
+      Dataset<Row> corruptRecordStash = null;
+      if (ErrorTableUtils.isErrorTableCorruptRecordColumnPresent(dataset)) {
+        corruptRecordStash = dataset.select(new Column(ERROR_TABLE_CURRUPT_RECORD_COL_NAME));
+        corruptRecordStash.cache();
+      }
+
       dataset = transformer.apply(jsc, sparkSession, dataset, transformerInfo.getProperties(properties, transformers));
-      // validate in every stage to ensure ErrorRecordColumn not dropped by one of the transformer and added by next transformer.
-      ErrorTableUtils.validate(dataset);
+
+      if (!ErrorTableUtils.isErrorTableCorruptRecordColumnPresent(dataset)) {
+        if (corruptRecordStash != null) {
+          // Restore original values via positional zip — works when the transformer
+          // only projects columns (row count and partition layout unchanged).
+          dataset = ErrorTableUtils.restoreCorruptRecordColumn(sparkSession, dataset, corruptRecordStash);
+        } else {
+          dataset = ErrorTableUtils.addNullValueErrorTableCorruptRecordColumn(dataset);
+        }
+      }
+
+      if (corruptRecordStash != null) {
+        corruptRecordStash.unpersist();
+      }
     }
     return dataset;
   }

hudi-utilities/src/test/java/org/apache/hudi/utilities/functional/TestErrorTableAwareChainedTransformer.java

@@ -21,7 +21,6 @@
 
 import org.apache.hudi.common.config.TypedProperties;
 import org.apache.hudi.common.util.Option;
-import org.apache.hudi.exception.HoodieValidationException;
 import org.apache.hudi.testutils.SparkClientFunctionalTestHarness;
 import org.apache.hudi.utilities.exception.HoodieTransformException;
 import org.apache.hudi.utilities.transform.ErrorTableAwareChainedTransformer;
@@ -76,17 +75,62 @@ public void testForErrorTableConfig() {
   }
 
   @Test
-  public void testForErrorRecordColumn() {
+  void testCorruptRecordPreservedAfterTransformerDropsIt() {
+    // Regression for ENG-41958: t1 populates _corrupt_record with "true", t2 drops the
+    // column via select("foo"). The chain must preserve the populated values, not re-inject
+    // as null — otherwise error rows silently flow into the main write path instead of the
+    // error table.
     Dataset<Row> original = getTestDataset();
 
     Transformer t1 = getErrorEventHandlerTransformer();
     Transformer t2 = getErrorRecordColumnDropTransformer();
-    Transformer t3 = (jsc, sparkSession, dataset, properties) -> dataset.withColumn("foo",
+    Transformer t3 = (jsc, sparkSession, dataset, props) -> dataset.withColumn("foo",
         dataset.col("foo").cast(IntegerType));
     TypedProperties properties = new TypedProperties();
     properties.setProperty(ERROR_TABLE_ENABLED.key(), "true");
     ErrorTableAwareChainedTransformer transformer = new ErrorTableAwareChainedTransformer(Arrays.asList(t1, t2, t3));
-    assertThrows(HoodieValidationException.class, () -> transformer.apply(jsc(), spark(), original, properties));
+    Dataset<Row> transformed = transformer.apply(jsc(), spark(), original, properties);
+
+    assertArrayEquals(new String[]{"foo", ERROR_TABLE_CURRUPT_RECORD_COL_NAME}, transformed.columns());
+    assertEquals(2, transformed.count());
+    // Values populated by t1 must survive t2 dropping the column
+    assertEquals(0, transformed.filter(new Column(ERROR_TABLE_CURRUPT_RECORD_COL_NAME).isNull()).count());
+    assertEquals(2, transformed.filter(new Column(ERROR_TABLE_CURRUPT_RECORD_COL_NAME).equalTo("true")).count());
+  }
+
+  @Test
+  void testCustomColumnProjectionPreservesCorruptRecord() {
+    // Single custom transformer doing column projection — drops _corrupt_record.
+    Dataset<Row> original = getTestDataset();
+
+    Transformer columnFilter = (jsc, sparkSession, dataset, props) -> dataset.select("foo");
+    TypedProperties properties = new TypedProperties();
+    properties.setProperty(ERROR_TABLE_ENABLED.key(), "true");
+    ErrorTableAwareChainedTransformer transformer =
+        new ErrorTableAwareChainedTransformer(Arrays.asList(columnFilter));
+    Dataset<Row> transformed = transformer.apply(jsc(), spark(), original, properties);
+
+    assertArrayEquals(new String[]{"foo", ERROR_TABLE_CURRUPT_RECORD_COL_NAME}, transformed.columns());
+    assertEquals(2, transformed.count());
+    assertEquals(2, transformed.filter(new Column(ERROR_TABLE_CURRUPT_RECORD_COL_NAME).isNull()).count());
+  }
+
+  @Test
+  void testTransformerPreservingCorruptRecordIsNoOp() {
+    // Transformer that keeps all columns — re-injection is a no-op.
+    Dataset<Row> original = getTestDataset();
+
+    Transformer keepAll = (jsc, sparkSession, dataset, props) -> dataset.withColumn("foo",
+        dataset.col("foo").cast(IntegerType));
+    TypedProperties properties = new TypedProperties();
+    properties.setProperty(ERROR_TABLE_ENABLED.key(), "true");
+    ErrorTableAwareChainedTransformer transformer =
+        new ErrorTableAwareChainedTransformer(Arrays.asList(keepAll));
+    Dataset<Row> transformed = transformer.apply(jsc(), spark(), original, properties);
+
+    assertArrayEquals(new String[]{"foo", ERROR_TABLE_CURRUPT_RECORD_COL_NAME}, transformed.columns());
+    assertEquals(2, transformed.count());
+    assertEquals(2, transformed.filter(new Column(ERROR_TABLE_CURRUPT_RECORD_COL_NAME).isNull()).count());
   }
 
   private Dataset<Row> getTestDataset() {