ENH: more datetime64 integration in core data algorithms per #996, close #1035

Author: wesm

pandas/core/common.py

@@ -66,7 +66,7 @@ def isnull(obj):
                 result = Series(result, index=obj.index, copy=False)
         elif obj.dtype == np.datetime64:
             # this is the NaT pattern
-            result = obj.ravel().view('i8') == 0x8000000000000000
+            result = obj.view('i8') == lib.NaT
         else:
             result = -np.isfinite(obj)
         return result
@@ -106,43 +106,59 @@ def _unpickle_array(bytes):
     arr = read_array(BytesIO(bytes))
     return arr
 
-def _take_1d_bool(arr, indexer, out, fill_value=np.nan):
-    view = arr.view(np.uint8)
-    outview = out.view(np.uint8)
+def _take_1d_datetime(arr, indexer, out, fill_value=np.nan):
+    view = arr.view(np.int64)
+    outview = out.view(np.int64)
     lib.take_1d_bool(view, indexer, outview, fill_value=fill_value)
 
-def _take_2d_axis0_bool(arr, indexer, out, fill_value=np.nan):
-    view = arr.view(np.uint8)
-    outview = out.view(np.uint8)
-    lib.take_2d_axis0_bool(view, indexer, outview, fill_value=fill_value)
+def _take_2d_axis0_datetime(arr, indexer, out, fill_value=np.nan):
+    view = arr.view(np.int64)
+    outview = out.view(np.int64)
+    lib.take_1d_bool(view, indexer, outview, fill_value=fill_value)
 
-def _take_2d_axis1_bool(arr, indexer, out, fill_value=np.nan):
+def _take_2d_axis1_datetime(arr, indexer, out, fill_value=np.nan):
     view = arr.view(np.uint8)
     outview = out.view(np.uint8)
-    lib.take_2d_axis1_bool(view, indexer, outview, fill_value=fill_value)
+    lib.take_1d_bool(view, indexer, outview, fill_value=fill_value)
+
+def _view_wrapper(f, wrap_dtype, na_override=None):
+    def wrapper(arr, indexer, out, fill_value=np.nan):
+        if na_override is not None and np.isnan(fill_value):
+            fill_value = na_override
+        view = arr.view(wrap_dtype)
+        outview = out.view(wrap_dtype)
+        f(view, indexer, outview, fill_value=fill_value)
+    return wrapper
+
 
 _take1d_dict = {
     'float64' : lib.take_1d_float64,
     'int32' : lib.take_1d_int32,
     'int64' : lib.take_1d_int64,
     'object' : lib.take_1d_object,
-    'bool' : _take_1d_bool
+    'bool' : _view_wrapper(lib.take_1d_bool, np.uint8),
+    'datetime64[us]' : _view_wrapper(lib.take_1d_int64, np.int64,
+                                     na_override=lib.NaT),
 }
 
 _take2d_axis0_dict = {
     'float64' : lib.take_2d_axis0_float64,
     'int32' : lib.take_2d_axis0_int32,
     'int64' : lib.take_2d_axis0_int64,
     'object' : lib.take_2d_axis0_object,
-    'bool' : _take_2d_axis0_bool
+    'bool' : _view_wrapper(lib.take_2d_axis0_bool, np.uint8),
+    'datetime64[us]' : _view_wrapper(lib.take_2d_axis0_int64, np.int64,
+                                     na_override=lib.NaT),
 }
 
 _take2d_axis1_dict = {
     'float64' : lib.take_2d_axis1_float64,
     'int32' : lib.take_2d_axis1_int32,
     'int64' : lib.take_2d_axis1_int64,
     'object' : lib.take_2d_axis1_object,
-    'bool' : _take_2d_axis1_bool
+    'bool' : _view_wrapper(lib.take_2d_axis1_bool, np.uint8),
+    'datetime64[us]' : _view_wrapper(lib.take_2d_axis1_int64, np.int64,
+                                     na_override=lib.NaT),
 }
 
 def _get_take2d_function(dtype_str, axis=0):
@@ -164,14 +180,13 @@ def take_1d(arr, indexer, out=None, fill_value=np.nan):
         indexer = np.array(indexer, dtype=np.int32)
 
     indexer = _ensure_int32(indexer)
-
     out_passed = out is not None
+    take_f = _take1d_dict.get(dtype_str)
 
     if dtype_str in ('int32', 'int64', 'bool'):
         try:
             if out is None:
                 out = np.empty(n, dtype=arr.dtype)
-            take_f = _take1d_dict[dtype_str]
             take_f(arr, indexer, out=out, fill_value=fill_value)
         except ValueError:
             mask = indexer == -1
@@ -186,10 +201,9 @@ def take_1d(arr, indexer, out=None, fill_value=np.nan):
                                     out.dtype)
                 out = _maybe_upcast(out)
                 np.putmask(out, mask, fill_value)
-    elif dtype_str in ('float64', 'object'):
+    elif dtype_str in ('float64', 'object', 'datetime64[us]'):
         if out is None:
             out = np.empty(n, dtype=arr.dtype)
-        take_f = _take1d_dict[dtype_str]
         take_f(arr, indexer, out=out, fill_value=fill_value)
     else:
         out = arr.take(indexer, out=out)
@@ -238,7 +252,7 @@ def take_2d(arr, indexer, out=None, mask=None, needs_masking=None, axis=0,
             take_f = _get_take2d_function(dtype_str, axis=axis)
             take_f(arr, indexer, out=out, fill_value=fill_value)
             return out
-    elif dtype_str in ('float64', 'object'):
+    elif dtype_str in ('float64', 'object', 'datetime64[us]'):
         if out is None:
             out = np.empty(out_shape, dtype=arr.dtype)
         take_f = _get_take2d_function(dtype_str, axis=axis)
@@ -301,9 +315,22 @@ def _need_upcast(values):
         return True
     return False
 
+def _interp_wrapper(f, wrap_dtype, na_override=None):
+    def wrapper(arr, mask, limit=None):
+        view = arr.view(wrap_dtype)
+        f(view, mask, limit=limit)
+    return wrapper
+
+_pad_1d_datetime = _interp_wrapper(lib.pad_inplace_int64, np.int64)
+_pad_2d_datetime = _interp_wrapper(lib.pad_2d_inplace_int64, np.int64)
+_backfill_1d_datetime = _interp_wrapper(lib.backfill_inplace_int64, np.int64)
+_backfill_2d_datetime = _interp_wrapper(lib.backfill_2d_inplace_int64, np.int64)
+
 def pad_1d(values, limit=None):
     if is_float_dtype(values):
         _method = lib.pad_inplace_float64
+    elif is_datetime64_dtype(values):
+        _method = _pad_1d_datetime
     elif values.dtype == np.object_:
         _method = lib.pad_inplace_object
     else: # pragma: no cover
@@ -314,21 +341,48 @@ def pad_1d(values, limit=None):
 def backfill_1d(values, limit=None):
     if is_float_dtype(values):
         _method = lib.backfill_inplace_float64
+    elif is_datetime64_dtype(values):
+        _method = _backfill_1d_datetime
     elif values.dtype == np.object_:
         _method = lib.backfill_inplace_object
     else: # pragma: no cover
         raise ValueError('Invalid dtype for padding')
 
     _method(values, isnull(values).view(np.uint8), limit=limit)
 
+def pad_2d(values, limit=None):
+    if is_float_dtype(values):
+        _method = lib.pad_2d_inplace_float64
+    elif is_datetime64_dtype(values):
+        _method = _pad_2d_datetime
+    elif values.dtype == np.object_:
+        _method = lib.pad_2d_inplace_object
+    else: # pragma: no cover
+        raise ValueError('Invalid dtype for padding')
+
+    _method(values, isnull(values).view(np.uint8), limit=limit)
+
+def backfill_2d(values, limit=None):
+    if is_float_dtype(values):
+        _method = lib.backfill_2d_inplace_float64
+    elif is_datetime64_dtype(values):
+        _method = _backfill_2d_datetime
+    elif values.dtype == np.object_:
+        _method = lib.backfill_2d_inplace_object
+    else: # pragma: no cover
+        raise ValueError('Invalid dtype for padding')
+
+    _method(values, isnull(values).view(np.uint8), limit=limit)
+
+
 def _consensus_name_attr(objs):
     name = objs[0].name
     for obj in objs[1:]:
         if obj.name != name:
             return None
     return name
 
-#-------------------------------------------------------------------------------
+#----------------------------------------------------------------------
 # Lots of little utilities
 
 def _infer_dtype(value):
@@ -579,6 +633,13 @@ def is_integer_dtype(arr_or_dtype):
     return (issubclass(tipo, np.integer) and not
             issubclass(tipo, np.datetime64))
 
+def is_datetime64_dtype(arr_or_dtype):
+    if isinstance(arr_or_dtype, np.dtype):
+        tipo = arr_or_dtype.type
+    else:
+        tipo = arr_or_dtype.dtype.type
+    return issubclass(tipo, np.datetime64)
+
 def is_float_dtype(arr_or_dtype):
     if isinstance(arr_or_dtype, np.dtype):
         tipo = arr_or_dtype.type

pandas/core/format.py

@@ -413,6 +413,8 @@ def format_array(values, formatter, float_format=None, na_rep='NaN',
         fmt_klass = FloatArrayFormatter
     elif com.is_integer_dtype(values.dtype):
         fmt_klass = IntArrayFormatter
+    elif com.is_datetime64_dtype(values.dtype):
+        fmt_klass = Datetime64Formatter
     else:
         fmt_klass = GenericArrayFormatter
 
@@ -549,6 +551,23 @@ def get_result(self):
         return _make_fixed_width(fmt_values, self.justify)
 
 
+class Datetime64Formatter(GenericArrayFormatter):
+
+    def get_result(self):
+        if self.formatter:
+            formatter = self.formatter
+        else:
+            def formatter(x):
+                if isnull(x):
+                    return 'NaT'
+                else:
+                    return str(x)
+
+        fmt_values = [formatter(x) for x in self.values]
+
+        return _make_fixed_width(fmt_values, self.justify)
+
+
 def _make_fixed_width(strings, justify='right'):
     if len(strings) == 0:
         return strings
@@ -609,7 +628,7 @@ def _has_names(index):
 
 def set_printoptions(precision=None, column_space=None, max_rows=None,
                      max_columns=None, colheader_justify='right',
-                     notebook_repr_html=None, 
+                     notebook_repr_html=None,
                      date_dayfirst=None, date_yearfirst=None):
     """
     Alter default behavior of DataFrame.toString

pandas/core/frame.py

@@ -2457,8 +2457,6 @@ def fillna(self, value=None, method='pad', axis=0, inplace=False,
         -------
         filled : DataFrame
         """
-        from pandas.core.internals import FloatBlock, ObjectBlock
-
         self._consolidate_inplace()
 
         if value is None:
@@ -2468,7 +2466,7 @@ def fillna(self, value=None, method='pad', axis=0, inplace=False,
             new_blocks = []
             method = com._clean_fill_method(method)
             for block in self._data.blocks:
-                if isinstance(block, (FloatBlock, ObjectBlock)):
+                if block._can_hold_na:
                     newb = block.interpolate(method, axis=axis,
                                              limit=limit, inplace=inplace)
                 else:

pandas/core/internals.py

@@ -225,33 +225,12 @@ def interpolate(self, method='pad', axis=0, inplace=False, limit=None):
         transf = (lambda x: x) if axis == 0 else (lambda x: x.T)
 
         if method == 'pad':
-            _pad(transf(values), limit=limit)
+            com.pad_2d(transf(values), limit=limit)
         else:
-            _backfill(transf(values), limit=limit)
+            com.backfill_2d(transf(values), limit=limit)
 
         return make_block(values, self.items, self.ref_items)
 
-def _pad(values, limit=None):
-    if com.is_float_dtype(values):
-        _method = lib.pad_2d_inplace_float64
-    elif values.dtype == np.object_:
-        _method = lib.pad_2d_inplace_object
-    else: # pragma: no cover
-        raise ValueError('Invalid dtype for padding')
-
-    _method(values, com.isnull(values).view(np.uint8),
-            limit=limit)
-
-def _backfill(values, limit=None):
-    if com.is_float_dtype(values):
-        _method = lib.backfill_2d_inplace_float64
-    elif values.dtype == np.object_:
-        _method = lib.backfill_2d_inplace_object
-    else: # pragma: no cover
-        raise ValueError('Invalid dtype for padding')
-
-    _method(values, com.isnull(values).view(np.uint8),
-            limit=limit)
 
 #-------------------------------------------------------------------------------
 # Is this even possible?
@@ -284,7 +263,7 @@ def should_store(self, value):
                               (np.integer, np.floating, np.bool_))
 
 class DatetimeBlock(IntBlock):
-    pass
+    _can_hold_na = True
 
 
 def make_block(values, items, ref_items, do_integrity_check=False):

pandas/core/series.py

@@ -1910,10 +1910,9 @@ def reindex(self, index=None, method=None, level=None, fill_value=np.nan,
         if len(self.index) == 0:
             return Series(nan, index=index, name=self.name)
 
-        new_index, fill_vec = self.index.reindex(index, method=method,
+        new_index, indexer = self.index.reindex(index, method=method,
                                                  level=level, limit=limit)
-        fill_vec = com._ensure_int32(fill_vec)
-        new_values = com.take_1d(self.values, fill_vec, fill_value=fill_value)
+        new_values = com.take_1d(self.values, indexer, fill_value=fill_value)
         return Series(new_values, index=new_index, name=self.name)
 
     def reindex_like(self, other, method=None, limit=None):

pandas/src/generate_code.py

@@ -282,13 +282,12 @@ def pad_%(name)s(ndarray[%(c_type)s] old, ndarray[%(c_type)s] new,
 @cython.wraparound(False)
 def pad_inplace_%(name)s(ndarray[%(c_type)s] values,
                          ndarray[uint8_t, cast=True] mask,
-                            limit=None):
+                         limit=None):
     cdef Py_ssize_t i, N
     cdef %(c_type)s val
     cdef int lim, fill_count = 0
 
     N = len(values)
-    val = np.nan
 
     if limit is None:
         lim = N
@@ -307,6 +306,7 @@ def pad_inplace_%(name)s(ndarray[%(c_type)s] values,
         else:
             fill_count = 0
             val = values[i]
+
 """
 
 pad_2d_template = """@cython.boundscheck(False)
@@ -320,8 +320,6 @@ def pad_2d_inplace_%(name)s(ndarray[%(c_type)s, ndim=2] values,
 
     K, N = (<object> values).shape
 
-    val = np.nan
-
     if limit is None:
         lim = N
     else:
@@ -795,6 +793,12 @@ def generate_from_template(template, ndim=1, exclude=None):
 templates_2d = [take_2d_axis0_template,
                 take_2d_axis1_template]
 
+
+# templates_1d_datetime = [take_1d_template]
+# templates_2d_datetime = [take_2d_axis0_template,
+#                          take_2d_axis1_template]
+
+
 def generate_take_cython_file(path='generated.pyx'):
     with open(path, 'w') as f:
         for template in templates_1d:
@@ -803,6 +807,12 @@ def generate_take_cython_file(path='generated.pyx'):
         for template in templates_2d:
             print >> f, generate_from_template(template, ndim=2)
 
+        # for template in templates_1d_datetime:
+        #     print >> f, generate_from_template_datetime(template)
+
+        # for template in templates_2d_datetime:
+        #     print >> f, generate_from_template_datetime(template, ndim=2)
+
         for template in nobool_1d_templates:
             print >> f, generate_from_template(template, exclude=['bool'])