Merge pull request #73 from static-frame/72/split-after-count

Author: flexatone

performance/__main__.py

@@ -25,6 +25,8 @@
 from performance.reference.util import isna_element as isna_element_ref
 from performance.reference.util import get_new_indexers_and_screen_ak
 from performance.reference.util import get_new_indexers_and_screen_ref
+from performance.reference.util import split_after_count as split_after_count_ref
+from performance.reference.util import count_iteration as count_iteration_ref
 
 from performance.reference.array_go import ArrayGO as ArrayGOREF
 
@@ -41,6 +43,8 @@
 from arraykit import array_deepcopy as array_deepcopy_ak
 from arraykit import delimited_to_arrays as delimited_to_arrays_ak
 from arraykit import isna_element as isna_element_ak
+from arraykit import split_after_count as split_after_count_ak
+from arraykit import count_iteration as count_iteration_ak
 
 from arraykit import ArrayGO as ArrayGOAK
 
@@ -695,6 +699,43 @@ class GetNewIndexersAndScreenPerfREF(GetNewIndexersAndScreenPerf):
     entry = staticmethod(get_new_indexers_and_screen_ref)
 
 
+
+
+#-------------------------------------------------------------------------------
+class SplitAfterCount(Perf):
+    NUMBER = 200_000
+
+    def __init__(self):
+        self.string = ''.join(['abcd,'] * 1000)
+
+    def main(self):
+        post = self.entry(self.string, ',', 20)
+
+class SplitAfterCountAK(SplitAfterCount):
+    entry = staticmethod(split_after_count_ak)
+
+class SplitAfterCountREF(SplitAfterCount):
+    entry = staticmethod(split_after_count_ref)
+
+
+#-------------------------------------------------------------------------------
+class CountIterations(Perf):
+    NUMBER = 10_000
+
+    def __init__(self):
+        self.strio = io.StringIO('\n'.join(['abcd'] * 10_000))
+
+    def main(self):
+        post = self.entry(self.strio)
+        self.strio.seek(0)
+
+class CountIterationsAK(CountIterations):
+    entry = staticmethod(count_iteration_ak)
+
+class CountIterationsREF(CountIterations):
+    entry = staticmethod(count_iteration_ref)
+
+
 #-------------------------------------------------------------------------------
 
 def get_arg_parser():

performance/reference/util.py

@@ -248,3 +248,18 @@ def get_new_indexers_and_screen_ak(
         return np.unique(indexers, return_inverse=True)
 
     return ak_routine(indexers, positions)
+
+
+def split_after_count(string: str, delimiter: str, count: int):
+    *left, right = string.split(delimiter, maxsplit=count)
+    return ','.join(left), right
+
+def count_iteration(iterable: tp.Iterable):
+    count = 0
+    for i in iterable:
+        count += 1
+    return count
+
+
+
+

src/__init__.py

@@ -19,3 +19,5 @@
 from ._arraykit import delimited_to_arrays as delimited_to_arrays
 from ._arraykit import iterable_str_to_array_1d as iterable_str_to_array_1d
 from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen
+from ._arraykit import split_after_count as split_after_count
+from ._arraykit import count_iteration as count_iteration

src/__init__.pyi

@@ -45,6 +45,15 @@ def delimited_to_arrays(
         decimalchar: str = '.',
         ) -> tp.List[np.array]: ...
 
+def split_after_count(
+        string: str,
+        *
+        delimiter: str,
+        count: int,
+        ) -> tp.Tuple[str, str]: ...
+
+def count_iteration(__iterable: tp.Iterable) -> int: ...
+
 def immutable_filter(__array: np.ndarray) -> np.ndarray: ...
 def mloc(__array: np.ndarray) -> int: ...
 def name_filter(__name: tp.Hashable) -> tp.Hashable: ...
@@ -57,4 +66,4 @@ 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: ...
-def get_new_indexers_and_screen(indexers: np.ndarray, positions: np.ndarray) -> tp.Tuple[np.ndarray, np.ndarray]: ...
+def get_new_indexers_and_screen(indexers: np.ndarray, positions: np.ndarray) -> tp.Tuple[np.ndarray, np.ndarray]: ...

src/_arraykit.c

@@ -1441,10 +1441,10 @@ AK_CPL_FromIterable(PyObject* iterable, bool type_parse, Py_UCS4 tsep, Py_UCS4 d
         }
         Py_DECREF(field);
     }
+    Py_DECREF(iter);
     if (PyErr_Occurred()) {
         return NULL;
     }
-    Py_DECREF(iter);
     return cpl;
 }
 
@@ -2955,6 +2955,97 @@ iterable_str_to_array_1d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwarg
     return AK_IterableStrToArray1D(iterable, dtype_specifier, tsep, decc);
 }
 
+
+static PyObject *
+split_after_count(PyObject *Py_UNUSED(m), PyObject *args)
+{
+    PyObject *string = NULL;
+    PyObject *delimiter = NULL;
+    int count = 0;
+
+    if (!PyArg_ParseTuple(args,
+            "OOi:split_after_count",
+            &string,
+            &delimiter,
+            &count)) {
+        return NULL;
+    }
+
+    if (!PyUnicode_Check(string)) {
+        PyErr_Format(PyExc_RuntimeError,
+                "a string is required, not %.200s",
+                Py_TYPE(string)->tp_name
+                );
+        return NULL;
+    }
+
+    if (count <= 0) {
+        PyErr_Format(PyExc_RuntimeError,
+                "count must be greater than zero, not %i",
+                count
+                );
+        return NULL;
+    }
+
+    Py_UCS4 delim_char;
+    if (AK_set_char(
+            "delimiter",
+            &delim_char,
+            delimiter,
+            '\0')) return NULL;
+
+    unsigned int kind = PyUnicode_KIND(string);
+    const void *data = PyUnicode_DATA(string);
+    Py_ssize_t pos = 0;
+    Py_ssize_t delim_count = 0;
+    Py_ssize_t linelen = PyUnicode_GET_LENGTH(string);
+    Py_UCS4 c;
+
+    while (pos < linelen) {
+        c = PyUnicode_READ(kind, data, pos);
+        if (c == delim_char) {
+            delim_count++;
+            if (delim_count == count) {
+                break; // to not include delim at transition
+                // do not increment pos so as to exclude in left
+            }
+        }
+        pos++;
+    }
+
+    PyObject* left = PyUnicode_Substring(string, 0, pos);
+    PyObject* right = PyUnicode_Substring(string, pos+1, linelen);
+    PyObject *result = PyTuple_Pack(2, left, right);
+    Py_DECREF(left);
+    Py_DECREF(right);
+    return result;
+}
+
+
+
+static PyObject *
+count_iteration(PyObject *Py_UNUSED(m), PyObject *iterable)
+{
+    PyObject *iter = PyObject_GetIter(iterable);
+    if (iter == NULL) return NULL;
+
+    int count = 0;
+    PyObject *v;
+
+    while ((v = PyIter_Next(iter))) {
+        count++;
+        Py_DECREF(v);
+    }
+    Py_DECREF(iter);
+    if (PyErr_Occurred()) {
+        return NULL;
+    }
+    PyObject* result = PyLong_FromLong(count);
+    if (result == NULL) return NULL;
+    return result;
+}
+
+
 //------------------------------------------------------------------------------
 
 // Return the integer version of the pointer to underlying data-buffer of array.
@@ -3087,7 +3178,7 @@ resolve_dtype(PyObject *Py_UNUSED(m), PyObject *args)
 {
     PyArray_Descr *d1, *d2;
     if (!PyArg_ParseTuple(args, "O!O!:resolve_dtype",
-                          &PyArrayDescr_Type, &d1, &PyArrayDescr_Type, &d2))
+            &PyArrayDescr_Type, &d1, &PyArrayDescr_Type, &d2))
     {
         return NULL;
     }
@@ -3817,6 +3908,8 @@ static PyMethodDef arraykit_methods[] =  {
             (PyCFunction)iterable_str_to_array_1d,
             METH_VARARGS | METH_KEYWORDS,
             NULL},
+    {"split_after_count", split_after_count, METH_VARARGS, NULL},
+    {"count_iteration", count_iteration, METH_O, NULL},
     {"isna_element", isna_element, METH_O, NULL},
     {"dtype_from_element", dtype_from_element, METH_O, NULL},
     {"get_new_indexers_and_screen",

test/test_util.py

@@ -3,7 +3,7 @@
 import datetime
 import unittest
 import warnings
-
+from io import StringIO
 import numpy as np  # type: ignore
 
 from arraykit import resolve_dtype
@@ -17,6 +17,9 @@
 from arraykit import array_deepcopy
 from arraykit import isna_element
 from arraykit import dtype_from_element
+from arraykit import split_after_count
+from arraykit import count_iteration
+
 from performance.reference.util import get_new_indexers_and_screen_ak as get_new_indexers_and_screen_full
 from arraykit import get_new_indexers_and_screen
 
@@ -442,5 +445,58 @@ def test_get_new_indexers_and_screen_b(self) -> None:
         assert tuple(map(list, postB)) == (list(indexersB), list(indexersB))
 
 
+    #---------------------------------------------------------------------------
+    def test_split_after_count_a(self) -> None:
+        post = split_after_count('a,b,c,d,e', ',', 2)
+        self.assertEqual(post[0], 'a,b')
+        self.assertEqual(post[1], 'c,d,e')
+
+    def test_split_after_count_b(self) -> None:
+        post = split_after_count('a,b,c,d,e', ',', 4)
+        self.assertEqual(post[0], 'a,b,c,d')
+        self.assertEqual(post[1], 'e')
+
+    def test_split_after_count_c(self) -> None:
+        post = split_after_count('a,b,c,d,e', ',', 5)
+        self.assertEqual(post[0], 'a,b,c,d,e')
+        self.assertEqual(post[1], '')
+
+    def test_split_after_count_d(self) -> None:
+        post = split_after_count('a', ',', 5)
+        self.assertEqual(post[0], 'a')
+        self.assertEqual(post[1], '')
+
+    def test_split_after_count_e(self) -> None:
+        with self.assertRaises(RuntimeError):
+            post = split_after_count('a,', ',', 0)
+
+    def test_split_after_count_f(self) -> None:
+        post = split_after_count('a,', ',', 1)
+        self.assertEqual(post[0], 'a')
+        self.assertEqual(post[1], '')
+
+    def test_split_after_count_g(self) -> None:
+        post = split_after_count(',', ',', 1)
+        self.assertEqual(post[0], '')
+        self.assertEqual(post[1], '')
+
+    def test_split_after_count_h(self) -> None:
+        post = split_after_count('a,b,c,d,e', '|', 5)
+        self.assertEqual(post[0], 'a,b,c,d,e')
+        self.assertEqual(post[1], '')
+
+
+    #---------------------------------------------------------------------------
+    def test_count_iteration_a(self) -> None:
+        post = count_iteration(('a', 'b', 'c', 'd'))
+        self.assertEqual(post, 4)
+
+    def test_count_iteration_b(self) -> None:
+        s1 = StringIO(',1,a,b\n-,1,43,54\nX,2,1,3\nY,1,8,10\n-,2,6,20')
+        post = count_iteration(s1)
+        self.assertEqual(post, 5)
+
+
+
 if __name__ == '__main__':
     unittest.main()