This is an automated email from the ASF dual-hosted git repository.
AlenkaF pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new b8079cd294 GH-54637: [Python] add `MemoryPool` support to
`pq.read_table` (#54919)
b8079cd294 is described below
commit b8079cd294ad60bdb81da8a0989e8fd27918c517
Author: Omkar Kabde <[email protected]>
AuthorDate: Wed Oct 28 20:27:19 2026 +0530
GH-54637: [Python] add `MemoryPool` support to `pq.read_table` (#54919)
### Rationale for this change
`pq.read_table` has no way to pass a memory pool. `ParquetReader` and
`Dataset.to_table` already accept one, it just wasn't added.
### What changes are included in this PR?
I added a `memory_pool` argument to `read_table`, `ParquetFile` and
`ParquetDataset`. It is passed to `ParquetReader` and to `Dataset.to_table`,
with docstrings updated.
### Are these changes tested?
Yes. I added new tests in `test_basic.py` - It reads a Parquet file with a
tracked memory pool and check that the pool was actually used. They cover
`read_table`, `ParquetFile` and `ParquetDataset`.
### Are there any user-facing changes?
New optional `memory_pool` keyword. No breaking changes.
### Was AI used for this PR?
**PR code and description written by:**
- [x] Human
- [x] AI
**Reviewed before submission by:**
- [x] Human
- [ ] AI
Github Issue - https://github.com/apache/arrow/issues/54637
* GitHub Issue: #54637
Authored-by: Omkar Kabde <[email protected]>
Signed-off-by: AlenkaF <[email protected]>
---
python/pyarrow/parquet/core.py | 28 ++++++++++++----
python/pyarrow/tests/parquet/test_basic.py | 51 ++++++++++++++++++++++++++++++
2 files changed, 73 insertions(+), 6 deletions(-)
diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py
index 253a8a3312..fa3b897ad2 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 @@ class ParquetFile:
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 @@ class ParquetFile:
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 @@ Examples
"""
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 @@ Examples
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 @@ Examples
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 @@ arrow_extensions_enabled : bool, default True
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 ce8c75f3c1..ba88118b82 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)
+
+
[email protected]
+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())
+
+
[email protected]("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
