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

Fixes pandas-dev#48700
Regressed in pandas-dev#9245
Regressed in 2f0ada3

Author: akx

doc/source/whatsnew/v1.6.0.rst

@@ -151,6 +151,8 @@ Performance improvements
 - Performance improvement in ``var`` for nullable dtypes (:issue:`48379`).
 - Performance improvements to :func:`read_sas` (:issue:`47403`, :issue:`47405`, :issue:`47656`, :issue:`48502`)
 - Memory improvement in :meth:`RangeIndex.sort_values` (:issue:`48801`)
+- Memory improvement in :class:`StataReader` when reading seekable files (:issue:`48922`)
+
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_160.bug_fixes:

pandas/io/stata.py

@@ -1164,15 +1164,25 @@ 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
+            self._copied_to_memory = False
+        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._copied_to_memory = True
 
         self._read_header()
         self._setup_dtype()
@@ -1192,7 +1202,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

@@ -2027,7 +2027,14 @@ 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:
+            # A zstd stream is not seekable, so the reader should have decompressed it.
+            # This assumption could change in the future.
+            assert reader._copied_to_memory == (compression == "zstd")
+
+            result = reader.read()
         tm.assert_frame_equal(result, df)