diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index eca83e8576d..bb8bacc7751 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -271,13 +271,11 @@ class FileReaderImpl : public FileReader { Status ReadColumn(int i, const std::vector& row_groups, ColumnReader* reader, std::shared_ptr* out) { BEGIN_PARQUET_CATCH_EXCEPTIONS - // TODO(wesm): This calculation doesn't make much sense when we have repeated - // schema nodes + // NextBatch()'s size is a number of records (rows), not leaf values, so use the + // row group's own row count directly rather than some column's num_values(). int64_t records_to_read = 0; for (auto row_group : row_groups) { - // Can throw exception - records_to_read += - reader_->metadata()->RowGroup(row_group)->ColumnChunk(i)->num_values(); + records_to_read += reader_->metadata()->RowGroup(row_group)->num_rows(); } #ifdef ARROW_WITH_OPENTELEMETRY std::string column_name = reader_->metadata()->schema()->Column(i)->name(); diff --git a/python/pyarrow/tests/parquet/test_encryption.py b/python/pyarrow/tests/parquet/test_encryption.py index 6a3842f3edf..5a34f32bc3f 100644 --- a/python/pyarrow/tests/parquet/test_encryption.py +++ b/python/pyarrow/tests/parquet/test_encryption.py @@ -501,45 +501,58 @@ def validate_kms_connection_config(kms_connection_config): validate_kms_connection_config(kms_connection_config_1) -@pytest.mark.xfail(reason="Plaintext footer - reading plaintext column subset" - " reads encrypted columns too") def test_encrypted_parquet_write_read_plain_footer_single_wrapping( tempdir, data_table): - """Write an encrypted parquet, with plaintext footer - and with single wrapping, - verify it's encrypted, and then read plaintext columns.""" + """ + Write an encrypted parquet, with plaintext footer and with single wrapping, + verify it's encrypted, and then read plaintext columns. Runs once with a + flat schema and once where the encrypted column `b` is itself a nested + (struct) field. + """ path = tempdir / PARQUET_NAME - # Encrypt the footer with the footer key, - # encrypt column `a` and column `b` with another key, - # keep `c` plaintext - encryption_config = pe.EncryptionConfiguration( - footer_key=FOOTER_KEY_NAME, - column_keys={ - COL_KEY_NAME: ["a", "b"], - }, - plaintext_footer=True, - double_wrapping=False) + for nested in [False, True]: + if nested: + table = pa.Table.from_pydict({ + 'a': pa.array([1, 2, 3]), + 'b': pa.array( + [{'x': 1, 'y': 2}, {'x': 3, 'y': 4}, {'x': 5, 'y': 6}], + type=pa.struct([('x', pa.int32()), ('y', pa.int32())])), + 'c': pa.array(['x', 'y', 'z']) + }) + else: + table = data_table + + # Encrypt the footer with the footer key, + # encrypt column `a` and column `b` with another key, keep `c` plaintext + encryption_config = pe.EncryptionConfiguration( + footer_key=FOOTER_KEY_NAME, + column_keys={ + COL_KEY_NAME: ["a", "b"], + }, + plaintext_footer=True, + double_wrapping=False) - kms_connection_config = pe.KmsConnectionConfig( - custom_kms_conf={ - FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), - COL_KEY_NAME: COL_KEY.decode("UTF-8"), - } - ) + kms_connection_config = pe.KmsConnectionConfig( + custom_kms_conf={ + FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"), + COL_KEY_NAME: COL_KEY.decode("UTF-8"), + } + ) - def kms_factory(kms_connection_configuration): - return InMemoryKmsClient(kms_connection_configuration) + def kms_factory(kms_connection_configuration): + return InMemoryKmsClient(kms_connection_configuration) - crypto_factory = pe.CryptoFactory(kms_factory) - # Write with encryption properties - write_encrypted_parquet(path, data_table, encryption_config, - kms_connection_config, crypto_factory) + crypto_factory = pe.CryptoFactory(kms_factory) + # Write with encryption properties + write_encrypted_parquet(path, table, encryption_config, + kms_connection_config, crypto_factory) - # # Read without decryption properties only the plaintext column - # result = pq.ParquetFile(path) - # result_table = result.read(columns='c', use_threads=False) - # assert table.num_rows == result_table.num_rows + # Read the plaintext column without decryption properties + with pq.ParquetFile(path) as result: + result_table = result.read(columns='c', use_threads=False) + assert table.num_rows == result_table.num_rows + assert table.select(['c']).equals(result_table) def test_encrypted_parquet_write_read_external(tempdir, data_table, diff --git a/python/pyarrow/tests/test_dataset_encryption.py b/python/pyarrow/tests/test_dataset_encryption.py index 0ef3931a4cf..ceb75ca603b 100644 --- a/python/pyarrow/tests/test_dataset_encryption.py +++ b/python/pyarrow/tests/test_dataset_encryption.py @@ -115,11 +115,9 @@ def do_test_dataset_encryption_decryption(table, extra_column_path=None): if extra_column_path: keys = dict(**KEYS, **{EXTRA_COL_KEY_NAME: EXTRA_COL_KEY}) column_keys = dict(**COLUMN_KEYS, **{EXTRA_COL_KEY_NAME: [extra_column_path]}) - extra_column_name = extra_column_path.split(".")[0] else: keys = KEYS column_keys = COLUMN_KEYS - extra_column_name = None # define the actual test def assert_decrypts( @@ -235,13 +233,10 @@ def assert_decrypts( for key_name, key in keys.items() if key_name in [FOOTER_KEY_NAME, column_key_name]} - # that one encrypted column can only be read - # if it is not a column path / nested field - plaintext_and_one_success = encrypted_column_name != extra_column_name plaintext_and_one = plaintext_column_names + [encrypted_column_name] assert_decrypts(read_keys, plaintext_column_names, True) - assert_decrypts(read_keys, plaintext_and_one, plaintext_and_one_success) + assert_decrypts(read_keys, plaintext_and_one, True) assert_decrypts(read_keys, encrypted_column_names, False) assert_decrypts(read_keys, all_column_names, False)