diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 253a8a33122..fa3b897ad23 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, @@ -1537,7 +1541,8 @@ def schema(self): """ return self._dataset.schema - def read(self, columns=None, use_threads=True, use_pandas_metadata=False): + def read(self, columns=None, use_threads=True, use_pandas_metadata=False, + memory_pool=None): """ Read (multiple) Parquet files as a single pyarrow.Table. @@ -1551,6 +1556,9 @@ def read(self, columns=None, use_threads=True, use_pandas_metadata=False): use_pandas_metadata : bool, default False If True and file has custom pandas schema metadata, ensure that index columns are also loaded. + memory_pool : pyarrow.MemoryPool, optional + Pool to allocate memory from. If not passed, the default memory + pool is used. Returns ------- @@ -1605,7 +1613,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=memory_pool ) # if the "pandas" metadata entry was retrieved from common_metadata, @@ -1813,6 +1821,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,8 +1920,9 @@ 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): + read_kwargs = {"memory_pool": memory_pool} try: dataset = ParquetDataset( source, @@ -1980,10 +1992,14 @@ 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, ) + # ParquetFile takes the pool when it is opened, not when it is read + read_kwargs = {} return dataset.read(columns=columns, use_threads=use_threads, - use_pandas_metadata=use_pandas_metadata) + use_pandas_metadata=use_pandas_metadata, + **read_kwargs) read_table.__doc__ = _read_table_docstring.format( diff --git a/python/pyarrow/tests/parquet/test_basic.py b/python/pyarrow/tests/parquet/test_basic.py index ce8c75f3c10..ba88118b82b 100644 --- a/python/pyarrow/tests/parquet/test_basic.py +++ b/python/pyarrow/tests/parquet/test_basic.py @@ -1089,3 +1089,54 @@ def test_filters_to_expression_deprecated_alias(): result = pq._filters_to_expression(filters) assert isinstance(result, ds.Expression) assert result.equals(expected) + + +@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).read( + use_threads=False, memory_pool=proxy_pool) + assert result.num_rows == table.num_rows + assert proxy_pool.bytes_allocated() > 0 + result = None + assert proxy_pool.bytes_allocated() == 0