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

Refs #48922

Author: akx

doc/source/whatsnew/v2.0.0.rst

@@ -197,6 +197,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

@@ -1156,7 +1156,9 @@ def __init__(
             raise ValueError("chunksize must be a positive integer when set.")
 
         # State variables for the file
+        self._entered = False
         self._path_or_buf = None
+        self._close_file = None
         self._has_string_data = False
         self._missing_values = False
         self._can_read_value_labels = False
@@ -1179,21 +1181,35 @@ 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,
+            )
+        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__(
@@ -1207,7 +1223,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])