You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "mbutrovich (via GitHub)" <gi...@apache.org> on 2025/09/23 19:34:37 UTC

[PR] feat: Parquet Modular Encryption support for native_datafusion and native_iceberg_compat readers [datafusion-comet]

mbutrovich opened a new pull request, #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447

   This is draft for now. I have some duplicate code to clean up and other minor refactoring to do, but it's ready to start playing in CI.
   
   I'll expand on this description soon.
   
   ## Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Closes #.
   
   ## Rationale for this change
   
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   
   We want to add Parquet Module Encryption support for the native readers. We use the encryption factory features added in DataFusion 50 to register an encryption factory that uses JNI to get decryption keys from Spark.
   
   ## What changes are included in this PR?
   
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   ## How are these changes tested?
   
   <!--
   We typically require tests for all PRs in order to:
   1. Prevent the code from being accidentally broken by subsequent changes
   2. Serve as another way to document the expected behavior of the code
   
   If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
   -->
   
   Existing PME tests with new readers added.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2408651442


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -191,6 +521,30 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     }
   }
 
+  /**
+   * Verify that the directory contains an encrypted parquet in plaintext footer mode by means of
+   * checking for all the parquet part files in the parquet directory that their magic string is
+   * PAR1, as defined in the spec:
+   * https://github.com/apache/parquet-format/blob/master/Encryption.md#55-plaintext-footer-mode
+   */
+  private def verifyParquetPlaintextFooter(parquetDir: String): Unit = {
+    val parquetPartitionFiles = getListOfParquetFiles(new File(parquetDir))
+    assert(parquetPartitionFiles.size >= 1)
+    parquetPartitionFiles.foreach { parquetFile =>
+      val magicString = "PAR1"
+      val magicStringLength = magicString.length()
+      val byteArray = new Array[Byte](magicStringLength)
+      val randomAccessFile = new RandomAccessFile(parquetFile, "r")
+      try {
+        randomAccessFile.read(byteArray, 0, magicStringLength)

Review Comment:
   Added.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2394775079


##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -201,6 +207,30 @@ abstract class CometNativeExec extends CometExec {
         // TODO: support native metrics for all operators.
         val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
+        // For each relation in a CometNativeScan generate a hadoopConf,
+        // for each file path in a relation associate with hadoopConf
+        val cometNativeScans: Seq[CometNativeScanExec] = this
+          .collectLeaves()
+          .filter(_.isInstanceOf[CometNativeScanExec])
+          .map(_.asInstanceOf[CometNativeScanExec])
+        val encryptedFilePaths = cometNativeScans.flatMap { scan =>
+          // This creates a hadoopConf that brings in any SQLConf "spark.hadoop.*" configs and
+          // per-relation configs since different tables might have different decryption
+          // properties.
+          val hadoopConf = scan.relation.sparkSession.sessionState
+            .newHadoopConfWithOptions(scan.relation.options)
+          val encryptionEnabled = CometParquetUtils.encryptionEnabled(hadoopConf)
+          if (encryptionEnabled) {
+            // hadoopConf isn't serializable, so we have to do a broadcasted config.
+            val broadcastedConf =
+              scan.relation.sparkSession.sparkContext

Review Comment:
   Right but the encryption options can either be in the SQLConf or the relation options, eventually it all gets stuff into the hadoopConf for the Parquet library. Seemed easiest to just broadcast the hadoopConf the same way the existing Comet Parquet scan does (see `CometParquetFileFormat`).



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391052576


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   Session hadoopConf is not what the scans use though. They add all the relation options (Parquet options like encryption keys) to the hadoopConf, so each scan can have a unique hadoopConf. Whether we could have a Comet plan with multiple Parquet scans is the real question.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373403929


##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."

Review Comment:
   (question) We don't support it now b/c it is on the write path (set by the Parquet writer)?



##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."

Review Comment:
   (question) We don't support it now b/c it is on the write path (set by Parquet writer)?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2393003156


##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()
+    }
+}
+
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+pub struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    pub fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",
+                )
+                .unwrap(),
+        })
+    }
+}
+
+impl KeyRetriever for CometKeyRetriever {
+    /// Get a data encryption key using the metadata stored in the Parquet file.
+    fn retrieve_key(&self, key_metadata: &[u8]) -> datafusion::parquet::errors::Result<Vec<u8>> {
+        use jni::{objects::JObject, signature::ReturnType};
+
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()
+            .map_err(|e| datafusion::parquet::errors::ParquetError::General(e.to_string()))?;
+
+        // Get the key unwrapper instance from GlobalRef
+        let unwrapper_instance = self.key_unwrapper.as_obj();
+
+        let instance: JObject = unsafe { JObject::from_raw(unwrapper_instance.as_raw()) };
+
+        // Convert file path to JString
+        let file_path_jstring = env.new_string(&self.file_path).unwrap();
+
+        // Convert key_metadata to JByteArray
+        let key_metadata_array = env.byte_array_from_slice(key_metadata).unwrap();
+
+        // Call instance method FileKeyUnwrapper.getKey(String, byte[]) -> byte[]
+        let result = unsafe {
+            env.call_method_unchecked(

Review Comment:
   This can throw an exception can it not?  Elsewhere we use `try_unwrap_or_throw` (errors.rs) to report the error.



##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   > Whether we could have a Comet plan with multiple Parquet scans is the real question.
   
   I don't know what you mean by this. What exactly are you calling a Parquet scan? 
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#issuecomment-3325352610

   ## [Codecov](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
   :x: Patch coverage is `0%` with `60 lines` in your changes missing coverage. Please review.
   :white_check_mark: Project coverage is 10.48%. Comparing base ([`f09f8af`](https://app.codecov.io/gh/apache/datafusion-comet/commit/f09f8af64c6599255e116a376f4f008f2fd63b43?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`8bac76a`](https://app.codecov.io/gh/apache/datafusion-comet/commit/8bac76a1cc2744c8b227decf32e71ecd3ddb2468?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
   :warning: Report is 536 commits behind head on main.
   
   | [Files with missing lines](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
   |---|---|---|
   | [...n/scala/org/apache/spark/sql/comet/operators.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fcomet%2Foperators.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9zcGFyay9zcWwvY29tZXQvb3BlcmF0b3JzLnNjYWxh) | 0.00% | [19 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...rg/apache/comet/parquet/CometFileKeyUnwrapper.java](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fcomet%2Fparquet%2FCometFileKeyUnwrapper.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-Y29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9jb21ldC9wYXJxdWV0L0NvbWV0RmlsZUtleVVud3JhcHBlci5qYXZh) | 0.00% | [17 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...ain/scala/org/apache/comet/CometExecIterator.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2FCometExecIterator.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9Db21ldEV4ZWNJdGVyYXRvci5zY2FsYQ==) | 0.00% | [11 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...va/org/apache/comet/parquet/NativeBatchReader.java](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fcomet%2Fparquet%2FNativeBatchReader.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-Y29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9jb21ldC9wYXJxdWV0L05hdGl2ZUJhdGNoUmVhZGVyLmphdmE=) | 0.00% | [6 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [.../scala/org/apache/comet/serde/QueryPlanSerde.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2Fserde%2FQueryPlanSerde.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9zZXJkZS9RdWVyeVBsYW5TZXJkZS5zY2FsYQ==) | 0.00% | [4 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...t/execution/shuffle/CometNativeShuffleWriter.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fcomet%2Fexecution%2Fshuffle%2FCometNativeShuffleWriter.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9zcGFyay9zcWwvY29tZXQvZXhlY3V0aW9uL3NodWZmbGUvQ29tZXROYXRpdmVTaHVmZmxlV3JpdGVyLnNjYWxh) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   | [...n/scala/org/apache/comet/rules/CometScanRule.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2Frules%2FCometScanRule.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9ydWxlcy9Db21ldFNjYW5SdWxlLnNjYWxh) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
   
   <details><summary>Additional details and impacted files</summary>
   
   
   
   ```diff
   @@              Coverage Diff              @@
   ##               main    #2447       +/-   ##
   =============================================
   - Coverage     56.12%   10.48%   -45.64%     
   + Complexity      976      232      -744     
   =============================================
     Files           119      147       +28     
     Lines         11743    13499     +1756     
     Branches       2251     2333       +82     
   =============================================
   - Hits           6591     1416     -5175     
   - Misses         4012    11906     +7894     
   + Partials       1140      177      -963     
   ```
   </details>
   
   [:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/datafusion-comet/pull/2447?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
   :loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
   <details><summary> :rocket: New features to boost your workflow: </summary>
   
   - :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
   </details>


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391082513


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =
+      new ConcurrentHashMap<>();
+
+  /**
+   * 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) {
+    // Use DecryptionPropertiesFactory.loadFactory to get the factory and then call
+    // getFileDecryptionProperties
+    DecryptionPropertiesFactory factory = factoryCache.get(hadoopConf);

Review Comment:
   The Parquet library decides what the fields it needs are. If we want to limit the versions we want to support I could start hard-coding the values, but that's creating future work for ourselves any time a new config is added.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391204172


##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()

Review Comment:
   This will only work for local files, not HDFS or S3. I'll need to sort that out today.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373388517


##########
common/src/main/java/org/apache/comet/parquet/NativeBatchReader.java:
##########
@@ -410,6 +410,16 @@ public void init() throws Throwable {
         }
       }
 
+      boolean encryptionEnabled =
+          !conf.get("parquet.crypto.factory.class").isEmpty()

Review Comment:
   `conf.get` returns `null` if no such property exists, so we might need a null check as well.



##########
common/src/main/java/org/apache/comet/parquet/NativeBatchReader.java:
##########
@@ -410,6 +410,16 @@ public void init() throws Throwable {
         }
       }
 
+      boolean encryptionEnabled =
+          !conf.get("parquet.crypto.factory.class").isEmpty()

Review Comment:
   (nit) Maybe we can use `DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME` and `KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373571298


##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."

Review Comment:
   Correct. Comet does not support write path at all. When the time comes, PME is already supported in DataFusion and Arrow-rs.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391075529


##########
common/src/main/scala/org/apache/comet/parquet/CometParquetUtils.scala:
##########
@@ -20,13 +20,30 @@
 package org.apache.comet.parquet
 
 import org.apache.hadoop.conf.Configuration
+import org.apache.parquet.crypto.DecryptionPropertiesFactory
+import org.apache.parquet.crypto.keytools.KeyToolkit
 import org.apache.spark.sql.internal.SQLConf
 
 object CometParquetUtils {
   private val PARQUET_FIELD_ID_WRITE_ENABLED = "spark.sql.parquet.fieldId.write.enabled"
   private val PARQUET_FIELD_ID_READ_ENABLED = "spark.sql.parquet.fieldId.read.enabled"
   private val IGNORE_MISSING_PARQUET_FIELD_ID = "spark.sql.parquet.fieldId.read.ignoreMissing"
 
+  // Map of unsupported encryption configuration key-value pairs
+  private val UNSUPPORTED_ENCRYPTION_CONFIGS: Map[String, Set[String]] = Map(
+    "parquet.encryption.algorithm" -> Set("AES_GCM_CTR_V1")

Review Comment:
   https://github.com/apache/parquet-java/blob/f3542ee639fc4dbdcc79f82c0d30f202e8a2338a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/keytools/PropertiesDrivenCryptoFactory.java#L60
   
   https://github.com/apache/arrow-rs/blob/8eca76df2986908bcfa175fa6b16c52d1019a874/parquet/src/file/metadata/parser.rs#L494
   
   I can add links.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2394600127


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   A scan node in a plan tree.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407818480


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -47,90 +50,399 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     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"
 
   test("SPARK-34990: Write and read an encrypted parquet") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            // Make sure encryption works with multiple Parquet files
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            verifyParquetEncrypted(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      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}") {
+
+        // Make sure encryption works with multiple Parquet files
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
         }
+      }
     }
   }
 
   test("SPARK-37117: Can't read files in Parquet encryption external key material mode") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.material.store.internally" -> "false",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        KeyToolkit.KEY_MATERIAL_INTERNAL_PROPERTY_NAME -> "false", // default is true
+        InMemoryKMS.KEY_LIST_PROPERTY_NAME ->
+          s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("SPARK-42114: Test of uniform parquet encryption") {
+
+    import testImplicits._
+
+    withTempDir { dir =>
+      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"key1: ${key1}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option("parquet.encryption.uniform.key", "key1")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Plain text footer mode") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        PropertiesDrivenCryptoFactory.PLAINTEXT_FOOTER_PROPERTY_NAME -> "true", // default is false
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetPlaintextFooter(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Change encryption algorithm") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        // default is AES_GCM_V1
+        PropertiesDrivenCryptoFactory.ENCRYPTION_ALGORITHM_PROPERTY_NAME -> "AES_GCM_CTR_V1",
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        // native_datafusion and native_iceberg_compat fall back due to Arrow-rs
+        // https://github.com/apache/arrow-rs/blob/main/parquet/src/file/metadata/parser.rs#L414

Review Comment:
   Correct. The library they use for decryption (I think it's ring) has support for it, but it's just not hooked up to arrow-rs yet.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407830899


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -47,90 +50,399 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     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"
 
   test("SPARK-34990: Write and read an encrypted parquet") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            // Make sure encryption works with multiple Parquet files
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            verifyParquetEncrypted(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      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}") {
+
+        // Make sure encryption works with multiple Parquet files
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
         }
+      }
     }
   }
 
   test("SPARK-37117: Can't read files in Parquet encryption external key material mode") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.material.store.internally" -> "false",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        KeyToolkit.KEY_MATERIAL_INTERNAL_PROPERTY_NAME -> "false", // default is true
+        InMemoryKMS.KEY_LIST_PROPERTY_NAME ->
+          s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("SPARK-42114: Test of uniform parquet encryption") {
+
+    import testImplicits._
+
+    withTempDir { dir =>
+      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"key1: ${key1}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option("parquet.encryption.uniform.key", "key1")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Plain text footer mode") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        PropertiesDrivenCryptoFactory.PLAINTEXT_FOOTER_PROPERTY_NAME -> "true", // default is false
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetPlaintextFooter(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Change encryption algorithm") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        // default is AES_GCM_V1
+        PropertiesDrivenCryptoFactory.ENCRYPTION_ALGORITHM_PROPERTY_NAME -> "AES_GCM_CTR_V1",
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        // native_datafusion and native_iceberg_compat fall back due to Arrow-rs
+        // https://github.com/apache/arrow-rs/blob/main/parquet/src/file/metadata/parser.rs#L414

Review Comment:
   Thanks @mbutrovich, then I think the line number to `arrow-rs` might be `L494` not `L414`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391079953


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =

Review Comment:
   Spark must also be keeping the same number in memory for its scans. Also, it's whatever subset of files this plan is responsible for.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391068852


##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -201,6 +207,30 @@ abstract class CometNativeExec extends CometExec {
         // TODO: support native metrics for all operators.
         val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
+        // For each relation in a CometNativeScan generate a hadoopConf,
+        // for each file path in a relation associate with hadoopConf
+        val cometNativeScans: Seq[CometNativeScanExec] = this
+          .collectLeaves()
+          .filter(_.isInstanceOf[CometNativeScanExec])
+          .map(_.asInstanceOf[CometNativeScanExec])
+        val encryptedFilePaths = cometNativeScans.flatMap { scan =>
+          // This creates a hadoopConf that brings in any SQLConf "spark.hadoop.*" configs and
+          // per-relation configs since different tables might have different decryption
+          // properties.
+          val hadoopConf = scan.relation.sparkSession.sessionState
+            .newHadoopConfWithOptions(scan.relation.options)
+          val encryptionEnabled = CometParquetUtils.encryptionEnabled(hadoopConf)
+          if (encryptionEnabled) {
+            // hadoopConf isn't serializable, so we have to do a broadcasted config.
+            val broadcastedConf =
+              scan.relation.sparkSession.sparkContext

Review Comment:
   It's only broadcasting one hadoopConf per relation (table). It is broadcasting a mapping from file to hadoopConf for each file.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391583917


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   Since each relation gets a hadoopConf, I could probably make the key the relation name to hash something a bit smaller, but this isn't really on the critical path of query execution so I don't expect it to be a huge win.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2395889210


##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -201,6 +207,30 @@ abstract class CometNativeExec extends CometExec {
         // TODO: support native metrics for all operators.
         val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
+        // For each relation in a CometNativeScan generate a hadoopConf,
+        // for each file path in a relation associate with hadoopConf
+        val cometNativeScans: Seq[CometNativeScanExec] = this
+          .collectLeaves()
+          .filter(_.isInstanceOf[CometNativeScanExec])
+          .map(_.asInstanceOf[CometNativeScanExec])
+        val encryptedFilePaths = cometNativeScans.flatMap { scan =>
+          // This creates a hadoopConf that brings in any SQLConf "spark.hadoop.*" configs and
+          // per-relation configs since different tables might have different decryption
+          // properties.
+          val hadoopConf = scan.relation.sparkSession.sessionState
+            .newHadoopConfWithOptions(scan.relation.options)
+          val encryptionEnabled = CometParquetUtils.encryptionEnabled(hadoopConf)
+          if (encryptionEnabled) {
+            // hadoopConf isn't serializable, so we have to do a broadcasted config.
+            val broadcastedConf =
+              scan.relation.sparkSession.sparkContext

Review Comment:
   Ok. I understand now what you're doing here. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2403518726


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   @mbutrovich I had this one open question about a plan with a union that may have more than one scan. Can you verify this will not be an issue.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407671215


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,146 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:

Review Comment:
   This diagram is super helpful, thanks a lot.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2395885679


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   I see what you mean now. What about a plan with a `union` ? 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2388718794


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   There is only one hadoop conf in a spark session so this may be overkill. 



##########
common/src/main/java/org/apache/comet/parquet/NativeBatchReader.java:
##########
@@ -410,6 +410,14 @@ public void init() throws Throwable {
         }
       }
 
+      boolean encryptionEnabled = CometParquetUtils.encryptionEnabled(conf);
+
+      // Create keyUnwrapper if encryption is enabled
+      CometFileKeyUnwrapper keyUnwrapper = encryptionEnabled ? new CometFileKeyUnwrapper() : null;
+      if (encryptionEnabled) {
+        keyUnwrapper.storeDecryptionKeyRetriever(file.filePath().toString(), conf);
+      }
+

Review Comment:
   ```suggestion
         CometFileKeyUnwrapper keyUnwrapper = null;
         if (encryptionEnabled) {
          keyUnwrapper = new CometFileKeyUnwrapper;
           keyUnwrapper.storeDecryptionKeyRetriever(file.filePath().toString(), conf);
         }
   ```



##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =

Review Comment:
   Every file path? This can get rather large when the number of files starts to reach 100K or more. 



##########
common/src/main/scala/org/apache/comet/parquet/CometParquetUtils.scala:
##########
@@ -20,13 +20,30 @@
 package org.apache.comet.parquet
 
 import org.apache.hadoop.conf.Configuration
+import org.apache.parquet.crypto.DecryptionPropertiesFactory
+import org.apache.parquet.crypto.keytools.KeyToolkit
 import org.apache.spark.sql.internal.SQLConf
 
 object CometParquetUtils {
   private val PARQUET_FIELD_ID_WRITE_ENABLED = "spark.sql.parquet.fieldId.write.enabled"
   private val PARQUET_FIELD_ID_READ_ENABLED = "spark.sql.parquet.fieldId.read.enabled"
   private val IGNORE_MISSING_PARQUET_FIELD_ID = "spark.sql.parquet.fieldId.read.ignoreMissing"
 
+  // Map of unsupported encryption configuration key-value pairs
+  private val UNSUPPORTED_ENCRYPTION_CONFIGS: Map[String, Set[String]] = Map(
+    "parquet.encryption.algorithm" -> Set("AES_GCM_CTR_V1")

Review Comment:
   Where do we get this from? 



##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =
+      new ConcurrentHashMap<>();
+
+  /**
+   * 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) {
+    // Use DecryptionPropertiesFactory.loadFactory to get the factory and then call
+    // getFileDecryptionProperties
+    DecryptionPropertiesFactory factory = factoryCache.get(hadoopConf);

Review Comment:
   Is this hadoop conf the entire hadoop configuration (which can have a thousand entries) or just the incremental properties specified for the session? Hashing this can become time comsuming. 



##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -201,6 +207,30 @@ abstract class CometNativeExec extends CometExec {
         // TODO: support native metrics for all operators.
         val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
+        // For each relation in a CometNativeScan generate a hadoopConf,
+        // for each file path in a relation associate with hadoopConf
+        val cometNativeScans: Seq[CometNativeScanExec] = this
+          .collectLeaves()
+          .filter(_.isInstanceOf[CometNativeScanExec])
+          .map(_.asInstanceOf[CometNativeScanExec])
+        val encryptedFilePaths = cometNativeScans.flatMap { scan =>
+          // This creates a hadoopConf that brings in any SQLConf "spark.hadoop.*" configs and
+          // per-relation configs since different tables might have different decryption
+          // properties.
+          val hadoopConf = scan.relation.sparkSession.sessionState
+            .newHadoopConfWithOptions(scan.relation.options)
+          val encryptionEnabled = CometParquetUtils.encryptionEnabled(hadoopConf)
+          if (encryptionEnabled) {
+            // hadoopConf isn't serializable, so we have to do a broadcasted config.
+            val broadcastedConf =
+              scan.relation.sparkSession.sparkContext

Review Comment:
   Can you explain a little what you are doing here? What is the additional hadoop conf information that needs to be broadcast per file path (as opposed to encryption properties that are defined once per table)? 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2398950970


##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()
+    }
+}
+
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+pub struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    pub fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",
+                )
+                .unwrap(),
+        })
+    }
+}
+
+impl KeyRetriever for CometKeyRetriever {
+    /// Get a data encryption key using the metadata stored in the Parquet file.
+    fn retrieve_key(&self, key_metadata: &[u8]) -> datafusion::parquet::errors::Result<Vec<u8>> {
+        use jni::{objects::JObject, signature::ReturnType};
+
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()
+            .map_err(|e| datafusion::parquet::errors::ParquetError::General(e.to_string()))?;
+
+        // Get the key unwrapper instance from GlobalRef
+        let unwrapper_instance = self.key_unwrapper.as_obj();
+
+        let instance: JObject = unsafe { JObject::from_raw(unwrapper_instance.as_raw()) };
+
+        // Convert file path to JString
+        let file_path_jstring = env.new_string(&self.file_path).unwrap();
+
+        // Convert key_metadata to JByteArray
+        let key_metadata_array = env.byte_array_from_slice(key_metadata).unwrap();
+
+        // Call instance method FileKeyUnwrapper.getKey(String, byte[]) -> byte[]
+        let result = unsafe {
+            env.call_method_unchecked(

Review Comment:
   I'm still wrapping my head around the JNI stuff, but I think `try_unwrap_or_throw` is applicable for JVM -> Native calls, but in this case we're going Native -> JNI. I think this patten (checking for exception below) maps to elsewhere: https://github.com/mbutrovich/datafusion-comet/blob/12f0fdc6d4161cd1eab548cc0376643940022f7d/native/core/src/jvm_bridge/mod.rs#L70



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2399364065


##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()
+    }
+}
+
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+pub struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    pub fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",
+                )
+                .unwrap(),
+        })
+    }
+}
+
+impl KeyRetriever for CometKeyRetriever {
+    /// Get a data encryption key using the metadata stored in the Parquet file.
+    fn retrieve_key(&self, key_metadata: &[u8]) -> datafusion::parquet::errors::Result<Vec<u8>> {
+        use jni::{objects::JObject, signature::ReturnType};
+
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()
+            .map_err(|e| datafusion::parquet::errors::ParquetError::General(e.to_string()))?;
+
+        // Get the key unwrapper instance from GlobalRef
+        let unwrapper_instance = self.key_unwrapper.as_obj();
+
+        let instance: JObject = unsafe { JObject::from_raw(unwrapper_instance.as_raw()) };
+
+        // Convert file path to JString
+        let file_path_jstring = env.new_string(&self.file_path).unwrap();
+
+        // Convert key_metadata to JByteArray
+        let key_metadata_array = env.byte_array_from_slice(key_metadata).unwrap();
+
+        // Call instance method FileKeyUnwrapper.getKey(String, byte[]) -> byte[]
+        let result = unsafe {
+            env.call_method_unchecked(

Review Comment:
   Ah, I think you are correct. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407733328


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -47,90 +50,399 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     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"
 
   test("SPARK-34990: Write and read an encrypted parquet") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            // Make sure encryption works with multiple Parquet files
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            verifyParquetEncrypted(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      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}") {
+
+        // Make sure encryption works with multiple Parquet files
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
         }
+      }
     }
   }
 
   test("SPARK-37117: Can't read files in Parquet encryption external key material mode") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.material.store.internally" -> "false",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        KeyToolkit.KEY_MATERIAL_INTERNAL_PROPERTY_NAME -> "false", // default is true
+        InMemoryKMS.KEY_LIST_PROPERTY_NAME ->
+          s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("SPARK-42114: Test of uniform parquet encryption") {
+
+    import testImplicits._
+
+    withTempDir { dir =>
+      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"key1: ${key1}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option("parquet.encryption.uniform.key", "key1")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Plain text footer mode") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        PropertiesDrivenCryptoFactory.PLAINTEXT_FOOTER_PROPERTY_NAME -> "true", // default is false
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetPlaintextFooter(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Change encryption algorithm") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        // default is AES_GCM_V1
+        PropertiesDrivenCryptoFactory.ENCRYPTION_ALGORITHM_PROPERTY_NAME -> "AES_GCM_CTR_V1",
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        // native_datafusion and native_iceberg_compat fall back due to Arrow-rs
+        // https://github.com/apache/arrow-rs/blob/main/parquet/src/file/metadata/parser.rs#L414

Review Comment:
   (question) Do you mean https://github.com/apache/arrow-rs/blob/main/parquet/src/file/metadata/parser.rs#L494 b/c arrow-rs doesn't support `AES_GCM_CTR_V1`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2408069317


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -47,90 +50,399 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     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"
 
   test("SPARK-34990: Write and read an encrypted parquet") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            // Make sure encryption works with multiple Parquet files
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            verifyParquetEncrypted(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      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}") {
+
+        // Make sure encryption works with multiple Parquet files
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
         }
+      }
     }
   }
 
   test("SPARK-37117: Can't read files in Parquet encryption external key material mode") {
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_DATAFUSION)
-    assume(CometConf.COMET_NATIVE_SCAN_IMPL.get() != CometConf.SCAN_NATIVE_ICEBERG_COMPAT)
 
     import testImplicits._
 
-    Seq("org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory").foreach {
-      factoryClass =>
-        withTempDir { dir =>
-          withSQLConf(
-            "parquet.crypto.factory.class" -> factoryClass,
-            "parquet.encryption.kms.client.class" ->
-              "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
-            "parquet.encryption.key.material.store.internally" -> "false",
-            "parquet.encryption.key.list" ->
-              s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
-
-            val inputDF = spark
-              .range(0, 2000)
-              .map(i => (i, i.toString, i.toFloat))
-              .repartition(10)
-              .toDF("a", "b", "c")
-            val parquetDir = new File(dir, "parquet").getCanonicalPath
-            inputDF.write
-              .option("parquet.encryption.column.keys", "key1: a, b; key2: c")
-              .option("parquet.encryption.footer.key", "footerKey")
-              .parquet(parquetDir)
-
-            val parquetDF = spark.read.parquet(parquetDir)
-            assert(parquetDF.inputFiles.nonEmpty)
-            val readDataset = parquetDF.select("a", "b", "c")
-
-            if (CometConf.COMET_ENABLED.get(conf)) {
-              checkSparkAnswerAndOperator(readDataset)
-            } else {
-              checkAnswer(readDataset, inputDF)
-            }
-          }
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        KeyToolkit.KEY_MATERIAL_INTERNAL_PROPERTY_NAME -> "false", // default is true
+        InMemoryKMS.KEY_LIST_PROPERTY_NAME ->
+          s"footerKey: ${footerKey}, key1: ${key1}, key2: ${key2}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("SPARK-42114: Test of uniform parquet encryption") {
+
+    import testImplicits._
+
+    withTempDir { dir =>
+      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"key1: ${key1}") {
+
+        val inputDF = spark
+          .range(0, 2000)
+          .map(i => (i, i.toString, i.toFloat))
+          .repartition(10)
+          .toDF("a", "b", "c")
+        val parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option("parquet.encryption.uniform.key", "key1")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Plain text footer mode") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        PropertiesDrivenCryptoFactory.PLAINTEXT_FOOTER_PROPERTY_NAME -> "true", // default is false
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetPlaintextFooter(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        if (CometConf.COMET_ENABLED.get(conf)) {
+          checkSparkAnswerAndOperator(readDataset)
+        } else {
+          checkAnswer(readDataset, inputDF)
+        }
+      }
+    }
+  }
+
+  test("Change encryption algorithm") {
+    import testImplicits._
+
+    withTempDir { dir =>
+      withSQLConf(
+        DecryptionPropertiesFactory.CRYPTO_FACTORY_CLASS_PROPERTY_NAME -> cryptoFactoryClass,
+        KeyToolkit.KMS_CLIENT_CLASS_PROPERTY_NAME ->
+          "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS",
+        // default is AES_GCM_V1
+        PropertiesDrivenCryptoFactory.ENCRYPTION_ALGORITHM_PROPERTY_NAME -> "AES_GCM_CTR_V1",
+        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 parquetDir = new File(dir, "parquet").getCanonicalPath
+        inputDF.write
+          .option(PropertiesDrivenCryptoFactory.COLUMN_KEYS_PROPERTY_NAME, "key1: a, b; key2: c")
+          .option(PropertiesDrivenCryptoFactory.FOOTER_KEY_PROPERTY_NAME, "footerKey")
+          .parquet(parquetDir)
+
+        verifyParquetEncrypted(parquetDir)
+
+        val parquetDF = spark.read.parquet(parquetDir)
+        assert(parquetDF.inputFiles.nonEmpty)
+        val readDataset = parquetDF.select("a", "b", "c")
+
+        // native_datafusion and native_iceberg_compat fall back due to Arrow-rs
+        // https://github.com/apache/arrow-rs/blob/main/parquet/src/file/metadata/parser.rs#L414

Review Comment:
   Yeah I shouldnt link against main since that can change
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2392911698


##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -201,6 +207,30 @@ abstract class CometNativeExec extends CometExec {
         // TODO: support native metrics for all operators.
         val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
+        // For each relation in a CometNativeScan generate a hadoopConf,
+        // for each file path in a relation associate with hadoopConf
+        val cometNativeScans: Seq[CometNativeScanExec] = this
+          .collectLeaves()
+          .filter(_.isInstanceOf[CometNativeScanExec])
+          .map(_.asInstanceOf[CometNativeScanExec])
+        val encryptedFilePaths = cometNativeScans.flatMap { scan =>
+          // This creates a hadoopConf that brings in any SQLConf "spark.hadoop.*" configs and
+          // per-relation configs since different tables might have different decryption
+          // properties.
+          val hadoopConf = scan.relation.sparkSession.sessionState
+            .newHadoopConfWithOptions(scan.relation.options)
+          val encryptionEnabled = CometParquetUtils.encryptionEnabled(hadoopConf)
+          if (encryptionEnabled) {
+            // hadoopConf isn't serializable, so we have to do a broadcasted config.
+            val broadcastedConf =
+              scan.relation.sparkSession.sparkContext

Review Comment:
   Why does this need to be broadcast? Won't each executor instance will have its own copy of the `scan.relation.options` already. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391204172


##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()

Review Comment:
   This will only work for local files, not HDFS or S3. I'll need to sort that out today.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2394600127


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   A scan node in a plan tree, specifically a stage that gets converted to a Comet native plan.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407544661


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   > Do you think I should go back to the map for `hadoopConf` to future proof a potential future `UNION` implementation?
   
   I think not. That's a pretty good explanation of why your assumption is valid, and the unit test verifies that. Thanks for adding the test!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2406118264


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   I added a test for `UNION`, and it didn't trigger the assertion. Here's what Claude says:
   ```
   Looking at line 700 in operators.scala:
   
     override def doExecuteColumnar(): RDD[ColumnarBatch] = {
       sparkContext.union(children.map(_.executeColumnar()))
     }
   
   CometUnionExec is NOT a native operator - it's a Comet operator that executes at the Spark/JVM level.
   
   Here's what's happening:
   
   1. Union stays in the same stage (it's a narrow transformation in Spark)
   2. But each child executes as a separate native plan:
     - Each CometNativeScan child calls executeColumnar() independently
     - Each creates its own isolated native plan (DataSourceExec)
     - The union itself happens in Spark by combining the RDDs
   3. This is why the assertion doesn't fire:
     - The assertion at operators.scala:218 checks collectLeaves() within a single native plan
     - But with CometUnionExec, each scan child is executed as a separate native plan
     - Each native plan only contains one scan, so the assertion passes
   
   Looking at your native plan output:
   - Plan #5055 - just DataSourceExec for parquet1
   - Plan #5057 - just DataSourceExec for parquet2
   - No UnionExec in the native layer - union happens in Spark
   
   So the encryption implementation is actually fine for unions because each scan gets its own execution context with its own hadoopConf. The concern about
   multiple scans with different encryption keys only applies if both scans were in the same native plan, which doesn't happen with the current CometUnion design.
   ```
   
   Do you think I should go back to the map for `hadoopConf` to future proof a potential future `UNION` implementation?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373651740


##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",

Review Comment:
   Yeah. `DecryptionKeyRetriever.getKey` returns `byte[]`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373409278


##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",

Review Comment:
   The return type of `getKey()` is `byte[]`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391052576


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   Session hadoopConf is not what the scans use though. They add all the relation options (Parquet options like keys) to the hadoopConf, so each scan can have a unique hadoopConf. Whether we could have a Comet plan with multiple Parquet scans is the real question.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2393772647


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =
+      new ConcurrentHashMap<>();
+
+  /**
+   * 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) {
+    // Use DecryptionPropertiesFactory.loadFactory to get the factory and then call
+    // getFileDecryptionProperties
+    DecryptionPropertiesFactory factory = factoryCache.get(hadoopConf);
+    if (factory == null) {
+      factory = DecryptionPropertiesFactory.loadFactory(hadoopConf);
+      factoryCache.put(hadoopConf, factory);

Review Comment:
   Since you use `ConcurrentMap` you probably want to use its `#computeIfAbsent()` method instead of `#get() + check for null + #put()`



##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()
+    }
+}
+
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+pub struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    pub fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;
+
+        Ok(CometKeyRetriever {
+            file_path: file_path.to_string(),
+            key_unwrapper,
+            get_key_method_id: env
+                .get_method_id(
+                    "org/apache/comet/parquet/CometFileKeyUnwrapper",
+                    "getKey",
+                    "(Ljava/lang/String;[B)[B",
+                )
+                .unwrap(),
+        })
+    }
+}
+
+impl KeyRetriever for CometKeyRetriever {
+    /// Get a data encryption key using the metadata stored in the Parquet file.
+    fn retrieve_key(&self, key_metadata: &[u8]) -> datafusion::parquet::errors::Result<Vec<u8>> {
+        use jni::{objects::JObject, signature::ReturnType};
+
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()
+            .map_err(|e| datafusion::parquet::errors::ParquetError::General(e.to_string()))?;
+
+        // Get the key unwrapper instance from GlobalRef
+        let unwrapper_instance = self.key_unwrapper.as_obj();
+
+        let instance: JObject = unsafe { JObject::from_raw(unwrapper_instance.as_raw()) };
+
+        // Convert file path to JString
+        let file_path_jstring = env.new_string(&self.file_path).unwrap();
+
+        // Convert key_metadata to JByteArray
+        let key_metadata_array = env.byte_array_from_slice(key_metadata).unwrap();
+
+        // Call instance method FileKeyUnwrapper.getKey(String, byte[]) -> byte[]
+        let result = unsafe {
+            env.call_method_unchecked(
+                instance,
+                self.get_key_method_id,
+                ReturnType::Array,
+                &[
+                    jni::objects::JValue::from(&file_path_jstring).as_jni(),
+                    jni::objects::JValue::from(&key_metadata_array).as_jni(),
+                ],
+            )
+        };
+
+        let result = result.unwrap();

Review Comment:
   Can you use `result?` instead ? To return an Err instead of panic-ing.
   Or use pattern match and return a custom error.
   
   Few more unwraps below.



##########
native/core/src/parquet/encryption_support.rs:
##########
@@ -0,0 +1,151 @@
+// 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.
+
+use crate::execution::operators::ExecutionError;
+use crate::jvm_bridge::JVMClasses;
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::common::extensions_options;
+use datafusion::config::EncryptionFactoryOptions;
+use datafusion::error::DataFusionError;
+use datafusion::execution::parquet_encryption::EncryptionFactory;
+use jni::objects::{GlobalRef, JMethodID};
+use object_store::path::Path;
+use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever};
+use parquet::encryption::encrypt::FileEncryptionProperties;
+use std::sync::Arc;
+
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    pub struct CometEncryptionConfig {
+        pub url_base: String, default = "file:///".into()
+    }
+}
+
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."
+                .parse()
+                .unwrap(),
+        ))
+    }
+
+    /// Generate file decryption properties to use when reading a Parquet file.
+    /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever`
+    /// that can determine the keys using the encryption metadata.
+    async fn get_file_decryption_properties(
+        &self,
+        options: &EncryptionFactoryOptions,
+        file_path: &Path,
+    ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+        let config: CometEncryptionConfig = options.to_extension_options()?;
+
+        let full_path: String = config.url_base + file_path.as_ref();
+        let key_retriever = CometKeyRetriever::new(&full_path, self.key_unwrapper.clone())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let decryption_properties =
+            FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+        Ok(Some(decryption_properties))
+    }
+}
+
+pub struct CometKeyRetriever {
+    file_path: String,
+    key_unwrapper: GlobalRef,
+    get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+    pub fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self, ExecutionError> {
+        // Get JNI environment
+        let mut env = JVMClasses::get_env()?;

Review Comment:
   Below you use
   ```rust
   let mut env = JVMClasses::get_env()
               .map_err(|e| datafusion::parquet::errors::ParquetError::General(e.to_string()))?;
   ```
   No need of error mapping here ?!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2394900454


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   I'll simplify this with a couple of assertions that a Comet plan should only have one scan node in it.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373403929


##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
     Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
 }
 
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+    struct CometEncryptionConfig {
+        url_base: String, default = "file:///".into()
+    }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+    pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+    async fn get_file_encryption_properties(
+        &self,
+        _options: &EncryptionFactoryOptions,
+        _schema: &SchemaRef,
+        _file_path: &Path,
+    ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+        Err(DataFusionError::NotImplemented(
+            "Comet does not support Parquet encryption yet."

Review Comment:
   (question) We don't support it b/c it is set by the Parquet writer?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#issuecomment-3353068211

   Results attached from the benchmark I added to `CometReadBenchmark`, and a small chart with highlights to see what the overhead of encryption is for the various readers.
   
   <img width="884" height="459" alt="decryption" src="https://github.com/user-attachments/assets/801e5f34-9c3e-4f3a-908d-cf3d6f3a1cb3" />
   
   [benchmark_decryption.txt](https://github.com/user-attachments/files/22624077/benchmark_decryption.txt)
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2391583917


##########
common/src/main/java/org/apache/comet/parquet/CometFileKeyUnwrapper.java:
##########
@@ -0,0 +1,145 @@
+/*
+ * 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.comet.parquet;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionKeyRetriever;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
+
+// spotless:off
+/*
+ * Architecture Overview:
+ *
+ *          JVM Side                           |                     Native Side
+ *   ┌─────────────────────────────────────┐   |   ┌─────────────────────────────────────┐
+ *   │     CometFileKeyUnwrapper           │   |   │       Parquet File Reading          │
+ *   │                                     │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      hadoopConf             │    │   |   │  │     file1.parquet           │    │
+ *   │  │   (Configuration)           │    │   |   │  │     file2.parquet           │    │
+ *   │  └─────────────────────────────┘    │   |   │  │     file3.parquet           │    │
+ *   │              │                      │   |   │  └─────────────────────────────┘    │
+ *   │              ▼                      │   |   │              │                      │
+ *   │  ┌─────────────────────────────┐    │   |   │              │                      │
+ *   │  │      factoryCache           │    │   |   │              ▼                      │
+ *   │  │   (many-to-one mapping)     │    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │                             │    │   |   │  │  Parse file metadata &      │    │
+ *   │  │ file1 ──┐                   │    │   |   │  │  extract keyMetadata        │    │
+ *   │  │ file2 ──┼─► DecryptionProps │    │   |   │  └─────────────────────────────┘    │
+ *   │  │ file3 ──┘      Factory      │    │   |   │              │                      │
+ *   │  └─────────────────────────────┘    │   |   │              │                      │
+ *   │              │                      │   |   │              ▼                      │
+ *   │              ▼                      │   |   │  ╔═════════════════════════════╗    │
+ *   │  ┌─────────────────────────────┐    │   |   │  ║        JNI CALL:            ║    │
+ *   │  │      retrieverCache         │    │   |   │  ║       getKey(filePath,      ║    │
+ *   │  │  filePath -> KeyRetriever   │◄───┼───┼───┼──║        keyMetadata)         ║    │
+ *   │  └─────────────────────────────┘    │   |   │  ╚═════════════════════════════╝    │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │                                     │
+ *   │  │  DecryptionKeyRetriever     │    │   |   │                                     │
+ *   │  │     .getKey(keyMetadata)    │    │   |   │                                     │
+ *   │  └─────────────────────────────┘    │   |   │                                     │
+ *   │              │                      │   |   │                                     │
+ *   │              ▼                      │   |   │                                     │
+ *   │  ┌─────────────────────────────┐    │   |   │  ┌─────────────────────────────┐    │
+ *   │  │      return key bytes       │────┼───┼───┼─►│   Use key for decryption    │    │
+ *   │  └─────────────────────────────┘    │   |   │  │    of parquet data          │    │
+ *   └─────────────────────────────────────┘   |   │  └─────────────────────────────┘    │
+ *                                             |   └─────────────────────────────────────┘
+ *                                             |
+ *                                    JNI Boundary
+ *
+ * Setup Phase (storeDecryptionKeyRetriever):
+ * 1. hadoopConf → DecryptionPropertiesFactory (cached in factoryCache)
+ * 2. Factory + filePath → DecryptionKeyRetriever (cached in retrieverCache)
+ *
+ * Runtime Phase (getKey):
+ * 3. Native code calls getKey(filePath, keyMetadata) ──► JVM
+ * 4. Retrieve cached DecryptionKeyRetriever for filePath
+ * 5. KeyRetriever.getKey(keyMetadata) → decrypted key bytes
+ * 6. Return key bytes ──► Native code for parquet decryption
+ */
+// spotless:on
+
+/**
+ * Helper class to access DecryptionKeyRetriever.getKey from native code via JNI. This class handles
+ * the complexity of creating and caching properly configured DecryptionKeyRetriever instances using
+ * DecryptionPropertiesFactory. The life of this object is meant to map to a single Comet plan, so
+ * associated with CometExecIterator.
+ */
+public class CometFileKeyUnwrapper {
+
+  // Each file path gets a unique DecryptionKeyRetriever
+  private final ConcurrentHashMap<String, DecryptionKeyRetriever> retrieverCache =
+      new ConcurrentHashMap<>();
+
+  // Each hadoopConf yields a unique DecryptionPropertiesFactory. While it's unlikely that
+  // this Comet plan contains more than one hadoopConf, we don't want to assume that. So we'll
+  // provide the ability to cache more than one Factory with a map.
+  private final ConcurrentHashMap<Configuration, DecryptionPropertiesFactory> factoryCache =

Review Comment:
   Since each relation gets a hadoopConf, I could probably make the key the relation name to hash something a bit smaller, but this isn't really on the critical path of query execution so I don't expect it to be a huge win.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich merged PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "hsiang-c (via GitHub)" <gi...@apache.org>.
hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2407800932


##########
spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala:
##########
@@ -191,6 +521,30 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
     }
   }
 
+  /**
+   * Verify that the directory contains an encrypted parquet in plaintext footer mode by means of
+   * checking for all the parquet part files in the parquet directory that their magic string is
+   * PAR1, as defined in the spec:
+   * https://github.com/apache/parquet-format/blob/master/Encryption.md#55-plaintext-footer-mode
+   */
+  private def verifyParquetPlaintextFooter(parquetDir: String): Unit = {
+    val parquetPartitionFiles = getListOfParquetFiles(new File(parquetDir))
+    assert(parquetPartitionFiles.size >= 1)
+    parquetPartitionFiles.foreach { parquetFile =>
+      val magicString = "PAR1"
+      val magicStringLength = magicString.length()
+      val byteArray = new Array[Byte](magicStringLength)
+      val randomAccessFile = new RandomAccessFile(parquetFile, "r")
+      try {
+        randomAccessFile.read(byteArray, 0, magicStringLength)

Review Comment:
   (nit) Do you think we should check the magic number at the tail?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] feat: Parquet Modular Encryption with Spark KMS for native readers [datafusion-comet]

Posted by "parthchandra (via GitHub)" <gi...@apache.org>.
parthchandra commented on PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#issuecomment-3330865362

   Also look at https://github.com/apache/parquet-java/blob/master/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestPropertiesDrivenEncryption.java to see if there are any tests that might be relevant here.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org