diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 4de88fe22997..4d246abc4ed8 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -272,6 +272,9 @@ class ParquetFile: If True, read Parquet logical types as Arrow extension types where possible (e.g., read JSON as the canonical `arrow.json` extension type or UUID as the canonical `arrow.uuid` extension type). + memory_pool : pyarrow.MemoryPool, optional + Pool to allocate memory from. If not passed, the default memory pool + is used. Examples -------- @@ -322,7 +325,8 @@ def __init__(self, source, *, metadata=None, common_metadata=None, decryption_properties=None, thrift_string_size_limit=None, thrift_container_size_limit=None, schema_depth_limit=None, filesystem=None, - page_checksum_verification=False, arrow_extensions_enabled=True): + page_checksum_verification=False, arrow_extensions_enabled=True, + memory_pool=None): self._close_source = getattr(source, 'closed', True) @@ -332,7 +336,7 @@ def __init__(self, source, *, metadata=None, common_metadata=None, source = filesystem.open_input_file(source) self._close_source = True # We opened it here, ensure we close it. - self.reader = ParquetReader() + self.reader = ParquetReader(memory_pool=memory_pool) self.reader.open( source, use_memory_map=memory_map, buffer_size=buffer_size, pre_buffer=pre_buffer, @@ -1392,6 +1396,9 @@ class ParquetDataset: If True, read Parquet logical types as Arrow extension types where possible, (e.g., read JSON as the canonical `arrow.json` extension type or UUID as the canonical `arrow.uuid` extension type). +memory_pool : pyarrow.MemoryPool, optional + Pool to allocate memory from. If not passed, the default memory pool + is used. Examples -------- @@ -1406,9 +1413,11 @@ def __init__(self, path_or_paths, filesystem=None, schema=None, *, filters=None, decryption_properties=None, thrift_string_size_limit=None, thrift_container_size_limit=None, schema_depth_limit=None, page_checksum_verification=False, - arrow_extensions_enabled=True): + arrow_extensions_enabled=True, memory_pool=None): import pyarrow.dataset as ds + self._memory_pool = memory_pool + # map format arguments read_options = { "pre_buffer": pre_buffer, @@ -1605,7 +1614,7 @@ def read(self, columns=None, use_threads=True, use_pandas_metadata=False): table = self._dataset.to_table( columns=columns, filter=self._filter_expression, - use_threads=use_threads + use_threads=use_threads, memory_pool=self._memory_pool ) # if the "pandas" metadata entry was retrieved from common_metadata, @@ -1813,6 +1822,9 @@ def partitioning(self): If True, read Parquet logical types as Arrow extension types where possible, (e.g., read JSON as the canonical `arrow.json` extension type or UUID as the canonical `arrow.uuid` extension type). +memory_pool : pyarrow.MemoryPool, optional + Pool to allocate memory from. If not passed, the default memory pool + is used. Returns ------- @@ -1909,7 +1921,7 @@ def read_table(source, *, columns=None, use_threads=True, decryption_properties=None, thrift_string_size_limit=None, thrift_container_size_limit=None, schema_depth_limit=None, page_checksum_verification=False, - arrow_extensions_enabled=True): + arrow_extensions_enabled=True, memory_pool=None): try: dataset = ParquetDataset( @@ -1932,6 +1944,7 @@ def read_table(source, *, columns=None, use_threads=True, schema_depth_limit=schema_depth_limit, page_checksum_verification=page_checksum_verification, arrow_extensions_enabled=arrow_extensions_enabled, + memory_pool=memory_pool, ) except ImportError: # fall back on ParquetFile for simple cases when pyarrow.dataset @@ -1980,6 +1993,7 @@ def read_table(source, *, columns=None, use_threads=True, thrift_container_size_limit=thrift_container_size_limit, schema_depth_limit=schema_depth_limit, page_checksum_verification=page_checksum_verification, + memory_pool=memory_pool, ) return dataset.read(columns=columns, use_threads=use_threads, diff --git a/python/pyarrow/tests/parquet/test_basic.py b/python/pyarrow/tests/parquet/test_basic.py index 8b91090989ab..d13e8834c1b6 100644 --- a/python/pyarrow/tests/parquet/test_basic.py +++ b/python/pyarrow/tests/parquet/test_basic.py @@ -1075,3 +1075,54 @@ def test_read_table_raises_value_error_when_ds_is_unavailable(monkeypatch, sourc with pytest.raises(ValueError, match="the 'source' argument"): pq.read_table(source=source) + + +@pytest.fixture +def proxy_pool(): + # Owned by the fixture so that it outlives every object created from it + # inside the test function (readers and tables keep a raw pool pointer). + return pa.proxy_memory_pool(pa.default_memory_pool()) + + +@pytest.mark.parametrize("use_dataset_path", [True, False]) +def test_read_table_memory_pool(tempdir, proxy_pool, use_dataset_path, + monkeypatch): + table = pa.table({'a': list(range(1000)), + 'b': [str(i) for i in range(1000)]}) + path = str(tempdir / 'data.parquet') + pq.write_table(table, path) + + if not use_dataset_path: + # Exercise the ParquetFile fallback used without pyarrow.dataset + monkeypatch.setitem(sys.modules, "pyarrow.dataset", None) + + result = pq.read_table(path, memory_pool=proxy_pool) + assert result.equals(table) + assert proxy_pool.bytes_allocated() >= table.nbytes + result = None + assert proxy_pool.bytes_allocated() == 0 + + +def test_parquet_file_memory_pool(tempdir, proxy_pool): + table = pa.table({'a': list(range(1000))}) + path = str(tempdir / 'data.parquet') + pq.write_table(table, path) + + result = pq.ParquetFile(path, memory_pool=proxy_pool).read() + assert result.equals(table) + assert proxy_pool.bytes_allocated() >= table.nbytes + result = None + assert proxy_pool.bytes_allocated() == 0 + + +def test_parquet_dataset_memory_pool(tempdir, proxy_pool): + table = pa.table({'a': list(range(1000)), 'part': ['x', 'y'] * 500}) + root = str(tempdir / 'dataset') + pq.write_to_dataset(table, root, partition_cols=['part']) + + result = pq.ParquetDataset(root, memory_pool=proxy_pool).read( + use_threads=False) + assert result.num_rows == table.num_rows + assert proxy_pool.bytes_allocated() > 0 + result = None + assert proxy_pool.bytes_allocated() == 0