fix split view boxing

Author: Ehsan Totoni

hpat/_str_ext.cpp

@@ -92,6 +92,8 @@ int64_t hash_str(std::string* in_str);
 void c_glob(uint32_t **offsets, char **data, uint8_t **null_bitmap, int64_t* num_strings, char* path);
 npy_intp array_size(PyArrayObject* arr);
 void* array_getptr1(PyArrayObject* arr, npy_intp ind);
+void array_setitem(PyArrayObject* arr, char* p, PyObject *s);
+
 
 PyMODINIT_FUNC PyInit_hstr_ext(void) {
     PyObject *m;
@@ -192,6 +194,8 @@ PyMODINIT_FUNC PyInit_hstr_ext(void) {
                             PyLong_FromVoidPtr((void*)(&array_size)));
     PyObject_SetAttrString(m, "array_getptr1",
                             PyLong_FromVoidPtr((void*)(&array_getptr1)));
+    PyObject_SetAttrString(m, "array_setitem",
+                            PyLong_FromVoidPtr((void*)(&array_setitem)));
     return m;
 }
 
@@ -249,11 +253,14 @@ void dtor_str_arr_split_view(str_arr_split_view_payload* in_str_arr, int64_t siz
     return;
 }
 
+// example: ['AB,CC', 'C,ABB,D', 'G', '', 'g,f']
+// offsets [0, 5, 12, 13, 13, 14, 17]
+// data_offsets [-1, 2, 5,   4, 6, 10, 12,  11, 13,   12, 13,   12, 14, 16]
+// index_offsets [0, 3, 7, 9, 11, 14]
 void str_arr_split_view_impl(str_arr_split_view_payload* out_view, int64_t n_strs, uint32_t* offsets, char* data, char sep)
 {
     uint32_t total_chars = offsets[n_strs];
-    printf("n_strs %d sep %c total chars:%d\n", n_strs, sep, total_chars);
-    //return;
+    // printf("n_strs %d sep %c total chars:%d\n", n_strs, sep, total_chars);
     uint32_t* index_offsets = new uint32_t[n_strs+1];
     std::vector<uint32_t> data_offs;
 
@@ -273,6 +280,8 @@ void str_arr_split_view_impl(str_arr_split_view_payload* out_view, int64_t n_str
             index_offsets[str_ind+1] = data_offs.size();
             str_ind++;
             if (str_ind == n_strs) break;  // all finished
+            // start new string
+            data_offs.push_back(data_ind-1);
             continue;  // stay on same data_ind for start of next string
         }
         if (data[data_ind] == sep)
@@ -286,14 +295,14 @@ void str_arr_split_view_impl(str_arr_split_view_payload* out_view, int64_t n_str
     // TODO: avoid copy
     std::copy(data_offs.cbegin(), data_offs.cend(), out_view->data_offsets);
 
-    printf("index_offsets: ");
-    for (int i=0; i<=n_strs; i++)
-        printf("%d ", index_offsets[i]);
-    printf("\n");
-    printf("data_offsets: ");
-    for (int i=0; i<data_offs.size(); i++)
-        printf("%d ", data_offs[i]);
-    printf("\n");
+    // printf("index_offsets: ");
+    // for (int i=0; i<=n_strs; i++)
+    //     printf("%d ", index_offsets[i]);
+    // printf("\n");
+    // printf("data_offsets: ");
+    // for (int i=0; i<data_offs.size(); i++)
+    //     printf("%d ", data_offs[i]);
+    // printf("\n");
     return;
 }
 
@@ -704,6 +713,16 @@ void* array_getptr1(PyArrayObject* arr, npy_intp ind)
     return PyArray_GETPTR1(arr, ind);
 }
 
+void array_setitem(PyArrayObject* arr, char* p, PyObject *s)
+{
+#define CHECK(expr, msg) if(!(expr)){std::cerr << msg << std::endl; return;}
+    // std::cout << "get array ptr " << ind << '\n';
+    int err = PyArray_SETITEM(arr, p, s);
+    CHECK(err==0, "setting item in numpy array failed");
+    return;
+#undef CHECK
+}
+
 // glob support
 void c_glob(uint32_t **offsets, char **data, uint8_t **null_bitmap, int64_t* num_strings, char* path)
 {

hpat/hiframes/split_impl.py

@@ -19,6 +19,10 @@
 import llvmlite.llvmpy.core as lc
 from llvmlite import ir as lir
 import llvmlite.binding as ll
+from llvmlite.llvmpy.core import Type as LLType
+from .. import hstr_ext
+ll.add_symbol('array_setitem', hstr_ext.array_setitem)
+ll.add_symbol('array_getptr1', hstr_ext.array_getptr1)
 
 from .. import hstr_ext
 ll.add_symbol('dtor_str_arr_split_view', hstr_ext.dtor_str_arr_split_view)
@@ -187,6 +191,15 @@ def box_str_arr_split_view(typ, val, c):
     out_arr = c.pyapi.call_method(
         np_class_obj, "ndarray", (num_items_obj, dtype))
 
+    # Array setitem call
+    arr_get_fnty = LLType.function(
+        lir.IntType(8).as_pointer(), [c.pyapi.pyobj, c.pyapi.py_ssize_t])
+    arr_get_fn = c.pyapi._get_function(arr_get_fnty, name="array_getptr1")
+    arr_setitem_fnty = LLType.function(
+        lir.VoidType(),
+        [c.pyapi.pyobj, lir.IntType(8).as_pointer(), c.pyapi.pyobj])
+    arr_setitem_fn = c.pyapi._get_function(
+        arr_setitem_fnty, name="array_setitem")
 
     # for each string
     with cgutils.for_range(builder, sp_view.num_items) as loop:
@@ -200,7 +213,8 @@ def box_str_arr_split_view(typ, val, c):
 
         # Build a new Python list
         nitems = builder.sub(list_end_offset, list_start_offset)
-        cgutils.printf(builder, "str %lld n %lld\n", str_ind, nitems)
+        nitems = builder.sub(nitems, nitems.type(1))
+        # cgutils.printf(builder, "str %lld n %lld\n", str_ind, nitems)
         list_obj = c.pyapi.list_new(nitems)
         with c.builder.if_then(cgutils.is_not_null(c.builder, list_obj),
                                likely=True):
@@ -211,10 +225,14 @@ def box_str_arr_split_view(typ, val, c):
                 # add 1 since starts from -1
                 data_start = builder.add(data_start, data_start.type(1))
                 data_end = builder.load(builder.gep(sp_view.data_offsets, [builder.add(start_index, start_index.type(1))]))
-                cgutils.printf(builder, "ind %lld %lld\n", data_start, data_end)
-                #itemobj =
-                #c.pyapi.list_setitem(obj, loop.index, itemobj)
-
+                # cgutils.printf(builder, "ind %lld %lld\n", data_start, data_end)
+                data_ptr = builder.gep(builder.extract_value(sp_view.data, 0), [data_start])
+                str_size = builder.sext(builder.sub(data_end, data_start), lir.IntType(64))
+                str_obj = c.pyapi.string_from_string_and_size(data_ptr, str_size)
+                c.pyapi.list_setitem(list_obj, loop.index, str_obj)
+
+        arr_ptr = builder.call(arr_get_fn, [out_arr, str_ind])
+        builder.call(arr_setitem_fn, [out_arr, arr_ptr, list_obj])
 
 
     c.pyapi.decref(np_class_obj)