API, Core, Spark 4.1: Add key_metadata and encrypted write support for StatisticsFile

api/src/main/java/org/apache/iceberg/StatisticsFile.java

@@ -18,6 +18,7 @@
  */
 package org.apache.iceberg;
 
+import java.nio.ByteBuffer;
 import java.util.List;
 
 /**
@@ -42,6 +43,13 @@ public interface StatisticsFile {
   /** Size of the Puffin footer. */
   long fileFooterSizeInBytes();
 
+  /**
+   * Returns metadata about how this file is encrypted, or null if the file is stored in plain text.
+   */
+  default ByteBuffer keyMetadata() {
+    return null;
+  }
+
   /** List of statistics contained in the file. Never null. */
   List<BlobMetadata> blobMetadata();
 }

core/src/main/java/org/apache/iceberg/GenericStatisticsFile.java

@@ -18,31 +18,51 @@
  */
 package org.apache.iceberg;
 
+import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.StringJoiner;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.util.ByteBuffers;
 
 public class GenericStatisticsFile implements StatisticsFile {
   private final long snapshotId;
   private final String path;
   private final long fileSizeInBytes;
   private final long fileFooterSizeInBytes;
+  private final byte[] keyMetadata;
   private final List<BlobMetadata> blobMetadata;
 
+  /**
+   * @deprecated will be removed in 1.13.0. Use {@link #GenericStatisticsFile(long, String, long,
+   *     long, ByteBuffer, List)} instead.
+   */
+  @Deprecated
   public GenericStatisticsFile(
       long snapshotId,
       String path,
       long fileSizeInBytes,
       long fileFooterSizeInBytes,
       List<BlobMetadata> blobMetadata) {
+    this(snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, null, blobMetadata);
+  }
+
+  public GenericStatisticsFile(
+      long snapshotId,
+      String path,
+      long fileSizeInBytes,
+      long fileFooterSizeInBytes,
+      ByteBuffer keyMetadata,
+      List<BlobMetadata> blobMetadata) {
     Preconditions.checkNotNull(path, "path is null");
     Preconditions.checkNotNull(blobMetadata, "blobMetadata is null");
     this.snapshotId = snapshotId;
     this.path = path;
     this.fileSizeInBytes = fileSizeInBytes;
     this.fileFooterSizeInBytes = fileFooterSizeInBytes;
+    this.keyMetadata = ByteBuffers.toByteArray(keyMetadata);
     this.blobMetadata = ImmutableList.copyOf(blobMetadata);
   }
 
@@ -66,6 +86,11 @@ public long fileFooterSizeInBytes() {
     return fileFooterSizeInBytes;
   }
 
+  @Override
+  public ByteBuffer keyMetadata() {
+    return keyMetadata != null ? ByteBuffer.wrap(keyMetadata) : null;
+  }
+
   @Override
   public List<BlobMetadata> blobMetadata() {
     return blobMetadata;
@@ -84,12 +109,19 @@ public boolean equals(Object o) {
         && fileSizeInBytes == that.fileSizeInBytes
         && fileFooterSizeInBytes == that.fileFooterSizeInBytes
         && Objects.equals(path, that.path)
+        && Arrays.equals(keyMetadata, that.keyMetadata)
         && Objects.equals(blobMetadata, that.blobMetadata);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, blobMetadata);
+    return Objects.hash(
+        snapshotId,
+        path,
+        fileSizeInBytes,
+        fileFooterSizeInBytes,
+        Arrays.hashCode(keyMetadata),
+        blobMetadata);
   }
 
   @Override
@@ -99,6 +131,7 @@ public String toString() {
         .add("path='" + path + "'")
         .add("fileSizeInBytes=" + fileSizeInBytes)
         .add("fileFooterSizeInBytes=" + fileFooterSizeInBytes)
+        .add("keyMetadata=" + (keyMetadata == null ? "null" : "(redacted)"))
         .add("blobMetadata=" + blobMetadata)
         .toString();
   }

core/src/main/java/org/apache/iceberg/RewriteTablePathUtil.java

@@ -181,6 +181,7 @@ private static List<StatisticsFile> updatePathInStatisticsFiles(
                     newPath(existing.path(), sourcePrefix, targetPrefix),
                     existing.fileSizeInBytes(),
                     existing.fileFooterSizeInBytes(),
+                    existing.keyMetadata(),
                     existing.blobMetadata()))
         .collect(Collectors.toList());
   }

core/src/main/java/org/apache/iceberg/StatisticsFileParser.java

@@ -21,6 +21,7 @@
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.databind.JsonNode;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -34,6 +35,7 @@ public class StatisticsFileParser {
   private static final String STATISTICS_PATH = "statistics-path";
   private static final String FILE_SIZE_IN_BYTES = "file-size-in-bytes";
   private static final String FILE_FOOTER_SIZE_IN_BYTES = "file-footer-size-in-bytes";
+  private static final String KEY_METADATA = "key-metadata";
   private static final String BLOB_METADATA = "blob-metadata";
   private static final String TYPE = "type";
   private static final String SEQUENCE_NUMBER = "sequence-number";
@@ -57,6 +59,11 @@ public static void toJson(StatisticsFile statisticsFile, JsonGenerator generator
     generator.writeStringField(STATISTICS_PATH, statisticsFile.path());
     generator.writeNumberField(FILE_SIZE_IN_BYTES, statisticsFile.fileSizeInBytes());
     generator.writeNumberField(FILE_FOOTER_SIZE_IN_BYTES, statisticsFile.fileFooterSizeInBytes());
+    if (statisticsFile.keyMetadata() != null) {
+      generator.writeFieldName(KEY_METADATA);
+      SingleValueParser.toJson(
+          DataFile.KEY_METADATA.type(), statisticsFile.keyMetadata(), generator);
+    }
     generator.writeArrayFieldStart(BLOB_METADATA);
     for (BlobMetadata blobMetadata : statisticsFile.blobMetadata()) {
       toJson(blobMetadata, generator);
@@ -70,6 +77,7 @@ static StatisticsFile fromJson(JsonNode node) {
     String path = JsonUtil.getString(STATISTICS_PATH, node);
     long fileSizeInBytes = JsonUtil.getLong(FILE_SIZE_IN_BYTES, node);
     long fileFooterSizeInBytes = JsonUtil.getLong(FILE_FOOTER_SIZE_IN_BYTES, node);
+    ByteBuffer keyMetadata = JsonUtil.getByteBufferOrNull(KEY_METADATA, node);
     ImmutableList.Builder<BlobMetadata> blobMetadata = ImmutableList.builder();
     JsonNode blobsJson = node.get(BLOB_METADATA);
     Preconditions.checkArgument(
@@ -80,7 +88,12 @@ static StatisticsFile fromJson(JsonNode node) {
       blobMetadata.add(blobMetadataFromJson(blobJson));
     }
     return new GenericStatisticsFile(
-        snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, blobMetadata.build());
+        snapshotId,
+        path,
+        fileSizeInBytes,
+        fileFooterSizeInBytes,
+        keyMetadata,
+        blobMetadata.build());
   }
 
   private static void toJson(BlobMetadata blobMetadata, JsonGenerator generator)

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/actions/ComputeTableStatsSparkAction.java

@@ -19,6 +19,7 @@
 package org.apache.iceberg.spark.actions;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
@@ -33,6 +34,10 @@
 import org.apache.iceberg.TableOperations;
 import org.apache.iceberg.actions.ComputeTableStats;
 import org.apache.iceberg.actions.ImmutableComputeTableStats;
+import org.apache.iceberg.encryption.EncryptedOutputFile;
+import org.apache.iceberg.encryption.EncryptingFileIO;
+import org.apache.iceberg.encryption.EncryptionKeyMetadata;
+import org.apache.iceberg.encryption.NativeEncryptionKeyMetadata;
 import org.apache.iceberg.exceptions.RuntimeIOException;
 import org.apache.iceberg.io.OutputFile;
 import org.apache.iceberg.puffin.Blob;
@@ -110,21 +115,35 @@ private Result doExecute() {
 
   private StatisticsFile writeStatsFile(List<Blob> blobs) {
     LOG.info("Writing stats for table {} for snapshot {}", table.name(), snapshotId());
-    OutputFile outputFile = table.io().newOutputFile(outputPath());
+    EncryptingFileIO io = EncryptingFileIO.combine(table.io(), table.encryption());
+    EncryptedOutputFile encryptedOutputFile = io.newEncryptingOutputFile(outputPath());
+    OutputFile outputFile = encryptedOutputFile.encryptingOutputFile();
+
     try (PuffinWriter writer = Puffin.write(outputFile).createdBy(appIdentifier()).build()) {
       blobs.forEach(writer::add);
       writer.finish();
+      long fileSize = writer.fileSize();
       return new GenericStatisticsFile(
           snapshotId(),
           outputFile.location(),
-          writer.fileSize(),
+          fileSize,
           writer.footerSize(),
+          keyMetadata(encryptedOutputFile, fileSize),
           GenericBlobMetadata.from(writer.writtenBlobsMetadata()));
     } catch (IOException e) {
       throw new RuntimeIOException(e);
     }
   }
 
+  private ByteBuffer keyMetadata(EncryptedOutputFile outputFile, long fileSizeInBytes) {
+    EncryptionKeyMetadata keyMetadata = outputFile.keyMetadata();
+    if (keyMetadata instanceof NativeEncryptionKeyMetadata nativeKeyMetadata) {
+      return nativeKeyMetadata.copyWithLength(fileSizeInBytes).buffer();
+    }
+
+    return keyMetadata.buffer();
+  }
+
   private List<Blob> generateNDVBlobs() {
     return NDVSketchUtil.generateBlobs(spark(), table, snapshot, columns());
   }

spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputeTableStatsEncryption.java

@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.iceberg.spark.actions;
+
+import static org.apache.iceberg.spark.actions.NDVSketchUtil.APACHE_DATASKETCHES_THETA_V1_NDV_PROPERTY;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.io.IOException;
+import org.apache.iceberg.CatalogProperties;
+import org.apache.iceberg.Parameters;
+import org.apache.iceberg.StatisticsFile;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.encryption.EncryptingFileIO;
+import org.apache.iceberg.encryption.UnitestKMS;
+import org.apache.iceberg.io.InputFile;
+import org.apache.iceberg.puffin.Puffin;
+import org.apache.iceberg.puffin.PuffinReader;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.spark.CatalogTestBase;
+import org.apache.iceberg.spark.SparkCatalogConfig;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+
+public class TestComputeTableStatsEncryption extends CatalogTestBase {
+
+  @Parameters(name = "catalogName = {0}, implementation = {1}, config = {2}")
+  protected static Object[][] parameters() {
+    return new Object[][] {
+      {
+        SparkCatalogConfig.HIVE.catalogName(),
+        SparkCatalogConfig.HIVE.implementation(),
+        ImmutableMap.builder()
+            .putAll(SparkCatalogConfig.HIVE.properties())
+            .put(CatalogProperties.ENCRYPTION_KMS_IMPL, UnitestKMS.class.getCanonicalName())
+            .build()
+      }
+    };
+  }
+
+  @BeforeEach
+  public void createTable() {
+    sql(
+        "CREATE TABLE %s (id bigint, data string) USING iceberg "
+            + "TBLPROPERTIES ('encryption.key-id'='%s', 'format-version'='3')",
+        tableName, UnitestKMS.MASTER_KEY_NAME1);
+    sql("INSERT INTO %s VALUES (1, 'a'), (1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')", tableName);
+  }
+
+  @AfterEach
+  public void removeTable() {
+    sql("DROP TABLE IF EXISTS %s", tableName);
+  }
+
+  @TestTemplate
+  public void writeAndReadEncryptedStatisticsFile() throws IOException {
+    // the validation catalog has to know about the KMS to be able to read the encrypted table
+    validationCatalog.initialize(catalogName, catalogConfig);
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    StatisticsFile statisticsFile =
+        SparkActions.get().computeTableStats(table).columns("id").execute().statisticsFile();
+
+    assertThat(statisticsFile.fileSizeInBytes()).isGreaterThan(0);
+    assertThat(statisticsFile.keyMetadata()).isNotNull();
+
+    table.refresh();
+    EncryptingFileIO io = EncryptingFileIO.combine(table.io(), table.encryption());
+    InputFile decryptingInputFile =
+        io.newDecryptingInputFile(
+            statisticsFile.path(), statisticsFile.fileSizeInBytes(), statisticsFile.keyMetadata());
+
+    try (PuffinReader reader = Puffin.read(decryptingInputFile).build()) {
+      assertThat(reader.fileMetadata().blobs())
+          .singleElement()
+          .satisfies(
+              blob ->
+                  assertThat(blob.properties())
+                      .containsEntry(APACHE_DATASKETCHES_THETA_V1_NDV_PROPERTY, "4"));
+    }
+  }
+
+  @TestTemplate
+  public void failToReadEncryptedStatisticsFile() throws IOException {
+    // the validation catalog has to know about the KMS to be able to read the encrypted table
+    validationCatalog.initialize(catalogName, catalogConfig);
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    StatisticsFile statisticsFile =
+        SparkActions.get().computeTableStats(table).columns("id").execute().statisticsFile();
+
+    table.refresh();
+
+    InputFile plainTextInputFile = table.io().newInputFile(statisticsFile.path());
+    try (PuffinReader reader = Puffin.read(plainTextInputFile).build()) {
+      assertThatThrownBy(reader::fileMetadata)
+          .isInstanceOf(IllegalStateException.class)
+          .hasMessageStartingWith("Invalid file: expected magic at offset");
+    }
+  }
+}