Make GetColumnMetaDecryptor return a Decryptor factory

EnricoMi · Feb 6, 2025 · 71f1b4c · 71f1b4c
1 parent 4595667
commit 71f1b4c
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 23 deletions.
diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc
@@ -338,16 +338,16 @@ class SerializedPageReader : public PageReader {
 
 void SerializedPageReader::InitDecryption() {
   // Prepare the AAD for quick update later.
-  if (crypto_ctx_.data_decryptor != nullptr) {
-    ARROW_DCHECK(!crypto_ctx_.data_decryptor->file_aad().empty());
-    data_decryptor_ = std::make_shared<Decryptor>(*crypto_ctx_.data_decryptor);
+  if (crypto_ctx_.data_decryptor) {
+    data_decryptor_ = crypto_ctx_.data_decryptor();
+    ARROW_DCHECK(!data_decryptor_->file_aad().empty());
     data_page_aad_ = encryption::CreateModuleAad(
         data_decryptor_->file_aad(), encryption::kDataPage, crypto_ctx_.row_group_ordinal,
         crypto_ctx_.column_ordinal, kNonPageOrdinal);
   }
-  if (crypto_ctx_.meta_decryptor != nullptr) {
-    ARROW_DCHECK(!crypto_ctx_.meta_decryptor->file_aad().empty());
-    meta_decryptor_ = std::make_shared<Decryptor>(*crypto_ctx_.meta_decryptor);
+  if (crypto_ctx_.meta_decryptor) {
+    meta_decryptor_ = crypto_ctx_.meta_decryptor();
+    ARROW_DCHECK(!meta_decryptor_->file_aad().empty());
     data_page_header_aad_ = encryption::CreateModuleAad(
         meta_decryptor_->file_aad(), encryption::kDataPageHeader,
         crypto_ctx_.row_group_ordinal, crypto_ctx_.column_ordinal, kNonPageOrdinal);

diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h
@@ -103,7 +103,8 @@ class PARQUET_EXPORT LevelDecoder {
 
 struct CryptoContext {
   CryptoContext(bool start_with_dictionary_page, int16_t rg_ordinal, int16_t col_ordinal,
-                std::shared_ptr<Decryptor> meta, std::shared_ptr<Decryptor> data)
+                std::function<std::shared_ptr<Decryptor>()> meta,
+                std::function<std::shared_ptr<Decryptor>()> data)
       : start_decrypt_with_dictionary_page(start_with_dictionary_page),
         row_group_ordinal(rg_ordinal),
         column_ordinal(col_ordinal),
@@ -114,8 +115,8 @@ struct CryptoContext {
   bool start_decrypt_with_dictionary_page = false;
   int16_t row_group_ordinal = -1;
   int16_t column_ordinal = -1;
-  std::shared_ptr<Decryptor> meta_decryptor;
-  std::shared_ptr<Decryptor> data_decryptor;
+  std::function<std::shared_ptr<Decryptor>()> meta_decryptor;
+  std::function<std::shared_ptr<Decryptor>()> data_decryptor;
 };
 
 // Abstract page iterator interface. This way, we can feed column pages to the

diff --git a/cpp/src/parquet/encryption/internal_file_decryptor.cc b/cpp/src/parquet/encryption/internal_file_decryptor.cc
@@ -194,18 +194,22 @@ std::shared_ptr<Decryptor> GetColumnDecryptor(
 
 }  // namespace
 
-std::shared_ptr<Decryptor> GetColumnMetaDecryptor(
+std::function<std::shared_ptr<Decryptor>()> GetColumnMetaDecryptor(
     const ColumnCryptoMetaData* crypto_metadata, InternalFileDecryptor* file_decryptor) {
-  return GetColumnDecryptor(crypto_metadata, file_decryptor,
-                            &InternalFileDecryptor::GetColumnMetaDecryptor,
-                            /*metadata=*/true);
+  return [&]() -> std::shared_ptr<Decryptor> {
+    return GetColumnDecryptor(crypto_metadata, file_decryptor,
+               &InternalFileDecryptor::GetColumnMetaDecryptor,
+               /*metadata=*/true);
+  };
 }
 
-std::shared_ptr<Decryptor> GetColumnDataDecryptor(
+std::function<std::shared_ptr<Decryptor>()> GetColumnDataDecryptor(
     const ColumnCryptoMetaData* crypto_metadata, InternalFileDecryptor* file_decryptor) {
-  return GetColumnDecryptor(crypto_metadata, file_decryptor,
-                            &InternalFileDecryptor::GetColumnDataDecryptor,
-                            /*metadata=*/false);
+  return [&]() -> std::shared_ptr<Decryptor> {
+    return GetColumnDecryptor(crypto_metadata, file_decryptor,
+                              &InternalFileDecryptor::GetColumnDataDecryptor,
+                              /*metadata=*/false);
+  };
 }
 
 void UpdateDecryptor(const std::shared_ptr<Decryptor>& decryptor,

diff --git a/cpp/src/parquet/encryption/internal_file_decryptor.h b/cpp/src/parquet/encryption/internal_file_decryptor.h
@@ -107,11 +107,11 @@ class InternalFileDecryptor {
 };
 
 /// Utility to get column meta decryptor of an encrypted column.
-std::shared_ptr<Decryptor> GetColumnMetaDecryptor(
+std::function<std::shared_ptr<Decryptor>()> GetColumnMetaDecryptor(
     const ColumnCryptoMetaData* crypto_metadata, InternalFileDecryptor* file_decryptor);
 
 /// Utility to get column data decryptor of an encrypted column.
-std::shared_ptr<Decryptor> GetColumnDataDecryptor(
+std::function<std::shared_ptr<Decryptor>()> GetColumnDataDecryptor(
     const ColumnCryptoMetaData* crypto_metadata, InternalFileDecryptor* file_decryptor);
 
 void UpdateDecryptor(const std::shared_ptr<Decryptor>& decryptor,

diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc
@@ -259,9 +259,9 @@ class SerializedRowGroup : public RowGroupReader::Contents {
     }
 
     // The column is encrypted
-    std::shared_ptr<Decryptor> meta_decryptor = GetColumnMetaDecryptor(
+    std::function<std::shared_ptr<Decryptor>()> meta_decryptor = GetColumnMetaDecryptor(
         crypto_metadata.get(), file_metadata_->file_decryptor().get());
-    std::shared_ptr<Decryptor> data_decryptor = GetColumnDataDecryptor(
+    std::function<std::shared_ptr<Decryptor>()> data_decryptor = GetColumnDataDecryptor(
         crypto_metadata.get(), file_metadata_->file_decryptor().get());
     ARROW_DCHECK_NE(meta_decryptor, nullptr);
     ARROW_DCHECK_NE(data_decryptor, nullptr);

diff --git a/cpp/src/parquet/page_index.cc b/cpp/src/parquet/page_index.cc
@@ -256,8 +256,9 @@ class RowGroupPageIndexReaderImpl : public RowGroupPageIndexReader {
     auto descr = row_group_metadata_->schema()->Column(i);
 
     // Get decryptor of column index if encrypted.
-    std::shared_ptr<Decryptor> decryptor = parquet::GetColumnMetaDecryptor(
+    auto decryptor_func = parquet::GetColumnMetaDecryptor(
         col_chunk->crypto_metadata().get(), file_decryptor_);
+    auto decryptor = decryptor_func();
     if (decryptor != nullptr) {
       UpdateDecryptor(decryptor, row_group_ordinal_, /*column_ordinal=*/i,
                       encryption::kColumnIndex);
@@ -295,8 +296,9 @@ class RowGroupPageIndexReaderImpl : public RowGroupPageIndexReader {
     uint32_t length = static_cast<uint32_t>(offset_index_location->length);
 
     // Get decryptor of offset index if encrypted.
-    std::shared_ptr<Decryptor> decryptor =
+    auto decryptor_func =
         GetColumnMetaDecryptor(col_chunk->crypto_metadata().get(), file_decryptor_);
+    auto decryptor = decryptor_func();
     if (decryptor != nullptr) {
       UpdateDecryptor(decryptor, row_group_ordinal_, /*column_ordinal=*/i,
                       encryption::kOffsetIndex);