REGR: be able to read Stata files without reading them fully into memory

Fixes pandas-dev#48700
Refs pandas-dev#9245
Refs pandas-dev#37639
Regressed in 6d1541e

Author: akx

doc/source/whatsnew/v1.6.0.rst

@@ -155,6 +155,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_160.bug_fixes:

pandas/io/stata.py

@@ -1164,15 +1164,23 @@ def __init__(
         self._lines_read = 0
 
         self._native_byteorder = _set_endianness(sys.byteorder)
-        with get_handle(
+
+        handles = get_handle(
             path_or_buf,
             "rb",
             storage_options=storage_options,
             is_text=False,
             compression=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()
@@ -1192,7 +1200,7 @@ def __exit__(
 
     def close(self) -> None:
         """close the handle if its open"""
-        self.path_or_buf.close()
+        self._close_file()
 
     def _set_encoding(self) -> None:
         """

pandas/tests/io/test_stata.py

@@ -1848,6 +1848,29 @@ 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 isinstance(reader.path_or_buf, io.BytesIO)
+        assert not reader.read().empty
+
+    # Test that we use a given fp exactly, if possible.
+    with open(file_path, "rb") as fp:
+        with StataReader(fp) as reader:
+            assert reader.path_or_buf is fp
+            assert not reader.read().empty
+
+    # 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 reader.path_or_buf is bio
+                assert not reader.read().empty
+
+
 @pytest.mark.parametrize("version", [114, 117, 118, 119, None])
 @pytest.mark.parametrize("use_dict", [True, False])
 @pytest.mark.parametrize("infer", [True, False])
@@ -2027,7 +2050,10 @@ def test_stata_compression(compression_only, read_infer, to_infer):
 
     with tm.ensure_clean(filename) as path:
         df.to_stata(path, compression=to_compression)
-        result = read_stata(path, compression=read_compression, index_col="index")
+        with StataReader(
+            path, compression=read_compression, index_col="index"
+        ) as reader:
+            result = reader.read()
         tm.assert_frame_equal(result, df)