Merge

Author: jbuyl

doc/source/whatsnew/v0.17.0.txt

@@ -613,3 +613,4 @@ Bug Fixes
 
 - Bug in ``pd.DataFrame`` when constructing an empty DataFrame with a string dtype (:issue:`9428`)
 - Bug in ``read_stata`` when reading a file with a different order set in ``columns`` (:issue:`10739`)
+- Bug in ``pd.unique`` for arrays with the ``datetime64`` or ``timedelta64`` dtype that meant an array with object dtype was returned instead the original dtype (:issue: `9431`)

pandas/core/algorithms.py

@@ -36,7 +36,7 @@ def match(to_match, values, na_sentinel=-1):
         values = np.array(values, dtype='O')
 
     f = lambda htype, caster: _match_generic(to_match, values, htype, caster)
-    result = _hashtable_algo(f, values.dtype)
+    result = _hashtable_algo(f, values.dtype, np.int64)
 
     if na_sentinel != -1:
 
@@ -66,14 +66,20 @@ def unique(values):
     return _hashtable_algo(f, values.dtype)
 
 
-def _hashtable_algo(f, dtype):
+def _hashtable_algo(f, dtype, return_dtype=None):
     """
     f(HashTable, type_caster) -> result
     """
     if com.is_float_dtype(dtype):
         return f(htable.Float64HashTable, com._ensure_float64)
     elif com.is_integer_dtype(dtype):
         return f(htable.Int64HashTable, com._ensure_int64)
+    elif com.is_datetime64_dtype(dtype):
+        return_dtype = return_dtype or 'M8[ns]'
+        return f(htable.Int64HashTable, com._ensure_int64).view(return_dtype)
+    elif com.is_timedelta64_dtype(dtype):
+        return_dtype = return_dtype or 'm8[ns]'
+        return f(htable.Int64HashTable, com._ensure_int64).view(return_dtype)
     else:
         return f(htable.PyObjectHashTable, com._ensure_object)
 

pandas/tests/test_algos.py

@@ -235,6 +235,50 @@ def test_on_index_object(self):
 
         tm.assert_almost_equal(result, expected)
 
+    def test_datetime64_dtype_array_returned(self):
+        # GH 9431
+        expected = np.array(['2015-01-03T00:00:00.000000000+0000',
+                             '2015-01-01T00:00:00.000000000+0000'], dtype='M8[ns]')
+
+        dt_index = pd.to_datetime(['2015-01-03T00:00:00.000000000+0000',
+                                   '2015-01-01T00:00:00.000000000+0000',
+                                   '2015-01-01T00:00:00.000000000+0000'])
+        result = algos.unique(dt_index)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+        s = pd.Series(dt_index)
+        result = algos.unique(s)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+        arr = s.values
+        result = algos.unique(arr)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+
+    def test_timedelta64_dtype_array_returned(self):
+        # GH 9431
+        expected = np.array([31200, 45678, 10000], dtype='m8[ns]')
+
+        td_index = pd.to_timedelta([31200, 45678, 31200, 10000, 45678])
+        result = algos.unique(td_index)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+        s = pd.Series(td_index)
+        result = algos.unique(s)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+        arr = s.values
+        result = algos.unique(arr)
+        tm.assert_numpy_array_equal(result, expected)
+        self.assertEqual(result.dtype, expected.dtype)
+
+
+
 class TestValueCounts(tm.TestCase):
     _multiprocess_can_split_ = True