GH-54637: [Python] add MemoryPool support to pq.read_table

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(

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