FIX: StataReader: don't buffer entire file into memory unless necessary

Refs #48922

Author: akx

doc/source/whatsnew/v2.0.0.rst

@@ -1161,6 +1161,8 @@ Performance improvements
 - Fixed a reference leak in :func:`read_hdf` (:issue:`37441`)
 - Fixed a memory leak in :meth:`DataFrame.to_json` and :meth:`Series.to_json` when serializing datetimes and timedeltas (:issue:`40443`)
 - Decreased memory usage in many :class:`DataFrameGroupBy` methods (:issue:`51090`)
+- Memory improvement in :class:`StataReader` when reading seekable files (:issue:`48922`)
+
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_200.bug_fixes:

pandas/io/stata.py

@@ -23,6 +23,7 @@
     TYPE_CHECKING,
     Any,
     AnyStr,
+    Callable,
     Final,
     Hashable,
     Sequence,
@@ -1148,6 +1149,7 @@ def __init__(
         self._encoding = ""
         self._chunksize = chunksize
         self._using_iterator = False
+        self._entered = False
         if self._chunksize is None:
             self._chunksize = 1
         elif not isinstance(chunksize, int) or chunksize <= 0:
@@ -1177,21 +1179,36 @@ def _open_file(self) -> None:
         """
         Open the file (with compression options, etc.), and read header information.
         """
-        with get_handle(
+        if not self._entered:
+            warnings.warn(
+                "StataReader is being used without using a context manager. "
+                "Using StataReader as a context manager is the only supported method.",
+                ResourceWarning,
+                stacklevel=find_stack_level(),
+            )
+        handles = get_handle(
             self._original_path_or_buf,
             "rb",
             storage_options=self._storage_options,
             is_text=False,
             compression=self._compression,
-        ) as handles:
-            # Copy to BytesIO, and ensure no encoding
-            self._path_or_buf = BytesIO(handles.handle.read())
+        )
+        if hasattr(handles.handle, "seekable") and handles.handle.seekable():
+            # If the handle is directly seekable, use it without an extra copy.
+            self._path_or_buf = handles.handle
+            self._close_file = handles.close
+        else:
+            # Copy to memory, and ensure no encoding.
+            with handles:
+                self._path_or_buf = BytesIO(handles.handle.read())
+            self._close_file = self._path_or_buf.close
 
         self._read_header()
         self._setup_dtype()
 
     def __enter__(self) -> StataReader:
         """enter context manager"""
+        self._entered = True
         return self
 
     def __exit__(
@@ -1200,12 +1217,26 @@ def __exit__(
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> None:
-        """exit context manager"""
-        self.close()
+        if self._close_file:
+            self._close_file()
 
     def close(self) -> None:
-        """close the handle if its open"""
-        self._path_or_buf.close()
+        """Close the handle if its open.
+
+        .. deprecated: 2.0.0
+
+           The close method is not part of the public API.
+           The only supported way to use StataReader is to use it as a context manager.
+        """
+        warnings.warn(
+            "The StataReader.close() method is not part of the public API and "
+            "will be removed in a future version without notice. "
+            "Using StataReader as a context manager is the only supported method.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        if self._close_file:
+            self._close_file()
 
     def _set_encoding(self) -> None:
         """

pandas/tests/io/test_stata.py

@@ -1889,6 +1889,44 @@ def test_backward_compat(version, datapath):
     tm.assert_frame_equal(old_dta, expected, check_dtype=False)
 
 
+def test_direct_read(datapath, monkeypatch):
+    file_path = datapath("io", "data", "stata", "stata-compat-118.dta")
+
+    # Test that opening a file path doesn't buffer the file.
+    with StataReader(file_path) as reader:
+        # Must not have been buffered to memory
+        assert not reader.read().empty
+        assert not isinstance(reader._path_or_buf, io.BytesIO)
+
+    # Test that we use a given fp exactly, if possible.
+    with open(file_path, "rb") as fp:
+        with StataReader(fp) as reader:
+            assert not reader.read().empty
+            assert reader._path_or_buf is fp
+
+    # Test that we use a given BytesIO exactly, if possible.
+    with open(file_path, "rb") as fp:
+        with io.BytesIO(fp.read()) as bio:
+            with StataReader(bio) as reader:
+                assert not reader.read().empty
+                assert reader._path_or_buf is bio
+
+
+def test_statareader_warns_when_used_without_context(datapath):
+    file_path = datapath("io", "data", "stata", "stata-compat-118.dta")
+    with tm.assert_produces_warning(
+        ResourceWarning,
+        match="without using a context manager",
+    ):
+        sr = StataReader(file_path)
+        sr.read()
+    with tm.assert_produces_warning(
+        FutureWarning,
+        match="is not part of the public API",
+    ):
+        sr.close()
+
+
 @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
 @pytest.mark.parametrize("use_dict", [True, False])
 @pytest.mark.parametrize("infer", [True, False])