ENH: move _ensure_{dtype} functions to Cython for speedup, close #1221

Author: wesm

pandas/core/common.py

@@ -711,36 +711,12 @@ def is_float_dtype(arr_or_dtype):
     return issubclass(tipo, np.floating)
 
 
-def _ensure_float64(arr):
-    if arr.dtype != np.float64:
-        arr = arr.astype(np.float64)
-    return arr
-
-def _ensure_int64(arr):
-    try:
-        if arr.dtype != np.int64:
-            arr = arr.astype(np.int64)
-        return arr
-    except AttributeError:
-        return np.array(arr, dtype=np.int64)
+_ensure_float64 = _algos.ensure_float64
+_ensure_int64 = _algos.ensure_int64
+_ensure_int32 = _algos.ensure_int32
+_ensure_platform_int = _algos.ensure_platform_int
+_ensure_object = _algos.ensure_object
 
-def _ensure_platform_int(labels):
-    try:
-        if labels.dtype != np.int_:  # pragma: no cover
-            labels = labels.astype(np.int_)
-        return labels
-    except AttributeError:
-        return np.array(labels, dtype=np.int_)
-
-def _ensure_int32(arr):
-    if arr.dtype != np.int32:
-        arr = arr.astype(np.int32)
-    return arr
-
-def _ensure_object(arr):
-    if arr.dtype != np.object_:
-        arr = arr.astype('O')
-    return arr
 
 def _astype_nansafe(arr, dtype):
     if (np.issubdtype(arr.dtype, np.floating) and

pandas/src/generate_code.py

@@ -810,6 +810,35 @@ def outer_join_indexer_%(name)s(ndarray[%(c_type)s] left,
 
 """
 
+# ensure_dtype functions
+
+ensure_dtype_template = """
+cpdef ensure_%(name)s(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_%(ctype)s:
+            return arr
+        else:
+            return arr.astype(np.%(dtype)s)
+    else:
+        return np.array(arr, dtype=np.%(dtype)s)
+
+"""
+
+ensure_functions = [
+    ('float64', 'FLOAT64', 'float64'),
+    ('int32', 'INT32', 'int32'),
+    ('int64', 'INT64', 'int64'),
+    ('platform_int', 'INT', 'int_'),
+    ('object', 'OBJECT', 'object_'),
+]
+
+def generate_ensure_dtypes():
+    output = StringIO()
+    for name, ctype, dtype in ensure_functions:
+        filled = ensure_dtype_template % locals()
+        output.write(filled)
+    return output.getvalue()
+
 #----------------------------------------------------------------------
 # Fast "put" logic for speeding up interleaving logic
 
@@ -916,6 +945,8 @@ def generate_take_cython_file(path='generated.pyx'):
         for template in nobool_1d_templates:
             print >> f, generate_from_template(template, exclude=['bool'])
 
+        print >> f, generate_ensure_dtypes()
+
         # print >> f, generate_put_functions()
 
 if __name__ == '__main__':

pandas/src/generated.pyx

@@ -3306,3 +3306,54 @@ def inner_join_indexer_int64(ndarray[int64_t] left,
     return result, lindexer, rindexer
 
 
+
+cpdef ensure_float64(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_FLOAT64:
+            return arr
+        else:
+            return arr.astype(np.float64)
+    else:
+        return np.array(arr, dtype=np.float64)
+
+
+cpdef ensure_int32(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_INT32:
+            return arr
+        else:
+            return arr.astype(np.int32)
+    else:
+        return np.array(arr, dtype=np.int32)
+
+
+cpdef ensure_int64(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_INT64:
+            return arr
+        else:
+            return arr.astype(np.int64)
+    else:
+        return np.array(arr, dtype=np.int64)
+
+
+cpdef ensure_platform_int(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_INT:
+            return arr
+        else:
+            return arr.astype(np.int_)
+    else:
+        return np.array(arr, dtype=np.int_)
+
+
+cpdef ensure_object(object arr):
+    if util.is_array(arr):
+        if (<ndarray> arr).descr.type_num == NPY_OBJECT:
+            return arr
+        else:
+            return arr.astype(np.object_)
+    else:
+        return np.array(arr, dtype=np.object_)
+
+

pandas/src/tseries.pyx

@@ -665,6 +665,7 @@ def value_count_int64(ndarray[int64_t] values):
 
     return result_keys, result_counts
 
+
 include "hashtable.pyx"
 include "datetime.pyx"
 include "skiplist.pyx"

pandas/src/util.pxd

@@ -60,4 +60,3 @@ cdef inline bint _checknull(object val):
 
 cdef inline bint _checknan(object val):
     return not cnp.PyArray_Check(val) and val != val
-