Skip to content

Commit 951663a

Browse files
milesgrangermilesgranger
authored
ARROW-13763: [Python] Close files in ParquetFile & ParquetDatasetPiece (#13821)
Will fix [ARROW-13763](https://issues.apache.org/jira/browse/ARROW-13763) A separate Jira issue will be made to address closing files in V2 ParquetDataset, which needs to be handled in the C++ layer. Adds context manager to `pq.ParquetFile` to close input file, and ensure reads within `pq.ParquetDataset` and `pq.read_table` are closed. ```python # user opened file-like object will not be closed with open('file.parquet', 'rb') as f: with pq.ParquetFile(f) as p: table = p.read() assert not f.closed # did not inadvertently close the open file assert not p.closed assert not f.closed # parquet context exit didn't close it assert not p.closed # references the input file status assert f.closed # normal context exit close assert p.closed # ... # path-like will be closed upon exit or `ParquetFile.close` with pq.ParquetFile('file.parquet') as p: table = p.read() assert not p.closed assert p.closed ``` Authored-by: Miles Granger <[email protected]> Signed-off-by: Antoine Pitrou <[email protected]>
1 parent b0422e5 commit 951663a

3 files changed

Lines changed: 90 additions & 6 deletions

File tree

‎python/pyarrow/_parquet.pyx‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1159,6 +1159,7 @@ cdef class ParquetReader(_Weakrefable):
11591159
CMemoryPool* pool
11601160
unique_ptr[FileReader] reader
11611161
FileMetaData _metadata
1162+
shared_ptr[CRandomAccessFile] rd_handle
11621163

11631164
cdef public:
11641165
_column_idx_map
@@ -1175,7 +1176,6 @@ cdef class ParquetReader(_Weakrefable):
11751176
thrift_string_size_limit=None,
11761177
thrift_container_size_limit=None):
11771178
cdef:
1178-
shared_ptr[CRandomAccessFile] rd_handle
11791179
shared_ptr[CFileMetaData] c_metadata
11801180
CReaderProperties properties = default_reader_properties()
11811181
ArrowReaderProperties arrow_props = (
@@ -1221,10 +1221,10 @@ cdef class ParquetReader(_Weakrefable):
12211221
string_to_timeunit(coerce_int96_timestamp_unit))
12221222

12231223
self.source = source
1224+
get_reader(source, use_memory_map, &self.rd_handle)
12241225

1225-
get_reader(source, use_memory_map, &rd_handle)
12261226
with nogil:
1227-
check_status(builder.Open(rd_handle, properties, c_metadata))
1227+
check_status(builder.Open(self.rd_handle, properties, c_metadata))
12281228

12291229
# Set up metadata
12301230
with nogil:
@@ -1435,6 +1435,19 @@ cdef class ParquetReader(_Weakrefable):
14351435
.ReadColumn(column_index, &out))
14361436
return pyarrow_wrap_chunked_array(out)
14371437

1438+
def close(self):
1439+
if not self.closed:
1440+
with nogil:
1441+
check_status(self.rd_handle.get().Close())
1442+
1443+
@property
1444+
def closed(self):
1445+
if self.rd_handle == NULL:
1446+
return True
1447+
with nogil:
1448+
closed = self.rd_handle.get().closed()
1449+
return closed
1450+
14381451

14391452
cdef shared_ptr[WriterProperties] _create_writer_properties(
14401453
use_dictionary=None,

‎python/pyarrow/parquet/core.py‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -294,9 +294,16 @@ def __init__(self, source, *, metadata=None, common_metadata=None,
294294
thrift_string_size_limit=thrift_string_size_limit,
295295
thrift_container_size_limit=thrift_container_size_limit,
296296
)
297+
self._close_source = getattr(source, 'closed', True)
297298
self.common_metadata = common_metadata
298299
self._nested_paths_by_prefix = self._build_nested_paths()
299300

301+
def __enter__(self):
302+
return self
303+
304+
def __exit__(self, *args, **kwargs):
305+
self.close()
306+
300307
def _build_nested_paths(self):
301308
paths = self.reader.column_paths
302309

@@ -376,6 +383,14 @@ def num_row_groups(self):
376383
"""
377384
return self.reader.num_row_groups
378385

386+
def close(self, force: bool = False):
387+
if self._close_source or force:
388+
self.reader.close()
389+
390+
@property
391+
def closed(self) -> bool:
392+
return self.reader.closed
393+
379394
def read_row_group(self, i, columns=None, use_threads=True,
380395
use_pandas_metadata=False):
381396
"""
@@ -1129,8 +1144,8 @@ def get_metadata(self):
11291144
-------
11301145
metadata : FileMetaData
11311146
"""
1132-
f = self.open()
1133-
return f.metadata
1147+
with self.open() as parquet:
1148+
return parquet.metadata
11341149

11351150
def open(self):
11361151
"""
@@ -1204,6 +1219,9 @@ def read(self, columns=None, use_threads=True, partitions=None,
12041219
arr = pa.DictionaryArray.from_arrays(indices, dictionary)
12051220
table = table.append_column(name, arr)
12061221

1222+
# To ParquetFile the source looked like it was already open, so won't
1223+
# actually close it without overriding.
1224+
reader.close(force=True)
12071225
return table
12081226

12091227

@@ -1890,7 +1908,8 @@ def read(self, columns=None, use_threads=True, use_pandas_metadata=False):
18901908
"""
18911909
tables = []
18921910
for piece in self._pieces:
1893-
table = piece.read(columns=columns, use_threads=use_threads,
1911+
table = piece.read(columns=columns,
1912+
use_threads=use_threads,
18941913
partitions=self._partitions,
18951914
use_pandas_metadata=use_pandas_metadata)
18961915
tables.append(table)

‎python/pyarrow/tests/parquet/test_parquet_file.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import io
1919
import os
20+
from unittest import mock
2021

2122
import pytest
2223

@@ -277,3 +278,54 @@ def test_pre_buffer(pre_buffer):
277278
buf.seek(0)
278279
pf = pq.ParquetFile(buf, pre_buffer=pre_buffer)
279280
assert pf.read().num_rows == N
281+
282+
283+
def test_parquet_file_explicitly_closed(tempdir):
284+
"""
285+
Unopened files should be closed explicitly after use,
286+
and previously opened files should be left open.
287+
Applies to read_table, ParquetDataset, and ParquetFile
288+
"""
289+
# create test parquet file
290+
fn = tempdir.joinpath('file.parquet')
291+
table = pa.table({'col1': [0, 1], 'col2': [0, 1]})
292+
pq.write_table(table, fn)
293+
294+
# read_table (legacy) with opened file (will leave open)
295+
with open(fn, 'rb') as f:
296+
pq.read_table(f, use_legacy_dataset=True)
297+
assert not f.closed # Didn't close it internally after read_table
298+
299+
# read_table (legacy) with unopened file (will close)
300+
with mock.patch.object(pq.ParquetFile, "close") as mock_close:
301+
pq.read_table(fn, use_legacy_dataset=True)
302+
mock_close.assert_called()
303+
304+
# ParquetDataset test (legacy) with unopened file (will close)
305+
with mock.patch.object(pq.ParquetFile, "close") as mock_close:
306+
pq.ParquetDataset(fn, use_legacy_dataset=True).read()
307+
mock_close.assert_called()
308+
309+
# ParquetDataset test (legacy) with opened file (will leave open)
310+
with open(fn, 'rb') as f:
311+
# ARROW-8075: support ParquetDataset from file-like, not just path-like
312+
with pytest.raises(TypeError, match='not a path-like object'):
313+
pq.ParquetDataset(f, use_legacy_dataset=True).read()
314+
assert not f.closed
315+
316+
# ParquetFile with opened file (will leave open)
317+
with open(fn, 'rb') as f:
318+
with pq.ParquetFile(f) as p:
319+
p.read()
320+
assert not f.closed
321+
assert not p.closed
322+
assert not f.closed # opened input file was not closed
323+
assert not p.closed # parquet file obj reports as not closed
324+
assert f.closed
325+
assert p.closed # parquet file being closed reflects underlying file
326+
327+
# ParquetFile with unopened file (will close)
328+
with pq.ParquetFile(fn) as p:
329+
p.read()
330+
assert not p.closed
331+
assert p.closed # parquet file obj reports as closed

0 commit comments

Comments
 (0)