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

Refs #48922

Author: akx

doc/source/whatsnew/v2.0.0.rst

@@ -231,6 +231,7 @@ Performance improvements
 - Performance improvements to :func:`read_sas` (:issue:`47403`, :issue:`47405`, :issue:`47656`, :issue:`48502`)
 - Memory improvement in :meth:`RangeIndex.sort_values` (:issue:`48801`)
 - Performance improvement in :class:`DataFrameGroupBy` and :class:`SeriesGroupBy` when ``by`` is a categorical type and ``sort=False`` (:issue:`48976`)
+- 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,
@@ -1150,6 +1151,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:
@@ -1159,6 +1161,7 @@ def __init__(
         # NB: _path_or_buf is mistyped on purpose, since the alternative is to placate
         #     mypy by having an assert before every read.
         self._path_or_buf: IO[bytes] = None  # type: ignore[assignment]
+        self._close_file: Callable[[], None] | None = None
         self._has_string_data = False
         self._missing_values = False
         self._can_read_value_labels = False
@@ -1181,21 +1184,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(
+                "Please use StataReader as a context manager to ensure that "
+                "the file is properly closed when reading is finished.",
+                ResourceWarning,
+                stacklevel=4,
+            )
+        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__(
@@ -1209,7 +1227,8 @@ def __exit__(
 
     def close(self) -> None:
         """close the handle if its open"""
-        self._path_or_buf.close()
+        if self._close_file:
+            self._close_file()
 
     def _set_encoding(self) -> None:
         """

pandas/tests/io/test_stata.py

@@ -1842,6 +1842,35 @@ 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):
+        StataReader(file_path).read()
+
+
 @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
 @pytest.mark.parametrize("use_dict", [True, False])
 @pytest.mark.parametrize("infer", [True, False])