[Kernel] Assign base row ID and default row commit version to AddFile (#3894)

qiyuandong-db · web-flow · commit 837cf22b5452 · 2025-01-08T08:56:04.000-08:00
#### Which Delta project/connector is this regarding?  - [ ] Spark - [ ] Standalone - [ ] Flink - [x] Kernel - [ ] Other (fill in here) ## Description  This PR implements the row tracking _support_ requirements in Delta Kernel, according to the [Delta Protocol](https://github.com/delta-io/delta/blob/master/PROTOCOL.md#row-tracking). Specifically, it includes: - assigning `baseRowId` and `defaultRowCommitVersion` to `AddFile` actions prior to committing them - maintaining the `rowIdHighWaterMark` of the `delta.rowTracking` metadata domain during the base row ID assignment, which is the highest assigned fresh row id for the table ## How was this patch tested?  Added tests in `RowTrackingSuite.scala`. This includes unit tests and integration tests with Delta-Spark. ## Does this PR introduce _any_ user-facing changes?  No.
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/DeltaErrors.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/DeltaErrors.java
@@ -288,6 +288,19 @@ public static ConcurrentWriteException concurrentDomainMetadataAction(
     return new ConcurrentWriteException(message);
   }
 
+  public static KernelException missingNumRecordsStatsForRowTracking() {
+    return new KernelException(
+        "Cannot write to a rowTracking-supported table without 'numRecords' statistics. "
+            + "Connectors are expected to populate the number of records statistics when "
+            + "writing to a Delta table with 'rowTracking' table feature supported.");
+  }
+
+  public static KernelException rowTrackingSupportedWithDomainMetadataUnsupported() {
+    return new KernelException(
+        "Feature 'rowTracking' is supported and depends on feature 'domainMetadata',"
+            + " but 'domainMetadata' is unsupported");
+  }
+
   /* ------------------------ HELPER METHODS ----------------------------- */
   private static String formatTimestamp(long millisSinceEpochUTC) {
     return new Timestamp(millisSinceEpochUTC).toInstant().toString();
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/InternalScanFileUtils.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/InternalScanFileUtils.java
@@ -32,6 +32,7 @@
 import java.net.URI;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * Utilities to extract information out of the scan file rows returned by {@link
@@ -42,6 +43,7 @@ private InternalScanFileUtils() {}
 
   private static final String TABLE_ROOT_COL_NAME = "tableRoot";
   private static final DataType TABLE_ROOT_DATA_TYPE = StringType.STRING;
+
   /** {@link Column} expression referring to the `partitionValues` in scan `add` file. */
   public static final Column ADD_FILE_PARTITION_COL_REF =
       new Column(new String[] {"add", "partitionValues"});
@@ -190,4 +192,14 @@ public static DeletionVectorDescriptor getDeletionVectorDescriptorFromRow(Row sc
   public static Column getPartitionValuesParsedRefInAddFile(String partitionColName) {
     return new Column(new String[] {"add", "partitionValues_parsed", partitionColName});
   }
+
+  public static Optional<Long> getBaseRowId(Row scanFile) {
+    Row addFileRow = getAddFileEntry(scanFile);
+    return new AddFile(addFileRow).getBaseRowId();
+  }
+
+  public static Optional<Long> getDefaultRowCommitVersion(Row scanFile) {
+    Row addFileRow = getAddFileEntry(scanFile);
+    return new AddFile(addFileRow).getDefaultRowCommitVersion();
+  }
 }
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/TableFeatures.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/TableFeatures.java
@@ -40,6 +40,7 @@ public class TableFeatures {
               add("typeWidening-preview");
               add("typeWidening");
               add(DOMAIN_METADATA_FEATURE_NAME);
+              add(ROW_TRACKING_FEATURE_NAME);
             }
           });
 
@@ -59,9 +60,10 @@ public class TableFeatures {
             }
           });
 
-  /** The feature name for domain metadata. */
   public static final String DOMAIN_METADATA_FEATURE_NAME = "domainMetadata";
 
+  public static final String ROW_TRACKING_FEATURE_NAME = "rowTracking";
+
   /** The minimum writer version required to support table features. */
   public static final int TABLE_FEATURES_MIN_WRITER_VERSION = 7;
 
@@ -101,7 +103,8 @@ public static void validateReadSupportedTable(
    *   <li>protocol writer version 1.
    *   <li>protocol writer version 2 only with appendOnly feature enabled.
    *   <li>protocol writer version 7 with {@code appendOnly}, {@code inCommitTimestamp}, {@code
-   *       columnMapping}, {@code typeWidening}, {@code domainMetadata} feature enabled.
+   *       columnMapping}, {@code typeWidening}, {@code domainMetadata}, {@code rowTracking} feature
+   *       enabled.
    * </ul>
    *
    * @param protocol Table protocol
@@ -137,6 +140,12 @@ public static void validateWriteSupportedTable(
             throw unsupportedWriterFeature(tablePath, writerFeature);
           }
         }
+        // Eventually we may have a way to declare and enforce dependencies between features.
+        // By putting this check for row tracking here, it makes it easier to spot that row
+        // tracking defines such a dependency that can be implicitly checked.
+        if (isRowTrackingSupported(protocol) && !isDomainMetadataSupported(protocol)) {
+          throw DeltaErrors.rowTrackingSupportedWithDomainMetadataUnsupported();
+        }
         break;
       default:
         throw unsupportedWriterProtocol(tablePath, minWriterVersion);
@@ -190,12 +199,17 @@ public static Set<String> extractAutomaticallyEnabledWriterFeatures(
    * @return true if the "domainMetadata" feature is supported, false otherwise
    */
   public static boolean isDomainMetadataSupported(Protocol protocol) {
-    List<String> writerFeatures = protocol.getWriterFeatures();
-    if (writerFeatures == null) {
-      return false;
-    }
-    return writerFeatures.contains(DOMAIN_METADATA_FEATURE_NAME)
-        && protocol.getMinWriterVersion() >= TABLE_FEATURES_MIN_WRITER_VERSION;
+    return isWriterFeatureSupported(protocol, DOMAIN_METADATA_FEATURE_NAME);
+  }
+
+  /**
+   * Check if the table protocol supports the "rowTracking" writer feature.
+   *
+   * @param protocol the protocol to check
+   * @return true if the protocol supports row tracking, false otherwise
+   */
+  public static boolean isRowTrackingSupported(Protocol protocol) {
+    return isWriterFeatureSupported(protocol, ROW_TRACKING_FEATURE_NAME);
   }
 
   /**
@@ -254,4 +268,13 @@ private static void validateNoInvariants(StructType tableSchema) {
       throw columnInvariantsNotSupported();
     }
   }
+
+  private static boolean isWriterFeatureSupported(Protocol protocol, String featureName) {
+    List<String> writerFeatures = protocol.getWriterFeatures();
+    if (writerFeatures == null) {
+      return false;
+    }
+    return writerFeatures.contains(featureName)
+        && protocol.getMinWriterVersion() >= TABLE_FEATURES_MIN_WRITER_VERSION;
+  }
 }
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/TransactionImpl.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/TransactionImpl.java
@@ -32,6 +32,7 @@
 import io.delta.kernel.internal.fs.Path;
 import io.delta.kernel.internal.replay.ConflictChecker;
 import io.delta.kernel.internal.replay.ConflictChecker.TransactionRebaseState;
+import io.delta.kernel.internal.rowtracking.RowTracking;
 import io.delta.kernel.internal.util.*;
 import io.delta.kernel.types.StructType;
 import io.delta.kernel.utils.CloseableIterable;
@@ -145,6 +146,18 @@ public TransactionCommitResult commit(Engine engine, CloseableIterable<Row> data
       CommitInfo attemptCommitInfo = generateCommitAction(engine);
       updateMetadataWithICTIfRequired(
           engine, attemptCommitInfo.getInCommitTimestamp(), readSnapshot.getVersion(engine));
+
+      // If row tracking is supported, assign base row IDs and default row commit versions to any
+      // AddFile actions that do not yet have them. If the row ID high watermark changes, emit a
+      // DomainMetadata action to update it.
+      if (TableFeatures.isRowTrackingSupported(protocol)) {
+        RowTracking.createNewHighWaterMarkIfNeeded(readSnapshot, dataActions)
+            .ifPresent(domainMetadatas::add);
+        dataActions =
+            RowTracking.assignBaseRowIdAndDefaultRowCommitVersion(
+                readSnapshot, commitAsVersion, dataActions);
+      }
+
       int numRetries = 0;
       do {
         logger.info("Committing transaction as version = {}.", commitAsVersion);
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/AddFile.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/AddFile.java
@@ -136,7 +136,7 @@ public static Row createAddFileRow(
 
   /** Constructs an {@link AddFile} action from the given 'AddFile' {@link Row}. */
   public AddFile(Row row) {
-    super(row, FULL_SCHEMA);
+    super(row);
   }
 
   public String getPath() {
@@ -181,10 +181,12 @@ public Optional<Long> getDefaultRowCommitVersion() {
   }
 
   public Optional<DataFileStatistics> getStats() {
-    int index = getFieldIndex("stats");
-    return row.isNullAt(index)
-        ? Optional.empty()
-        : DataFileStatistics.deserializeFromJson(row.getString(index));
+    return getFieldIndexOpt("stats")
+        .flatMap(
+            index ->
+                row.isNullAt(index)
+                    ? Optional.empty()
+                    : DataFileStatistics.deserializeFromJson(row.getString(index)));
   }
 
   public Optional<Long> getNumRecords() {
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/RowBackedAction.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/RowBackedAction.java
@@ -22,6 +22,7 @@
 import io.delta.kernel.types.*;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * An abstract base class for Delta Log actions that are backed by a {@link Row}. This design is to
@@ -34,26 +35,30 @@ public abstract class RowBackedAction {
   /** The underlying {@link Row} that represents an action and contains all its field values. */
   protected final Row row;
 
-  protected RowBackedAction(Row row, StructType expectedSchema) {
-    checkArgument(
-        row.getSchema().equals(expectedSchema),
-        "Expected row schema: %s, found: %s",
-        expectedSchema,
-        row.getSchema());
-
+  protected RowBackedAction(Row row) {
     this.row = row;
   }
 
   /**
-   * Returns the index of the field with the given name in the full schema of the row. Throws an
-   * {@link IllegalArgumentException} if the field is not found.
+   * Returns the index of the field with the given name in the schema of the row. Throws an {@link
+   * IllegalArgumentException} if the field is not found.
    */
   protected int getFieldIndex(String fieldName) {
     int index = row.getSchema().indexOf(fieldName);
     checkArgument(index >= 0, "Field '%s' not found in schema: %s", fieldName, row.getSchema());
     return index;
   }
 
+  /**
+   * Returns the index of the field with the given name in the schema of the row, or {@link
+   * Optional#empty()} if the field is not found. This should be used when the underlying row may or
+   * may not contain that field.
+   */
+  protected Optional<Integer> getFieldIndexOpt(String fieldName) {
+    int index = row.getSchema().indexOf(fieldName);
+    return index >= 0 ? Optional.of(index) : Optional.empty();
+  }
+
   /**
    * Returns a new {@link Row} with the same schema and values as the row backing this action, but
    * with the value of the field with the given name overridden by the given value.
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/SingleAction.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/actions/SingleAction.java
@@ -72,7 +72,7 @@ public class SingleAction {
   // schema for those fields here.
 
   private static final int TXN_ORDINAL = FULL_SCHEMA.indexOf("txn");
-  private static final int ADD_FILE_ORDINAL = FULL_SCHEMA.indexOf("add");
+  public static final int ADD_FILE_ORDINAL = FULL_SCHEMA.indexOf("add");
   private static final int REMOVE_FILE_ORDINAL = FULL_SCHEMA.indexOf("remove");
   private static final int METADATA_ORDINAL = FULL_SCHEMA.indexOf("metaData");
   private static final int PROTOCOL_ORDINAL = FULL_SCHEMA.indexOf("protocol");
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/rowtracking/RowTracking.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/rowtracking/RowTracking.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright (2024) The Delta Lake Project Authors.
+ *
+ * Licensed 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 io.delta.kernel.internal.rowtracking;
+
+import static io.delta.kernel.internal.util.Preconditions.checkArgument;
+
+import io.delta.kernel.data.Row;
+import io.delta.kernel.internal.DeltaErrors;
+import io.delta.kernel.internal.SnapshotImpl;
+import io.delta.kernel.internal.TableFeatures;
+import io.delta.kernel.internal.actions.*;
+import io.delta.kernel.utils.CloseableIterable;
+import io.delta.kernel.utils.CloseableIterator;
+import java.io.IOException;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicLong;
+
+/** A collection of helper methods for working with row tracking. */
+public class RowTracking {
+  private RowTracking() {
+    // Empty constructor to prevent instantiation of this class
+  }
+
+  /**
+   * Assigns base row IDs and default row commit versions to {@link AddFile} actions in the provided
+   * {@code dataActions}. This method should be called when processing data actions during commit
+   * preparation and before the actual commit.
+   *
+   * <p>This method should be called exactly once per transaction during commit preparation, i.e.,
+   * not for each commit attempt. And it should only be called when the 'rowTracking' feature is
+   * supported.
+   *
+   * <p>For {@link AddFile} actions missing a base row ID, assigns the current row ID high watermark
+   * plus 1. The high watermark is then incremented by the number of records in the file. For
+   * actions missing a default row commit version, assigns the specified commit version.
+   *
+   * @param snapshot the snapshot of the table that this transaction is reading at
+   * @param commitVersion the version of the commit for default row commit version assignment
+   * @param dataActions the {@link CloseableIterable} of data actions to process
+   * @return an {@link CloseableIterable} of data actions with base row IDs and default row commit
+   *     versions assigned
+   */
+  public static CloseableIterable<Row> assignBaseRowIdAndDefaultRowCommitVersion(
+      SnapshotImpl snapshot, long commitVersion, CloseableIterable<Row> dataActions) {
+    checkArgument(
+        TableFeatures.isRowTrackingSupported(snapshot.getProtocol()),
+        "Base row ID and default row commit version are assigned "
+            + "only when feature 'rowTracking' is supported.");
+
+    return new CloseableIterable<Row>() {
+      @Override
+      public void close() throws IOException {
+        dataActions.close();
+      }
+
+      @Override
+      public CloseableIterator<Row> iterator() {
+        // Used to keep track of the current high watermark as we iterate through the data actions.
+        // Use an AtomicLong to allow for updating the high watermark in the lambda.
+        final AtomicLong currRowIdHighWatermark = new AtomicLong(readRowIdHighWaterMark(snapshot));
+        return dataActions
+            .iterator()
+            .map(
+                row -> {
+                  // Non-AddFile actions are returned unchanged
+                  if (row.isNullAt(SingleAction.ADD_FILE_ORDINAL)) {
+                    return row;
+                  }
+
+                  AddFile addFile = new AddFile(row.getStruct(SingleAction.ADD_FILE_ORDINAL));
+
+                  // Assign base row ID if missing
+                  if (!addFile.getBaseRowId().isPresent()) {
+                    final long numRecords = getNumRecordsOrThrow(addFile);
+                    addFile = addFile.withNewBaseRowId(currRowIdHighWatermark.get() + 1L);
+                    currRowIdHighWatermark.addAndGet(numRecords);
+                  }
+
+                  // Assign default row commit version if missing
+                  if (!addFile.getDefaultRowCommitVersion().isPresent()) {
+                    addFile = addFile.withNewDefaultRowCommitVersion(commitVersion);
+                  }
+
+                  return SingleAction.createAddFileSingleAction(addFile.toRow());
+                });
+      }
+    };
+  }
+
+  /**
+   * Returns a {@link DomainMetadata} action if the row ID high watermark has changed due to newly
+   * processed {@link AddFile} actions.
+   *
+   * <p>This method should be called during commit preparation to prepare the domain metadata
+   * actions for commit. It should be called only when the 'rowTracking' feature is supported.
+   *
+   * @param snapshot the snapshot of the table that this transaction is reading at
+   * @param dataActions the iterable of data actions that may update the high watermark
+   */
+  public static Optional<DomainMetadata> createNewHighWaterMarkIfNeeded(
+      SnapshotImpl snapshot, CloseableIterable<Row> dataActions) {
+    checkArgument(
+        TableFeatures.isRowTrackingSupported(snapshot.getProtocol()),
+        "Row ID high watermark is updated only when feature 'rowTracking' is supported.");
+
+    final long prevRowIdHighWatermark = readRowIdHighWaterMark(snapshot);
+    // Use an AtomicLong to allow for updating the high watermark in the lambda
+    final AtomicLong newRowIdHighWatermark = new AtomicLong(prevRowIdHighWatermark);
+
+    dataActions.forEach(
+        row -> {
+          if (!row.isNullAt(SingleAction.ADD_FILE_ORDINAL)) {
+            AddFile addFile = new AddFile(row.getStruct(SingleAction.ADD_FILE_ORDINAL));
+            if (!addFile.getBaseRowId().isPresent()) {
+              newRowIdHighWatermark.addAndGet(getNumRecordsOrThrow(addFile));
+            }
+          }
+        });
+
+    return (newRowIdHighWatermark.get() != prevRowIdHighWatermark)
+        ? Optional.of(new RowTrackingMetadataDomain(newRowIdHighWatermark.get()).toDomainMetadata())
+        : Optional.empty();
+  }
+
+  /**
+   * Reads the current row ID high watermark from the snapshot, or returns a default value if
+   * missing.
+   */
+  private static long readRowIdHighWaterMark(SnapshotImpl snapshot) {
+    return RowTrackingMetadataDomain.fromSnapshot(snapshot)
+        .map(RowTrackingMetadataDomain::getRowIdHighWaterMark)
+        .orElse(RowTrackingMetadataDomain.MISSING_ROW_ID_HIGH_WATERMARK);
+  }
+
+  /**
+   * Get the number of records from the AddFile's statistics. It errors out if statistics are
+   * missing.
+   */
+  private static long getNumRecordsOrThrow(AddFile addFile) {
+    return addFile.getNumRecords().orElseThrow(DeltaErrors::missingNumRecordsStatsForRowTracking);
+  }
+}
diff --git a/kernel/kernel-api/src/main/java/io/delta/kernel/internal/rowtracking/RowTrackingMetadataDomain.java b/kernel/kernel-api/src/main/java/io/delta/kernel/internal/rowtracking/RowTrackingMetadataDomain.java
diff --git a/kernel/kernel-api/src/test/scala/io/delta/kernel/internal/TableFeaturesSuite.scala b/kernel/kernel-api/src/test/scala/io/delta/kernel/internal/TableFeaturesSuite.scala
diff --git a/kernel/kernel-defaults/src/test/scala/io/delta/kernel/defaults/RowTrackingSuite.scala b/kernel/kernel-defaults/src/test/scala/io/delta/kernel/defaults/RowTrackingSuite.scala
