Skip to content

Commit 70f5c26

Browse files
EnricoMiEnricoMi
authored
GH-51361: [C++][Parquet] Derive records_to_read in FileReaderImpl::ReadColumn from RowGroup (#51362)
### Rationale for this change Fixes #51361. ### What changes are included in this PR? Method `FileReaderImpl::ReadColumn` should derive `records_to_read` from the RowGroup rather than ColumnChunk's `num_values`. Deriving the right ColumnChunk index in `FileReaderImpl::DecodeRowGroups` is not trivial for nested schemas. This simplifies `FileReaderImpl::ReadColumn` and fixes #51361. This was silently masked for full-schema reads and for columns with identical num_values(), but surfaces as a hard failure when an earlier, unselected column requires decryption: reading only a trailing plaintext column of a partially column-key-encrypted, plaintext-footer Parquet file threw "Cannot decrypt ColumnMetadata" even though the requested column was never encrypted. This never corrupts data on unencrypted files: ReadColumn's wrong index is only ever used to look up ColumnChunk(i)->num_values(), a count fed into the *already-correct* reader as an upper bound on how many records to decode. Every row contributes at least one definition/repetition-level entry, so num_values() for any column is always >= that row group's true row count, and every column in a row group shares the same row count. ### Are these changes tested? Yes, in the context of reading a plaintext column of a partially encrypted Parquet file. This cannot be tested with non-encrypted files. ### Are there any user-facing changes? No. ### Was AI used for this PR? In accordance to the [AI generation guidelines](https://arrow.apache.org/docs/dev/developers/overview.html#ai-generated-code), please disclose below whether and how AI was used in this PR. **PR code and description written by:** - [X] Human - [X] AI **Reviewed before submission by:** - [X] Human - [ ] AI - [ ] Not reviewed * GitHub Issue: #51361 Lead-authored-by: Enrico Minack <enrico.minack@insightsoftmax.com> Co-authored-by: Enrico Minack <github@enrico.minack.dev> Signed-off-by: Adam Reeve <adreeve@gmail.com>
1 parent c07de67 commit 70f5c26

3 files changed

Lines changed: 48 additions & 42 deletions

File tree

‎cpp/src/parquet/arrow/reader.cc‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -271,13 +271,11 @@ class FileReaderImpl : public FileReader {
271271
Status ReadColumn(int i, const std::vector<int>& row_groups, ColumnReader* reader,
272272
std::shared_ptr<ChunkedArray>* out) {
273273
BEGIN_PARQUET_CATCH_EXCEPTIONS
274-
// TODO(wesm): This calculation doesn't make much sense when we have repeated
275-
// schema nodes
274+
// NextBatch()'s size is a number of records (rows), not leaf values, so use the
275+
// row group's own row count directly rather than some column's num_values().
276276
int64_t records_to_read = 0;
277277
for (auto row_group : row_groups) {
278-
// Can throw exception
279-
records_to_read +=
280-
reader_->metadata()->RowGroup(row_group)->ColumnChunk(i)->num_values();
278+
records_to_read += reader_->metadata()->RowGroup(row_group)->num_rows();
281279
}
282280
#ifdef ARROW_WITH_OPENTELEMETRY
283281
std::string column_name = reader_->metadata()->schema()->Column(i)->name();

‎python/pyarrow/tests/parquet/test_encryption.py‎

Lines changed: 44 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -501,45 +501,58 @@ def validate_kms_connection_config(kms_connection_config):
501501
validate_kms_connection_config(kms_connection_config_1)
502502

503503

504-
@pytest.mark.xfail(reason="Plaintext footer - reading plaintext column subset"
505-
" reads encrypted columns too")
506504
def test_encrypted_parquet_write_read_plain_footer_single_wrapping(
507505
tempdir, data_table):
508-
"""Write an encrypted parquet, with plaintext footer
509-
and with single wrapping,
510-
verify it's encrypted, and then read plaintext columns."""
506+
"""
507+
Write an encrypted parquet, with plaintext footer and with single wrapping,
508+
verify it's encrypted, and then read plaintext columns. Runs once with a
509+
flat schema and once where the encrypted column `b` is itself a nested
510+
(struct) field.
511+
"""
511512
path = tempdir / PARQUET_NAME
512513

513-
# Encrypt the footer with the footer key,
514-
# encrypt column `a` and column `b` with another key,
515-
# keep `c` plaintext
516-
encryption_config = pe.EncryptionConfiguration(
517-
footer_key=FOOTER_KEY_NAME,
518-
column_keys={
519-
COL_KEY_NAME: ["a", "b"],
520-
},
521-
plaintext_footer=True,
522-
double_wrapping=False)
514+
for nested in [False, True]:
515+
if nested:
516+
table = pa.Table.from_pydict({
517+
'a': pa.array([1, 2, 3]),
518+
'b': pa.array(
519+
[{'x': 1, 'y': 2}, {'x': 3, 'y': 4}, {'x': 5, 'y': 6}],
520+
type=pa.struct([('x', pa.int32()), ('y', pa.int32())])),
521+
'c': pa.array(['x', 'y', 'z'])
522+
})
523+
else:
524+
table = data_table
525+
526+
# Encrypt the footer with the footer key,
527+
# encrypt column `a` and column `b` with another key, keep `c` plaintext
528+
encryption_config = pe.EncryptionConfiguration(
529+
footer_key=FOOTER_KEY_NAME,
530+
column_keys={
531+
COL_KEY_NAME: ["a", "b"],
532+
},
533+
plaintext_footer=True,
534+
double_wrapping=False)
523535

524-
kms_connection_config = pe.KmsConnectionConfig(
525-
custom_kms_conf={
526-
FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"),
527-
COL_KEY_NAME: COL_KEY.decode("UTF-8"),
528-
}
529-
)
536+
kms_connection_config = pe.KmsConnectionConfig(
537+
custom_kms_conf={
538+
FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"),
539+
COL_KEY_NAME: COL_KEY.decode("UTF-8"),
540+
}
541+
)
530542

531-
def kms_factory(kms_connection_configuration):
532-
return InMemoryKmsClient(kms_connection_configuration)
543+
def kms_factory(kms_connection_configuration):
544+
return InMemoryKmsClient(kms_connection_configuration)
533545

534-
crypto_factory = pe.CryptoFactory(kms_factory)
535-
# Write with encryption properties
536-
write_encrypted_parquet(path, data_table, encryption_config,
537-
kms_connection_config, crypto_factory)
546+
crypto_factory = pe.CryptoFactory(kms_factory)
547+
# Write with encryption properties
548+
write_encrypted_parquet(path, table, encryption_config,
549+
kms_connection_config, crypto_factory)
538550

539-
# # Read without decryption properties only the plaintext column
540-
# result = pq.ParquetFile(path)
541-
# result_table = result.read(columns='c', use_threads=False)
542-
# assert table.num_rows == result_table.num_rows
551+
# Read the plaintext column without decryption properties
552+
with pq.ParquetFile(path) as result:
553+
result_table = result.read(columns='c', use_threads=False)
554+
assert table.num_rows == result_table.num_rows
555+
assert table.select(['c']).equals(result_table)
543556

544557

545558
def test_encrypted_parquet_write_read_external(tempdir, data_table,

‎python/pyarrow/tests/test_dataset_encryption.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -115,11 +115,9 @@ def do_test_dataset_encryption_decryption(table, extra_column_path=None):
115115
if extra_column_path:
116116
keys = dict(**KEYS, **{EXTRA_COL_KEY_NAME: EXTRA_COL_KEY})
117117
column_keys = dict(**COLUMN_KEYS, **{EXTRA_COL_KEY_NAME: [extra_column_path]})
118-
extra_column_name = extra_column_path.split(".")[0]
119118
else:
120119
keys = KEYS
121120
column_keys = COLUMN_KEYS
122-
extra_column_name = None
123121

124122
# define the actual test
125123
def assert_decrypts(
@@ -235,13 +233,10 @@ def assert_decrypts(
235233
for key_name, key in keys.items()
236234
if key_name in [FOOTER_KEY_NAME, column_key_name]}
237235

238-
# that one encrypted column can only be read
239-
# if it is not a column path / nested field
240-
plaintext_and_one_success = encrypted_column_name != extra_column_name
241236
plaintext_and_one = plaintext_column_names + [encrypted_column_name]
242237

243238
assert_decrypts(read_keys, plaintext_column_names, True)
244-
assert_decrypts(read_keys, plaintext_and_one, plaintext_and_one_success)
239+
assert_decrypts(read_keys, plaintext_and_one, True)
245240
assert_decrypts(read_keys, encrypted_column_names, False)
246241
assert_decrypts(read_keys, all_column_names, False)
247242

0 commit comments

Comments
 (0)