use fast zip with a placeholder value just for np.nan

Author: Chang She

pandas/core/frame.py

@@ -2344,7 +2344,7 @@ def dropna(self, axis=0, how='any', thresh=None, subset=None):
         new_labels = labels[mask]
         return self.reindex(**{axis_name: new_labels})
 
-    def drop_duplicates(self, cols=None, take_last=False):
+    def drop_duplicates(self, cols=None, take_last=False, skipna=True):
         """
         Return DataFrame with duplicate rows removed, optionally only
         considering certain columns
@@ -2356,15 +2356,17 @@ def drop_duplicates(self, cols=None, take_last=False):
             default use all of the columns
         take_last : boolean, default False
             Take the last observed row in a row. Defaults to the first row
+        skipna : boolean, default True
+            If True then keep NaN
 
         Returns
         -------
         deduplicated : DataFrame
         """
-        duplicated = self.duplicated(cols, take_last=take_last)
+        duplicated = self.duplicated(cols, take_last=take_last, skipna=skipna)
         return self[-duplicated]
 
-    def duplicated(self, cols=None, take_last=False):
+    def duplicated(self, cols=None, take_last=False, skipna=True):
         """
         Return boolean Series denoting duplicate rows, optionally only
         considering certain columns
@@ -2376,20 +2378,29 @@ def duplicated(self, cols=None, take_last=False):
             default use all of the columns
         take_last : boolean, default False
             Take the last observed row in a row. Defaults to the first row
+        skipna : boolean, default True
+            If True then NaN are not marked as duplicates
 
         Returns
         -------
         duplicated : Series
         """
+        zip_func = lib.fast_zip if skipna else lib.fast_zip_fillna
+
         if cols is not None:
             if isinstance(cols, list):
-                keys = zip(*[self[x] for x in cols])
+                values = [self[x].values for x in cols]
+                keys = zip_func(values)
+                dup_func = lib.duplicated_skipna
             else:
-                keys = list(self[cols])
+                keys = self[cols]
+                dup_func = lib.duplicated_skipna if skipna else lib.duplicated
         else:
-            keys = zip(*self.values.T)
+            values = list(self.values.T)
+            keys = zip_func(values)
+            dup_func = lib.duplicated_skipna
 
-        duplicated = lib.duplicated(keys, take_last=take_last)
+        duplicated = dup_func(list(keys), take_last=take_last)
         return Series(duplicated, index=self.index)
 
     #----------------------------------------------------------------------
@@ -4527,7 +4538,6 @@ def _homogenize(data, index, columns, dtype=None):
 def _put_str(s, space):
     return ('%s' % s)[:space].ljust(space)
 
-
 def _is_sequence(x):
     try:
         iter(x)

pandas/src/groupby.pyx

@@ -1301,12 +1301,39 @@ def count_level_2d(ndarray[uint8_t, ndim=2, cast=True] mask,
 
     return counts
 
+def duplicated_skipna(list values, take_last=False):
+    cdef:
+        Py_ssize_t i, n
+        dict seen = {}
+        object row
+
+    n = len(values)
+    cdef ndarray[uint8_t] result = np.zeros(n, dtype=np.uint8)
+
+    if take_last:
+        for i from n > i >= 0:
+            row = values[i]
+            if row in seen:
+                result[i] = 1
+            else:
+                seen[row] = None
+                result[i] = 0
+    else:
+        for i from 0 <= i < n:
+            row = values[i]
+            if row in seen:
+                result[i] = 1
+            else:
+                seen[row] = None
+                result[i] = 0
+
+    return result.view(np.bool_)
 
 def duplicated(list values, take_last=False):
     cdef:
         Py_ssize_t i, n
         dict seen = {}
-        int has_nan = 0
+        bint has_nan = 0
         object row
 
     n = len(values)
@@ -1318,7 +1345,7 @@ def duplicated(list values, take_last=False):
             if row in seen:
                 result[i] = 1
             elif row != row:
-                if has_nan == 1:
+                if has_nan:
                     result[i] = 1
                 else:
                     has_nan = 1
@@ -1332,7 +1359,7 @@ def duplicated(list values, take_last=False):
             if row in seen:
                 result[i] = 1
             elif row != row:
-                if has_nan == 1:
+                if has_nan:
                     result[i] = 1
                 else:
                     has_nan = 1
@@ -1343,7 +1370,6 @@ def duplicated(list values, take_last=False):
 
     return result.view(np.bool_)
 
-
 def generate_slices(ndarray[int64_t] labels, Py_ssize_t ngroups):
     cdef:
         Py_ssize_t i, group_size, n, lab, start

pandas/src/tseries.pyx

@@ -404,6 +404,57 @@ def fast_zip(list ndarrays):
 
     return result
 
+cdef class _PandasNull:
+    pass
+
+pandas_null = _PandasNull()
+
+def fast_zip_fillna(list ndarrays, fill_value=pandas_null):
+    '''
+    For zipping multiple ndarrays into an ndarray of tuples
+    '''
+    cdef:
+        Py_ssize_t i, j, k, n
+        ndarray[object] result
+        flatiter it
+        object val, tup
+
+    k = len(ndarrays)
+    n = len(ndarrays[0])
+
+    result = np.empty(n, dtype=object)
+
+    # initialize tuples on first pass
+    arr = ndarrays[0]
+    it = <flatiter> PyArray_IterNew(arr)
+    for i in range(n):
+        val = PyArray_GETITEM(arr, PyArray_ITER_DATA(it))
+        tup = PyTuple_New(k)
+
+        if val != val:
+            val = fill_value
+
+        PyTuple_SET_ITEM(tup, 0, val)
+        Py_INCREF(val)
+        result[i] = tup
+        PyArray_ITER_NEXT(it)
+
+    for j in range(1, k):
+        arr = ndarrays[j]
+        it = <flatiter> PyArray_IterNew(arr)
+        if len(arr) != n:
+            raise ValueError('all arrays must be same length')
+
+        for i in range(n):
+            val = PyArray_GETITEM(arr, PyArray_ITER_DATA(it))
+            if val != val:
+                val = fill_value
+
+            PyTuple_SET_ITEM(result[i], j, val)
+            Py_INCREF(val)
+            PyArray_ITER_NEXT(it)
+
+    return result
 
 def get_reverse_indexer(ndarray[int64_t] indexer, Py_ssize_t length):
     cdef:

pandas/tests/test_frame.py

@@ -1,3 +1,4 @@
+
 # pylint: disable-msg=W0612,E1101
 from copy import deepcopy
 from datetime import datetime, timedelta
@@ -3322,20 +3323,20 @@ def test_drop_duplicates_NA(self):
                         'D' : range(8)})
 
         # single column
-        result = df.drop_duplicates('C')
+        result = df.drop_duplicates('C', skipna=False)
         expected = df[:2]
         assert_frame_equal(result, expected)
 
-        result = df.drop_duplicates('C', take_last=True)
+        result = df.drop_duplicates('C', take_last=True, skipna=False)
         expected = df.ix[[3, 7]]
         assert_frame_equal(result, expected)
 
         # multi column
-        result = df.drop_duplicates(['C', 'B'])
+        result = df.drop_duplicates(['C', 'B'], skipna=False)
         expected = df.ix[[0, 1, 2, 4]]
         assert_frame_equal(result, expected)
 
-        result = df.drop_duplicates(['C', 'B'], take_last=True)
+        result = df.drop_duplicates(['C', 'B'], take_last=True, skipna=False)
         expected = df.ix[[1, 3, 6, 7]]
         assert_frame_equal(result, expected)