Merge pull request #33 from chaburkland/2/dtype_from_element

2/dtype from element

Author: chaburkland

performance/__main__.py

@@ -1,3 +1,4 @@
+import collections
 import datetime
 import timeit
 import argparse
@@ -13,6 +14,7 @@
 from performance.reference.util import row_1d_filter as row_1d_filter_ref
 from performance.reference.util import resolve_dtype as resolve_dtype_ref
 from performance.reference.util import resolve_dtype_iter as resolve_dtype_iter_ref
+from performance.reference.util import dtype_from_element as dtype_from_element_ref
 from performance.reference.util import array_deepcopy as array_deepcopy_ref
 from performance.reference.util import isna_element as isna_element_ref
 
@@ -27,6 +29,7 @@
 from arraykit import row_1d_filter as row_1d_filter_ak
 from arraykit import resolve_dtype as resolve_dtype_ak
 from arraykit import resolve_dtype_iter as resolve_dtype_iter_ak
+from arraykit import dtype_from_element as dtype_from_element_ak
 from arraykit import array_deepcopy as array_deepcopy_ak
 from arraykit import isna_element as isna_element_ak
 
@@ -250,6 +253,45 @@ class ArrayGOPerfREF(ArrayGOPerf):
     entry = staticmethod(ArrayGOREF)
 
 
+#-------------------------------------------------------------------------------
+class DtypeFromElementPerf(Perf):
+    NUMBER = 1000
+
+    def pre(self):
+        NT = collections.namedtuple('NT', tuple('abc'))
+
+        self.values = [
+                np.longlong(-1), np.int_(-1), np.intc(-1), np.short(-1), np.byte(-1),
+                np.ubyte(1), np.ushort(1), np.uintc(1), np.uint(1), np.ulonglong(1),
+                np.half(1.0), np.single(1.0), np.float_(1.0), np.longfloat(1.0),
+                np.csingle(1.0j), np.complex_(1.0j), np.clongfloat(1.0j),
+                np.bool_(0), np.str_('1'), np.unicode_('1'), np.void(1),
+                np.object(), np.datetime64('NaT'), np.timedelta64('NaT'), np.nan,
+                12, 12.0, True, None, float('NaN'), object(), (1, 2, 3),
+                NT(1, 2, 3), datetime.date(2020, 12, 31), datetime.timedelta(14),
+        ]
+
+        # Datetime & Timedelta
+        for precision in ['ns', 'us', 'ms', 's', 'm', 'h', 'D', 'M', 'Y']:
+            for kind, ctor in (('m', np.timedelta64), ('M', np.datetime64)):
+                self.values.append(ctor(12, precision))
+
+        for size in (1, 8, 16, 32, 64, 128, 256, 512):
+            self.values.append(bytes(size))
+            self.values.append('x' * size)
+
+    def main(self):
+        for _ in range(40):
+            for val in self.values:
+                self.entry(val)
+
+class DtypeFromElementPerfAK(DtypeFromElementPerf):
+    entry = staticmethod(dtype_from_element_ak)
+
+class DtypeFromElementPerfREF(DtypeFromElementPerf):
+    entry = staticmethod(dtype_from_element_ref)
+
+
 #-------------------------------------------------------------------------------
 class IsNaElementPerf(Perf):
     NUMBER = 1000
@@ -336,7 +378,6 @@ def main():
 
     records = [('cls', 'func', 'ak', 'ref', 'ref/ak')]
     for cls_perf in Perf.__subclasses__(): # only get one level
-        print(cls_perf)
         cls_map = {}
         if match and cls_perf.__name__ not in match:
             continue

performance/reference/util.py

@@ -199,3 +199,20 @@ def isna_element(value: tp.Any) -> bool:
         return np.isnat(value) #type: ignore
 
     return value is None
+
+
+def dtype_from_element(value: tp.Optional[tp.Hashable]) -> np.dtype:
+    '''Given an arbitrary hashable to be treated as an element, return the appropriate dtype. This was created to avoid using np.array(value).dtype, which for a Tuple does not return object.
+    '''
+    if value is np.nan:
+        # NOTE: this will not catch all NaN instances, but will catch any default NaNs in function signatures that reference the same NaN object found on the NP root namespace
+        return DTYPE_FLOAT_DEFAULT
+    if value is None:
+        return DTYPE_OBJECT
+    if isinstance(value, tuple):
+        return DTYPE_OBJECT
+    if hasattr(value, 'dtype'):
+        return value.dtype #type: ignore
+    # NOTE: calling array and getting dtype on np.nan is faster than combining isinstance, isnan calls
+    return np.array(value).dtype
+

src/__init__.py

@@ -15,3 +15,4 @@
 from ._arraykit import resolve_dtype as resolve_dtype
 from ._arraykit import resolve_dtype_iter as resolve_dtype_iter
 from ._arraykit import isna_element as isna_element
+from ._arraykit import dtype_from_element as dtype_from_element

src/__init__.pyi

@@ -31,3 +31,5 @@ def array_deepcopy(__array: np.array, memo: tp.Dict[int, tp.Any]) -> np.ndarray:
 def resolve_dtype(__d1: np.dtype, __d2: np.dtype) -> np.dtype: ...
 def resolve_dtype_iter(__dtypes: tp.Iterable[np.dtype]) -> np.dtype: ...
 def isna_element(__value: tp.Any) -> bool: ...
+def dtype_from_element(__value: tp.Optional[tp.Hashable]) -> np.dtype: ...
+

src/_arraykit.c

@@ -365,6 +365,71 @@ resolve_dtype_iter(PyObject *Py_UNUSED(m), PyObject *arg)
 //------------------------------------------------------------------------------
 // general utility
 
+static PyObject *
+dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg)
+{
+    // -------------------------------------------------------------------------
+    // 1. Handle fast, exact type checks first.
+
+    // None
+    if (arg == Py_None) {
+        return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
+    }
+
+    // Float
+    if (PyFloat_CheckExact(arg)) {
+        return (PyObject*)PyArray_DescrFromType(NPY_DOUBLE);
+    }
+
+    // Integers
+    if (PyLong_CheckExact(arg)) {
+        return (PyObject*)PyArray_DescrFromType(NPY_LONG);
+    }
+
+    // Bool
+    if (PyBool_Check(arg)) {
+        return (PyObject*)PyArray_DescrFromType(NPY_BOOL);
+    }
+
+    PyObject* dtype = NULL;
+
+    // String
+    if (PyUnicode_CheckExact(arg)) {
+        PyArray_Descr* descr = PyArray_DescrFromType(NPY_UNICODE);
+        if (descr == NULL) {
+            return NULL;
+        }
+        dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
+        Py_DECREF(descr);
+        return dtype;
+    }
+
+    // Bytes
+    if (PyBytes_CheckExact(arg)) {
+        PyArray_Descr* descr = PyArray_DescrFromType(NPY_STRING);
+        if (descr == NULL) {
+            return NULL;
+        }
+        dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
+        Py_DECREF(descr);
+        return dtype;
+    }
+
+    // -------------------------------------------------------------------------
+    // 2. Construct dtype (slightly more complicated)
+
+    // Already known
+    dtype = PyObject_GetAttrString(arg, "dtype");
+    if (dtype) {
+        return dtype;
+    }
+    PyErr_Clear();
+
+    // -------------------------------------------------------------------------
+    // 3. Handles everything else.
+    return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
+}
+
 static PyObject *
 isna_element(PyObject *Py_UNUSED(m), PyObject *arg)
 {
@@ -706,6 +771,7 @@ static PyMethodDef arraykit_methods[] =  {
     {"resolve_dtype", resolve_dtype, METH_VARARGS, NULL},
     {"resolve_dtype_iter", resolve_dtype_iter, METH_O, NULL},
     {"isna_element", isna_element, METH_O, NULL},
+    {"dtype_from_element", dtype_from_element, METH_O, NULL},
     {NULL},
 };
 

test/test_util.py

@@ -1,3 +1,4 @@
+import collections
 import datetime
 import unittest
 import itertools
@@ -14,6 +15,7 @@
 from arraykit import immutable_filter
 from arraykit import array_deepcopy
 from arraykit import isna_element
+from arraykit import dtype_from_element
 
 from performance.reference.util import mloc as mloc_ref
 
@@ -299,7 +301,73 @@ def test_isna_element_false(self) -> None:
         self.assertFalse(isna_element(datetime.date(2020, 12, 31)))
         self.assertFalse(isna_element(False))
 
+    def test_dtype_from_element_core_dtypes(self) -> None:
+        dtypes = [
+                np.longlong,
+                np.int_,
+                np.intc,
+                np.short,
+                np.byte,
+                np.ubyte,
+                np.ushort,
+                np.uintc,
+                np.uint,
+                np.ulonglong,
+                np.half,
+                np.single,
+                np.float_,
+                np.longfloat,
+                np.csingle,
+                np.complex_,
+                np.clongfloat,
+                np.bool_,
+        ]
+        for dtype in dtypes:
+            self.assertEqual(dtype, dtype_from_element(dtype()))
+
+    def test_dtype_from_element_str_and_misc_dtypes(self) -> None:
+        dtype_obj_pairs = [
+                (np.dtype('<U1'), np.str_('1')),
+                (np.dtype('<U1'), np.unicode_('1')),
+                (np.dtype('V1'), np.void(1)),
+                (np.dtype('O'), np.object()),
+                (np.dtype('<M8'), np.datetime64('NaT')),
+                (np.dtype('<m8'), np.timedelta64('NaT')),
+                (np.float_, np.nan),
+        ]
+        for dtype, obj in dtype_obj_pairs:
+            self.assertEqual(dtype, dtype_from_element(obj))
+
+    def test_dtype_from_element_obj_dtypes(self) -> None:
+        NT = collections.namedtuple('NT', tuple('abc'))
+
+        dtype_obj_pairs = [
+                (np.int_, 12),
+                (np.float_, 12.0),
+                (np.bool_, True),
+                (np.dtype('O'), None),
+                (np.float_, float('NaN')),
+                (np.dtype('O'), object()),
+                (np.dtype('O'), (1, 2, 3)),
+                (np.dtype('O'), NT(1, 2, 3)),
+                (np.dtype('O'), datetime.date(2020, 12, 31)),
+                (np.dtype('O'), datetime.timedelta(14)),
+        ]
+        for dtype, obj in dtype_obj_pairs:
+            self.assertEqual(dtype, dtype_from_element(obj))
+
+    def test_dtype_from_element_time_dtypes(self) -> None:
+        # Datetime & Timedelta
+        for precision in ['ns', 'us', 'ms', 's', 'm', 'h', 'D', 'M', 'Y']:
+            for kind, ctor in (('m', np.timedelta64), ('M', np.datetime64)):
+                obj = ctor(12, precision)
+                self.assertEqual(np.dtype(f'<{kind}8[{precision}]'), dtype_from_element(obj))
+
+    def test_dtype_from_element_str_and_bytes_dtypes(self) -> None:
+        for size in (1, 8, 16, 32, 64, 128, 256, 512):
+            self.assertEqual(np.dtype(f'|S{size}'), dtype_from_element(bytes(size)))
+            self.assertEqual(np.dtype(f'<U{size}'), dtype_from_element('x' * size))
+
 
 if __name__ == '__main__':
     unittest.main()
-