fix: Normalize s3 paths for PME key retriever (#2874)

Author: mbutrovich

common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java

@@ -101,13 +101,35 @@ public class CometFileKeyUnwrapper {
   // Cache the hadoopConf just to assert the assumption above.
   private Configuration conf = null;
 
+  /**
+   * Normalizes S3 URI schemes to a canonical form. S3 can be accessed via multiple schemes (s3://,
+   * s3a://, s3n://) that refer to the same logical filesystem. This method ensures consistent cache
+   * lookups regardless of which scheme is used.
+   *
+   * @param filePath The file path that may contain an S3 URI
+   * @return The file path with normalized S3 scheme (s3a://)
+   */
+  private String normalizeS3Scheme(final String filePath) {
+    // Normalize s3:// and s3n:// to s3a:// for consistent cache lookups
+    // This handles the case where ObjectStoreUrl uses s3:// but Spark uses s3a://
+    String s3Prefix = "s3://";
+    String s3nPrefix = "s3n://";
+    if (filePath.startsWith(s3Prefix)) {
+      return "s3a://" + filePath.substring(s3Prefix.length());
+    } else if (filePath.startsWith(s3nPrefix)) {
+      return "s3a://" + filePath.substring(s3nPrefix.length());
+    }
+    return filePath;
+  }
+
   /**
    * Creates and stores a DecryptionKeyRetriever instance for the given file path.
    *
    * @param filePath The path to the Parquet file
    * @param hadoopConf The Hadoop Configuration to use for this file path
    */
   public void storeDecryptionKeyRetriever(final String filePath, final Configuration hadoopConf) {
+    final String normalizedPath = normalizeS3Scheme(filePath);
     // Use DecryptionPropertiesFactory.loadFactory to get the factory and then call
     // getFileDecryptionProperties
     if (factory == null) {
@@ -122,7 +144,7 @@ public void storeDecryptionKeyRetriever(final String filePath, final Configurati
         factory.getFileDecryptionProperties(hadoopConf, path);
 
     DecryptionKeyRetriever keyRetriever = decryptionProperties.getKeyRetriever();
-    retrieverCache.put(filePath, keyRetriever);
+    retrieverCache.put(normalizedPath, keyRetriever);
   }
 
   /**
@@ -136,7 +158,8 @@ public void storeDecryptionKeyRetriever(final String filePath, final Configurati
    */
   public byte[] getKey(final String filePath, final byte[] keyMetadata)
       throws ParquetCryptoRuntimeException {
-    DecryptionKeyRetriever keyRetriever = retrieverCache.get(filePath);
+    final String normalizedPath = normalizeS3Scheme(filePath);
+    DecryptionKeyRetriever keyRetriever = retrieverCache.get(normalizedPath);
     if (keyRetriever == null) {
       throw new ParquetCryptoRuntimeException(
           "Failed to find DecryptionKeyRetriever for path: " + filePath);

spark/src/test/scala/org/apache/comet/parquet/ParquetReadFromS3Suite.scala

@@ -19,6 +19,12 @@
 
 package org.apache.comet.parquet
 
+import java.nio.charset.StandardCharsets
+import java.util.Base64
+
+import org.apache.parquet.crypto.DecryptionPropertiesFactory
+import org.apache.parquet.crypto.keytools.{KeyToolkit, PropertiesDrivenCryptoFactory}
+import org.apache.parquet.crypto.keytools.mocks.InMemoryKMS
 import org.apache.spark.sql.{DataFrame, SaveMode}
 import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -30,6 +36,15 @@ class ParquetReadFromS3Suite extends CometS3TestBase with AdaptiveSparkPlanHelpe
 
   override protected val testBucketName = "test-bucket"
 
+  // Encryption keys for testing parquet encryption
+  private val encoder = Base64.getEncoder
+  private val footerKey =
+    encoder.encodeToString("0123456789012345".getBytes(StandardCharsets.UTF_8))
+  private val key1 = encoder.encodeToString("1234567890123450".getBytes(StandardCharsets.UTF_8))
+  private val key2 = encoder.encodeToString("1234567890123451".getBytes(StandardCharsets.UTF_8))
+  private val cryptoFactoryClass =
+    "org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory"
+
   private def writeTestParquetFile(filePath: String): Unit = {
     val df = spark.range(0, 1000)
     df.write.format("parquet").mode(SaveMode.Overwrite).save(filePath)
@@ -76,4 +91,32 @@ class ParquetReadFromS3Suite extends CometS3TestBase with AdaptiveSparkPlanHelpe
     assertCometScan(df)
     assert(df.first().getLong(0) == 499500)
   }
+
+  test("write and read encrypted parquet from S3") {
+    import testImplicits._
+
+    withSQLConf(
+      DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+      KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+        "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+      InMemoryKMS.KEY_LIST_PROPERTY_NAME ->
+        s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
+
+      val inputDF = spark
+        .range(0, 1000)
+        .map(i => (i, i.toString, i.toFloat))
+        .repartition(5)
+        .toDF("a", "b", "c")
+
+      val testFilePath = s"s3a://$testBucketName/data/encrypted-test.parquet"
+      inputDF.write
+        .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+        .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+        .parquet(testFilePath)
+
+      val df = spark.read.parquet(testFilePath).agg(sum(col("a")))
+      assertCometScan(df)
+      assert(df.first().getLong(0) == 499500)
+    }
+  }
 }