[omkar-334]

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:

• Human
• AI

Reviewed before submission by:

• Human
• AI

Github Issue - #54637

• GitHub Issue: [Parquet][Python] pq.read_table should allow passing MemoryPool #54637

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[github-actions]

⚠️ GitHub issue #54637 has been automatically assigned in GitHub to PR creator.

[omkar-334]

My understanding before implementing/opening the PR -

Function read_table supports 2 types of datasets

1. dataset = ParquetDataset

• self._dataset = ds.dataset
  This function returns multiple types, not so straightforward. The class Dataset (_dataset.pyx) has three types of datasets -

  classes = {
        'union': UnionDataset,
        'filesystem': FileSystemDataset,
        'in-memory': InMemoryDataset,
    }
  

  • _filesystem_dataset -> FileSystemDatasetFactory
  • _union_dataset -> UnionDataset
  • _in_memory_dataset -> InMemoryDataset
    We don;t need to change code for all the classes, the below unction makes it easy for us

• table = self._dataset.to_table -> to_table (_dataset.pyx)
  This function in pyx supports MemoryPool -> we just need to edit this.

2. dataset = ParquetFile

• self.reader = ParquetReader()
  • python/pyarrow/_parquet.pyx - ParquetReader supports CMemoryPool* pool
• self.reader.read_all
• return dataset.read

[omkar-334]

Tests passing

[omkar-334]

cc @pitrou

[pitrou]

Thanks for doing this @omkar-334 . This looks good on the principle, just a small suggestion for improving tests.

[pitrou]

How about we exercise this a bit more specifically, such as:

Suggested change

	assert proxy_pool.max_memory() > 0
	assert proxy_pool.bytes_allocated() >= table.nbytes
	result = None
	assert proxy_pool.bytes_allocated() == 0