split view opt WIP

Author: Ehsan Totoni

hpat/_str_ext.cpp

@@ -29,13 +29,22 @@ struct str_arr_payload {
     uint8_t* null_bitmap;
 };
 
+// XXX: equivalent to payload data model in split_impl.py
+struct str_arr_split_view_payload {
+    uint32_t *index_offsets;
+    uint32_t *data_offsets;
+    // uint8_t* null_bitmap;
+};
+
 // taken from Arrow bin-util.h
 static constexpr uint8_t kBitmask[] = {1, 2, 4, 8, 16, 32, 64, 128};
 
 void* init_string(char*, int64_t);
 void* init_string_const(char* in_str);
 void dtor_string(std::string** in_str, int64_t size, void* in);
 void dtor_string_array(str_arr_payload* in_str, int64_t size, void* in);
+void dtor_str_arr_split_view(str_arr_split_view_payload* in_str_arr, int64_t size, void* in);
+void str_arr_split_view_impl(str_arr_split_view_payload* out_view, int64_t n_strs, uint32_t* offsets, char* data, char sep);
 const char* get_c_str(std::string* s);
 const char* get_char_ptr(char c);
 void* str_concat(std::string* s1, std::string* s2);
@@ -103,6 +112,10 @@ PyMODINIT_FUNC PyInit_hstr_ext(void) {
                             PyLong_FromVoidPtr((void*)(&dtor_string)));
     PyObject_SetAttrString(m, "dtor_string_array",
                             PyLong_FromVoidPtr((void*)(&dtor_string_array)));
+    PyObject_SetAttrString(m, "dtor_str_arr_split_view",
+                            PyLong_FromVoidPtr((void*)(&dtor_str_arr_split_view)));
+    PyObject_SetAttrString(m, "str_arr_split_view_impl",
+                            PyLong_FromVoidPtr((void*)(&str_arr_split_view_impl)));
     PyObject_SetAttrString(m, "get_c_str",
                             PyLong_FromVoidPtr((void*)(&get_c_str)));
     PyObject_SetAttrString(m, "get_char_ptr",
@@ -225,6 +238,65 @@ void dtor_string_array(str_arr_payload* in_str_arr, int64_t size, void* in)
     return;
 }
 
+void dtor_str_arr_split_view(str_arr_split_view_payload* in_str_arr, int64_t size, void* in)
+{
+    // printf("str arr dtor size: %lld\n", in_str_arr->size);
+    // printf("num chars: %d\n", in_str_arr->offsets[in_str_arr->size]);
+    delete[] in_str_arr->index_offsets;
+    delete[] in_str_arr->data_offsets;
+    // if (in_str_arr->null_bitmap != nullptr)
+    //     delete[] in_str_arr->null_bitmap;
+    return;
+}
+
+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;
+    uint32_t* index_offsets = new uint32_t[n_strs+1];
+    std::vector<uint32_t> data_offs;
+
+    data_offs.push_back(-1);
+    index_offsets[0] = 0;
+    // uint32_t curr_data_off = 0;
+
+    int data_ind = offsets[0];
+    int str_ind = 0;
+    // while there are chars to consume, equal since the first if will consume it
+    while (data_ind <= total_chars)
+    {
+        // string has finished
+        if (data_ind == offsets[str_ind+1])
+        {
+            data_offs.push_back(data_ind);
+            index_offsets[str_ind+1] = data_offs.size();
+            str_ind++;
+            if (str_ind == n_strs) break;  // all finished
+            continue;  // stay on same data_ind for start of next string
+        }
+        if (data[data_ind] == sep)
+        {
+            data_offs.push_back(data_ind);
+        }
+        data_ind++;
+    }
+    out_view->index_offsets = index_offsets;
+    out_view->data_offsets = new uint32_t[data_offs.size()];
+    // 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");
+    return;
+}
+
 const char* get_c_str(std::string* s)
 {
     // printf("in get %s\n", s->c_str());
@@ -507,7 +579,7 @@ void string_array_from_sequence(PyObject * obj, int64_t * no_strings, uint32_t *
         PyGILState_Release(gilstate);
         return;
     }
- 
+
     *no_strings = -1;
     *offset_table = NULL;
     *buffer = NULL;

hpat/hiframes/boxing.py

@@ -22,6 +22,8 @@
 from hpat.hiframes.pd_categorical_ext import (PDCategoricalDtype,
     box_categorical_array, unbox_categorical_array)
 from hpat.hiframes.pd_series_ext import SeriesType, arr_to_series_type
+from hpat.hiframes.split_impl import (string_array_split_view_type,
+    box_str_arr_split_view)
 
 from .. import hstr_ext
 import llvmlite.binding as ll
@@ -318,6 +320,8 @@ def _box_series_data(dtype, data_typ, val, c):
         arr = box_datetime_date_array(data_typ, val, c)
     elif isinstance(dtype, PDCategoricalDtype):
         arr = box_categorical_array(data_typ, val, c)
+    elif data_typ == string_array_split_view_type:
+        arr = box_str_arr_split_view(data_typ, val, c)
     elif dtype == types.List(string_type):
         arr = box_list(list_string_array_type, val, c)
     else:

hpat/hiframes/hiframes_typed.py

@@ -34,7 +34,7 @@
 from hpat.io.pio_api import h5dataset_type
 from hpat.hiframes.rolling import get_rolling_setup_args
 from hpat.hiframes.aggregate import Aggregate
-from hpat.hiframes import series_kernels
+from hpat.hiframes import series_kernels, split_impl
 from hpat.hiframes.series_kernels import series_replace_funcs
 
 
@@ -1687,6 +1687,7 @@ def _run_series_str_replace(self, assign, lhs, arr, rhs, nodes):
 
     def _run_series_str_split(self, assign, lhs, arr, rhs, nodes):
         sep = rhs.args[0]  # TODO: support default whitespace separator
+        sep_typ = self.typemap[sep.name]
 
         def _str_split_impl(str_arr, sep):
             numba.parfor.init_prange()
@@ -1698,6 +1699,13 @@ def _str_split_impl(str_arr, sep):
 
             return hpat.hiframes.api.init_series(out_arr)
 
+
+        if isinstance(sep_typ, types.StringLiteral) and len(sep_typ.literal_value) == 1:
+            def _str_split_impl(str_arr, sep):
+                out_arr = hpat.hiframes.split_impl.compute_split_view(
+                    str_arr, sep)
+                return hpat.hiframes.api.init_series(out_arr)
+
         return self._replace_func(_str_split_impl, [arr, sep], pre_nodes=nodes)
 
     def _run_series_str_get(self, assign, lhs, arr, rhs, nodes):

hpat/hiframes/pd_series_ext.py

@@ -21,7 +21,7 @@
     CategoricalArray)
 from hpat.hiframes.rolling import supported_rolling_funcs
 import datetime
-
+from hpat.hiframes.split_impl import string_array_split_view_type
 
 class SeriesType(types.IterableType):
     """Temporary type class for Series objects.
@@ -562,7 +562,11 @@ def resolve_replace(self, ary, args, kws):
 
     @bound_function("strmethod.split")
     def resolve_split(self, ary, args, kws):
-        return signature(SeriesType(types.List(string_type)), *args)
+        out = SeriesType(types.List(string_type))
+        if (len(args) == 1 and isinstance(args[0], types.StringLiteral)
+                and len(args[0].literal_value) == 1):
+            out = SeriesType(types.List(string_type), string_array_split_view_type)
+        return signature(out, *args)
 
     @bound_function("strmethod.get")
     def resolve_get(self, ary, args, kws):