Repository: pandas-dev/pandas

Issue:
[pandas/io/stata.py] Importing option for importing specified columns from stata .dta files
It would be great if a `columns` option could be added to the Stata Reader. Some `dta` files can be large and I would like to import a few specified columns. Current I am importing the entire file and returning a selection such as:

``` python
data = pd.read_stata(fn)[['year', 'exports']]
```

I would propose `pd.read_stata(fn, columns=['year', 'exports'])`



Discussion:
cc @bashtage

should be called names for compat with other readers

This will be very difficult since dta files are stored in binary as rows and the reader sequentially reads each row, parses the binary and keeps the data.  A `names` input would realistically only be able to drop columns once all data had been imported, and so would not be incredibly useful, since this can be trivially done from any `DataFrame`.  

Rewriting the reader to skip columns would be a substantial undertaking, imo.  This said, there may be one way to do it that wouldn't be too onerous, but I won't have any time for a while.

As a workaround I would recommend doing a once and for all conversion from a DTA to a more native pandas format (e.g. hdf) and then not worrying about it.

@sanguineturtle 

Recent changes to master allow 100,000 observations with 5 variables to be read in around 0.05 seconds. Hopefully this performance means reading select columns is not really a performance issue any more.  See #8073 for benchmarks.

@jreback 

This probably isn't remotely possible with the modifications I made to the reader since it uses `np.frombuffer` for performance.  However, the performance improvements are large enough that this isn't an issue except possibly in memory constrained situations.  Moreover, reading only a subset of columns to minimize memory impact would probably be so slow (relative to master) that it wouldn't be be useful in situations where the dataset was very big.

yep in fact in hdf you read entire rows with every column and then just reindex at the end 
but good for now
thanks!

@bashtage This is great news. Thanks for letting me know. My dataset has ~27 million observations with 8 variables. The main reason I was looking for column selection is that 4 are redundant identifiers. As you have suggested I looked into `HDF` files and have written a conversion method (for the class) and moved the source to using `HDF` files (via pandas) with year keys which has proven to have very fast retrieval. I am still learning about `HDF` but it's in built compression delivers some pretty amazing results on file size etc. 


Resolving patch:
diff --git a/pandas/io/stata.py b/pandas/io/stata.py
--- a/pandas/io/stata.py
+++ b/pandas/io/stata.py
@@ -29,7 +29,7 @@
 
 def read_stata(filepath_or_buffer, convert_dates=True,
                convert_categoricals=True, encoding=None, index=None,
-               convert_missing=False, preserve_dtypes=True):
+               convert_missing=False, preserve_dtypes=True, columns=None):
     """
     Read Stata file into DataFrame
 
@@ -55,11 +55,14 @@ def read_stata(filepath_or_buffer, convert_dates=True,
     preserve_dtypes : boolean, defaults to True
         Preserve Stata datatypes. If False, numeric data are upcast to pandas
         default types for foreign data (float64 or int64)
+    columns : list or None
+        Columns to retain.  Columns will be returned in the given order.  None
+        returns all columns
     """
     reader = StataReader(filepath_or_buffer, encoding)
 
     return reader.data(convert_dates, convert_categoricals, index,
-                       convert_missing, preserve_dtypes)
+                       convert_missing, preserve_dtypes, columns)
 
 _date_formats = ["%tc", "%tC", "%td", "%d", "%tw", "%tm", "%tq", "%th", "%ty"]
 
@@ -977,7 +980,7 @@ def _read_strls(self):
             self.path_or_buf.read(1)  # zero-termination
 
     def data(self, convert_dates=True, convert_categoricals=True, index=None,
-             convert_missing=False, preserve_dtypes=True):
+             convert_missing=False, preserve_dtypes=True, columns=None):
         """
         Reads observations from Stata file, converting them into a dataframe
 
@@ -999,6 +1002,10 @@ def data(self, convert_dates=True, convert_categoricals=True, index=None,
         preserve_dtypes : boolean, defaults to True
             Preserve Stata datatypes. If False, numeric data are upcast to
             pandas default types for foreign data (float64 or int64)
+        columns : list or None
+            Columns to retain.  Columns will be returned in the given order.
+            None returns all columns
+
         Returns
         -------
         y : DataFrame instance
@@ -1034,6 +1041,35 @@ def data(self, convert_dates=True, convert_categoricals=True, index=None,
             data = DataFrame.from_records(data, index=index)
             data.columns = self.varlist
 
+        if columns is not None:
+            column_set = set(columns)
+            if len(column_set) != len(columns):
+                raise ValueError('columns contains duplicate entries')
+            unmatched = column_set.difference(data.columns)
+            if unmatched:
+                raise ValueError('The following columns were not found in the '
+                                 'Stata data set: ' +
+                                 ', '.join(list(unmatched)))
+            # Copy information for retained columns for later processing
+            dtyplist = []
+            typlist = []
+            fmtlist = []
+            lbllist = []
+            matched = set()
+            for i, col in enumerate(data.columns):
+                if col in column_set:
+                    matched.update([col])
+                    dtyplist.append(self.dtyplist[i])
+                    typlist.append(self.typlist[i])
+                    fmtlist.append(self.fmtlist[i])
+                    lbllist.append(self.lbllist[i])
+
+            data = data[columns]
+            self.dtyplist = dtyplist
+            self.typlist = typlist
+            self.fmtlist = fmtlist
+            self.lbllist = lbllist
+
         for col, typ in zip(data, self.typlist):
             if type(typ) is int:
                 data[col] = data[col].apply(self._null_terminate, convert_dtype=True,)


Test patch:
diff --git a/pandas/io/tests/test_stata.py b/pandas/io/tests/test_stata.py
--- a/pandas/io/tests/test_stata.py
+++ b/pandas/io/tests/test_stata.py
@@ -720,12 +720,29 @@ def test_dtype_conversion(self):
 
         tm.assert_frame_equal(expected, conversion)
 
+    def test_drop_column(self):
+        expected = self.read_csv(self.csv15)
+        expected['byte_'] = expected['byte_'].astype(np.int8)
+        expected['int_'] = expected['int_'].astype(np.int16)
+        expected['long_'] = expected['long_'].astype(np.int32)
+        expected['float_'] = expected['float_'].astype(np.float32)
+        expected['double_'] = expected['double_'].astype(np.float64)
+        expected['date_td'] = expected['date_td'].apply(datetime.strptime,
+                                                        args=('%Y-%m-%d',))
 
+        columns = ['byte_', 'int_', 'long_']
+        expected = expected[columns]
+        dropped = read_stata(self.dta15_117, convert_dates=True,
+                             columns=columns)
 
+        tm.assert_frame_equal(expected, dropped)
+        with tm.assertRaises(ValueError):
+            columns = ['byte_', 'byte_']
+            read_stata(self.dta15_117, convert_dates=True, columns=columns)
 
-
-
-
+        with tm.assertRaises(ValueError):
+            columns = ['byte_', 'int_', 'long_', 'not_found']
+            read_stata(self.dta15_117, convert_dates=True, columns=columns)
 
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],