ARROW-11644: [Python][Parquet] Low-level Parquet decryption in Python

ci/scripts/python_wheel_manylinux_build.sh

@@ -56,6 +56,7 @@ echo "=== (${PYTHON_VERSION}) Building Arrow C++ libraries ==="
 : ${ARROW_MIMALLOC:=ON}
 : ${ARROW_ORC:=ON}
 : ${ARROW_PARQUET:=ON}
+: ${PARQUET_REQUIRE_ENCRYPTION:=ON}
 : ${ARROW_PLASMA:=ON}
 : ${ARROW_S3:=ON}
 : ${ARROW_TENSORFLOW:=ON}
@@ -89,6 +90,7 @@ cmake \
     -DARROW_ORC=${ARROW_ORC} \
     -DARROW_PACKAGE_KIND="python-wheel-manylinux${MANYLINUX_VERSION}" \
     -DARROW_PARQUET=${ARROW_PARQUET} \
+    -DPARQUET_REQUIRE_ENCRYPTION=ON \
     -DARROW_PLASMA=${ARROW_PLASMA} \
     -DARROW_PYTHON=ON \
     -DARROW_RPATH_ORIGIN=ON \

python/pyarrow/_parquet.pxd

@@ -18,6 +18,7 @@
 # distutils: language = c++
 # cython: language_level = 3
 
+from libcpp.map cimport map as c_map
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport (CChunkedArray, CSchema, CStatus,
                                         CTable, CMemoryPool, CBuffer,
@@ -336,7 +337,8 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil:
         void disable_buffered_stream()
         void set_buffer_size(int64_t buf_size)
         int64_t buffer_size() const
-
+        void file_decryption_properties(
+            shared_ptr[FileDecryptionProperties] decryption)
     CReaderProperties default_reader_properties()
 
     cdef cppclass ArrowReaderProperties:
@@ -376,6 +378,9 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
             Builder* encoding(const c_string& path,
                               ParquetEncoding encoding)
             Builder* write_batch_size(int64_t batch_size)
+            Builder* encryption(
+                shared_ptr[FileEncryptionProperties] file_encryption_properties
+            )
             shared_ptr[WriterProperties] build()
 
     cdef cppclass ArrowWriterProperties:
@@ -393,6 +398,55 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
             shared_ptr[ArrowWriterProperties] build()
         c_bool support_deprecated_int96_timestamps()
 
+    cdef enum ParquetCipherType:
+        AES_GCM_V1 "parquet::ParquetCipher::AES_GCM_V1"
+        AES_GCM_CTR_V1 "parquet::ParquetCipher::AES_GCM_CTR_V1"
+
+    cdef cppclass ColumnEncryptionProperties:
+        cppclass Builder:
+            Builder(const c_string& name)
+            Builder(const shared_ptr[ColumnPath]& path)
+            Builder* key(c_string column_key)
+            Builder* key_metadata(const c_string& key_metadata)
+            Builder* key_id(const c_string& key_id)
+            shared_ptr[ColumnEncryptionProperties] build()
+        c_string column_path()
+        c_bool is_encrypted()
+        c_bool is_encrypted_with_footer_key()
+        c_string key()
+        c_string key_metadata()
+        WipeOutEncryptionKey()
+        c_bool is_utilized()
+        shared_ptr[ColumnEncryptionProperties] DeepClone()
+
+    cdef cppclass FileEncryptionProperties:
+        cppclass Builder:
+            Builder(const c_string& footer_key)
+            Builder* set_plaintext_footer()
+            Builder* algorithm(ParquetCipherType parquet_cipher)
+            Builder* footer_key_id(const c_string& key_id)
+            Builder* footer_key_metadata(const c_string& footer_key_metadata)
+            Builder* aad_prefix(const c_string& aad_prefix)
+            Builder* disable_aad_prefix_storage()
+            Builder* encrypted_columns(
+                const c_map[c_string, shared_ptr[
+                    ColumnEncryptionProperties]]& encrypted_columns)
+            shared_ptr[FileEncryptionProperties] build()
+        c_bool encrypted_footer()
+        # EncryptionAlgorithm algorithm()
+        c_string footer_key()
+        c_string footer_key_metadata()
+        c_string file_aad()
+        shared_ptr[c_map[c_string, shared_ptr[
+            ColumnEncryptionProperties]]] column_encryption_properties(
+                const c_string& column_path)
+        c_bool is_utilized()
+        set_utilized()
+        WipeOutEncryptionKeys()
+        shared_ptr[FileEncryptionProperties] DeepClone(c_string new_aad_prefix)
+        c_map[c_string, shared_ptr[
+            ColumnEncryptionProperties]] encrypted_columns()
+
 
 cdef extern from "parquet/arrow/reader.h" namespace "parquet::arrow" nogil:
     cdef cppclass FileReader:
@@ -450,6 +504,58 @@ cdef extern from "parquet/arrow/reader.h" namespace "parquet::arrow" nogil:
         const shared_ptr[const CKeyValueMetadata]& key_value_metadata,
         shared_ptr[CSchema]* out)
 
+cdef extern from "parquet/encryption/encryption.h" namespace "parquet" nogil:
+    cdef cppclass DecryptionKeyRetriever:
+        c_string GetKey(const c_string& key_metadata)
+
+    cdef cppclass IntegerKeyIdReceiver:
+        PutKey(uint32_t key_id, const c_string& key)
+        c_string GetKey(const c_string& key_metadata)
+
+    cdef cppclass StringKeyIdReceiver:
+        PutKey(const c_string& key_id, const c_string& key)
+        c_string GetKey(const c_string& key_metadata)
+
+    cdef cppclass ColumnDecryptionProperties:
+        cppclass Builder:
+            Builder(const c_string& name)
+            Builder(const shared_ptr[ColumnPath]& path)
+            Builder* key(const c_string& key)
+            shared_ptr[ColumnDecryptionProperties] build()
+        c_string column_path()
+        c_string key()
+        c_bool is_utilized()
+        c_bool set_utilized()
+        WipeOutDecryptionKey()
+        shared_ptr[ColumnDecryptionProperties] DeepClone()
+
+    cdef cppclass FileDecryptionProperties:
+        cppclass Builder:
+            Builder()
+            Builder* footer_key(const c_string footer_key)
+            Builder* column_keys(const c_map[c_string, shared_ptr[
+                ColumnDecryptionProperties]]& column_decryption_properties)
+            Builder* key_retriever(
+                const shared_ptr[DecryptionKeyRetriever]& key_retriever)
+            Builder* disable_footer_signature_verification()
+            Builder* aad_prefix(const c_string& aad_prefix)
+            # Builder* aad_prefix_verifier(shared_ptr[
+            #    AADPrefixVerifier] aad_prefix_verifier]
+            Builder* plaintext_files_allowed()
+            shared_ptr[FileDecryptionProperties] build()
+        c_string column_key()
+        c_string footer_key()
+        c_string aad_prefix()
+        const shared_ptr[DecryptionKeyRetriever]& key_retriever()
+        c_bool check_plaintext_footer_integrity()
+        c_bool plaintext_files_allowed()
+        # shared_ptr[AADPrefixVerifier]& aad_prefix_verifier()
+        WipeOutDecryptionKeys()
+        c_bool is_utilized()
+        set_utilized()
+        shared_ptr[FileDecryptionProperties] DeepClone()
+
+
 cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil:
 
     CStatus ToParquetSchema(
@@ -501,7 +607,8 @@ cdef shared_ptr[WriterProperties] _create_writer_properties(
     data_page_size=*,
     compression_level=*,
     use_byte_stream_split=*,
-    data_page_version=*) except *
+    data_page_version=*,
+    lowlevel_encryption_properties=*) except *
 
 
 cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties(

python/pyarrow/_parquet.pyx

@@ -913,6 +913,68 @@ cdef ParquetCompression compression_from_name(name):
         return ParquetCompression_UNCOMPRESSED
 
 
+cdef cppclass PythonKeyRetriever(DecryptionKeyRetriever):
+    # Retrieve a decryption key from a Python function.
+    object pyobject
+
+    __init__(object pyobject):
+        this.pyobject = pyobject
+
+    c_string GetKey(const c_string& key_metadata) nogil:
+        with gil:
+            return this.pyobject(key_metadata)
+
+
+cdef class LowLevelDecryptionProperties(_Weakrefable):
+    """Wrapper for FileDecryptionProperties."""
+    cdef:
+        shared_ptr[FileDecryptionProperties] decryption_properties
+
+    def __cinit__(self,
+                  footer_key=None,
+                  # Map column name (str) to key (bytes):
+                  column_keys=None,
+                  # Callable that takes key metadata (bytes), returns the key
+                  # (bytes); equivalent to DecryptionKeyRetriever in C++ API.
+                  retrieve_key=None,
+                  disable_footer_signature_verification=False,
+                  plaintext_files_allowed=False):
+        cdef c_map[c_string, shared_ptr[
+            ColumnDecryptionProperties]] c_column_keys
+        cdef FileDecryptionProperties.Builder builder
+        cdef c_string column, key
+        cdef shared_ptr[DecryptionKeyRetriever] key_retriever
+
+        builder = FileDecryptionProperties.Builder()
+
+        if retrieve_key is not None:
+            if not callable(retrieve_key):
+                raise ValueError(
+                    "retrieve_key %s is not callable" % (retrieve_key,))
+            key_retriever = shared_ptr[DecryptionKeyRetriever](
+                new PythonKeyRetriever(retrieve_key)
+            )
+            builder.key_retriever(key_retriever)
+
+        if footer_key is not None:
+            builder.footer_key(footer_key)
+
+        if column_keys is not None:
+            for pycolumn, key in column_keys.items():
+                column = pycolumn.encode("utf-8")
+                c_column_keys[column] = ColumnDecryptionProperties.Builder(
+                    column).key(key).build()
+            builder.column_keys(c_column_keys)
+
+        if disable_footer_signature_verification:
+            builder.disable_footer_signature_verification()
+
+        if plaintext_files_allowed:
+            builder.plaintext_files_allowed()
+
+        self.decryption_properties = builder.build()
+
+
 cdef class ParquetReader(_Weakrefable):
     cdef:
         object source
@@ -929,7 +991,8 @@ cdef class ParquetReader(_Weakrefable):
 
     def open(self, object source not None, bint use_memory_map=True,
              read_dictionary=None, FileMetaData metadata=None,
-             int buffer_size=0, bint pre_buffer=False):
+             int buffer_size=0, bint pre_buffer=False,
+             LowLevelDecryptionProperties lldecrypt=None):
         cdef:
             shared_ptr[CRandomAccessFile] rd_handle
             shared_ptr[CFileMetaData] c_metadata
@@ -938,6 +1001,7 @@ cdef class ParquetReader(_Weakrefable):
                 default_arrow_reader_properties())
             c_string path
             FileReaderBuilder builder
+            shared_ptr[FileDecryptionProperties] decryption_properties
 
         if metadata is not None:
             c_metadata = metadata.sp_metadata
@@ -950,6 +1014,10 @@ cdef class ParquetReader(_Weakrefable):
         else:
             raise ValueError('Buffer size must be larger than zero')
 
+        if lldecrypt is not None:
+            decryption_properties = lldecrypt.decryption_properties
+            properties.file_decryption_properties(decryption_properties)
+
         arrow_props.set_pre_buffer(pre_buffer)
 
         self.source = source
@@ -1175,19 +1243,61 @@ cdef class ParquetReader(_Weakrefable):
         return pyarrow_wrap_chunked_array(out)
 
 
+cdef class LowLevelEncryptionProperties(_Weakrefable):
+    """Wrapper for FileEncryptionProperties.
+
+    Not intended for direct use by Python users, currently, just wrapped to
+    enable testing of low-level decryption.
+    """
+    cdef:
+        shared_ptr[FileEncryptionProperties] encryption_properties
+
+    def __cinit__(self, footer_key, column_keys=None, plaintext_footer=False,
+                  footer_key_metadata=None, column_keys_metadata=None):
+        cdef FileEncryptionProperties.Builder* builder
+        cdef c_string column
+        cdef c_map[c_string, shared_ptr[
+            ColumnEncryptionProperties]] c_column_keys
+
+        builder = new FileEncryptionProperties.Builder(footer_key)
+        if footer_key_metadata is not None:
+            builder.footer_key_metadata(footer_key_metadata)
+
+        if column_keys is not None:
+            for pycolumn, key in column_keys.items():
+                column = pycolumn.encode("utf-8")
+                if column_keys_metadata is not None:
+                    key_id = column_keys_metadata[pycolumn]
+                else:
+                    key_id = column
+                c_column_keys[column] = ColumnEncryptionProperties.Builder(
+                    column).key(key).key_metadata(key_id).build()
+            builder.encrypted_columns(c_column_keys)
+
+        if plaintext_footer:
+            builder.set_plaintext_footer()
+
+        self.encryption_properties = builder.build()
+        del builder
+
+
 cdef shared_ptr[WriterProperties] _create_writer_properties(
-        use_dictionary=None,
-        compression=None,
-        version=None,
-        write_statistics=None,
-        data_page_size=None,
-        compression_level=None,
-        use_byte_stream_split=False,
-        data_page_version=None) except *:
+    use_dictionary=None,
+    compression=None,
+    version=None,
+    write_statistics=None,
+    data_page_size=None,
+    compression_level=None,
+    use_byte_stream_split=False,
+    data_page_version=None,
+    lowlevel_encryption_properties=None
+) except *:
     """General writer properties"""
     cdef:
         shared_ptr[WriterProperties] properties
         WriterProperties.Builder props
+        LowLevelEncryptionProperties llep
+        shared_ptr[FileEncryptionProperties] encryption_properties
 
     # data_page_version
 
@@ -1266,6 +1376,12 @@ cdef shared_ptr[WriterProperties] _create_writer_properties(
     if data_page_size is not None:
         props.data_pagesize(data_page_size)
 
+    # encryption
+    if lowlevel_encryption_properties is not None:
+        llep = lowlevel_encryption_properties
+        encryption_properties = llep.encryption_properties
+        props.encryption(encryption_properties)
+
     properties = props.build()
 
     return properties
@@ -1365,6 +1481,7 @@ cdef class ParquetWriter(_Weakrefable):
                   use_byte_stream_split=False,
                   writer_engine_version=None,
                   data_page_version=None,
+                  lowlevel_encryption_properties=None,
                   use_compliant_nested_type=False):
         cdef:
             shared_ptr[WriterProperties] properties
@@ -1391,7 +1508,8 @@ cdef class ParquetWriter(_Weakrefable):
             data_page_size=data_page_size,
             compression_level=compression_level,
             use_byte_stream_split=use_byte_stream_split,
-            data_page_version=data_page_version
+            data_page_version=data_page_version,
+            lowlevel_encryption_properties=lowlevel_encryption_properties,
         )
         arrow_properties = _create_arrow_writer_properties(
             use_deprecated_int96_timestamps=use_deprecated_int96_timestamps,

python/pyarrow/parquet.py

@@ -35,7 +35,8 @@
 from pyarrow._parquet import (ParquetReader, Statistics,  # noqa
                               FileMetaData, RowGroupMetaData,
                               ColumnChunkMetaData,
-                              ParquetSchema, ColumnSchema)
+                              ParquetSchema, ColumnSchema,
+                              LowLevelDecryptionProperties)
 from pyarrow.fs import (LocalFileSystem, FileSystem,
                         _resolve_filesystem_and_path, _ensure_filesystem)
 from pyarrow import filesystem as legacyfs
@@ -213,15 +214,19 @@ class ParquetFile:
         Coalesce and issue file reads in parallel to improve performance on
         high-latency filesystems (e.g. S3). If True, Arrow will use a
         background I/O thread pool.
+    low_level_decryption: LowLevelDecryptionProperties, default None
+        Low-level decryption properties. Note that a nicer high-level API
+        should be available eventually.
     """
 
     def __init__(self, source, metadata=None, common_metadata=None,
                  read_dictionary=None, memory_map=False, buffer_size=0,
-                 pre_buffer=False):
+                 pre_buffer=False, low_level_decryption=None):
         self.reader = ParquetReader()
         self.reader.open(source, use_memory_map=memory_map,
                          buffer_size=buffer_size, pre_buffer=pre_buffer,
-                         read_dictionary=read_dictionary, metadata=metadata)
+                         read_dictionary=read_dictionary, metadata=metadata,
+                         lldecrypt=low_level_decryption)
         self.common_metadata = common_metadata
         self._nested_paths_by_prefix = self._build_nested_paths()