optimized to_numeric() using string array view

Author: Ehsan Totoni

hpat/_str_ext.cpp

@@ -16,6 +16,8 @@ using std::regex;
 using std::regex_search;
 #endif
 
+#include <boost/lexical_cast.hpp>
+
 #ifndef _WIN32
 #include <glob.h>
 #endif
@@ -76,6 +78,8 @@ void* getitem_string_array_std(uint32_t *offsets, char *data, int64_t index);
 void print_str(std::string* str);
 void print_char(char c);
 void print_int(int64_t val);
+int str_arr_to_int64(int64_t* out, uint32_t *offsets, char *data, int64_t index);
+int str_arr_to_float64(double* out, uint32_t *offsets, char *data, int64_t index);
 void* compile_regex(std::string* pat);
 bool str_contains_regex(std::string* str, regex* e);
 bool str_contains_noregex(std::string* str, std::string* pat);
@@ -167,6 +171,10 @@ PyMODINIT_FUNC PyInit_hstr_ext(void) {
                             PyLong_FromVoidPtr((void*)(&print_char)));
     PyObject_SetAttrString(m, "print_int",
                             PyLong_FromVoidPtr((void*)(&print_int)));
+    PyObject_SetAttrString(m, "str_arr_to_int64",
+                            PyLong_FromVoidPtr((void*)(&str_arr_to_int64)));
+    PyObject_SetAttrString(m, "str_arr_to_float64",
+                            PyLong_FromVoidPtr((void*)(&str_arr_to_float64)));
     PyObject_SetAttrString(m, "compile_regex",
                             PyLong_FromVoidPtr((void*)(&compile_regex)));
     PyObject_SetAttrString(m, "str_contains_noregex",
@@ -493,6 +501,42 @@ void* getitem_string_array_std(uint32_t *offsets, char *data, int64_t index)
     return new std::string(&data[start], size);
 }
 
+
+int str_arr_to_int64(int64_t* out, uint32_t *offsets, char *data, int64_t index)
+{
+    uint32_t size = offsets[index+1]-offsets[index];
+    uint32_t start = offsets[index];
+    try
+    {
+        *out = boost::lexical_cast<int64_t>(data+start, (std::size_t)size);
+        return 0;
+    }
+    catch(const boost::bad_lexical_cast &)
+    {
+        *out = 0;
+        return -1;
+    }
+    return -1;
+}
+
+int str_arr_to_float64(double* out, uint32_t *offsets, char *data, int64_t index)
+{
+    uint32_t size = offsets[index+1]-offsets[index];
+    uint32_t start = offsets[index];
+    try
+    {
+        *out = boost::lexical_cast<double>(data+start, (std::size_t)size);
+        return 0;
+    }
+    catch(const boost::bad_lexical_cast &)
+    {
+        *out = std::nan("");  // TODO: numpy NaN
+        return -1;
+    }
+    return -1;
+}
+
+
 void* compile_regex(std::string* pat)
 {
     // printf("compiling\n");

hpat/hiframes/hiframes_typed.py

@@ -762,28 +762,30 @@ def _flatten_impl(A):
 
         if func_name == 'to_numeric':
             out_dtype = self.typemap[lhs.name].dtype
-            conv_func = int
+            conv_func = hpat.str_arr_ext.str_arr_item_to_int64
             if out_dtype == types.float64:
-                conv_func = float
+                conv_func = hpat.str_arr_ext.str_arr_item_to_float64
             else:
                 assert out_dtype == types.int64
 
             # TODO: handle non-Series input
 
             def _to_numeric_impl(A):
+                # TODO: fix distributed
                 numba.parfor.init_prange()
                 n = len(A)
                 B = np.empty(n, out_dtype)
                 for i in numba.parfor.internal_prange(n):
-                    B[i] = conv_func(A[i])
+                    conv_func(get_c_arr_ptr(B.ctypes, i), A, i)
 
                 return hpat.hiframes.api.init_series(B)
 
             nodes = []
             data = self._get_series_data(rhs.args[0], nodes)
             return self._replace_func(_to_numeric_impl, [data],
                 pre_nodes=nodes,
-                extra_globals={'out_dtype': out_dtype, 'conv_func': conv_func})
+                extra_globals={'out_dtype': out_dtype, 'conv_func': conv_func,
+                    'get_c_arr_ptr': hpat.hiframes.split_impl.get_c_arr_ptr})
 
         if func_name == 'parse_datetimes_from_strings':
             nodes = []

hpat/hiframes/split_impl.py

@@ -293,8 +293,12 @@ def codegen(context, builder, sig, args):
 
 @intrinsic
 def get_c_arr_ptr(typingctx, c_arr, ind_t=None):
+    assert isinstance(c_arr, (types.CPointer, types.ArrayCTypes))
     def codegen(context, builder, sig, args):
         in_arr, ind = args
+        if isinstance(sig.args[0], types.ArrayCTypes):
+            in_arr = builder.extract_value(in_arr, 0)
+
         return builder.bitcast(
             builder.gep(in_arr, [ind]), lir.IntType(8).as_pointer())
 

hpat/str_arr_ext.py

@@ -593,6 +593,8 @@ def str_arr_len(str_arr):
 ll.add_symbol('print_int', hstr_ext.print_int)
 ll.add_symbol('convert_len_arr_to_offset', hstr_ext.convert_len_arr_to_offset)
 ll.add_symbol('set_string_array_range', hstr_ext.set_string_array_range)
+ll.add_symbol('str_arr_to_int64', hstr_ext.str_arr_to_int64)
+ll.add_symbol('str_arr_to_float64', hstr_ext.str_arr_to_float64)
 
 convert_len_arr_to_offset = types.ExternalFunction("convert_len_arr_to_offset", types.void(types.voidptr, types.intp))
 
@@ -1084,6 +1086,56 @@ def str_arr_arr_impl(str_arr, ind_arr):
     res = context.compile_internal(builder, str_arr_arr_impl, sig, args)
     return res
 
+@intrinsic
+def str_arr_item_to_int64(typingctx, out_ptr_t, str_arr_t, ind_t=None):
+    assert str_arr_t == string_array_type
+    assert ind_t == types.int64
+
+    def codegen(context, builder, sig, args):
+        out_ptr, arr, ind = args
+        out_ptr = builder.bitcast(out_ptr, lir.IntType(64).as_pointer())
+        string_array = context.make_helper(builder, string_array_type, arr)
+        fnty = lir.FunctionType(
+            lir.IntType(32),
+            [lir.IntType(64).as_pointer(),
+             lir.IntType(32).as_pointer(),
+             lir.IntType(8).as_pointer(),
+             lir.IntType(64)])
+        fn_setitem = builder.module.get_or_insert_function(
+            fnty, name="str_arr_to_int64")
+        return builder.call(
+            fn_setitem,
+            [out_ptr, string_array.offsets, string_array.data, ind])
+
+    return types.int32(
+        out_ptr_t, string_array_type, types.int64), codegen
+
+
+@intrinsic
+def str_arr_item_to_float64(typingctx, out_ptr_t, str_arr_t, ind_t=None):
+    assert str_arr_t == string_array_type
+    assert ind_t == types.int64
+
+    def codegen(context, builder, sig, args):
+        out_ptr, arr, ind = args
+        out_ptr = builder.bitcast(out_ptr, lir.DoubleType().as_pointer())
+        string_array = context.make_helper(builder, string_array_type, arr)
+        fnty = lir.FunctionType(
+            lir.IntType(32),
+            [lir.DoubleType().as_pointer(),
+             lir.IntType(32).as_pointer(),
+             lir.IntType(8).as_pointer(),
+             lir.IntType(64)])
+        fn_setitem = builder.module.get_or_insert_function(
+            fnty, name="str_arr_to_float64")
+        return builder.call(
+            fn_setitem,
+            [out_ptr, string_array.offsets, string_array.data, ind])
+
+    return types.int32(
+        out_ptr_t, string_array_type, types.int64), codegen
+
+
 
 # TODO: support array of strings
 # @typeof_impl.register(np.ndarray)

hpat/tests/test_hiframes.py

@@ -453,6 +453,17 @@ def test_impl(df):
         self.assertEqual(count_array_REPs(), 3)
         self.assertEqual(count_parfor_REPs(), 0)
 
+    def test_str_get_to_numeric(self):
+        def test_impl(df):
+            B = df.A.str.split(',')
+            C = pd.to_numeric(B.str.get(1), errors='coerce')
+            return C
+
+        df = pd.DataFrame({'A': ['AB,12', 'C,321,D']})
+        hpat_func = hpat.jit(locals={'C': hpat.int64[:]})(test_impl)
+        pd.testing.assert_series_equal(
+            hpat_func(df), test_impl(df), check_names=False)
+
     def test_str_flatten(self):
         def test_impl(df):
             A = df.A.str.split(',')