Skip to content

Commit cd1811b

Browse files
authored
GH-48254: [Python][Parquet] Support extension types in read_schema (#48255)
### Rationale for this change pq.read_schema drops extension types (UUID comes back as fixed_size_binary[16]), while ParquetFile.schema_arrow and read_table preserve them. Schema inspection via metadata should match table/extension behavior. ### What changes are included in this PR? - Plumb arrow_extensions_enabled into read_schema and return schema_arrow when enabled so extension types are preserved. - Add regression test ensuring UUID extension types are retained by read_schema and downgraded to binary(16) when extensions are disabled. ### Are these changes tested? - Yes: added unit test test_read_schema_uuid_extension_type ### Are there any user-facing changes? - Behavior improvement: read_schema now preserves extension types (e.g., UUID) when extensions are enabled; no API break Notes: - I don't know if the fact the column types being returned are now extension<arrow.uuid> instead of fixed_size_binary[16], is considered a breaking change. - This PR patch was AI generated, but I personally reviewed it, the scope is small, and it looks fine to me. * GitHub Issue: #48254 Authored-by: Nicolas Vandeginste <n.vandeginste@abc-arbitrage.com> Signed-off-by: Raúl Cumplido <raulcumplido@gmail.com>
1 parent 4ea1ad2 commit cd1811b

2 files changed

Lines changed: 46 additions & 10 deletions

File tree

python/pyarrow/parquet/core.py

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -265,9 +265,9 @@ class ParquetFile:
265265
page_checksum_verification : bool, default False
266266
If True, verify the checksum for each page read from the file.
267267
arrow_extensions_enabled : bool, default True
268-
If True, read Parquet logical types as Arrow extension types where possible,
269-
(e.g., read JSON as the canonical `arrow.json` extension type or UUID as
270-
the canonical `arrow.uuid` extension type).
268+
If True, read Parquet logical types as Arrow extension types where
269+
possible (e.g., read JSON as the canonical `arrow.json` extension type
270+
or UUID as the canonical `arrow.uuid` extension type).
271271
272272
Examples
273273
--------
@@ -2372,7 +2372,7 @@ def write_metadata(schema, where, metadata_collector=None, filesystem=None,
23722372

23732373

23742374
def read_metadata(where, memory_map=False, decryption_properties=None,
2375-
filesystem=None):
2375+
filesystem=None, arrow_extensions_enabled=True):
23762376
"""
23772377
Read FileMetaData from footer of a single Parquet file.
23782378
@@ -2387,6 +2387,10 @@ def read_metadata(where, memory_map=False, decryption_properties=None,
23872387
If nothing passed, will be inferred based on path.
23882388
Path will try to be found in the local on-disk filesystem otherwise
23892389
it will be parsed as an URI to determine the filesystem.
2390+
arrow_extensions_enabled : bool, default True
2391+
If True, read Parquet logical types as Arrow extension types where
2392+
possible (e.g. UUID as the canonical `arrow.uuid` extension type).
2393+
If False, use the underlying storage types instead.
23902394
23912395
Returns
23922396
-------
@@ -2416,13 +2420,17 @@ def read_metadata(where, memory_map=False, decryption_properties=None,
24162420
file_ctx = where = filesystem.open_input_file(where)
24172421

24182422
with file_ctx:
2419-
file = ParquetFile(where, memory_map=memory_map,
2420-
decryption_properties=decryption_properties)
2423+
file = ParquetFile(
2424+
where,
2425+
memory_map=memory_map,
2426+
decryption_properties=decryption_properties,
2427+
arrow_extensions_enabled=arrow_extensions_enabled,
2428+
)
24212429
return file.metadata
24222430

24232431

24242432
def read_schema(where, memory_map=False, decryption_properties=None,
2425-
filesystem=None):
2433+
filesystem=None, arrow_extensions_enabled=True):
24262434
"""
24272435
Read effective Arrow schema from Parquet file metadata.
24282436
@@ -2437,6 +2445,9 @@ def read_schema(where, memory_map=False, decryption_properties=None,
24372445
If nothing passed, will be inferred based on path.
24382446
Path will try to be found in the local on-disk filesystem otherwise
24392447
it will be parsed as an URI to determine the filesystem.
2448+
arrow_extensions_enabled : bool, default True
2449+
If True, read Parquet logical types as Arrow extension types where
2450+
possible (e.g., UUID as the canonical `arrow.uuid` extension type).
24402451
24412452
Returns
24422453
-------
@@ -2462,9 +2473,12 @@ def read_schema(where, memory_map=False, decryption_properties=None,
24622473

24632474
with file_ctx:
24642475
file = ParquetFile(
2465-
where, memory_map=memory_map,
2466-
decryption_properties=decryption_properties)
2467-
return file.schema.to_arrow_schema()
2476+
where,
2477+
memory_map=memory_map,
2478+
decryption_properties=decryption_properties,
2479+
arrow_extensions_enabled=arrow_extensions_enabled,
2480+
)
2481+
return file.schema_arrow
24682482

24692483

24702484
__all__ = (

python/pyarrow/tests/parquet/test_metadata.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -849,3 +849,25 @@ def msg(c):
849849

850850
with pytest.raises(TypeError, match=msg("FileMetaData")):
851851
pq.FileMetaData()
852+
853+
854+
def test_read_schema_uuid_extension_type(tmp_path):
855+
# These are the raw 16-byte payloads for
856+
# UUID("e460f970-8351-474e-ac7f-a4673e4ba8cb").bytes and
857+
# UUID("1e741495-eed5-43ea-9bd7-73dc91424baf").bytes.
858+
data = [
859+
b'\xe4`\xf9p\x83QGN\xac\x7f\xa4g>K\xa8\xcb',
860+
b'\x1et\x14\x95\xee\xd5C\xea\x9b\xd7s\xdc\x91BK\xaf',
861+
None,
862+
]
863+
table = pa.table([pa.array(data, type=pa.uuid())], names=["ext"])
864+
865+
file_path = tmp_path / "uuid.parquet"
866+
file_path_str = str(file_path)
867+
pq.write_table(table, file_path_str, store_schema=False)
868+
869+
schema_default = pq.read_schema(file_path_str)
870+
assert schema_default.field("ext").type == pa.uuid()
871+
872+
schema_disabled = pq.read_schema(file_path_str, arrow_extensions_enabled=False)
873+
assert schema_disabled.field("ext").type == pa.binary(16)

0 commit comments

Comments
 (0)