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,
@@ -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,

python/pyarrow/tests/parquet/test_basic.py

@@ -1075,3 +1075,47 @@ 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.max_memory() > 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.max_memory() > 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()
+    assert result.num_rows == table.num_rows
+    assert proxy_pool.max_memory() > 0