[Kernel] Read the fileSizeHistogram from CRCInfo and propagate it through transaction metrics (delta-io#4403)

<!--
Thanks for sending a pull request!  Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://github.com/delta-io/delta/blob/master/CONTRIBUTING.md
2. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP]
Your PR title ...'.
  3. Be sure to keep the PR description updated to reflect all changes.
  4. Please write your PR title to summarize what this PR proposes.
5. If possible, provide a concise example to reproduce the issue for a
faster review.
6. If applicable, include the corresponding issue number in the PR title
and link it in the body.
-->

#### Which Delta project/connector is this regarding?
<!--
Please add the component selected below to the beginning of the pull
request title
For example: [Spark] Title of my pull request
-->

- [ ] Spark
- [ ] Standalone
- [ ] Flink
- [x] Kernel
- [ ] Other (fill in here)

## Description

Reads the fileSizeHistogram from CRC and add it as part of transaction
metrics. Updates it as the transaction is executed.

Note, this is a little weird since it only adds read support and no
write support, so we only actually have histograms available for new
tables in E2E tests. But delta-io#4328
will add write support and fill out the full feature afterward.
Documented throughout the changes needed when adding write support.

## How was this patch tested?

Updates tests.

## Does this PR introduce _any_ user-facing changes?

No

Author: allisonport-db

kernel/kernel-api/src/main/java/io/delta/kernel/internal/TransactionImpl.java

@@ -257,7 +257,13 @@ public Protocol getProtocol() {
   public TransactionCommitResult commit(Engine engine, CloseableIterable<Row> dataActions)
       throws ConcurrentWriteException {
     checkState(!closed, "Transaction is already attempted to commit. Create a new transaction.");
-    TransactionMetrics transactionMetrics = new TransactionMetrics();
+    // For a new table or when fileSizeHistogram is available in the CRC of the readSnapshot
+    // we update it in the commit. When it is not available we do nothing.
+    TransactionMetrics transactionMetrics =
+        isNewTable
+            ? TransactionMetrics.forNewTable()
+            : TransactionMetrics.withExistingTableFileSizeHistogram(
+                readSnapshot.getCurrentCrcInfo().flatMap(CRCInfo::getFileSizeHistogram));
     try {
       long committedVersion =
           transactionMetrics.totalCommitTimer.time(
@@ -332,6 +338,9 @@ private long commitWithRetry(
             commitAsVersion = rebaseState.getLatestVersion() + 1;
             dataActions = rebaseState.getUpdatedDataActions();
             domainMetadatas = Optional.of(rebaseState.getUpdatedDomainMetadatas());
+            // Action counters may be partially incremented from previous tries, reset the counters
+            // to 0 and drop fileSizeHistogram
+            transactionMetrics.resetActionMetricsForRetry();
           }
         }
         numTries++;
@@ -448,8 +457,6 @@ private long doCommit(
         }
       }
 
-      // Action counters may be partially incremented from previous tries, reset the counters to 0
-      transactionMetrics.resetCounters();
       boolean isAppendOnlyTable = APPEND_ONLY_ENABLED.fromMetadata(metadata);
 
       // Write the staged data to a delta file
@@ -488,17 +495,12 @@ private long doCommit(
   private void incrementMetricsForFileActionRow(TransactionMetrics txnMetrics, Row fileActionRow) {
     txnMetrics.totalActionsCounter.increment();
     if (!fileActionRow.isNullAt(ADD_FILE_ORDINAL)) {
-      txnMetrics.addFilesCounter.increment();
-      txnMetrics.addFilesSizeInBytesCounter.increment(
-          new AddFile(fileActionRow.getStruct(ADD_FILE_ORDINAL)).getSize());
-      // TODO increment fileSizeHistogram
+      txnMetrics.updateForAddFile(new AddFile(fileActionRow.getStruct(ADD_FILE_ORDINAL)).getSize());
     } else if (!fileActionRow.isNullAt(REMOVE_FILE_ORDINAL)) {
-      txnMetrics.removeFilesCounter.increment();
       RemoveFile removeFile = new RemoveFile(fileActionRow.getStruct(REMOVE_FILE_ORDINAL));
-      long fileSize =
+      long removeFileSize =
           removeFile.getSize().orElseThrow(DeltaErrorsInternal::missingRemoveFileSizeDuringCommit);
-      txnMetrics.removeFilesSizeInBytesCounter.increment(fileSize);
-      // TODO decrement fileSizeHistogram
+      txnMetrics.updateForRemoveFile(removeFileSize);
     }
   }
 
@@ -585,14 +587,18 @@ private void recordTransactionReport(
   private Optional<CRCInfo> buildPostCommitCrcInfoIfCurrentCrcAvailable(
       long commitAtVersion, TransactionMetricsResult metricsResult) {
     if (isNewTable) {
+      // We don't need to worry about conflicting transaction here since new tables always commit
+      // metadata (and thus fail any conflicts)
       return Optional.of(
           new CRCInfo(
               commitAtVersion,
               metadata,
               protocol,
               metricsResult.getTotalAddFilesSizeInBytes(),
               metricsResult.getNumAddFiles(),
-              Optional.of(txnId.toString())));
+              Optional.of(txnId.toString()),
+              Optional.empty() // once we support writing CRC populate here
+              ));
     }
 
     return currentCrcInfo
@@ -607,7 +613,9 @@ private Optional<CRCInfo> buildPostCommitCrcInfoIfCurrentCrcAvailable(
                     // TODO: handle RemoveFiles for calculating table size and num of files.
                     lastCrcInfo.getTableSizeBytes() + metricsResult.getTotalAddFilesSizeInBytes(),
                     lastCrcInfo.getNumFiles() + metricsResult.getNumAddFiles(),
-                    Optional.of(txnId.toString())));
+                    Optional.of(txnId.toString()),
+                    Optional.empty() // once we support writing CRC populate here
+                    ));
   }
 
   /**

kernel/kernel-api/src/main/java/io/delta/kernel/internal/checksum/CRCInfo.java

@@ -24,6 +24,7 @@
 import io.delta.kernel.internal.actions.Metadata;
 import io.delta.kernel.internal.actions.Protocol;
 import io.delta.kernel.internal.data.GenericRow;
+import io.delta.kernel.internal.stats.FileSizeHistogram;
 import io.delta.kernel.internal.util.InternalUtils;
 import io.delta.kernel.types.LongType;
 import io.delta.kernel.types.StringType;
@@ -46,6 +47,7 @@ public class CRCInfo {
   private static final String METADATA = "metadata";
   private static final String PROTOCOL = "protocol";
   private static final String TXN_ID = "txnId";
+  private static final String FILE_SIZE_HISTOGRAM = "fileSizeHistogram";
 
   public static final StructType CRC_FILE_SCHEMA =
       new StructType()
@@ -55,7 +57,8 @@ public class CRCInfo {
           .add(NUM_PROTOCOL, LongType.LONG)
           .add(METADATA, Metadata.FULL_SCHEMA)
           .add(PROTOCOL, Protocol.FULL_SCHEMA)
-          .add(TXN_ID, StringType.STRING, /*nullable*/ true);
+          .add(TXN_ID, StringType.STRING, /*nullable*/ true)
+          .add(FILE_SIZE_HISTOGRAM, FileSizeHistogram.FULL_SCHEMA, /*nullable*/ true);
 
   public static Optional<CRCInfo> fromColumnarBatch(
       long version, ColumnarBatch batch, int rowId, String crcFilePath) {
@@ -79,13 +82,18 @@ public static Optional<CRCInfo> fromColumnarBatch(
         txnIdColumnVector.isNullAt(rowId)
             ? Optional.empty()
             : Optional.of(txnIdColumnVector.getString(rowId));
+    Optional<FileSizeHistogram> fileSizeHistogram =
+        FileSizeHistogram.fromColumnVector(
+            batch.getColumnVector(getSchemaIndex(FILE_SIZE_HISTOGRAM)), rowId);
 
     //  protocol and metadata are nullable per fromColumnVector's implementation.
     if (protocol == null || metadata == null) {
       logger.warn("Invalid checksum file missing protocol and/or metadata: {}", crcFilePath);
       return Optional.empty();
     }
-    return Optional.of(new CRCInfo(version, metadata, protocol, tableSizeBytes, numFiles, txnId));
+    return Optional.of(
+        new CRCInfo(
+            version, metadata, protocol, tableSizeBytes, numFiles, txnId, fileSizeHistogram));
   }
 
   private final long version;
@@ -94,14 +102,16 @@ public static Optional<CRCInfo> fromColumnarBatch(
   private final long tableSizeBytes;
   private final long numFiles;
   private final Optional<String> txnId;
+  private final Optional<FileSizeHistogram> fileSizeHistogram;
 
   public CRCInfo(
       long version,
       Metadata metadata,
       Protocol protocol,
       long tableSizeBytes,
       long numFiles,
-      Optional<String> txnId) {
+      Optional<String> txnId,
+      Optional<FileSizeHistogram> fileSizeHistogram) {
     checkArgument(tableSizeBytes >= 0);
     checkArgument(numFiles >= 0);
     this.version = version;
@@ -110,6 +120,7 @@ public CRCInfo(
     this.tableSizeBytes = tableSizeBytes;
     this.numFiles = numFiles;
     this.txnId = requireNonNull(txnId);
+    this.fileSizeHistogram = requireNonNull(fileSizeHistogram);
   }
 
   /** The version of the Delta table that this CRCInfo represents. */
@@ -139,6 +150,11 @@ public Optional<String> getTxnId() {
     return txnId;
   }
 
+  /** The {@link FileSizeHistogram} stored in this CRCInfo. */
+  public Optional<FileSizeHistogram> getFileSizeHistogram() {
+    return fileSizeHistogram;
+  }
+
   /**
    * Encode as a {@link Row} object with the schema {@link CRCInfo#CRC_FILE_SCHEMA}.
    *
@@ -156,12 +172,14 @@ public Row toRow() {
 
     // Add optional fields
     txnId.ifPresent(txn -> values.put(getSchemaIndex(TXN_ID), txn));
+    // TODO write fileSizeHistogram here
     return new GenericRow(CRC_FILE_SCHEMA, values);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(version, metadata, protocol, tableSizeBytes, numFiles, txnId);
+    return Objects.hash(
+        version, metadata, protocol, tableSizeBytes, numFiles, txnId, fileSizeHistogram);
   }
 
   @Override
@@ -175,7 +193,8 @@ public boolean equals(Object o) {
         && numFiles == other.numFiles
         && metadata.equals(other.metadata)
         && protocol.equals(other.protocol)
-        && txnId.equals(other.txnId);
+        && txnId.equals(other.txnId)
+        && fileSizeHistogram.equals(other.fileSizeHistogram);
   }
 
   private static int getSchemaIndex(String fieldName) {

kernel/kernel-api/src/main/java/io/delta/kernel/internal/metrics/TransactionMetrics.java

@@ -15,7 +15,12 @@
  */
 package io.delta.kernel.internal.metrics;
 
+import static io.delta.kernel.internal.util.Preconditions.checkArgument;
+
+import io.delta.kernel.internal.stats.FileSizeHistogram;
+import io.delta.kernel.metrics.FileSizeHistogramResult;
 import io.delta.kernel.metrics.TransactionMetricsResult;
+import java.util.Optional;
 
 /**
  * Stores the metrics for an ongoing transaction. These metrics are updated and recorded throughout
@@ -27,32 +32,92 @@
  */
 public class TransactionMetrics {
 
+  /**
+   * @return a fresh TransactionMetrics object with a default tableFileSizeHistogram (with 0 counts)
+   */
+  public static TransactionMetrics forNewTable() {
+    return new TransactionMetrics(Optional.of(FileSizeHistogram.createDefaultHistogram()));
+  }
+
+  /**
+   * @return a fresh TransactionMetrics object with an initial tableFileSizeHistogram as provided
+   */
+  public static TransactionMetrics withExistingTableFileSizeHistogram(
+      Optional<FileSizeHistogram> tableFileSizeHistogram) {
+    return new TransactionMetrics(tableFileSizeHistogram);
+  }
+
   public final Timer totalCommitTimer = new Timer();
 
   public final Counter commitAttemptsCounter = new Counter();
 
-  public final Counter addFilesCounter = new Counter();
+  private final Counter addFilesCounter = new Counter();
 
-  public final Counter removeFilesCounter = new Counter();
+  private final Counter removeFilesCounter = new Counter();
 
   public final Counter totalActionsCounter = new Counter();
 
-  public final Counter addFilesSizeInBytesCounter = new Counter();
+  private final Counter addFilesSizeInBytesCounter = new Counter();
 
-  public final Counter removeFilesSizeInBytesCounter = new Counter();
+  private final Counter removeFilesSizeInBytesCounter = new Counter();
+
+  private Optional<FileSizeHistogram> tableFileSizeHistogram;
+
+  private TransactionMetrics(Optional<FileSizeHistogram> tableFileSizeHistogram) {
+    this.tableFileSizeHistogram = tableFileSizeHistogram;
+  }
 
   /**
-   * Resets the action counters (addFilesCounter, removeFilesCounter and totalActionsCounter) to 0.
-   * Action counters may be partially incremented if an action iterator is not read to completion
-   * (i.e. if an exception interrupts a file write). This allows us to reset the counters so that we
-   * can increment them correctly from 0 on a retry.
+   * Updates the metrics for a seen AddFile with size {@code addFileSize}. Specifically, updates
+   * addFilesCounter, addFilesSizeInBytesCounter, and tableFileSizeHistogram. Note, it does NOT
+   * increment totalActionsCounter, this needs to be done separately.
+   *
+   * @param addFileSize the size of the add file to update the metrics for
    */
-  public void resetCounters() {
+  public void updateForAddFile(long addFileSize) {
+    checkArgument(addFileSize >= 0, "File size must be non-negative, got %s", addFileSize);
+    addFilesCounter.increment();
+    addFilesSizeInBytesCounter.increment(addFileSize);
+    tableFileSizeHistogram.ifPresent(histogram -> histogram.insert(addFileSize));
+  }
+
+  /**
+   * Updates the metrics for a seen RemoveFile with size {@code removeFileSize}. Specifically,
+   * updates removeFilesCounter, removeFilesSizeInBytesCounter, and tableFileSizeHistogram. Note, it
+   * does NOT increment totalActionsCounter, this needs to be done separately.
+   *
+   * @param removeFileSize the size of the remove file to update the metrics for
+   */
+  public void updateForRemoveFile(long removeFileSize) {
+    checkArgument(removeFileSize >= 0, "File size must be non-negative, got %s", removeFileSize);
+    removeFilesCounter.increment();
+    removeFilesSizeInBytesCounter.increment(removeFileSize);
+    tableFileSizeHistogram.ifPresent(histogram -> histogram.remove(removeFileSize));
+  }
+
+  /**
+   * Resets any action metrics for a failed commit to prepare them for retrying. Specifically,
+   *
+   * <ul>
+   *   <li>Resets addFilesCounter, removeFilesCounter, totalActionsCounter,
+   *       addFilesSizeInBytesCounter, and removeFilesSizeInBytesCounter to 0
+   *   <li>Sets tableFileSizeHistogram to be empty since we don't know the updated distribution
+   *       after the conflicting txn committed
+   * </ul>
+   *
+   * Action counters / tableFileSizeHistogram may be partially incremented if an action iterator is
+   * not read to completion (i.e. if an exception interrupts a file write). This allows us to reset
+   * the counters so that we can increment them correctly from 0 on a retry.
+   */
+  public void resetActionMetricsForRetry() {
     addFilesCounter.reset();
     addFilesSizeInBytesCounter.reset();
     removeFilesCounter.reset();
     totalActionsCounter.reset();
     removeFilesSizeInBytesCounter.reset();
+    // For now, on retry we set tableFileSizeHistogram = Optional.empty() because we don't know the
+    // correct state of tableFileSizeHistogram after conflicting transaction has committed
+    tableFileSizeHistogram = Optional.empty();
   }
 
   public TransactionMetricsResult captureTransactionMetricsResult() {
@@ -65,6 +130,8 @@ public TransactionMetricsResult captureTransactionMetricsResult() {
       final long numRemoveFiles = removeFilesCounter.value();
       final long numTotalActions = totalActionsCounter.value();
       final long totalRemoveFileSizeInBytes = removeFilesSizeInBytesCounter.value();
+      final Optional<FileSizeHistogramResult> tableFileSizeHistogramResult =
+          tableFileSizeHistogram.map(FileSizeHistogram::captureFileSizeHistogramResult);
 
       @Override
       public long getTotalCommitDurationNs() {
@@ -100,6 +167,11 @@ public long getTotalAddFilesSizeInBytes() {
       public long getTotalRemoveFilesSizeInBytes() {
         return totalRemoveFileSizeInBytes;
       }
+
+      @Override
+      public Optional<FileSizeHistogramResult> getTableFileSizeHistogram() {
+        return tableFileSizeHistogramResult;
+      }
     };
   }
 
@@ -108,13 +180,14 @@ public String toString() {
     return String.format(
         "TransactionMetrics(totalCommitTimer=%s, commitAttemptsCounter=%s, addFilesCounter=%s, "
             + "removeFilesCounter=%s, totalActionsCounter=%s, totalAddFilesSizeInBytes=%s,"
-            + "totalRemoveFilesSizeInBytes=%s)",
+            + "totalRemoveFilesSizeInBytes=%s, tableFileSizeHistogram=%s)",
         totalCommitTimer,
         commitAttemptsCounter,
         addFilesCounter,
         removeFilesCounter,
         totalActionsCounter,
         addFilesSizeInBytesCounter,
-        removeFilesSizeInBytesCounter);
+        removeFilesSizeInBytesCounter,
+        tableFileSizeHistogram);
   }
 }

kernel/kernel-api/src/main/java/io/delta/kernel/internal/stats/FileSizeHistogram.java

@@ -23,6 +23,7 @@
 import io.delta.kernel.internal.data.GenericRow;
 import io.delta.kernel.internal.util.InternalUtils;
 import io.delta.kernel.internal.util.VectorUtils;
+import io.delta.kernel.metrics.FileSizeHistogramResult;
 import io.delta.kernel.types.ArrayType;
 import io.delta.kernel.types.LongType;
 import io.delta.kernel.types.StructType;
@@ -159,6 +160,8 @@ private static long[] createDefaultBinBoundaries() {
     return boundaries;
   }
 
+  // TODO factory from FileSizeHistogramResult
+
   ////////////////////////////////////
   // Member variables and methods  //
   ////////////////////////////////////
@@ -257,6 +260,30 @@ public Row toRow() {
     return new GenericRow(FULL_SCHEMA, value);
   }
 
+  public FileSizeHistogramResult captureFileSizeHistogramResult() {
+    return new FileSizeHistogramResult() {
+      final long[] copiedSortedBinBoundaries =
+          Arrays.copyOf(sortedBinBoundaries, sortedBinBoundaries.length);
+      final long[] copiedFileCounts = Arrays.copyOf(fileCounts, fileCounts.length);
+      final long[] copiedTotalBytes = Arrays.copyOf(totalBytes, totalBytes.length);
+
+      @Override
+      public long[] getSortedBinBoundaries() {
+        return copiedSortedBinBoundaries;
+      }
+
+      @Override
+      public long[] getFileCounts() {
+        return copiedFileCounts;
+      }
+
+      @Override
+      public long[] getTotalBytes() {
+        return copiedTotalBytes;
+      }
+    };
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;