From db9353914f615a141bbad4303f0168ae0f35f58f Mon Sep 17 00:00:00 2001 From: samukweku Date: Sat, 22 Aug 2026 08:23:42 +1000 Subject: [PATCH 1/3] [PERF] Simplify the conditional_join Python/Rust boundary Remove a duplicate _sum_starts_ends definition, replace ~87 repeated per-call dtype-dispatch dicts in _agg_functions.py/_compare.py/ _binary_search.py with one cached dispatcher (_rs_func), and swap repeat_index calls for np.repeat, which benchmarks faster at every scale tested and is more immediately readable. Issue #1649 @samukweku Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01V3KAbkp6JV96KYXNc6EJmN --- CHANGELOG.md | 1 + .../_conditional_join/_agg_functions.py | 1111 +---------------- .../_conditional_join/_binary_search.py | 183 +-- .../functions/_conditional_join/_compare.py | 291 +---- .../_conditional_join/_dtype_dispatch.py | 36 + .../_greater_than_indices.py | 7 +- .../functions/_conditional_join/_helpers.py | 22 +- .../_conditional_join/_less_than_indices.py | 7 +- .../_conditional_join/_range_indices.py | 7 +- pixi.lock | 4 +- tests/functions/test_conditional_join.py | 37 + 11 files changed, 175 insertions(+), 1531 deletions(-) create mode 100644 janitor/functions/_conditional_join/_dtype_dispatch.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..461e72d17 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Changelog ## [Unreleased] +- [PERF] Simplify the `conditional_join` Python/Rust boundary: removed a duplicate `_sum_starts_ends` definition, replaced ~87 repeated per-call dtype-dispatch dicts with one cached dispatcher, and switched `repeat_index` calls to `np.repeat`. - Issue #1649 @samukweku - [ENH] Avoid copying column data during `conditional_join` input validation. - Issue #1645, PR #1642 @samukweku - [ENH] Speed up `conditional_join` with an unsorted right join key and diff --git a/janitor/functions/_conditional_join/_agg_functions.py b/janitor/functions/_conditional_join/_agg_functions.py index 96fa0b7a2..7282f6fba 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,6 +1,8 @@ import janitor_rs import numpy as np +from ._dtype_dispatch import _rs_func + def _sum_starts( arr: np.ndarray, @@ -10,23 +12,7 @@ def _sum_starts( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_start_int64, - "int32": janitor_rs.compute_sum_start_int32, - "int16": janitor_rs.compute_sum_start_int16, - "int8": janitor_rs.compute_sum_start_int8, - "uint64": janitor_rs.compute_sum_start_uint64, - "uint32": janitor_rs.compute_sum_start_uint32, - "uint16": janitor_rs.compute_sum_start_uint16, - "uint8": janitor_rs.compute_sum_start_uint8, - "float64": janitor_rs.compute_sum_start_f64, - "float32": janitor_rs.compute_sum_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_start", arr.dtype.name) return func(arr=arr, starts=starts, booleans=booleans) @@ -38,23 +24,7 @@ def _sum_ends( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_end_int64, - "int32": janitor_rs.compute_sum_end_int32, - "int16": janitor_rs.compute_sum_end_int16, - "int8": janitor_rs.compute_sum_end_int8, - "uint64": janitor_rs.compute_sum_end_uint64, - "uint32": janitor_rs.compute_sum_end_uint32, - "uint16": janitor_rs.compute_sum_end_uint16, - "uint8": janitor_rs.compute_sum_end_uint8, - "float64": janitor_rs.compute_sum_end_f64, - "float32": janitor_rs.compute_sum_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_end", arr.dtype.name) return func(arr=arr, ends=ends, booleans=booleans) @@ -164,23 +134,7 @@ def _min_starts( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_start_int64, - "int32": janitor_rs.compute_min_start_int32, - "int16": janitor_rs.compute_min_start_int16, - "int8": janitor_rs.compute_min_start_int8, - "uint64": janitor_rs.compute_min_start_uint64, - "uint32": janitor_rs.compute_min_start_uint32, - "uint16": janitor_rs.compute_min_start_uint16, - "uint8": janitor_rs.compute_min_start_uint8, - "float64": janitor_rs.compute_min_start_f64, - "float32": janitor_rs.compute_min_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_start", arr.dtype.name) return func(arr=arr, starts=starts, booleans=booleans) @@ -192,23 +146,7 @@ def _min_ends( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_end_int64, - "int32": janitor_rs.compute_min_end_int32, - "int16": janitor_rs.compute_min_end_int16, - "int8": janitor_rs.compute_min_end_int8, - "uint64": janitor_rs.compute_min_end_uint64, - "uint32": janitor_rs.compute_min_end_uint32, - "uint16": janitor_rs.compute_min_end_uint16, - "uint8": janitor_rs.compute_min_end_uint8, - "float64": janitor_rs.compute_min_end_f64, - "float32": janitor_rs.compute_min_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_end", arr.dtype.name) return func(arr=arr, ends=ends, booleans=booleans) @@ -220,23 +158,7 @@ def _max_starts( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_start_int64, - "int32": janitor_rs.compute_max_start_int32, - "int16": janitor_rs.compute_max_start_int16, - "int8": janitor_rs.compute_max_start_int8, - "uint64": janitor_rs.compute_max_start_uint64, - "uint32": janitor_rs.compute_max_start_uint32, - "uint16": janitor_rs.compute_max_start_uint16, - "uint8": janitor_rs.compute_max_start_uint8, - "float64": janitor_rs.compute_max_start_f64, - "float32": janitor_rs.compute_max_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_start", arr.dtype.name) return func(arr=arr, starts=starts, booleans=booleans) @@ -248,23 +170,7 @@ def _max_ends( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_end_int64, - "int32": janitor_rs.compute_max_end_int32, - "int16": janitor_rs.compute_max_end_int16, - "int8": janitor_rs.compute_max_end_int8, - "uint64": janitor_rs.compute_max_end_uint64, - "uint32": janitor_rs.compute_max_end_uint32, - "uint16": janitor_rs.compute_max_end_uint16, - "uint8": janitor_rs.compute_max_end_uint8, - "float64": janitor_rs.compute_max_end_f64, - "float32": janitor_rs.compute_max_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_end", arr.dtype.name) return func(arr=arr, ends=ends, booleans=booleans) @@ -276,23 +182,7 @@ def _prod_starts( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_start_int64, - "int32": janitor_rs.compute_prod_start_int32, - "int16": janitor_rs.compute_prod_start_int16, - "int8": janitor_rs.compute_prod_start_int8, - "uint64": janitor_rs.compute_prod_start_uint64, - "uint32": janitor_rs.compute_prod_start_uint32, - "uint16": janitor_rs.compute_prod_start_uint16, - "uint8": janitor_rs.compute_prod_start_uint8, - "float64": janitor_rs.compute_prod_start_f64, - "float32": janitor_rs.compute_prod_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_start", arr.dtype.name) return func(arr=arr, starts=starts, booleans=booleans) @@ -304,23 +194,7 @@ def _prod_ends( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_end_int64, - "int32": janitor_rs.compute_prod_end_int32, - "int16": janitor_rs.compute_prod_end_int16, - "int8": janitor_rs.compute_prod_end_int8, - "uint64": janitor_rs.compute_prod_end_uint64, - "uint32": janitor_rs.compute_prod_end_uint32, - "uint16": janitor_rs.compute_prod_end_uint16, - "uint8": janitor_rs.compute_prod_end_uint8, - "float64": janitor_rs.compute_prod_end_f64, - "float32": janitor_rs.compute_prod_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_end", arr.dtype.name) return func(arr=arr, ends=ends, booleans=booleans) @@ -334,23 +208,7 @@ def _sum_starts_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_start_match_int64, - "int32": janitor_rs.compute_sum_start_match_int32, - "int16": janitor_rs.compute_sum_start_match_int16, - "int8": janitor_rs.compute_sum_start_match_int8, - "uint64": janitor_rs.compute_sum_start_match_uint64, - "uint32": janitor_rs.compute_sum_start_match_uint32, - "uint16": janitor_rs.compute_sum_start_match_uint16, - "uint8": janitor_rs.compute_sum_start_match_uint8, - "float64": janitor_rs.compute_sum_start_match_f64, - "float32": janitor_rs.compute_sum_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -370,23 +228,7 @@ def _sum_ends_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_end_match_int64, - "int32": janitor_rs.compute_sum_end_match_int32, - "int16": janitor_rs.compute_sum_end_match_int16, - "int8": janitor_rs.compute_sum_end_match_int8, - "uint64": janitor_rs.compute_sum_end_match_uint64, - "uint32": janitor_rs.compute_sum_end_match_uint32, - "uint16": janitor_rs.compute_sum_end_match_uint16, - "uint8": janitor_rs.compute_sum_end_match_uint8, - "float64": janitor_rs.compute_sum_end_match_f64, - "float32": janitor_rs.compute_sum_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_end_match", arr.dtype.name) return func(arr=arr, ends=ends, counts=counts, matches=matches, booleans=booleans) @@ -400,23 +242,7 @@ def _max_starts_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_start_match_int64, - "int32": janitor_rs.compute_max_start_match_int32, - "int16": janitor_rs.compute_max_start_match_int16, - "int8": janitor_rs.compute_max_start_match_int8, - "uint64": janitor_rs.compute_max_start_match_uint64, - "uint32": janitor_rs.compute_max_start_match_uint32, - "uint16": janitor_rs.compute_max_start_match_uint16, - "uint8": janitor_rs.compute_max_start_match_uint8, - "float64": janitor_rs.compute_max_start_match_f64, - "float32": janitor_rs.compute_max_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -436,23 +262,7 @@ def _max_ends_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_end_match_int64, - "int32": janitor_rs.compute_max_end_match_int32, - "int16": janitor_rs.compute_max_end_match_int16, - "int8": janitor_rs.compute_max_end_match_int8, - "uint64": janitor_rs.compute_max_end_match_uint64, - "uint32": janitor_rs.compute_max_end_match_uint32, - "uint16": janitor_rs.compute_max_end_match_uint16, - "uint8": janitor_rs.compute_max_end_match_uint8, - "float64": janitor_rs.compute_max_end_match_f64, - "float32": janitor_rs.compute_max_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_end_match", arr.dtype.name) return func(arr=arr, ends=ends, counts=counts, matches=matches, booleans=booleans) @@ -466,23 +276,7 @@ def _min_starts_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_start_match_int64, - "int32": janitor_rs.compute_min_start_match_int32, - "int16": janitor_rs.compute_min_start_match_int16, - "int8": janitor_rs.compute_min_start_match_int8, - "uint64": janitor_rs.compute_min_start_match_uint64, - "uint32": janitor_rs.compute_min_start_match_uint32, - "uint16": janitor_rs.compute_min_start_match_uint16, - "uint8": janitor_rs.compute_min_start_match_uint8, - "float64": janitor_rs.compute_min_start_match_f64, - "float32": janitor_rs.compute_min_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -502,23 +296,7 @@ def _min_ends_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_end_match_int64, - "int32": janitor_rs.compute_min_end_match_int32, - "int16": janitor_rs.compute_min_end_match_int16, - "int8": janitor_rs.compute_min_end_match_int8, - "uint64": janitor_rs.compute_min_end_match_uint64, - "uint32": janitor_rs.compute_min_end_match_uint32, - "uint16": janitor_rs.compute_min_end_match_uint16, - "uint8": janitor_rs.compute_min_end_match_uint8, - "float64": janitor_rs.compute_min_end_match_f64, - "float32": janitor_rs.compute_min_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_end_match", arr.dtype.name) return func(arr=arr, ends=ends, counts=counts, matches=matches, booleans=booleans) @@ -532,23 +310,7 @@ def _sum_positions( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_positions_int64, - "int32": janitor_rs.compute_sum_positions_int32, - "int16": janitor_rs.compute_sum_positions_int16, - "int8": janitor_rs.compute_sum_positions_int8, - "uint64": janitor_rs.compute_sum_positions_uint64, - "uint32": janitor_rs.compute_sum_positions_uint32, - "uint16": janitor_rs.compute_sum_positions_uint16, - "uint8": janitor_rs.compute_sum_positions_uint8, - "float64": janitor_rs.compute_sum_positions_f64, - "float32": janitor_rs.compute_sum_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -568,23 +330,7 @@ def _prod_starts_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_start_match_int64, - "int32": janitor_rs.compute_prod_start_match_int32, - "int16": janitor_rs.compute_prod_start_match_int16, - "int8": janitor_rs.compute_prod_start_match_int8, - "uint64": janitor_rs.compute_prod_start_match_uint64, - "uint32": janitor_rs.compute_prod_start_match_uint32, - "uint16": janitor_rs.compute_prod_start_match_uint16, - "uint8": janitor_rs.compute_prod_start_match_uint8, - "float64": janitor_rs.compute_prod_start_match_f64, - "float32": janitor_rs.compute_prod_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -604,23 +350,7 @@ def _prod_ends_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_end_match_int64, - "int32": janitor_rs.compute_prod_end_match_int32, - "int16": janitor_rs.compute_prod_end_match_int16, - "int8": janitor_rs.compute_prod_end_match_int8, - "uint64": janitor_rs.compute_prod_end_match_uint64, - "uint32": janitor_rs.compute_prod_end_match_uint32, - "uint16": janitor_rs.compute_prod_end_match_uint16, - "uint8": janitor_rs.compute_prod_end_match_uint8, - "float64": janitor_rs.compute_prod_end_match_f64, - "float32": janitor_rs.compute_prod_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_end_match", arr.dtype.name) return func(arr=arr, ends=ends, counts=counts, matches=matches, booleans=booleans) @@ -634,23 +364,7 @@ def _prod_positions( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_positions_int64, - "int32": janitor_rs.compute_prod_positions_int32, - "int16": janitor_rs.compute_prod_positions_int16, - "int8": janitor_rs.compute_prod_positions_int8, - "uint64": janitor_rs.compute_prod_positions_uint64, - "uint32": janitor_rs.compute_prod_positions_uint32, - "uint16": janitor_rs.compute_prod_positions_uint16, - "uint8": janitor_rs.compute_prod_positions_uint8, - "float64": janitor_rs.compute_prod_positions_f64, - "float32": janitor_rs.compute_prod_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -670,23 +384,7 @@ def _min_positions( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_positions_int64, - "int32": janitor_rs.compute_min_positions_int32, - "int16": janitor_rs.compute_min_positions_int16, - "int8": janitor_rs.compute_min_positions_int8, - "uint64": janitor_rs.compute_min_positions_uint64, - "uint32": janitor_rs.compute_min_positions_uint32, - "uint16": janitor_rs.compute_min_positions_uint16, - "uint8": janitor_rs.compute_min_positions_uint8, - "float64": janitor_rs.compute_min_positions_f64, - "float32": janitor_rs.compute_min_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -706,23 +404,7 @@ def _max_positions( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_positions_int64, - "int32": janitor_rs.compute_max_positions_int32, - "int16": janitor_rs.compute_max_positions_int16, - "int8": janitor_rs.compute_max_positions_int8, - "uint64": janitor_rs.compute_max_positions_uint64, - "uint32": janitor_rs.compute_max_positions_uint32, - "uint16": janitor_rs.compute_max_positions_uint16, - "uint8": janitor_rs.compute_max_positions_uint8, - "float64": janitor_rs.compute_max_positions_f64, - "float32": janitor_rs.compute_max_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -741,23 +423,7 @@ def _max_starts_ends( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_start_end_int64, - "int32": janitor_rs.compute_max_start_end_int32, - "int16": janitor_rs.compute_max_start_end_int16, - "int8": janitor_rs.compute_max_start_end_int8, - "uint64": janitor_rs.compute_max_start_end_uint64, - "uint32": janitor_rs.compute_max_start_end_uint32, - "uint16": janitor_rs.compute_max_start_end_uint16, - "uint8": janitor_rs.compute_max_start_end_uint8, - "float64": janitor_rs.compute_max_start_end_f64, - "float32": janitor_rs.compute_max_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_start_end", arr.dtype.name) return func(arr=arr, starts=starts, ends=ends, booleans=booleans) @@ -770,23 +436,7 @@ def _min_starts_ends( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_start_end_int64, - "int32": janitor_rs.compute_min_start_end_int32, - "int16": janitor_rs.compute_min_start_end_int16, - "int8": janitor_rs.compute_min_start_end_int8, - "uint64": janitor_rs.compute_min_start_end_uint64, - "uint32": janitor_rs.compute_min_start_end_uint32, - "uint16": janitor_rs.compute_min_start_end_uint16, - "uint8": janitor_rs.compute_min_start_end_uint8, - "float64": janitor_rs.compute_min_start_end_f64, - "float32": janitor_rs.compute_min_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_start_end", arr.dtype.name) return func(arr=arr, starts=starts, ends=ends, booleans=booleans) @@ -799,23 +449,7 @@ def _sum_starts_ends( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_start_end_int64, - "int32": janitor_rs.compute_sum_start_end_int32, - "int16": janitor_rs.compute_sum_start_end_int16, - "int8": janitor_rs.compute_sum_start_end_int8, - "uint64": janitor_rs.compute_sum_start_end_uint64, - "uint32": janitor_rs.compute_sum_start_end_uint32, - "uint16": janitor_rs.compute_sum_start_end_uint16, - "uint8": janitor_rs.compute_sum_start_end_uint8, - "float64": janitor_rs.compute_sum_start_end_f64, - "float32": janitor_rs.compute_sum_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_start_end", arr.dtype.name) return func(arr=arr, starts=starts, ends=ends, booleans=booleans) @@ -828,23 +462,7 @@ def _prod_starts_ends( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_start_end_int64, - "int32": janitor_rs.compute_prod_start_end_int32, - "int16": janitor_rs.compute_prod_start_end_int16, - "int8": janitor_rs.compute_prod_start_end_int8, - "uint64": janitor_rs.compute_prod_start_end_uint64, - "uint32": janitor_rs.compute_prod_start_end_uint32, - "uint16": janitor_rs.compute_prod_start_end_uint16, - "uint8": janitor_rs.compute_prod_start_end_uint8, - "float64": janitor_rs.compute_prod_start_end_f64, - "float32": janitor_rs.compute_prod_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_start_end", arr.dtype.name) return func(arr=arr, starts=starts, ends=ends, booleans=booleans) @@ -859,23 +477,7 @@ def _prod_starts_ends_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_start_end_match_int64, - "int32": janitor_rs.compute_prod_start_end_match_int32, - "int16": janitor_rs.compute_prod_start_end_match_int16, - "int8": janitor_rs.compute_prod_start_end_match_int8, - "uint64": janitor_rs.compute_prod_start_end_match_uint64, - "uint32": janitor_rs.compute_prod_start_end_match_uint32, - "uint16": janitor_rs.compute_prod_start_end_match_uint16, - "uint8": janitor_rs.compute_prod_start_end_match_uint8, - "float64": janitor_rs.compute_prod_start_end_match_f64, - "float32": janitor_rs.compute_prod_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -886,35 +488,6 @@ def _prod_starts_ends_matches( ) -def _sum_starts_ends( - arr: np.ndarray, - starts: np.ndarray, - ends: np.ndarray, - booleans: np.ndarray, -) -> tuple: - """ - Compute sum - """ - mapping = { - "int64": janitor_rs.compute_sum_start_end_int64, - "int32": janitor_rs.compute_sum_start_end_int32, - "int16": janitor_rs.compute_sum_start_end_int16, - "int8": janitor_rs.compute_sum_start_end_int8, - "uint64": janitor_rs.compute_sum_start_end_uint64, - "uint32": janitor_rs.compute_sum_start_end_uint32, - "uint16": janitor_rs.compute_sum_start_end_uint16, - "uint8": janitor_rs.compute_sum_start_end_uint8, - "float64": janitor_rs.compute_sum_start_end_f64, - "float32": janitor_rs.compute_sum_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") - return func(arr=arr, starts=starts, ends=ends, booleans=booleans) - - def _sum_starts_ends_matches( arr: np.ndarray, starts: np.ndarray, @@ -926,23 +499,7 @@ def _sum_starts_ends_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_start_end_match_int64, - "int32": janitor_rs.compute_sum_start_end_match_int32, - "int16": janitor_rs.compute_sum_start_end_match_int16, - "int8": janitor_rs.compute_sum_start_end_match_int8, - "uint64": janitor_rs.compute_sum_start_end_match_uint64, - "uint32": janitor_rs.compute_sum_start_end_match_uint32, - "uint16": janitor_rs.compute_sum_start_end_match_uint16, - "uint8": janitor_rs.compute_sum_start_end_match_uint8, - "float64": janitor_rs.compute_sum_start_end_match_f64, - "float32": janitor_rs.compute_sum_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -964,23 +521,7 @@ def _min_starts_ends_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_start_end_match_int64, - "int32": janitor_rs.compute_min_start_end_match_int32, - "int16": janitor_rs.compute_min_start_end_match_int16, - "int8": janitor_rs.compute_min_start_end_match_int8, - "uint64": janitor_rs.compute_min_start_end_match_uint64, - "uint32": janitor_rs.compute_min_start_end_match_uint32, - "uint16": janitor_rs.compute_min_start_end_match_uint16, - "uint8": janitor_rs.compute_min_start_end_match_uint8, - "float64": janitor_rs.compute_min_start_end_match_f64, - "float32": janitor_rs.compute_min_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1002,23 +543,7 @@ def _max_starts_ends_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_start_end_match_int64, - "int32": janitor_rs.compute_max_start_end_match_int32, - "int16": janitor_rs.compute_max_start_end_match_int16, - "int8": janitor_rs.compute_max_start_end_match_int8, - "uint64": janitor_rs.compute_max_start_end_match_uint64, - "uint32": janitor_rs.compute_max_start_end_match_uint32, - "uint16": janitor_rs.compute_max_start_end_match_uint16, - "uint8": janitor_rs.compute_max_start_end_match_uint8, - "float64": janitor_rs.compute_max_start_end_match_f64, - "float32": janitor_rs.compute_max_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1039,23 +564,7 @@ def _prod_rev_starts( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_start_int64, - "int32": janitor_rs.compute_prod_rev_start_int32, - "int16": janitor_rs.compute_prod_rev_start_int16, - "int8": janitor_rs.compute_prod_rev_start_int8, - "uint64": janitor_rs.compute_prod_rev_start_uint64, - "uint32": janitor_rs.compute_prod_rev_start_uint32, - "uint16": janitor_rs.compute_prod_rev_start_uint16, - "uint8": janitor_rs.compute_prod_rev_start_uint8, - "float64": janitor_rs.compute_prod_rev_start_f64, - "float32": janitor_rs.compute_prod_rev_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_start", arr.dtype.name) return func(arr=arr, starts=starts, index=index, booleans=booleans, length=length) @@ -1069,23 +578,7 @@ def _prod_rev_ends( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_end_int64, - "int32": janitor_rs.compute_prod_rev_end_int32, - "int16": janitor_rs.compute_prod_rev_end_int16, - "int8": janitor_rs.compute_prod_rev_end_int8, - "uint64": janitor_rs.compute_prod_rev_end_uint64, - "uint32": janitor_rs.compute_prod_rev_end_uint32, - "uint16": janitor_rs.compute_prod_rev_end_uint16, - "uint8": janitor_rs.compute_prod_rev_end_uint8, - "float64": janitor_rs.compute_prod_rev_end_f64, - "float32": janitor_rs.compute_prod_rev_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_end", arr.dtype.name) return func(arr=arr, ends=ends, index=index, booleans=booleans, length=length) @@ -1101,23 +594,7 @@ def _prod_rev_starts_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_start_match_int64, - "int32": janitor_rs.compute_prod_rev_start_match_int32, - "int16": janitor_rs.compute_prod_rev_start_match_int16, - "int8": janitor_rs.compute_prod_rev_start_match_int8, - "uint64": janitor_rs.compute_prod_rev_start_match_uint64, - "uint32": janitor_rs.compute_prod_rev_start_match_uint32, - "uint16": janitor_rs.compute_prod_rev_start_match_uint16, - "uint8": janitor_rs.compute_prod_rev_start_match_uint8, - "float64": janitor_rs.compute_prod_rev_start_match_f64, - "float32": janitor_rs.compute_prod_rev_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1141,23 +618,7 @@ def _prod_rev_ends_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_end_match_int64, - "int32": janitor_rs.compute_prod_rev_end_match_int32, - "int16": janitor_rs.compute_prod_rev_end_match_int16, - "int8": janitor_rs.compute_prod_rev_end_match_int8, - "uint64": janitor_rs.compute_prod_rev_end_match_uint64, - "uint32": janitor_rs.compute_prod_rev_end_match_uint32, - "uint16": janitor_rs.compute_prod_rev_end_match_uint16, - "uint8": janitor_rs.compute_prod_rev_end_match_uint8, - "float64": janitor_rs.compute_prod_rev_end_match_f64, - "float32": janitor_rs.compute_prod_rev_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_end_match", arr.dtype.name) return func( arr=arr, index=index, @@ -1181,23 +642,7 @@ def _prod_rev_positions( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_positions_int64, - "int32": janitor_rs.compute_prod_rev_positions_int32, - "int16": janitor_rs.compute_prod_rev_positions_int16, - "int8": janitor_rs.compute_prod_rev_positions_int8, - "uint64": janitor_rs.compute_prod_rev_positions_uint64, - "uint32": janitor_rs.compute_prod_rev_positions_uint32, - "uint16": janitor_rs.compute_prod_rev_positions_uint16, - "uint8": janitor_rs.compute_prod_rev_positions_uint8, - "float64": janitor_rs.compute_prod_rev_positions_f64, - "float32": janitor_rs.compute_prod_rev_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1220,23 +665,7 @@ def _prod_rev_starts_ends( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_start_end_int64, - "int32": janitor_rs.compute_prod_rev_start_end_int32, - "int16": janitor_rs.compute_prod_rev_start_end_int16, - "int8": janitor_rs.compute_prod_rev_start_end_int8, - "uint64": janitor_rs.compute_prod_rev_start_end_uint64, - "uint32": janitor_rs.compute_prod_rev_start_end_uint32, - "uint16": janitor_rs.compute_prod_rev_start_end_uint16, - "uint8": janitor_rs.compute_prod_rev_start_end_uint8, - "float64": janitor_rs.compute_prod_rev_start_end_f64, - "float32": janitor_rs.compute_prod_rev_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_start_end", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1260,23 +689,7 @@ def _prod_rev_starts_ends_matches( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_start_end_match_int64, - "int32": janitor_rs.compute_prod_rev_start_end_match_int32, - "int16": janitor_rs.compute_prod_rev_start_end_match_int16, - "int8": janitor_rs.compute_prod_rev_start_end_match_int8, - "uint64": janitor_rs.compute_prod_rev_start_end_match_uint64, - "uint32": janitor_rs.compute_prod_rev_start_end_match_uint32, - "uint16": janitor_rs.compute_prod_rev_start_end_match_uint16, - "uint8": janitor_rs.compute_prod_rev_start_end_match_uint8, - "float64": janitor_rs.compute_prod_rev_start_end_match_f64, - "float32": janitor_rs.compute_prod_rev_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1299,23 +712,7 @@ def _min_rev_starts( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_start_int64, - "int32": janitor_rs.compute_min_rev_start_int32, - "int16": janitor_rs.compute_min_rev_start_int16, - "int8": janitor_rs.compute_min_rev_start_int8, - "uint64": janitor_rs.compute_min_rev_start_uint64, - "uint32": janitor_rs.compute_min_rev_start_uint32, - "uint16": janitor_rs.compute_min_rev_start_uint16, - "uint8": janitor_rs.compute_min_rev_start_uint8, - "float64": janitor_rs.compute_min_rev_start_f64, - "float32": janitor_rs.compute_min_rev_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_start", arr.dtype.name) return func(arr=arr, starts=starts, index=index, booleans=booleans, length=length) @@ -1329,23 +726,7 @@ def _min_rev_ends( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_end_int64, - "int32": janitor_rs.compute_min_rev_end_int32, - "int16": janitor_rs.compute_min_rev_end_int16, - "int8": janitor_rs.compute_min_rev_end_int8, - "uint64": janitor_rs.compute_min_rev_end_uint64, - "uint32": janitor_rs.compute_min_rev_end_uint32, - "uint16": janitor_rs.compute_min_rev_end_uint16, - "uint8": janitor_rs.compute_min_rev_end_uint8, - "float64": janitor_rs.compute_min_rev_end_f64, - "float32": janitor_rs.compute_min_rev_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_end", arr.dtype.name) return func(arr=arr, ends=ends, index=index, booleans=booleans, length=length) @@ -1361,23 +742,7 @@ def _min_rev_starts_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_start_match_int64, - "int32": janitor_rs.compute_min_rev_start_match_int32, - "int16": janitor_rs.compute_min_rev_start_match_int16, - "int8": janitor_rs.compute_min_rev_start_match_int8, - "uint64": janitor_rs.compute_min_rev_start_match_uint64, - "uint32": janitor_rs.compute_min_rev_start_match_uint32, - "uint16": janitor_rs.compute_min_rev_start_match_uint16, - "uint8": janitor_rs.compute_min_rev_start_match_uint8, - "float64": janitor_rs.compute_min_rev_start_match_f64, - "float32": janitor_rs.compute_min_rev_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1401,23 +766,7 @@ def _min_rev_ends_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_end_match_int64, - "int32": janitor_rs.compute_min_rev_end_match_int32, - "int16": janitor_rs.compute_min_rev_end_match_int16, - "int8": janitor_rs.compute_min_rev_end_match_int8, - "uint64": janitor_rs.compute_min_rev_end_match_uint64, - "uint32": janitor_rs.compute_min_rev_end_match_uint32, - "uint16": janitor_rs.compute_min_rev_end_match_uint16, - "uint8": janitor_rs.compute_min_rev_end_match_uint8, - "float64": janitor_rs.compute_min_rev_end_match_f64, - "float32": janitor_rs.compute_min_rev_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_end_match", arr.dtype.name) return func( arr=arr, index=index, @@ -1441,23 +790,7 @@ def _min_rev_positions( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_positions_int64, - "int32": janitor_rs.compute_min_rev_positions_int32, - "int16": janitor_rs.compute_min_rev_positions_int16, - "int8": janitor_rs.compute_min_rev_positions_int8, - "uint64": janitor_rs.compute_min_rev_positions_uint64, - "uint32": janitor_rs.compute_min_rev_positions_uint32, - "uint16": janitor_rs.compute_min_rev_positions_uint16, - "uint8": janitor_rs.compute_min_rev_positions_uint8, - "float64": janitor_rs.compute_min_rev_positions_f64, - "float32": janitor_rs.compute_min_rev_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1480,23 +813,7 @@ def _min_rev_starts_ends( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_start_end_int64, - "int32": janitor_rs.compute_min_rev_start_end_int32, - "int16": janitor_rs.compute_min_rev_start_end_int16, - "int8": janitor_rs.compute_min_rev_start_end_int8, - "uint64": janitor_rs.compute_min_rev_start_end_uint64, - "uint32": janitor_rs.compute_min_rev_start_end_uint32, - "uint16": janitor_rs.compute_min_rev_start_end_uint16, - "uint8": janitor_rs.compute_min_rev_start_end_uint8, - "float64": janitor_rs.compute_min_rev_start_end_f64, - "float32": janitor_rs.compute_min_rev_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_start_end", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1520,23 +837,7 @@ def _min_rev_starts_ends_matches( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_start_end_match_int64, - "int32": janitor_rs.compute_min_rev_start_end_match_int32, - "int16": janitor_rs.compute_min_rev_start_end_match_int16, - "int8": janitor_rs.compute_min_rev_start_end_match_int8, - "uint64": janitor_rs.compute_min_rev_start_end_match_uint64, - "uint32": janitor_rs.compute_min_rev_start_end_match_uint32, - "uint16": janitor_rs.compute_min_rev_start_end_match_uint16, - "uint8": janitor_rs.compute_min_rev_start_end_match_uint8, - "float64": janitor_rs.compute_min_rev_start_end_match_f64, - "float32": janitor_rs.compute_min_rev_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1559,23 +860,7 @@ def _max_rev_starts( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_start_int64, - "int32": janitor_rs.compute_max_rev_start_int32, - "int16": janitor_rs.compute_max_rev_start_int16, - "int8": janitor_rs.compute_max_rev_start_int8, - "uint64": janitor_rs.compute_max_rev_start_uint64, - "uint32": janitor_rs.compute_max_rev_start_uint32, - "uint16": janitor_rs.compute_max_rev_start_uint16, - "uint8": janitor_rs.compute_max_rev_start_uint8, - "float64": janitor_rs.compute_max_rev_start_f64, - "float32": janitor_rs.compute_max_rev_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_start", arr.dtype.name) return func(arr=arr, starts=starts, index=index, booleans=booleans, length=length) @@ -1589,23 +874,7 @@ def _max_rev_ends( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_end_int64, - "int32": janitor_rs.compute_max_rev_end_int32, - "int16": janitor_rs.compute_max_rev_end_int16, - "int8": janitor_rs.compute_max_rev_end_int8, - "uint64": janitor_rs.compute_max_rev_end_uint64, - "uint32": janitor_rs.compute_max_rev_end_uint32, - "uint16": janitor_rs.compute_max_rev_end_uint16, - "uint8": janitor_rs.compute_max_rev_end_uint8, - "float64": janitor_rs.compute_max_rev_end_f64, - "float32": janitor_rs.compute_max_rev_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_end", arr.dtype.name) return func(arr=arr, ends=ends, index=index, booleans=booleans, length=length) @@ -1621,23 +890,7 @@ def _max_rev_starts_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_start_match_int64, - "int32": janitor_rs.compute_max_rev_start_match_int32, - "int16": janitor_rs.compute_max_rev_start_match_int16, - "int8": janitor_rs.compute_max_rev_start_match_int8, - "uint64": janitor_rs.compute_max_rev_start_match_uint64, - "uint32": janitor_rs.compute_max_rev_start_match_uint32, - "uint16": janitor_rs.compute_max_rev_start_match_uint16, - "uint8": janitor_rs.compute_max_rev_start_match_uint8, - "float64": janitor_rs.compute_max_rev_start_match_f64, - "float32": janitor_rs.compute_max_rev_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1661,23 +914,7 @@ def _max_rev_ends_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_end_match_int64, - "int32": janitor_rs.compute_max_rev_end_match_int32, - "int16": janitor_rs.compute_max_rev_end_match_int16, - "int8": janitor_rs.compute_max_rev_end_match_int8, - "uint64": janitor_rs.compute_max_rev_end_match_uint64, - "uint32": janitor_rs.compute_max_rev_end_match_uint32, - "uint16": janitor_rs.compute_max_rev_end_match_uint16, - "uint8": janitor_rs.compute_max_rev_end_match_uint8, - "float64": janitor_rs.compute_max_rev_end_match_f64, - "float32": janitor_rs.compute_max_rev_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_end_match", arr.dtype.name) return func( arr=arr, index=index, @@ -1701,23 +938,7 @@ def _max_rev_positions( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_positions_int64, - "int32": janitor_rs.compute_max_rev_positions_int32, - "int16": janitor_rs.compute_max_rev_positions_int16, - "int8": janitor_rs.compute_max_rev_positions_int8, - "uint64": janitor_rs.compute_max_rev_positions_uint64, - "uint32": janitor_rs.compute_max_rev_positions_uint32, - "uint16": janitor_rs.compute_max_rev_positions_uint16, - "uint8": janitor_rs.compute_max_rev_positions_uint8, - "float64": janitor_rs.compute_max_rev_positions_f64, - "float32": janitor_rs.compute_max_rev_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1740,23 +961,7 @@ def _max_rev_starts_ends( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_start_end_int64, - "int32": janitor_rs.compute_max_rev_start_end_int32, - "int16": janitor_rs.compute_max_rev_start_end_int16, - "int8": janitor_rs.compute_max_rev_start_end_int8, - "uint64": janitor_rs.compute_max_rev_start_end_uint64, - "uint32": janitor_rs.compute_max_rev_start_end_uint32, - "uint16": janitor_rs.compute_max_rev_start_end_uint16, - "uint8": janitor_rs.compute_max_rev_start_end_uint8, - "float64": janitor_rs.compute_max_rev_start_end_f64, - "float32": janitor_rs.compute_max_rev_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_start_end", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1780,23 +985,7 @@ def _max_rev_starts_ends_matches( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_start_end_match_int64, - "int32": janitor_rs.compute_max_rev_start_end_match_int32, - "int16": janitor_rs.compute_max_rev_start_end_match_int16, - "int8": janitor_rs.compute_max_rev_start_end_match_int8, - "uint64": janitor_rs.compute_max_rev_start_end_match_uint64, - "uint32": janitor_rs.compute_max_rev_start_end_match_uint32, - "uint16": janitor_rs.compute_max_rev_start_end_match_uint16, - "uint8": janitor_rs.compute_max_rev_start_end_match_uint8, - "float64": janitor_rs.compute_max_rev_start_end_match_f64, - "float32": janitor_rs.compute_max_rev_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -1819,23 +1008,7 @@ def _prod_rev_no_ranges( """ Compute prod """ - mapping = { - "int64": janitor_rs.compute_prod_rev_no_range_int64, - "int32": janitor_rs.compute_prod_rev_no_range_int32, - "int16": janitor_rs.compute_prod_rev_no_range_int16, - "int8": janitor_rs.compute_prod_rev_no_range_int8, - "uint64": janitor_rs.compute_prod_rev_no_range_uint64, - "uint32": janitor_rs.compute_prod_rev_no_range_uint32, - "uint16": janitor_rs.compute_prod_rev_no_range_uint16, - "uint8": janitor_rs.compute_prod_rev_no_range_uint8, - "float64": janitor_rs.compute_prod_rev_no_range_f64, - "float32": janitor_rs.compute_prod_rev_no_range_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_prod_rev_no_range", arr.dtype.name) return func( arr=arr, left_index=left_index, @@ -1855,23 +1028,7 @@ def _max_rev_no_ranges( """ Compute max """ - mapping = { - "int64": janitor_rs.compute_max_rev_no_range_int64, - "int32": janitor_rs.compute_max_rev_no_range_int32, - "int16": janitor_rs.compute_max_rev_no_range_int16, - "int8": janitor_rs.compute_max_rev_no_range_int8, - "uint64": janitor_rs.compute_max_rev_no_range_uint64, - "uint32": janitor_rs.compute_max_rev_no_range_uint32, - "uint16": janitor_rs.compute_max_rev_no_range_uint16, - "uint8": janitor_rs.compute_max_rev_no_range_uint8, - "float64": janitor_rs.compute_max_rev_no_range_f64, - "float32": janitor_rs.compute_max_rev_no_range_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_max_rev_no_range", arr.dtype.name) return func( arr=arr, left_index=left_index, @@ -1891,23 +1048,7 @@ def _min_rev_no_ranges( """ Compute min """ - mapping = { - "int64": janitor_rs.compute_min_rev_no_range_int64, - "int32": janitor_rs.compute_min_rev_no_range_int32, - "int16": janitor_rs.compute_min_rev_no_range_int16, - "int8": janitor_rs.compute_min_rev_no_range_int8, - "uint64": janitor_rs.compute_min_rev_no_range_uint64, - "uint32": janitor_rs.compute_min_rev_no_range_uint32, - "uint16": janitor_rs.compute_min_rev_no_range_uint16, - "uint8": janitor_rs.compute_min_rev_no_range_uint8, - "float64": janitor_rs.compute_min_rev_no_range_f64, - "float32": janitor_rs.compute_min_rev_no_range_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_min_rev_no_range", arr.dtype.name) return func( arr=arr, left_index=left_index, @@ -1927,23 +1068,7 @@ def _sum_rev_no_ranges( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_no_range_int64, - "int32": janitor_rs.compute_sum_rev_no_range_int32, - "int16": janitor_rs.compute_sum_rev_no_range_int16, - "int8": janitor_rs.compute_sum_rev_no_range_int8, - "uint64": janitor_rs.compute_sum_rev_no_range_uint64, - "uint32": janitor_rs.compute_sum_rev_no_range_uint32, - "uint16": janitor_rs.compute_sum_rev_no_range_uint16, - "uint8": janitor_rs.compute_sum_rev_no_range_uint8, - "float64": janitor_rs.compute_sum_rev_no_range_f64, - "float32": janitor_rs.compute_sum_rev_no_range_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_no_range", arr.dtype.name) return func( arr=arr, left_index=left_index, @@ -1963,23 +1088,7 @@ def _sum_rev_starts( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_start_int64, - "int32": janitor_rs.compute_sum_rev_start_int32, - "int16": janitor_rs.compute_sum_rev_start_int16, - "int8": janitor_rs.compute_sum_rev_start_int8, - "uint64": janitor_rs.compute_sum_rev_start_uint64, - "uint32": janitor_rs.compute_sum_rev_start_uint32, - "uint16": janitor_rs.compute_sum_rev_start_uint16, - "uint8": janitor_rs.compute_sum_rev_start_uint8, - "float64": janitor_rs.compute_sum_rev_start_f64, - "float32": janitor_rs.compute_sum_rev_start_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_start", arr.dtype.name) return func(arr=arr, starts=starts, index=index, booleans=booleans, length=length) @@ -1993,23 +1102,7 @@ def _sum_rev_ends( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_end_int64, - "int32": janitor_rs.compute_sum_rev_end_int32, - "int16": janitor_rs.compute_sum_rev_end_int16, - "int8": janitor_rs.compute_sum_rev_end_int8, - "uint64": janitor_rs.compute_sum_rev_end_uint64, - "uint32": janitor_rs.compute_sum_rev_end_uint32, - "uint16": janitor_rs.compute_sum_rev_end_uint16, - "uint8": janitor_rs.compute_sum_rev_end_uint8, - "float64": janitor_rs.compute_sum_rev_end_f64, - "float32": janitor_rs.compute_sum_rev_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_end", arr.dtype.name) return func(arr=arr, ends=ends, index=index, booleans=booleans, length=length) @@ -2025,23 +1118,7 @@ def _sum_rev_starts_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_start_match_int64, - "int32": janitor_rs.compute_sum_rev_start_match_int32, - "int16": janitor_rs.compute_sum_rev_start_match_int16, - "int8": janitor_rs.compute_sum_rev_start_match_int8, - "uint64": janitor_rs.compute_sum_rev_start_match_uint64, - "uint32": janitor_rs.compute_sum_rev_start_match_uint32, - "uint16": janitor_rs.compute_sum_rev_start_match_uint16, - "uint8": janitor_rs.compute_sum_rev_start_match_uint8, - "float64": janitor_rs.compute_sum_rev_start_match_f64, - "float32": janitor_rs.compute_sum_rev_start_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_start_match", arr.dtype.name) return func( arr=arr, starts=starts, @@ -2065,23 +1142,7 @@ def _sum_rev_ends_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_end_match_int64, - "int32": janitor_rs.compute_sum_rev_end_match_int32, - "int16": janitor_rs.compute_sum_rev_end_match_int16, - "int8": janitor_rs.compute_sum_rev_end_match_int8, - "uint64": janitor_rs.compute_sum_rev_end_match_uint64, - "uint32": janitor_rs.compute_sum_rev_end_match_uint32, - "uint16": janitor_rs.compute_sum_rev_end_match_uint16, - "uint8": janitor_rs.compute_sum_rev_end_match_uint8, - "float64": janitor_rs.compute_sum_rev_end_match_f64, - "float32": janitor_rs.compute_sum_rev_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_end_match", arr.dtype.name) return func( arr=arr, index=index, @@ -2105,23 +1166,7 @@ def _sum_rev_positions( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_positions_int64, - "int32": janitor_rs.compute_sum_rev_positions_int32, - "int16": janitor_rs.compute_sum_rev_positions_int16, - "int8": janitor_rs.compute_sum_rev_positions_int8, - "uint64": janitor_rs.compute_sum_rev_positions_uint64, - "uint32": janitor_rs.compute_sum_rev_positions_uint32, - "uint16": janitor_rs.compute_sum_rev_positions_uint16, - "uint8": janitor_rs.compute_sum_rev_positions_uint8, - "float64": janitor_rs.compute_sum_rev_positions_f64, - "float32": janitor_rs.compute_sum_rev_positions_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_positions", arr.dtype.name) return func( arr=arr, starts=starts, @@ -2144,23 +1189,7 @@ def _sum_rev_starts_ends( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_start_end_int64, - "int32": janitor_rs.compute_sum_rev_start_end_int32, - "int16": janitor_rs.compute_sum_rev_start_end_int16, - "int8": janitor_rs.compute_sum_rev_start_end_int8, - "uint64": janitor_rs.compute_sum_rev_start_end_uint64, - "uint32": janitor_rs.compute_sum_rev_start_end_uint32, - "uint16": janitor_rs.compute_sum_rev_start_end_uint16, - "uint8": janitor_rs.compute_sum_rev_start_end_uint8, - "float64": janitor_rs.compute_sum_rev_start_end_f64, - "float32": janitor_rs.compute_sum_rev_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_start_end", arr.dtype.name) return func( arr=arr, starts=starts, @@ -2184,23 +1213,7 @@ def _sum_rev_starts_ends_matches( """ Compute sum """ - mapping = { - "int64": janitor_rs.compute_sum_rev_start_end_match_int64, - "int32": janitor_rs.compute_sum_rev_start_end_match_int32, - "int16": janitor_rs.compute_sum_rev_start_end_match_int16, - "int8": janitor_rs.compute_sum_rev_start_end_match_int8, - "uint64": janitor_rs.compute_sum_rev_start_end_match_uint64, - "uint32": janitor_rs.compute_sum_rev_start_end_match_uint32, - "uint16": janitor_rs.compute_sum_rev_start_end_match_uint16, - "uint8": janitor_rs.compute_sum_rev_start_end_match_uint8, - "float64": janitor_rs.compute_sum_rev_start_end_match_f64, - "float32": janitor_rs.compute_sum_rev_start_end_match_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compute_sum_rev_start_end_match", arr.dtype.name) return func( arr=arr, starts=starts, diff --git a/janitor/functions/_conditional_join/_binary_search.py b/janitor/functions/_conditional_join/_binary_search.py index 1c514455c..df957de47 100644 --- a/janitor/functions/_conditional_join/_binary_search.py +++ b/janitor/functions/_conditional_join/_binary_search.py @@ -1,6 +1,7 @@ -import janitor_rs import numpy as np +from ._dtype_dispatch import _rs_func + def _binary_search_lt( left: np.ndarray, @@ -11,23 +12,7 @@ def _binary_search_lt( """ Get starts for < joins """ - mapping = { - "int64": janitor_rs.binary_search_lt_int64, - "int32": janitor_rs.binary_search_lt_int32, - "int16": janitor_rs.binary_search_lt_int16, - "int8": janitor_rs.binary_search_lt_int8, - "uint64": janitor_rs.binary_search_lt_uint64, - "uint32": janitor_rs.binary_search_lt_uint32, - "uint16": janitor_rs.binary_search_lt_uint16, - "uint8": janitor_rs.binary_search_lt_uint8, - "float64": janitor_rs.binary_search_lt_f64, - "float32": janitor_rs.binary_search_lt_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_lt", left.dtype.name) return func(left, right, starts, ends) @@ -40,23 +25,7 @@ def _binary_search_le( """ Get starts for <= joins """ - mapping = { - "int64": janitor_rs.binary_search_le_int64, - "int32": janitor_rs.binary_search_le_int32, - "int16": janitor_rs.binary_search_le_int16, - "int8": janitor_rs.binary_search_le_int8, - "uint64": janitor_rs.binary_search_le_uint64, - "uint32": janitor_rs.binary_search_le_uint32, - "uint16": janitor_rs.binary_search_le_uint16, - "uint8": janitor_rs.binary_search_le_uint8, - "float64": janitor_rs.binary_search_le_f64, - "float32": janitor_rs.binary_search_le_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_le", left.dtype.name) return func(left, right, starts, ends) @@ -69,23 +38,7 @@ def _binary_search_gt( """ Get ends for > joins """ - mapping = { - "int64": janitor_rs.binary_search_gt_int64, - "int32": janitor_rs.binary_search_gt_int32, - "int16": janitor_rs.binary_search_gt_int16, - "int8": janitor_rs.binary_search_gt_int8, - "uint64": janitor_rs.binary_search_gt_uint64, - "uint32": janitor_rs.binary_search_gt_uint32, - "uint16": janitor_rs.binary_search_gt_uint16, - "uint8": janitor_rs.binary_search_gt_uint8, - "float64": janitor_rs.binary_search_gt_f64, - "float32": janitor_rs.binary_search_gt_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_gt", left.dtype.name) return func(left, right, starts, ends) @@ -98,23 +51,7 @@ def _binary_search_ge( """ Get ends for >= joins """ - mapping = { - "int64": janitor_rs.binary_search_ge_int64, - "int32": janitor_rs.binary_search_ge_int32, - "int16": janitor_rs.binary_search_ge_int16, - "int8": janitor_rs.binary_search_ge_int8, - "uint64": janitor_rs.binary_search_ge_uint64, - "uint32": janitor_rs.binary_search_ge_uint32, - "uint16": janitor_rs.binary_search_ge_uint16, - "uint8": janitor_rs.binary_search_ge_uint8, - "float64": janitor_rs.binary_search_ge_f64, - "float32": janitor_rs.binary_search_ge_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_ge", left.dtype.name) return func(left, right, starts, ends) @@ -126,23 +63,7 @@ def _binary_search_lt_first( """ Get starts for < joins """ - mapping = { - "int64": janitor_rs.binary_search_lt_first_int64, - "int32": janitor_rs.binary_search_lt_first_int32, - "int16": janitor_rs.binary_search_lt_first_int16, - "int8": janitor_rs.binary_search_lt_first_int8, - "uint64": janitor_rs.binary_search_lt_first_uint64, - "uint32": janitor_rs.binary_search_lt_first_uint32, - "uint16": janitor_rs.binary_search_lt_first_uint16, - "uint8": janitor_rs.binary_search_lt_first_uint8, - "float64": janitor_rs.binary_search_lt_first_f64, - "float32": janitor_rs.binary_search_lt_first_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_lt_first", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None @@ -155,23 +76,7 @@ def _binary_search_le_first( """ Get starts for <= joins """ - mapping = { - "int64": janitor_rs.binary_search_le_first_int64, - "int32": janitor_rs.binary_search_le_first_int32, - "int16": janitor_rs.binary_search_le_first_int16, - "int8": janitor_rs.binary_search_le_first_int8, - "uint64": janitor_rs.binary_search_le_first_uint64, - "uint32": janitor_rs.binary_search_le_first_uint32, - "uint16": janitor_rs.binary_search_le_first_uint16, - "uint8": janitor_rs.binary_search_le_first_uint8, - "float64": janitor_rs.binary_search_le_first_f64, - "float32": janitor_rs.binary_search_le_first_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_le_first", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None @@ -184,23 +89,7 @@ def _binary_search_gt_first( """ Get ends for > joins """ - mapping = { - "int64": janitor_rs.binary_search_gt_first_int64, - "int32": janitor_rs.binary_search_gt_first_int32, - "int16": janitor_rs.binary_search_gt_first_int16, - "int8": janitor_rs.binary_search_gt_first_int8, - "uint64": janitor_rs.binary_search_gt_first_uint64, - "uint32": janitor_rs.binary_search_gt_first_uint32, - "uint16": janitor_rs.binary_search_gt_first_uint16, - "uint8": janitor_rs.binary_search_gt_first_uint8, - "float64": janitor_rs.binary_search_gt_first_f64, - "float32": janitor_rs.binary_search_gt_first_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_gt_first", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None @@ -213,23 +102,7 @@ def _binary_search_ge_first( """ Get ends for >= joins """ - mapping = { - "int64": janitor_rs.binary_search_ge_first_int64, - "int32": janitor_rs.binary_search_ge_first_int32, - "int16": janitor_rs.binary_search_ge_first_int16, - "int8": janitor_rs.binary_search_ge_first_int8, - "uint64": janitor_rs.binary_search_ge_first_uint64, - "uint32": janitor_rs.binary_search_ge_first_uint32, - "uint16": janitor_rs.binary_search_ge_first_uint16, - "uint8": janitor_rs.binary_search_ge_first_uint8, - "float64": janitor_rs.binary_search_ge_first_f64, - "float32": janitor_rs.binary_search_ge_first_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_ge_first", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None @@ -242,23 +115,7 @@ def _binary_search_gt_regions( """ Get ends for > joins """ - mapping = { - "int64": janitor_rs.binary_search_gt_regions_int64, - "int32": janitor_rs.binary_search_gt_regions_int32, - "int16": janitor_rs.binary_search_gt_regions_int16, - "int8": janitor_rs.binary_search_gt_regions_int8, - "uint64": janitor_rs.binary_search_gt_regions_uint64, - "uint32": janitor_rs.binary_search_gt_regions_uint32, - "uint16": janitor_rs.binary_search_gt_regions_uint16, - "uint8": janitor_rs.binary_search_gt_regions_uint8, - "float64": janitor_rs.binary_search_gt_regions_f64, - "float32": janitor_rs.binary_search_gt_regions_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_gt_regions", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None @@ -271,23 +128,7 @@ def _binary_search_ge_regions( """ Get ends for >= joins """ - mapping = { - "int64": janitor_rs.binary_search_ge_regions_int64, - "int32": janitor_rs.binary_search_ge_regions_int32, - "int16": janitor_rs.binary_search_ge_regions_int16, - "int8": janitor_rs.binary_search_ge_regions_int8, - "uint64": janitor_rs.binary_search_ge_regions_uint64, - "uint32": janitor_rs.binary_search_ge_regions_uint32, - "uint16": janitor_rs.binary_search_ge_regions_uint16, - "uint8": janitor_rs.binary_search_ge_regions_uint8, - "float64": janitor_rs.binary_search_ge_regions_f64, - "float32": janitor_rs.binary_search_ge_regions_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("binary_search_ge_regions", left.dtype.name) search_indices, left_index, total = func(left, right, left_index) if not total: return None diff --git a/janitor/functions/_conditional_join/_compare.py b/janitor/functions/_conditional_join/_compare.py index 7b05f6af9..fcb72b7f2 100644 --- a/janitor/functions/_conditional_join/_compare.py +++ b/janitor/functions/_conditional_join/_compare.py @@ -1,6 +1,7 @@ -import janitor_rs import numpy as np +from ._dtype_dispatch import _rs_func + def _compare_ne_no_ranges( left: np.ndarray, @@ -14,23 +15,7 @@ def _compare_ne_no_ranges( """ Compute for no ranges (no starts/ends) for != operator """ - mapping = { - "int64": janitor_rs.compare_no_range_ne_int64, - "int32": janitor_rs.compare_no_range_ne_int32, - "int16": janitor_rs.compare_no_range_ne_int16, - "int8": janitor_rs.compare_no_range_ne_int8, - "uint64": janitor_rs.compare_no_range_ne_uint64, - "uint32": janitor_rs.compare_no_range_ne_uint32, - "uint16": janitor_rs.compare_no_range_ne_uint16, - "uint8": janitor_rs.compare_no_range_ne_uint8, - "float64": janitor_rs.compare_no_range_ne_f64, - "float32": janitor_rs.compare_no_range_ne_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_no_range_ne", left.dtype.name) return func( left, right, @@ -51,23 +36,7 @@ def _compare_no_ranges( """ Compute for no ranges (no starts/ends) """ - mapping = { - "int64": janitor_rs.compare_no_range_int64, - "int32": janitor_rs.compare_no_range_int32, - "int16": janitor_rs.compare_no_range_int16, - "int8": janitor_rs.compare_no_range_int8, - "uint64": janitor_rs.compare_no_range_uint64, - "uint32": janitor_rs.compare_no_range_uint32, - "uint16": janitor_rs.compare_no_range_uint16, - "uint8": janitor_rs.compare_no_range_uint8, - "float64": janitor_rs.compare_no_range_f64, - "float32": janitor_rs.compare_no_range_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_no_range", left.dtype.name) return func( left, right, @@ -88,23 +57,7 @@ def _compare_ne_first_run_starts_only( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_start_ne_1st_int64, - "int32": janitor_rs.compare_start_ne_1st_int32, - "int16": janitor_rs.compare_start_ne_1st_int16, - "int8": janitor_rs.compare_start_ne_1st_int8, - "uint64": janitor_rs.compare_start_ne_1st_uint64, - "uint32": janitor_rs.compare_start_ne_1st_uint32, - "uint16": janitor_rs.compare_start_ne_1st_uint16, - "uint8": janitor_rs.compare_start_ne_1st_uint8, - "float64": janitor_rs.compare_start_ne_1st_f64, - "float32": janitor_rs.compare_start_ne_1st_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start_ne_1st", left.dtype.name) return func( left, right, @@ -130,23 +83,7 @@ def _compare_ne_starts_only( """ compute for starts """ - mapping = { - "int64": janitor_rs.compare_start_ne_int64, - "int32": janitor_rs.compare_start_ne_int32, - "int16": janitor_rs.compare_start_ne_int16, - "int8": janitor_rs.compare_start_ne_int8, - "uint64": janitor_rs.compare_start_ne_uint64, - "uint32": janitor_rs.compare_start_ne_uint32, - "uint16": janitor_rs.compare_start_ne_uint16, - "uint8": janitor_rs.compare_start_ne_uint8, - "float64": janitor_rs.compare_start_ne_f64, - "float32": janitor_rs.compare_start_ne_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start_ne", left.dtype.name) return func( left, right, @@ -172,23 +109,7 @@ def _compare_ne_first_run_ends_only( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_end_ne_1st_int64, - "int32": janitor_rs.compare_end_ne_1st_int32, - "int16": janitor_rs.compare_end_ne_1st_int16, - "int8": janitor_rs.compare_end_ne_1st_int8, - "uint64": janitor_rs.compare_end_ne_1st_uint64, - "uint32": janitor_rs.compare_end_ne_1st_uint32, - "uint16": janitor_rs.compare_end_ne_1st_uint16, - "uint8": janitor_rs.compare_end_ne_1st_uint8, - "float64": janitor_rs.compare_end_ne_1st_f64, - "float32": janitor_rs.compare_end_ne_1st_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_end_ne_1st", left.dtype.name) return func( left, right, @@ -214,23 +135,7 @@ def _compare_ne_ends_only( """ compute for ends """ - mapping = { - "int64": janitor_rs.compare_end_ne_int64, - "int32": janitor_rs.compare_end_ne_int32, - "int16": janitor_rs.compare_end_ne_int16, - "int8": janitor_rs.compare_end_ne_int8, - "uint64": janitor_rs.compare_end_ne_uint64, - "uint32": janitor_rs.compare_end_ne_uint32, - "uint16": janitor_rs.compare_end_ne_uint16, - "uint8": janitor_rs.compare_end_ne_uint8, - "float64": janitor_rs.compare_end_ne_f64, - "float32": janitor_rs.compare_end_ne_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_end_ne", left.dtype.name) return func( left, right, @@ -257,23 +162,7 @@ def _compare_ne_first_run_starts_ends( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_start_end_ne_1st_int64, - "int32": janitor_rs.compare_start_end_ne_1st_int32, - "int16": janitor_rs.compare_start_end_ne_1st_int16, - "int8": janitor_rs.compare_start_end_ne_1st_int8, - "uint64": janitor_rs.compare_start_end_ne_1st_uint64, - "uint32": janitor_rs.compare_start_end_ne_1st_uint32, - "uint16": janitor_rs.compare_start_end_ne_1st_uint16, - "uint8": janitor_rs.compare_start_end_ne_1st_uint8, - "float64": janitor_rs.compare_start_end_ne_1st_f64, - "float32": janitor_rs.compare_start_end_ne_1st_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start_end_ne_1st", left.dtype.name) return func( left, right, @@ -300,23 +189,7 @@ def _compare_ne_starts_ends( """ compute for starts and ends """ - mapping = { - "int64": janitor_rs.compare_start_end_ne_int64, - "int32": janitor_rs.compare_start_end_ne_int32, - "int16": janitor_rs.compare_start_end_ne_int16, - "int8": janitor_rs.compare_start_end_ne_int8, - "uint64": janitor_rs.compare_start_end_ne_uint64, - "uint32": janitor_rs.compare_start_end_ne_uint32, - "uint16": janitor_rs.compare_start_end_ne_uint16, - "uint8": janitor_rs.compare_start_end_ne_uint8, - "float64": janitor_rs.compare_start_end_ne_f64, - "float32": janitor_rs.compare_start_end_ne_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start_end_ne", left.dtype.name) return func( left, right, @@ -339,23 +212,7 @@ def _compare_first_run_starts_only( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_first_start_int64, - "int32": janitor_rs.compare_first_start_int32, - "int16": janitor_rs.compare_first_start_int16, - "int8": janitor_rs.compare_first_start_int8, - "uint64": janitor_rs.compare_first_start_uint64, - "uint32": janitor_rs.compare_first_start_uint32, - "uint16": janitor_rs.compare_first_start_uint16, - "uint8": janitor_rs.compare_first_start_uint8, - "float64": janitor_rs.compare_first_start_f64, - "float32": janitor_rs.compare_first_start_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_first_start", left.dtype.name) return func(left, right, starts, op) @@ -370,23 +227,7 @@ def _compare_starts_only( """ compute for starts """ - mapping = { - "int64": janitor_rs.compare_start_int64, - "int32": janitor_rs.compare_start_int32, - "int16": janitor_rs.compare_start_int16, - "int8": janitor_rs.compare_start_int8, - "uint64": janitor_rs.compare_start_uint64, - "uint32": janitor_rs.compare_start_uint32, - "uint16": janitor_rs.compare_start_uint16, - "uint8": janitor_rs.compare_start_uint8, - "float64": janitor_rs.compare_start_f64, - "float32": janitor_rs.compare_start_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start", left.dtype.name) return func(left, right, starts, counts_array, matches, op) @@ -396,23 +237,7 @@ def _compare_first_run_ends_only( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_first_end_int64, - "int32": janitor_rs.compare_first_end_int32, - "int16": janitor_rs.compare_first_end_int16, - "int8": janitor_rs.compare_first_end_int8, - "uint64": janitor_rs.compare_first_end_uint64, - "uint32": janitor_rs.compare_first_end_uint32, - "uint16": janitor_rs.compare_first_end_uint16, - "uint8": janitor_rs.compare_first_end_uint8, - "float64": janitor_rs.compare_first_end_f64, - "float32": janitor_rs.compare_first_end_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_first_end", left.dtype.name) return func(left, right, ends, op) @@ -427,23 +252,7 @@ def _compare_ends_only( """ compute for ends """ - mapping = { - "int64": janitor_rs.compare_end_int64, - "int32": janitor_rs.compare_end_int32, - "int16": janitor_rs.compare_end_int16, - "int8": janitor_rs.compare_end_int8, - "uint64": janitor_rs.compare_end_uint64, - "uint32": janitor_rs.compare_end_uint32, - "uint16": janitor_rs.compare_end_uint16, - "uint8": janitor_rs.compare_end_uint8, - "float64": janitor_rs.compare_end_f64, - "float32": janitor_rs.compare_end_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_end", left.dtype.name) return func(left, right, ends, counts_array, matches, op) @@ -457,23 +266,7 @@ def _compare_first_run_starts_ends( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_first_start_end_int64, - "int32": janitor_rs.compare_first_start_end_int32, - "int16": janitor_rs.compare_first_start_end_int16, - "int8": janitor_rs.compare_first_start_end_int8, - "uint64": janitor_rs.compare_first_start_end_uint64, - "uint32": janitor_rs.compare_first_start_end_uint32, - "uint16": janitor_rs.compare_first_start_end_uint16, - "uint8": janitor_rs.compare_first_start_end_uint8, - "float64": janitor_rs.compare_first_start_end_f64, - "float32": janitor_rs.compare_first_start_end_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_first_start_end", left.dtype.name) return func(left, right, starts, ends, op) @@ -488,23 +281,7 @@ def _compare_starts_ends( """ compute for starts and ends """ - mapping = { - "int64": janitor_rs.compare_start_end_int64, - "int32": janitor_rs.compare_start_end_int32, - "int16": janitor_rs.compare_start_end_int16, - "int8": janitor_rs.compare_start_end_int8, - "uint64": janitor_rs.compare_start_end_uint64, - "uint32": janitor_rs.compare_start_end_uint32, - "uint16": janitor_rs.compare_start_end_uint16, - "uint8": janitor_rs.compare_start_end_uint8, - "float64": janitor_rs.compare_start_end_f64, - "float32": janitor_rs.compare_start_end_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_start_end", left.dtype.name) return func(left, right, starts, ends, matches, op) @@ -519,23 +296,7 @@ def _compare_positions( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_posns_int64, - "int32": janitor_rs.compare_posns_int32, - "int16": janitor_rs.compare_posns_int16, - "int8": janitor_rs.compare_posns_int8, - "uint64": janitor_rs.compare_posns_uint64, - "uint32": janitor_rs.compare_posns_uint32, - "uint16": janitor_rs.compare_posns_uint16, - "uint8": janitor_rs.compare_posns_uint8, - "float64": janitor_rs.compare_posns_f64, - "float32": janitor_rs.compare_posns_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_posns", left.dtype.name) return func( left=left, right=right, @@ -560,23 +321,7 @@ def _compare_positions_ne( """ compute for first run """ - mapping = { - "int64": janitor_rs.compare_posns_ne_int64, - "int32": janitor_rs.compare_posns_ne_int32, - "int16": janitor_rs.compare_posns_ne_int16, - "int8": janitor_rs.compare_posns_ne_int8, - "uint64": janitor_rs.compare_posns_ne_uint64, - "uint32": janitor_rs.compare_posns_ne_uint32, - "uint16": janitor_rs.compare_posns_ne_uint16, - "uint8": janitor_rs.compare_posns_ne_uint8, - "float64": janitor_rs.compare_posns_ne_f64, - "float32": janitor_rs.compare_posns_ne_f32, - } - dtype_name = left.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") + func = _rs_func("compare_posns_ne", left.dtype.name) return func( left=left, right=right, diff --git a/janitor/functions/_conditional_join/_dtype_dispatch.py b/janitor/functions/_conditional_join/_dtype_dispatch.py new file mode 100644 index 000000000..575e64265 --- /dev/null +++ b/janitor/functions/_conditional_join/_dtype_dispatch.py @@ -0,0 +1,36 @@ +"""Shared dtype -> janitor_rs function dispatch for conditional_join.""" + +from functools import lru_cache +from typing import Callable + +import janitor_rs + +# Maps a numpy dtype name to the suffix used in the corresponding +# `janitor_rs` function name (float64/float32 shorten to f64/f32; every +# other dtype keeps its numpy name). +_DTYPE_SUFFIXES = { + "int64": "int64", + "int32": "int32", + "int16": "int16", + "int8": "int8", + "uint64": "uint64", + "uint32": "uint32", + "uint16": "uint16", + "uint8": "uint8", + "float64": "f64", + "float32": "f32", +} + + +@lru_cache(maxsize=None) +def _rs_func(family: str, dtype_name: str) -> Callable: + """ + ELI5: for a given operation `family` (e.g. "compute_sum_start"), look up + the `janitor_rs` function built for `dtype_name` once, then remember it + so later calls for the same (family, dtype) pair skip the lookup. + """ + try: + suffix = _DTYPE_SUFFIXES[dtype_name] + except KeyError: + raise KeyError(f"Unsupported data type -> {dtype_name}") from None + return getattr(janitor_rs, f"{family}_{suffix}") diff --git a/janitor/functions/_conditional_join/_greater_than_indices.py b/janitor/functions/_conditional_join/_greater_than_indices.py index ba5afcc8b..408840df5 100644 --- a/janitor/functions/_conditional_join/_greater_than_indices.py +++ b/janitor/functions/_conditional_join/_greater_than_indices.py @@ -1,5 +1,4 @@ # helper functions for >/>= -import janitor_rs import numpy as np import pandas as pd @@ -87,9 +86,5 @@ def _greater_than_indices( ) right = [right_index[:ind] for ind in search_indices] right = np.concatenate(right) - left = janitor_rs.repeat_index( - index=left_index, - counts=search_indices, - length=search_indices.sum(), - ) + left = np.repeat(left_index, search_indices) return {"left_index": left, "right_index": right} diff --git a/janitor/functions/_conditional_join/_helpers.py b/janitor/functions/_conditional_join/_helpers.py index 281111175..91cd5765b 100644 --- a/janitor/functions/_conditional_join/_helpers.py +++ b/janitor/functions/_conditional_join/_helpers.py @@ -560,9 +560,7 @@ def _build_indices_positions( Build indices for multiple joins """ if keep == "all": - left_index = janitor_rs.repeat_index( - index=left_index, counts=counts_array, length=total - ) + left_index = np.repeat(left_index, counts_array) right_index = janitor_rs.build_positional_index( index=right_index, positions=positions, length=total ) @@ -609,29 +607,17 @@ def build_indices_matches( Build indices for multiple joins, where `matches` exist """ if (keep == "all") and (starts is not None) and (ends is None): - left = janitor_rs.repeat_index( - index=left_index, - counts=counts_array, - length=total, - ) + left = np.repeat(left_index, counts_array) right = janitor_rs.index_starts_only( index=right_index, starts=starts, matches=matches, length=total ) elif (keep == "all") and (starts is None) and (ends is not None): - left = janitor_rs.repeat_index( - index=left_index, - counts=counts_array, - length=total, - ) + left = np.repeat(left_index, counts_array) right = janitor_rs.index_ends_only( index=right_index, ends=ends, matches=matches, length=total ) elif (keep == "all") and (starts is not None) and (ends is not None): - left = janitor_rs.repeat_index( - index=left_index, - counts=counts_array, - length=total, - ) + left = np.repeat(left_index, counts_array) right = janitor_rs.index_starts_and_ends( index=right_index, starts=starts, diff --git a/janitor/functions/_conditional_join/_less_than_indices.py b/janitor/functions/_conditional_join/_less_than_indices.py index 60ec532f4..cbb2ca635 100644 --- a/janitor/functions/_conditional_join/_less_than_indices.py +++ b/janitor/functions/_conditional_join/_less_than_indices.py @@ -1,5 +1,4 @@ # helper functions for =3.0.0 - natsort>=8.4.0,<9 diff --git a/tests/functions/test_conditional_join.py b/tests/functions/test_conditional_join.py index 94fb70ac5..4115f6d3d 100644 --- a/tests/functions/test_conditional_join.py +++ b/tests/functions/test_conditional_join.py @@ -1,5 +1,6 @@ import operator +import janitor_rs import numpy as np import pandas as pd import pytest @@ -8,6 +9,7 @@ from pandas.testing import assert_frame_equal import janitor as jn +from janitor.functions._conditional_join._dtype_dispatch import _rs_func from janitor.testing_utils.strategies import ( conditional_df, conditional_right, @@ -6591,3 +6593,38 @@ def test_equi_le_ge_ge_ne_agg_rev(df, right): ).sort_index() actual = actual.loc[expected.index] assert_frame_equal(expected, actual) + + +def test_rs_func_dispatch_returns_correct_function(): + """`_rs_func` resolves the right janitor_rs export for family/dtype.""" + assert ( + _rs_func("compute_sum_start_end", "int64") + is janitor_rs.compute_sum_start_end_int64 + ) + assert _rs_func("binary_search_lt", "float32") is janitor_rs.binary_search_lt_f32 + assert _rs_func("binary_search_lt", "float64") is janitor_rs.binary_search_lt_f64 + + +def test_rs_func_unsupported_dtype_raises_keyerror(): + """Unsupported dtypes raise the same KeyError message as before.""" + with pytest.raises(KeyError, match="Unsupported data type -> bool"): + _rs_func("compute_sum_start_end", "bool") + + +@pytest.mark.parametrize( + "index, counts", + [ + (np.array([], dtype=np.int64), np.array([], dtype=np.int64)), + (np.array([1, 2, 3], dtype=np.int64), np.array([0, 0, 0], dtype=np.int64)), + (np.array([1, 2, 3], dtype=np.int64), np.array([2, 0, 3], dtype=np.int64)), + (np.arange(1000, dtype=np.int64), np.full(1000, 5, dtype=np.int64)), + ], +) +def test_np_repeat_matches_janitor_rs_repeat_index(index, counts): + """`np.repeat` is a drop-in replacement for `janitor_rs.repeat_index`.""" + expected = janitor_rs.repeat_index( + index=index, counts=counts, length=int(counts.sum()) + ) + actual = np.repeat(index, counts) + np.testing.assert_array_equal(actual, expected) + assert actual.dtype == expected.dtype == np.int64 From 3e895e85b61616e88bae34afb83cba37346a7691 Mon Sep 17 00:00:00 2001 From: samukweku Date: Wed, 26 Aug 2026 13:08:07 +1000 Subject: [PATCH 2/3] perf: coordinate reverse positions aggregation APIs --- .../_conditional_join/_agg_functions.py | 10 --------- .../_conditional_join/_get_join_aggs.py | 21 ++++++++++--------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/janitor/functions/_conditional_join/_agg_functions.py b/janitor/functions/_conditional_join/_agg_functions.py index 7282f6fba..227c9e7cc 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -112,7 +112,6 @@ def _size_rev_positions( ends: np.ndarray, index: np.ndarray, positions: np.ndarray, - length: int, ) -> tuple: """ Compute size_rev @@ -122,7 +121,6 @@ def _size_rev_positions( ends=ends, index=index, positions=positions, - length=length, ) @@ -637,7 +635,6 @@ def _prod_rev_positions( index: np.ndarray, positions: np.ndarray, booleans: np.ndarray, - length: int, ) -> tuple: """ Compute prod @@ -650,7 +647,6 @@ def _prod_rev_positions( index=index, positions=positions, booleans=booleans, - length=length, ) @@ -785,7 +781,6 @@ def _min_rev_positions( index: np.ndarray, positions: np.ndarray, booleans: np.ndarray, - length: int, ) -> tuple: """ Compute min @@ -798,7 +793,6 @@ def _min_rev_positions( index=index, positions=positions, booleans=booleans, - length=length, ) @@ -933,7 +927,6 @@ def _max_rev_positions( index: np.ndarray, positions: np.ndarray, booleans: np.ndarray, - length: int, ) -> tuple: """ Compute max @@ -946,7 +939,6 @@ def _max_rev_positions( index=index, positions=positions, booleans=booleans, - length=length, ) @@ -1161,7 +1153,6 @@ def _sum_rev_positions( index: np.ndarray, positions: np.ndarray, booleans: np.ndarray, - length: int, ) -> tuple: """ Compute sum @@ -1174,7 +1165,6 @@ def _sum_rev_positions( index=index, positions=positions, booleans=booleans, - length=length, ) diff --git a/janitor/functions/_conditional_join/_get_join_aggs.py b/janitor/functions/_conditional_join/_get_join_aggs.py index d757c6f41..c3c102d9b 100644 --- a/janitor/functions/_conditional_join/_get_join_aggs.py +++ b/janitor/functions/_conditional_join/_get_join_aggs.py @@ -268,7 +268,6 @@ def _agg_join_left(df: pd.DataFrame, aggfunc: list, indices: dict) -> pd.DataFra ends=indices["ends"], positions=indices["positions"], index=indices["right_index"], - length=indices["right_index"].size, ) else: ser = df.loc[indices["left_index"], column_name] @@ -276,15 +275,17 @@ def _agg_join_left(df: pd.DataFrame, aggfunc: list, indices: dict) -> pd.DataFra booleans = pd.isna(arr) arr = _helpers._convert_array_to_numpy(array=arr) func = mapping[agg] - _index, out = func( - arr=arr, - starts=indices["starts"], - ends=indices["ends"], - positions=indices["positions"], - index=indices["right_index"], - booleans=booleans, - length=indices["right_index"].size, - ) + kwargs = { + "arr": arr, + "starts": indices["starts"], + "ends": indices["ends"], + "positions": indices["positions"], + "index": indices["right_index"], + "booleans": booleans, + } + if agg not in {"sum", "min"}: + kwargs["length"] = indices["right_index"].size + _index, out = func(**kwargs) if agg in { "sum", "prod", From ba38ce7799f5b83b6f028396f561b65219114cba Mon Sep 17 00:00:00 2001 From: samukweku Date: Thu, 27 Aug 2026 14:35:46 +1000 Subject: [PATCH 3/3] docs: explain reverse aggregation contract --- AGENTS.md | 17 +++++++++++++++++ .../_conditional_join/_agg_functions.py | 13 +++++++++++++ 2 files changed, 30 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index 34a7f1921..92362cbc2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -548,6 +548,23 @@ CLI. include MkDocs. The documentation task is available in the `docs` environment. **Recommendation**: Run `pixi run -e docs build-docs` to build documentation. +### [2026-08-27] Document the Reverse Aggregation Boundary Contract + +**Context**: Coordinating pyjanitor with janitor-rs reverse aggregation kernels. +**Learning**: Reverse match kernels consume a flattened candidate tape. Direct +Rust callers must provide a non-empty tape with the expected shape. Pyjanitor +owns the producer invariant that match values are 0 or 1. A batch whose every +candidate range is zero-width is valid at the Python level, so pyjanitor must +return the typed empty aggregation result before calling Rust. Individual +zero-width rows remain valid when the overall tape is non-empty. +**Recommendation**: Keep this contract documented and tested at both +boundaries: Rust rejects malformed direct inputs, while pyjanitor short-circuits +legitimate no-candidate batches. Treat `length` as a legacy compatibility +argument only when supporting an older janitor-rs wheel. +**ELI5**: Python builds one roll of candidate tickets and Rust checks its shape. +If there are no tickets, Python returns the empty answer instead of handing an +empty roll to Rust. + --- ## Version History diff --git a/janitor/functions/_conditional_join/_agg_functions.py b/janitor/functions/_conditional_join/_agg_functions.py index 227c9e7cc..b58914103 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,3 +1,16 @@ +"""Adapters for conditional-join aggregation kernels. + +Reverse match kernels receive a flattened candidate tape. Rust intentionally +rejects an empty tape, while pyjanitor pre-filters an all-zero-width batch and +returns the normal typed empty result before dispatch. Individual zero-width +ranges are valid when the overall tape is non-empty. The comparison stage +owns the invariant that match values are 0 or 1; Rust validates tape shape. + +ELI5: Python builds one long roll of candidate tickets and Rust checks that +the roll has the right shape. If there are no tickets at all, Python answers +with an empty result instead of handing an empty roll to Rust. +""" + import janitor_rs import numpy as np