Allow retrieving Parquet decryption keys using the key metadata (#7286)

* Allow retrieving decryption keys from key metadata

* Refactor to disallow setting keys and a key retriever

* Remove unnecessary clone

* Implement PartialEq explicitly for DecryptionKeys rather than FileDecryptionProperties

* Add key retrieval methods to FileDecryptionProperties

* Get column path from ColumnCryptoMetaData

* Tidy up

* Revert rename of parameter

Author: adamreeve

parquet/src/arrow/arrow_reader/mod.rs

@@ -700,36 +700,19 @@ impl<T: ChunkReader + 'static> Iterator for ReaderPageIterator<T> {
 
         #[cfg(feature = "encryption")]
         let crypto_context = if let Some(file_decryptor) = self.metadata.file_decryptor() {
-            let column_name = self
-                .metadata
-                .file_metadata()
-                .schema_descr()
-                .column(self.column_idx);
-
-            if file_decryptor.is_column_encrypted(column_name.name()) {
-                let data_decryptor = file_decryptor.get_column_data_decryptor(column_name.name());
-                let data_decryptor = match data_decryptor {
-                    Ok(data_decryptor) => data_decryptor,
-                    Err(err) => return Some(Err(err)),
-                };
-
-                let metadata_decryptor =
-                    file_decryptor.get_column_metadata_decryptor(column_name.name());
-                let metadata_decryptor = match metadata_decryptor {
-                    Ok(metadata_decryptor) => metadata_decryptor,
-                    Err(err) => return Some(Err(err)),
-                };
-
-                let crypto_context = CryptoContext::new(
-                    rg_idx,
-                    self.column_idx,
-                    data_decryptor,
-                    metadata_decryptor,
-                    file_decryptor.file_aad().clone(),
-                );
-                Some(Arc::new(crypto_context))
-            } else {
-                None
+            match meta.crypto_metadata() {
+                Some(crypto_metadata) => {
+                    match CryptoContext::for_column(
+                        file_decryptor,
+                        crypto_metadata,
+                        rg_idx,
+                        self.column_idx,
+                    ) {
+                        Ok(context) => Some(Arc::new(context)),
+                        Err(err) => return Some(Err(err)),
+                    }
+                }
+                None => None,
             }
         } else {
             None

parquet/src/arrow/async_reader/mod.rs

@@ -1027,36 +1027,6 @@ impl RowGroups for InMemoryRowGroup<'_> {
     }
 
     fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>> {
-        #[cfg(feature = "encryption")]
-        let crypto_context = if let Some(file_decryptor) = self.metadata.clone().file_decryptor() {
-            let column_name = &self
-                .metadata
-                .clone()
-                .file_metadata()
-                .schema_descr()
-                .column(i);
-
-            if file_decryptor.is_column_encrypted(column_name.name()) {
-                let data_decryptor =
-                    file_decryptor.get_column_data_decryptor(column_name.name())?;
-                let metadata_decryptor =
-                    file_decryptor.get_column_metadata_decryptor(column_name.name())?;
-
-                let crypto_context = CryptoContext::new(
-                    self.row_group_idx,
-                    i,
-                    data_decryptor,
-                    metadata_decryptor,
-                    file_decryptor.file_aad().clone(),
-                );
-                Some(Arc::new(crypto_context))
-            } else {
-                None
-            }
-        } else {
-            None
-        };
-
         match &self.column_chunks[i] {
             None => Err(ParquetError::General(format!(
                 "Invalid column index {i}, column was not fetched"
@@ -1067,14 +1037,29 @@ impl RowGroups for InMemoryRowGroup<'_> {
                     // filter out empty offset indexes (old versions specified Some(vec![]) when no present)
                     .filter(|index| !index.is_empty())
                     .map(|index| index[i].page_locations.clone());
-                let metadata = self.metadata.row_group(self.row_group_idx);
+                let column_metadata = self.metadata.row_group(self.row_group_idx).column(i);
                 let page_reader = SerializedPageReader::new(
                     data.clone(),
-                    metadata.column(i),
+                    column_metadata,
                     self.row_count,
                     page_locations,
                 )?;
 
+                #[cfg(feature = "encryption")]
+                let crypto_context = if let Some(file_decryptor) = self.metadata.file_decryptor() {
+                    match column_metadata.crypto_metadata() {
+                        Some(crypto_metadata) => Some(Arc::new(CryptoContext::for_column(
+                            file_decryptor,
+                            crypto_metadata,
+                            self.row_group_idx,
+                            i,
+                        )?)),
+                        None => None,
+                    }
+                } else {
+                    None
+                };
+
                 #[cfg(feature = "encryption")]
                 let page_reader = page_reader.with_crypto_context(crypto_context);