Spaces:
Running
Running
| # 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. | |
| import shutil | |
| import os | |
| from datetime import timedelta | |
| import pyarrow as pa | |
| import pyarrow.dataset as ds | |
| import pyarrow.parquet.encryption as pe | |
| from pyarrow.tests.parquet.encryption import InMemoryKmsClient | |
| """ A sample to demonstrate parquet dataset encryption and decryption""" | |
| # create a list of dictionaries that will represent our dataset | |
| table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], | |
| 'n_legs': [2, 2, 4, 4, 5, 100], | |
| 'animal': ["Flamingo", "Parrot", "Dog", "Horse", | |
| "Brittle stars", "Centipede"]}) | |
| # create a PyArrow dataset from the table | |
| dataset = ds.dataset(table) | |
| FOOTER_KEY = b"0123456789112345" | |
| FOOTER_KEY_NAME = "footer_key" | |
| COL_KEY = b"1234567890123450" | |
| COL_KEY_NAME = "col_key" | |
| encryption_config = pe.EncryptionConfiguration( | |
| footer_key=FOOTER_KEY_NAME, | |
| plaintext_footer=False, | |
| # Use COL_KEY_NAME to encrypt `n_legs` and `animal` columns. | |
| column_keys={ | |
| COL_KEY_NAME: ["n_legs", "animal"], | |
| }, | |
| encryption_algorithm="AES_GCM_V1", | |
| # requires timedelta or an assertion is raised | |
| cache_lifetime=timedelta(minutes=5.0), | |
| data_key_length_bits=256) | |
| kms_connection_config = pe.KmsConnectionConfig( | |
| custom_kms_conf={ | |
| FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), | |
| COL_KEY_NAME: COL_KEY.decode("UTF-8"), | |
| } | |
| ) | |
| decryption_config = pe.DecryptionConfiguration(cache_lifetime=300) | |
| def kms_factory(kms_connection_configuration): | |
| return InMemoryKmsClient(kms_connection_configuration) | |
| crypto_factory = pe.CryptoFactory(kms_factory) | |
| parquet_encryption_cfg = ds.ParquetEncryptionConfig( | |
| crypto_factory, kms_connection_config, encryption_config) | |
| parquet_decryption_cfg = ds.ParquetDecryptionConfig(crypto_factory, | |
| kms_connection_config, | |
| decryption_config) | |
| # set encryption config for parquet fragment scan options | |
| pq_scan_opts = ds.ParquetFragmentScanOptions() | |
| pq_scan_opts.parquet_decryption_config = parquet_decryption_cfg | |
| pformat = pa.dataset.ParquetFileFormat(default_fragment_scan_options=pq_scan_opts) | |
| if os.path.exists('sample_dataset'): | |
| shutil.rmtree('sample_dataset') | |
| write_options = pformat.make_write_options( | |
| encryption_config=parquet_encryption_cfg) | |
| ds.write_dataset(data=dataset, base_dir="sample_dataset", | |
| partitioning=['year'], format=pformat, file_options=write_options) | |
| # read the dataset back | |
| dataset = ds.dataset('sample_dataset', format=pformat) | |
| # print the dataset | |
| print(dataset.to_table()) | |