Merge pull request #63 from chaburkland/burkland/new_indexers_from_indexer_subset

Add get_new_indexers_and_screen function

Author: chaburkland

performance/__main__.py

@@ -1,7 +1,9 @@
-import collections
+from collections import namedtuple
 import datetime
 import timeit
 import argparse
+import typing as tp
+from enum import Enum
 
 import numpy as np
 
@@ -17,6 +19,8 @@
 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
+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.array_go import ArrayGO as ArrayGOREF
 
@@ -258,7 +262,7 @@ class DtypeFromElementPerf(Perf):
     NUMBER = 1000
 
     def __init__(self):
-        NT = collections.namedtuple('NT', tuple('abc'))
+        NT = namedtuple('NT', tuple('abc'))
 
         self.values = [
                 np.longlong(-1), np.int_(-1), np.intc(-1), np.short(-1), np.byte(-1),
@@ -359,6 +363,103 @@ class IsNaElementPerfREF(IsNaElementPerf):
     entry = staticmethod(isna_element_ref)
 
 
+#-------------------------------------------------------------------------------
+class GetNewIndexersAndScreenPerf(Perf):
+    FUNCTIONS = (
+        "ordered",
+        "unordered",
+        "tiled",
+        "repeat",
+        "quick_exit",
+        "late_exit",
+        "small",
+        "large",
+    )
+    NUMBER = 5
+
+    TILED = "tiled"
+    REPEATED = "repeated"
+    ORDERED = "ordered"
+    UNORDERED = "unordered"
+
+    class Key(tp.NamedTuple):
+        type1: str
+        type2: str
+        increment: int
+        scale: int
+
+    def __init__(self):
+        NUMBERS = np.arange(500_000, dtype=np.int64)
+        POSITIONS = np.arange(500_000, dtype=np.int64)
+
+        np.random.seed(0)
+
+        self.cases: tp.Dict[self.Key, tp.Tuple[np.ndarray, np.ndarray]] = {}
+
+        for scale in (5, 50, 500, 5_000, 50_000):
+            tiled_ordered = np.tile(NUMBERS[:scale], len(NUMBERS) // scale)
+            repeated_ordered = np.repeat(NUMBERS[:scale], len(NUMBERS) // scale)
+            tiled_unordered = tiled_ordered.copy()
+            repeated_unordered = repeated_ordered.copy()
+            np.random.shuffle(tiled_unordered)
+            np.random.shuffle(repeated_unordered)
+
+            increment = scale
+            while increment <= len(NUMBERS):
+                positions = POSITIONS[:increment]
+                key_kwargs = dict(increment=increment, scale=scale)
+                self.cases[
+                    self.Key(type1=self.TILED, type2=self.ORDERED, **key_kwargs)
+                ] = (tiled_ordered, positions)
+                self.cases[
+                    self.Key(type1=self.REPEATED, type2=self.ORDERED, **key_kwargs)
+                ] = (repeated_ordered, positions)
+                self.cases[
+                    self.Key(type1=self.TILED, type2=self.UNORDERED, **key_kwargs)
+                ] = (tiled_unordered, positions)
+                self.cases[
+                    self.Key(type1=self.REPEATED, type2=self.UNORDERED, **key_kwargs)
+                ] = (repeated_unordered, positions)
+                increment *= 10
+
+    def evaluate_cases_by_condition(self, condition):
+        for key, (indexers, positions) in self.cases.items():
+            if condition(key):
+                self.entry(indexers=indexers, positions=positions)
+
+    def ordered(self):
+        self.evaluate_cases_by_condition(lambda key: key.type2 == self.ORDERED)
+
+    def unordered(self):
+        self.evaluate_cases_by_condition(lambda key: key.type2 == self.UNORDERED)
+
+    def tiled(self):
+        self.evaluate_cases_by_condition(lambda key: key.type1 == self.TILED)
+
+    def repeat(self):
+        self.evaluate_cases_by_condition(lambda key: key.type1 == self.REPEATED)
+
+    def quick_exit(self):
+        self.evaluate_cases_by_condition(lambda key: key.increment == key.scale)
+
+    def late_exit(self):
+        self.evaluate_cases_by_condition(lambda key: key.increment > key.scale)
+
+    def small(self):
+        self.evaluate_cases_by_condition(lambda key: key.scale <= 500)
+
+    def large(self):
+        self.evaluate_cases_by_condition(lambda key: key.scale > 500)
+
+
+class GetNewIndexersAndScreenPerfAK(GetNewIndexersAndScreenPerf):
+    entry = staticmethod(get_new_indexers_and_screen_ak)
+
+
+class GetNewIndexersAndScreenPerfREF(GetNewIndexersAndScreenPerf):
+    entry = staticmethod(get_new_indexers_and_screen_ref)
+
+
 #-------------------------------------------------------------------------------
 
 def get_arg_parser():
@@ -399,7 +500,7 @@ def main():
                         number=cls_runner.NUMBER)
             records.append((cls_perf.__name__, func_attr, results['ak'], results['ref'], results['ref'] / results['ak']))
 
-    width = 24
+    width = 32
     for record in records:
         print(''.join(
             (r.ljust(width) if isinstance(r, str) else str(round(r, 8)).ljust(width)) for r in record

performance/reference/util.py

@@ -216,3 +216,35 @@ def dtype_from_element(value: tp.Optional[tp.Hashable]) -> np.dtype:
     # NOTE: calling array and getting dtype on np.nan is faster than combining isinstance, isnan calls
     return np.array(value).dtype
 
+
+def get_new_indexers_and_screen_ref(
+        indexers: np.ndarray,
+        positions: np.ndarray,
+    ) -> tp.Tuple[np.ndarray, np.ndarray]:
+
+    positions = indexers.argsort()
+
+    # get the sorted indexers
+    indexers = indexers[positions]
+
+    mask = np.empty(indexers.shape, dtype=DTYPE_BOOL)
+    mask[0] = True
+    mask[1:] = indexers[1:] != indexers[:-1]
+
+    new_indexers = np.empty(mask.shape, dtype=DTYPE_INT_DEFAULT)
+    new_indexers[positions] = np.cumsum(mask) - 1
+    new_indexers.flags.writeable = False
+
+    return new_indexers, indexers[mask]
+
+
+def get_new_indexers_and_screen_ak(
+        indexers: np.ndarray,
+        positions: np.ndarray,
+    ) -> tp.Tuple[np.ndarray, np.ndarray]:
+    from arraykit import get_new_indexers_and_screen as ak_routine
+
+    if len(positions) > len(indexers):
+        return np.unique(indexers, return_inverse=True)
+
+    return ak_routine(indexers, positions)

src/__init__.py

@@ -16,3 +16,4 @@
 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
+from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen

src/__init__.pyi

@@ -8,7 +8,7 @@ __version__: str
 
 class ArrayGO:
 
-    values: np.array
+    values: np.ndarray
     def __init__(
         self, iterable: tp.Iterable[object], *, own_iterable: bool = ...
     ) -> None: ...
@@ -20,16 +20,16 @@ class ArrayGO:
     def copy(self: _T) -> _T: ...
     def extend(self, __values: tp.Iterable[object]) -> None: ...
 
-def immutable_filter(__array: np.array) -> np.array: ...
-def mloc(__array: np.array) -> int: ...
+def immutable_filter(__array: np.ndarray) -> np.ndarray: ...
+def mloc(__array: np.ndarray) -> int: ...
 def name_filter(__name: tp.Hashable) -> tp.Hashable: ...
-def shape_filter(__array: np.array) -> np.ndarray: ...
-def column_2d_filter(__array: np.array) -> np.ndarray: ...
-def column_1d_filter(__array: np.array) -> np.ndarray: ...
-def row_1d_filter(__array: np.array) -> np.ndarray: ...
-def array_deepcopy(__array: np.array, memo: tp.Dict[int, tp.Any]) -> np.ndarray: ...
+def shape_filter(__array: np.ndarray) -> np.ndarray: ...
+def column_2d_filter(__array: np.ndarray) -> np.ndarray: ...
+def column_1d_filter(__array: np.ndarray) -> np.ndarray: ...
+def row_1d_filter(__array: np.ndarray) -> np.ndarray: ...
+def array_deepcopy(__array: np.ndarray, 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: ...
-
+def get_new_indexers_and_screen(__indexers: np.ndarray, __positions: np.ndarray) -> tp.Tuple[np.ndarray, np.ndarray]: ...