8000 FIX: StataReader: don't buffer entire file into memory unless necessary · akx/pandas@2accab5 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2accab5

Browse files
akxakx
committed
FIX: StataReader: don't buffer entire file into memory unless necessary
Refs pandas-dev#48922
1 parent 024442a commit 2accab5

File tree

3 files changed

+78
-8
lines changed

3 files changed

+78
-8
lines changed

‎doc/source/whatsnew/v2.0.0.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1061,6 +1061,7 @@ Performance improvements
10611061
- Fixed a reference leak in :func:`read_hdf` (:issue:`37441`)
10621062
- Fixed a memory leak in :meth:`DataFrame.to_json` and :meth:`Series.to_json` when serializing datetimes and timedeltas (:issue:`40443`)
10631063
- Decreased memory usage in many :class:`DataFrameGroupBy` methods (:issue:`51090`)
1064+
- Memory improvement in :class:`StataReader` when reading seekable files (:issue:`48922`)
10641065
-
10651066

10661067
.. ---------------------------------------------------------------------------

‎pandas/io/stata.py‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
TYPE_CHECKING,
2424
Any,
2525
AnyStr,
26+
Callable,
2627
Final,
2728
Hashable,
2829
Sequence,
@@ -1152,6 +1153,7 @@ def __init__(
11521153
self._encoding = ""
11531154
self._chunksize = chunksize
11541155
self._using_iterator = False
1156+
self._entered = False
11551157
if self._chunksize is None:
11561158
self._chunksize = 1
11571159
elif not isinstance(chunksize, int) or chunksize <= 0:
@@ -1181,21 +1183,36 @@ def _open_file(self) -> None:
11811183
"""
11821184
Open the file (with compression options, etc.), and read header information.
11831185
"""
1184-
with get_handle(
1186+
if not self._entered:
1187+
warnings.warn(
1188+
"StataReader is being used without using a context manager. "
1189+
"Using StataReader as a context manager is the only supported method.",
1190+
ResourceWarning,
1191+
stacklevel=4,
1192+
)
1193+
handles = get_handle(
11851194
self._original_path_or_buf,
11861195
"rb",
11871196
storage_options=self._storage_options,
11881197
is_text=False,
11891198
compression=self._compression,
1190-
) as handles:
1191-
# Copy to BytesIO, and ensure no encoding
1192-
self._path_or_buf = BytesIO(handles.handle.read())
1199+
)
1200+
if hasattr(handles.handle, "seekable") and handles.handle.seekable():
1201+
# If the handle is directly seekable, use it without an extra copy.
1202+
self._path_or_buf = handles.handle
1203+
self._close_file = handles.close
1204+
else:
1205+
# Copy to memory, and ensure no encoding.
1206+
with handles:
1207+
self._path_or_buf = BytesIO(handles.handle.read())
1208+
self._close_file = self._path_or_buf.close
11931209

11941210
self._read_header()
11951211
self._setup_dtype()
11961212

11971213
def __enter__(self) -> StataReader:
11981214
"""enter context manager"""
1215+
self._entered = True
11991216
return self
12001217

12011218
def __exit__(
@@ -1204,12 +1221,26 @@ def __exit__(
12041221
exc_value: BaseException | None,
12051222
traceback: TracebackType | None,
12061223
) -> None:
1207-
"""exit context manager"""
1208-
self.close()
1224+
if self._close_file:
1225+
self._close_file()
12091226

12101227
def close(self) -> None:
1211-
"""close the handle if its open"""
1212-
self._path_or_buf.close()
1228+
"""Close the handle if its open.
1229+
1230+
.. deprecated: 2.0.0
1231+
1232+
The close method is not part of the public API.
1233+
The only supported way to use StataReader is to use it as a context manager.
1234+
"""
1235+
warnings.warn(
1236+
"The StataReader.close() method is not part of the public API and "
1237+
"will be removed in a future version without notice. "
1238+
"Using StataReader as a context manager is the only supported method.",
1239+
FutureWarning,
1240+
stacklevel=2,
1241+
)
1242+
if self._close_file:
1243+
self._close_file()
12131244

12141245
def _set_encoding(self) -> None:
12151246
"""

‎pandas/tests/io/test_stata.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1895,6 +1895,44 @@ def test_backward_compat(version, datapath):
18951895
tm.assert_frame_equal(old_dta, expected, check_dtype=False)
18961896

18971897

1898+
def test_direct_read(datapath, monkeypatch):
1899+
file_path = datapath("io", "data", "stata", "stata-compat-118.dta")
1900+
1901+
# Test that opening a file path doesn't buffer the file.
1902+
with StataReader(file_path) as reader:
1903+
# Must not have been buffered to memory
1904+
assert not reader.read().empty
1905+
assert not isinstance(reader._path_or_buf, io.BytesIO)
1906+
1907+
# Test that we use a given fp exactly, if possible.
1908+
with open(file_path, "rb") as fp:
1909+
with StataReader(fp) as reader:
1910+
assert not reader.read().empty
1911+
assert reader._path_or_buf is fp
1912+
1913+
# Test that we use a given BytesIO exactly, if possible.
1914+
with open(file_path, "rb") as fp:
1915+
with io.BytesIO(fp.read()) as bio:
1916+
with StataReader(bio) as reader:
1917+
assert not reader.read().empty
1918+
assert reader._path_or_buf is bio
1919+
1920+
1921+
def test_statareader_warns_when_used_without_context(datapath):
1922+
file_path = datapath("io", "data", "stata", "stata-compat-118.dta")
1923+
with tm.assert_produces_warning(
1924+
ResourceWarning,
1925+
match="without using a context manager",
1926+
):
1927+
sr = StataReader(file_path)
1928+
sr.read()
1929+
with tm.assert_produces_warning(
1930+
FutureWarning,
1931+
match="is not part of the public API",
1932+
):
1933+
sr.close()
1934+
1935+
18981936
@pytest.mark.parametrize("version", [114, 117, 118, 119, None])
18991937
@pytest.mark.parametrize("use_dict", [True, False])
19001938
@pytest.mark.parametrize("infer", [True, False])

0 commit comments

Comments
 (0)
0