ENH: better impl for na_map in string methods. close #2602

Author: wesm

pandas/core/strings.py

@@ -86,18 +86,30 @@ def _length_check(others):
 
 def _na_map(f, arr, na_result=np.nan):
     # should really _check_ for NA
-    def g(x):
-        try:
-            return f(x)
-        except (TypeError, AttributeError):
-            return na_result
-    return _map(g, arr)
+    return _map(f, arr, na_mask=True, na_value=na_result)
 
 
-def _map(f, arr):
+def _map(f, arr, na_mask=False, na_value=np.nan):
     if not isinstance(arr, np.ndarray):
         arr = np.asarray(arr, dtype=object)
-    return lib.map_infer(arr, f)
+    if na_mask:
+        mask = isnull(arr)
+        try:
+            result = lib.map_infer_mask(arr, f, mask.view(np.uint8))
+        except (TypeError, AttributeError):
+            def g(x):
+                try:
+                    return f(x)
+                except (TypeError, AttributeError):
+                    return na_value
+            return _map(g, arr)
+        if na_value is not np.nan:
+            np.putmask(result, mask, na_value)
+            if result.dtype == object:
+                result = lib.maybe_convert_objects(result)
+        return result
+    else:
+        return lib.map_infer(arr, f)
 
 
 def str_count(arr, pat, flags=0):

pandas/io/pytables.py

@@ -248,7 +248,7 @@ def __repr__(self):
                 # a table
                 if _is_table_type(n):
                     values.append(str(create_table(self, n)))
-                
+
                 # a group
                 elif kind is None:
                     values.append('unknown type')
@@ -400,7 +400,7 @@ def remove(self, key, where=None):
             if where is None:
                 group = self.get_node(key)
                 group._f_remove(recursive=True)
-            
+
             # delete from the table
             else:
                 if not _is_table_type(group):
@@ -962,7 +962,7 @@ def infer(self, table):
     def convert(self, values):
         """ set the values from this selection """
         self.values = Index(_maybe_convert(values[self.cname], self.kind))
-   
+
     @property
     def attrs(self):
         return self.table._v_attrs
@@ -1362,7 +1362,7 @@ def infer_axes(self):
               return a boolean indicating if we have a valid table or not """
 
         table = self.table
-        if table is None: 
+        if table is None:
             return False
 
         self.index_axes, self.values_axes = [ a.infer(self.table) for a in self.indexables if a.is_indexable ], [ a.infer(self.table) for a in self.indexables if not a.is_indexable ]
@@ -1412,13 +1412,13 @@ def create_axes(self, axes, obj, validate = True, min_itemsize = None):
                     indexer = len(self.non_index_axes)
                     exist_axis = existing_table.non_index_axes[indexer][1]
                     if append_axis != exist_axis:
-                        
+
                         # ahah! -> reindex
                         if sorted(append_axis) == sorted(exist_axis):
                             append_axis = exist_axis
 
                 self.non_index_axes.append((i,append_axis))
-        
+
         # set axis positions (based on the axes)
         self.index_axes = [ index_axes_map[a].set_pos(j) for j, a in enumerate(axes) ]
         j = len(self.index_axes)
@@ -1444,7 +1444,7 @@ def create_axes(self, axes, obj, validate = True, min_itemsize = None):
 
             # a string column
             if b.dtype.name == 'object':
-                
+
                 # itemsize is the maximum length of a string (along any dimension)
                 itemsize = _itemsize_string_array(values)
 
@@ -1489,7 +1489,7 @@ def reindex(obj, axis, filt, ordered):
             ordd = ordered & filt
             ordd = sorted(ordered.get_indexer(ordd))
             return obj.reindex_axis(ordered.take(ordd), axis = obj._get_axis_number(axis_name), copy = False)
-            
+
         # apply the selection filters (but keep in the same order)
         if self.selection.filter:
             for axis, filt in self.selection.filter:
@@ -1558,7 +1558,7 @@ class LegacyTable(Table):
 
         """
     _indexables = [IndexCol(name = 'index',  axis = 1, pos = 0),
-                   IndexCol(name = 'column', axis = 2, pos = 1, index_kind = 'columns_kind'), 
+                   IndexCol(name = 'column', axis = 2, pos = 1, index_kind = 'columns_kind'),
                    DataCol( name = 'fields', cname = 'values', kind_attr = 'fields', pos = 2) ]
     table_type = 'legacy'
     ndim       = 3
@@ -1569,7 +1569,7 @@ def write(self, **kwargs):
     def read(self, where=None):
         """ we have n indexable columns, with an arbitrary number of data axes """
 
-        
+
         if not self.read_axes(where): return None
 
         factors  = [ Categorical.from_array(a.values) for a in self.index_axes ]
@@ -1591,7 +1591,7 @@ def read(self, where=None):
 
                 # the data need to be sorted
                 sorted_values = c.take_data().take(sorter, axis=0)
-                
+
                 take_labels   = [ l.take(sorter) for l in labels ]
                 items         = Index(c.values)
                 block         = block2d_to_blocknd(sorted_values, items, tuple(N), take_labels)
@@ -1767,7 +1767,7 @@ def delete(self, where = None):
             # final element
             if groups[-1] != ln:
                 groups.append(ln)
-            
+
             # initial element
             if groups[0] != 0:
                 groups.insert(0,0)
@@ -1893,7 +1893,7 @@ def create_table(parent, group, typ = None, **kwargs):
 
 def _itemsize_string_array(arr):
     """ return the maximum size of elements in a strnig array """
-    return max([ str_len(arr[v]).max() for v in range(arr.shape[0]) ])
+    return max([ str_len(arr[v].ravel()).max() for v in range(arr.shape[0]) ])
 
 def _convert_index(index):
     if isinstance(index, DatetimeIndex):
@@ -2289,4 +2289,3 @@ def f(values, freq=None, tz=None):
                                              tz=tz)
         return f
     return klass
-

pandas/tests/test_strings.py

@@ -64,7 +64,7 @@ def test_count(self):
         self.assert_(isinstance(result, Series))
         tm.assert_almost_equal(result, exp)
 
-        #mixed
+        # mixed
         mixed = ['a', NA, 'b', True, datetime.today(), 'foo', None, 1, 2.]
         rs = strings.str_count(mixed, 'a')
         xp = [1, NA, 0, NA, NA, 0, NA, NA, NA]

vb_suite/frame_methods.py

@@ -113,3 +113,13 @@ def f(K=500):
 
 frame_insert_500_columns = Benchmark('f()', setup,
                                      start_date=datetime(2011, 1, 1))
+
+#----------------------------------------------------------------------
+# strings methods, #2602
+
+setup = common_setup + """
+s = Series(['abcdefg', np.nan]*500000)
+"""
+
+series_string_vector_slice = Benchmark('s.str[:5]', setup,
+                                       start_date=datetime(2012, 8, 1))