ENH: Allow usecols to accept callable (GH14154)

Author: WillAyd

doc/source/whatsnew/v0.19.0.txt

@@ -515,6 +515,7 @@ Other enhancements
 - ``astype()`` will now accept a dict of column name to data types mapping as the ``dtype`` argument. (:issue:`12086`)
 - The ``pd.read_json`` and ``DataFrame.to_json`` has gained support for reading and writing json lines with ``lines`` option see :ref:`Line delimited json <io.jsonl>` (:issue:`9180`)
 - :func:``read_excel`` now supports the true_values and false_values keyword arguments (:issue:`13347`)
+- The ``usecols`` argument now accepts a callable function as a value  (:issue:`14154`)
 
 .. _whatsnew_0190.api:
 

pandas/io/parsers.py

@@ -33,6 +33,7 @@
 from pandas.util.decorators import Appender
 
 import pandas.lib as lib
+import pandas.core.common as com
 import pandas.parser as _parser
 
 
@@ -85,8 +86,8 @@
     MultiIndex is used. If you have a malformed file with delimiters at the end
     of each line, you might consider index_col=False to force pandas to _not_
     use the first column as the index (row names)
-usecols : array-like, default None
-    Return a subset of the columns. All elements in this array must either
+usecols : array-like or callable, default None
+    Return a subset of the columns. If array-like, all elements must either
     be positional (i.e. integer indices into the document columns) or strings
     that correspond to column names provided either by the user in `names` or
     inferred from the document header row(s). For example, a valid `usecols`
@@ -975,17 +976,27 @@ def _is_index_col(col):
     return col is not None and col is not False
 
 
+def _evaluate_usecols(usecols, names):
+    if callable(usecols):
+        return set([i for i, name in enumerate(names)
+                    if com._apply_if_callable(usecols, name)])
+    else:
+        return usecols
+
+
 def _validate_usecols_arg(usecols):
     """
     Check whether or not the 'usecols' parameter
     contains all integers (column selection by index)
     or strings (column by name). Raises a ValueError
     if that is not the case.
     """
-    msg = ("The elements of 'usecols' must "
-           "either be all strings, all unicode, or all integers")
+    msg = ("'usecols' must either be all strings, all unicode, "
+           "all integers or callable")
 
     if usecols is not None:
+        if callable(usecols):
+            return usecols
         usecols_dtype = lib.infer_dtype(usecols)
         if usecols_dtype not in ('empty', 'integer',
                                  'string', 'unicode'):
@@ -1425,11 +1436,12 @@ def __init__(self, src, **kwds):
         self.orig_names = self.names[:]
 
         if self.usecols:
-            if len(self.names) > len(self.usecols):
+            usecols = _evaluate_usecols(self.usecols, self.orig_names)
+            if len(self.names) > len(usecols):
                 self.names = [n for i, n in enumerate(self.names)
-                              if (i in self.usecols or n in self.usecols)]
+                              if (i in usecols or n in usecols)]
 
-            if len(self.names) < len(self.usecols):
+            if len(self.names) < len(usecols):
                 raise ValueError("Usecols do not match names.")
 
         self._set_noconvert_columns()
@@ -1588,9 +1600,10 @@ def read(self, nrows=None):
 
     def _filter_usecols(self, names):
         # hackish
-        if self.usecols is not None and len(names) != len(self.usecols):
+        usecols = _evaluate_usecols(self.usecols, names)
+        if usecols is not None and len(names) != len(usecols):
             names = [name for i, name in enumerate(names)
-                     if i in self.usecols or name in self.usecols]
+                     if i in usecols or name in usecols]
         return names
 
     def _get_index_names(self):
@@ -2191,7 +2204,9 @@ def _handle_usecols(self, columns, usecols_key):
         usecols_key is used if there are string usecols.
         """
         if self.usecols is not None:
-            if any([isinstance(u, string_types) for u in self.usecols]):
+            if callable(self.usecols):
+                col_indices = _evaluate_usecols(self.usecols, usecols_key)
+            elif any([isinstance(u, string_types) for u in self.usecols]):
                 if len(columns) > 1:
                     raise ValueError("If using multiple headers, usecols must "
                                      "be integers.")

pandas/io/tests/parser/usecols.py

@@ -23,8 +23,9 @@ def test_raise_on_mixed_dtype_usecols(self):
         1000,2000,3000
         4000,5000,6000
         """
-        msg = ("The elements of 'usecols' must "
-               "either be all strings, all unicode, or all integers")
+
+        msg = ("'usecols' must either be all strings, all unicode, "
+               "all integers or callable")
         usecols = [0, 'b', 2]
 
         with tm.assertRaisesRegexp(ValueError, msg):
@@ -302,8 +303,8 @@ def test_usecols_with_mixed_encoding_strings(self):
         3.568935038,7,False,a
         '''
 
-        msg = ("The elements of 'usecols' must "
-               "either be all strings, all unicode, or all integers")
+        msg = ("'usecols' must either be all strings, all unicode, "
+               "all integers or callable")
 
         with tm.assertRaisesRegexp(ValueError, msg):
             self.read_csv(StringIO(s), usecols=[u'AAA', b'BBB'])
@@ -366,3 +367,25 @@ def test_np_array_usecols(self):
         expected = DataFrame([[1, 2]], columns=usecols)
         result = self.read_csv(StringIO(data), usecols=usecols)
         tm.assert_frame_equal(result, expected)
+
+    def test_callable_usecols(self):
+        s = '''AaA,bBb,CCC,ddd
+        0.056674973,8,True,a
+        2.613230982,2,False,b
+        3.568935038,7,False,a
+        '''
+
+        data = {
+            'AaA': {
+                0: 0.056674972999999997,
+                1: 2.6132309819999997,
+                2: 3.5689350380000002
+            },
+            'bBb': {0: 8, 1: 2, 2: 7},
+            'ddd': {0: 'a', 1: 'b', 2: 'a'}
+        }
+        expected = DataFrame(data)
+
+        df = self.read_csv(StringIO(s), usecols=lambda x:
+                           x.upper() in ['AAA', 'BBB', 'DDD'])
+        tm.assert_frame_equal(df, expected)

pandas/parser.pyx

@@ -35,6 +35,7 @@ cimport util
 
 import pandas.lib as lib
 import pandas.compat as compat
+import pandas.core.common as com
 from pandas.types.common import (is_categorical_dtype, CategoricalDtype,
                                  is_integer_dtype, is_float_dtype,
                                  is_bool_dtype, is_object_dtype,
@@ -273,7 +274,7 @@ cdef class TextReader:
         object file_handle, na_fvalues
         object true_values, false_values
         object dsource
-        bint na_filter, verbose, has_usecols, has_mi_columns
+        bint na_filter, verbose, has_usecols, has_mi_columns, callable_usecols
         int parser_start
         list clocks
         char *c_encoding
@@ -297,8 +298,10 @@ cdef class TextReader:
         object compression
         object mangle_dupe_cols
         object tupleize_cols
+        object usecols
         list dtype_cast_order
-        set noconvert, usecols
+        set noconvert
+
 
     def __cinit__(self, source,
                   delimiter=b',',
@@ -434,7 +437,11 @@ cdef class TextReader:
         # suboptimal
         if usecols is not None:
             self.has_usecols = 1
-            self.usecols = set(usecols)
+            if callable(usecols):
+                self.callable_usecols = 1
+                self.usecols = usecols
+            else:
+                self.usecols = set(usecols)
 
         # XXX
         if skipfooter > 0:
@@ -698,7 +705,6 @@ cdef class TextReader:
             cdef StringPath path = _string_path(self.c_encoding)
 
         header = []
-
         if self.parser.header_start >= 0:
 
             # Header is in the file
@@ -816,7 +822,7 @@ cdef class TextReader:
             #                        'data has %d fields'
             #                        % (passed_count, field_count))
 
-            if self.has_usecols and self.allow_leading_cols:
+            if self.has_usecols and self.allow_leading_cols and not self.callable_usecols:
                 nuse = len(self.usecols)
                 if nuse == passed_count:
                     self.leading_cols = 0
@@ -1010,17 +1016,25 @@ cdef class TextReader:
 
         results = {}
         nused = 0
+        
         for i in range(self.table_width):
+                
             if i < self.leading_cols:
                 # Pass through leading columns always
                 name = i
-            elif self.usecols and nused == len(self.usecols):
+            elif self.usecols and not self.callable_usecols and nused == len(self.usecols):
                 # Once we've gathered all requested columns, stop. GH5766
                 break
             else:
                 name = self._get_column_name(i, nused)
-                if self.has_usecols and not (i in self.usecols or
-                                             name in self.usecols):
+                usecols = set()
+                if self.callable_usecols:
+                    if com._apply_if_callable(self.usecols, name):
+                        usecols = set([i])
+                else:
+                    usecols = self.usecols
+                if self.has_usecols and not (i in usecols or
+                                             name in usecols):
                     continue
                 nused += 1
 
@@ -1336,6 +1350,7 @@ def _maybe_upcast(arr):
 
     return arr
 
+
 cdef enum StringPath:
     CSTRING
     UTF8