diff --git a/AGENTS.md b/AGENTS.md index 34a7f1921..d983bffe8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,7 +37,14 @@ these into the main sections. task. - **Test-Driven**: Always run tests after making code changes. - **Document**: Keep docstrings up-to-date using Google-style format. -- **Lint Markdown**: Always run `markdownlint` on markdown files after editing. +- **ELI5 Comments**: Use ELI5-style explanations generously in new or + modified code -- comments and docstrings should explain the mechanism + (why it works, what trade-off it makes), not just restate what the code + does. + +The `pixi run`, `markdownlint`, and `uvx marimo convert` rules are each +stated once, in their own section below (Package Manager, Markdown +Linting, Notebook Commands) -- not repeated throughout this file. --- @@ -413,17 +420,13 @@ pixi run pytest -m "not turtle" -v ## Common Anti-Patterns to Avoid -### ❌ DON'T - -1. **Don't run Python/pytest without pixi** +The `pixi run`, notebook-conversion, and markdownlint rules live in their +own sections above (Package Manager, Notebook Commands, Markdown Linting) +-- not repeated here. This list covers anti-patterns with no other home. - ```bash - # Wrong - python script.py - pytest tests/ - ``` +### ❌ DON'T -2. **Don't mutate input DataFrames** +1. **Don't mutate input DataFrames** ```python # Wrong @@ -432,32 +435,15 @@ pixi run pytest -m "not turtle" -v return df ``` -3. **Don't manually convert notebooks** - - ```bash - # Wrong - don't write custom conversion scripts - python convert_notebook.py - ``` - -4. **Don't forget to add tests** +2. **Don't forget to add tests** - Every new function needs corresponding tests -5. **Don't skip docstrings** +3. **Don't skip docstrings** - Interrogate enforces >55% docstring coverage -6. **Don't forget to lint markdown** - - Always run `markdownlint` on markdown files after editing - ### ✅ DO -1. **Always use pixi run** - - ```bash - pixi run pytest tests/ - pixi run python script.py - ``` - -2. **Work on copies** +1. **Work on copies** ```python def my_func(df): @@ -466,22 +452,9 @@ pixi run pytest -m "not turtle" -v return df ``` -3. **Use uvx marimo for notebooks** - - ```bash - uvx marimo convert notebook.ipynb -o notebook.py - ``` - -4. **Write tests alongside code** - -5. **Write Google-style docstrings with examples** +2. **Write tests alongside code** -6. **Run markdownlint on markdown files** - - ```bash - markdownlint AGENTS.md - # Install if not on PATH: pixi global install markdownlint-cli - ``` +3. **Write Google-style docstrings with examples** --- @@ -526,14 +499,6 @@ Add entries in the format: **Recommendation**: How to apply this learning --> -### [2025-12-19] Always Run markdownlint - -**Context**: Editing AGENTS.md file -**Learning**: Markdown files should be linted with `markdownlint` to ensure -consistent formatting and catch issues like long lines. -**Recommendation**: After editing any markdown file, run `markdownlint `. -If not installed, use `pixi global install markdownlint-cli`. - ### [2026-02-07] Open PRs with GitHub CLI **Context**: User requested opening a PR after pushing changes. @@ -548,6 +513,41 @@ 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-22] Always Use a New Worktree for a New Issue/Branch + +**Context**: Working on issue #1653 in the shared `~/github/pyjanitor` clone +while another concurrent session had issue #1648 checked out in a sibling +worktree (`~/github/pyjanitor-1648`); the shared clone's HEAD ended up +detached mid-session from that other work, which was only caught by +inspecting `git status`/`git reflog` before committing. +**Learning**: The shared main clone directory gets reused across sessions +and branches, so starting a new issue there risks colliding with another +in-progress checkout, or committing on top of the wrong branch/detached +HEAD. +**Recommendation**: For every new issue or branch, create a dedicated +`git worktree add ~/github/pyjanitor- -b origin/dev` +(or the equivalent for the current base branch) instead of checking out a +new branch in the shared clone. Never leave uncommitted work or run +destructive git operations in the shared clone without first checking +`git status`/`git branch --show-current` for signs of concurrent use. + +### [2026-08-22] Use ELI5 Comments Generously + +**Context**: Requested while adding new NumPy kernels to +`_agg_functions.py` (Issue #1653) with a non-obvious tie/null/NaN +contract to preserve. +**Learning**: Terse docstrings (e.g. the file's existing "Compute min") +are fine for simple pass-throughs, but new or non-obvious logic -- +especially anything with a subtle invariant, a workaround, or a +trade-off a reader wouldn't guess -- benefits from an explicit ELI5 +explanation of the mechanism, not just a restatement of what the code +does. +**Recommendation**: When writing or substantially modifying a function, +default to adding a short ELI5 paragraph explaining *why* it works that +way, particularly for dispatch/threshold logic, non-obvious algorithms, +and anything replicating an existing implementation's exact quirky +behavior. + --- ## Version History diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..8ba72069e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +- [PERF] Use an O(n + m) NumPy prefix/suffix running-argmin/argmax for + forward `conditional_join` range `min`/`max` aggregations + (`join_agg(..., aggfunc=[(col, "min"|"max")])`), replacing the O(sum + of interval widths) Rust kernels once query density crosses a + benchmarked work-factor threshold; sparse queries stay on the Rust + kernels. Preserves exact first-occurrence tie-breaking, null-skipping, + and the existing float NaN-comparison quirk. - Issue #1653, PR #1674 + @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..495194a4a 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,6 +1,286 @@ import janitor_rs import numpy as np +_ARGEXT_WORK_FACTOR = 20 + + +def _use_argext(arr_size: int, total_width: int) -> bool: + """ + Use the O(n) prefix/suffix precompute only when repeated Rust range + scans would cost more. + + ELI5: same idea as the prefix-sum work-factor check for range sums -- + only pay for building the running-best array once the Rust kernel + would otherwise re-scan the array more than `_ARGEXT_WORK_FACTOR` + times over. Benchmarked empirically (see PR description): this + kernel does more per-element work than a plain running sum, so its + break-even point is higher. + """ + return total_width > (_ARGEXT_WORK_FACTOR * arr_size) + + +def _running_argext_compact(vals: np.ndarray, is_max: bool, strict: bool) -> np.ndarray: + """ + Positions (into `vals`) of the running best value, scanning `vals` + left to right. `vals` must already have every null/NaN entry removed + -- that's what makes this safe to vectorize with no sentinel value: + there's nothing to accidentally collide with a real data point at a + narrow dtype's extreme (e.g. 255 for uint8). + + ELI5: walk once, remembering the position of the best value seen so + far. `strict=True` means a tie does NOT replace the earlier record + (used for prefixes); `strict=False` means a tie DOES replace it + (used for suffixes, scanning a reversed array, so "replace on tie" + ends up keeping the smaller/earlier original index). + """ + m = vals.size + # ELI5: `running_val[i]` = the best value seen anywhere in vals[:i+1]. + # `fmax`/`fmin` are the same as `maximum`/`minimum` for these already- + # cleaned values, computed once for the whole array in one call + # (that's the "O(n) instead of O(n) per row" trick). + running_val = (np.fmax if is_max else np.fmin).accumulate(vals) + + # ELI5: a position is a "new record" exactly when it's the reason + # `running_val` changed (or didn't change but ties count -- see + # `strict` below). Index 0 is always a record: there's nothing + # before it to compare against, so it's automatically the best seen + # so far. + is_new_record = np.empty(m, dtype=bool) + is_new_record[0] = True + if strict: + # Prefix rule: only a STRICT improvement counts as a new record, + # so an earlier tie keeps its spot -- `running_val` only changes + # on a genuine improvement, so comparing consecutive running + # values directly tells us exactly where that happened. + is_new_record[1:] = ( + (running_val[1:] > running_val[:-1]) + if is_max + else (running_val[1:] < running_val[:-1]) + ) + else: + # Suffix rule: a TIE also counts as a new record (see the + # `strict=False` note above -- this scans a reversed array, so + # "record on tie" is what makes the earlier original index win). + # `vals[i] == running_val[i]` is true both when vals[i] set a new + # best AND when it merely matched the existing best. + is_new_record[1:] = vals[1:] == running_val[1:] + + # ELI5: turn "is this position a record?" into "what's the latest + # record position seen so far?" -- put each record's own index at + # its position, -1 everywhere else, then let `maximum.accumulate` + # carry the largest (i.e. most recent) index forward over the -1s. + candidate = np.where(is_new_record, np.arange(m), -1) + return np.maximum.accumulate(candidate) + + +def _isnan_if_float(arr: np.ndarray) -> np.ndarray: + """`np.isnan` raises on integer dtypes, which never have NaN anyway.""" + if np.issubdtype(arr.dtype, np.floating): + return np.isnan(arr) + return np.zeros(arr.shape, dtype=bool) + + +def _prefix_argext(arr: np.ndarray, booleans: np.ndarray, is_max: bool) -> np.ndarray: + """ + result[e] = the Rust `compute_{min,max}_end_*` answer for that `e`, + for every e in 0..n, computed once in O(n) instead of once per row. + + ELI5: the Rust kernel restarts its scan from index 0 every single + time, no matter what `end` is, so every row is just a different-length + prefix of the *same* walk. Walk it once and remember the running best + position; every row can then just read off the answer for its own + `end`. + + Floats get one extra wrinkle: the Rust kernel's `<`/`>` comparisons + are IEEE754, where anything compared with NaN is always false. So + once the scan's very first non-null value is a NaN, nothing -- + real or NaN -- can ever replace it (`0 < NaN` and `NaN < 0` are both + false), and that NaN's position "freezes" as the answer for every + later `end` too. A NaN encountered *after* a real anchor, on the + other hand, is just silently never selected, exactly like a null. + """ + n = arr.size + # result[e] is the answer for range [0, e); -1 is "no valid element + # yet" (covers e == 0, and the whole array being null). + result = np.full(n + 1, -1, dtype=np.int64) + if n == 0: + return result + not_null = ~booleans + if not not_null.any(): + return result # every position is null -> every answer is -1 + + # ELI5: the very first non-null element is special -- it's the one + # every row's scan eventually reaches first (see the docstring: all + # prefixes share the same walk from index 0). `np.argmax` on a bool + # array returns the index of the first True. + first_valid = int(np.argmax(not_null)) + isnan = _isnan_if_float(arr) + + if isnan[first_valid]: + # Frozen case: nothing can ever beat/tie a NaN anchor (IEEE754), + # so every `end` that reaches past `first_valid` is stuck + # reporting that NaN's position; every `end` before it hasn't + # found any real value yet, so it's still -1 (already the + # default fill above). + result[first_valid + 1 :] = first_valid + return result + + # Anchor is real: NaN can now never win (it can't beat a real value, + # and a real value can't lose to one either), so treat NaN exactly + # like a null for the rest of this scan and compact both away. + extended_null = booleans | isnan + # original positions of the real, non-null values + valid_idx = np.flatnonzero(~extended_null) + valid_vals = arr[valid_idx] # ...and their values, same order + # Run the O(m) compact scan (m = count of real values), then map its + # answers (positions *within* `valid_vals`) back to real array + # positions via `valid_idx`. + running_pos = valid_idx[_running_argext_compact(valid_vals, is_max, strict=True)] + + # ELI5: `result[e]` should reuse the compact scan's answer for + # whichever real value was the LAST one inside [0, e). Count how + # many real values exist in [0, e) (that's `valid_count_prefix[e]`, + # a running count with a leading 0 so index 0 means "none yet"), and + # that count doubles as a 1-indexed position into `running_pos` -- + # subtract 1 to make it a normal 0-indexed lookup. + # dtype=np.int64 pins the accumulator explicitly -- np.cumsum's + # platform-default int (int32 on 64-bit Windows) would silently + # overflow past ~2.1 billion elements otherwise. + valid_count_prefix = np.concatenate( + ([0], np.cumsum(~extended_null, dtype=np.int64)) + ) + has_any = valid_count_prefix > 0 + k = np.clip(valid_count_prefix - 1, 0, max(running_pos.size - 1, 0)) + result[:] = np.where(has_any, running_pos[k], -1) + return result + + +def _suffix_argext(arr: np.ndarray, booleans: np.ndarray, is_max: bool) -> np.ndarray: + """ + result[s] = the Rust `compute_{min,max}_start_*` answer for that `s`, + for every s in 0..n, computed once in O(n) instead of once per row. + + ELI5: unlike the prefix case, each row's `start` restarts its own + scan from a different position, so which value gets "frozen" by the + NaN quirk (see `_prefix_argext`) can differ row to row -- a single + shared backward scan over the whole array isn't enough on its own. + So this splits the work: `next_valid` finds each `s`'s own anchor + position in O(n); if that anchor is a NaN the answer is just that + position (frozen); otherwise NaN behaves like null for the rest of + that row's range, and the answer comes from one shared "clean" + backward scan that skips nulls and NaNs alike. + """ + n = arr.size + # result[s] is the answer for range [s, n); -1 is "no valid element + # from s onward" (also serves as the s == n empty-range answer). + result = np.full(n + 1, -1, dtype=np.int64) + if n == 0: + return result + idx = np.arange(n) + isnan = _isnan_if_float(arr) + + # ELI5: `next_valid[s]` = the smallest index >= s that isn't null, + # or n if there isn't one. Built with a classic "fill backward" trick: + # put each non-null position's own index in `raw`, n (a value larger + # than any real index) everywhere else, then `minimum.accumulate` on + # the REVERSED array carries the smallest real index found so far + # back toward the front; reversing again undoes the flip so `raw[s]` + # lines up with position s again. This is the "each row re-anchors + # at its own start" position from the docstring, found for every s + # at once. + raw = np.where(~booleans, idx, n) + next_valid = np.minimum.accumulate(raw[::-1])[::-1] + + # Once we know a row's anchor isn't a NaN (checked below via + # `next_valid`), NaN behaves exactly like null for the rest of that + # row's scan -- so compact both away together, same as the prefix + # case, just scanning right-to-left. + extended_null = booleans | isnan + # original positions of the real, non-null values + valid_idx = np.flatnonzero(~extended_null) + if valid_idx.size: + valid_vals = arr[valid_idx] + # Reverse both arrays so "left to right on the reversed array" + # is the same walk as "right to left on the real array" -- + # `_running_argext_compact(..., strict=False)` then applies the + # tie-keeps-the-later-record rule, which (because everything is + # reversed) works out to "ties keep the smaller original index". + rev_pos_in_compact = _running_argext_compact( + valid_vals[::-1], is_max, strict=False + ) + # Map compact-array positions back to real array positions, in + # the same reversed order used above. + running_pos_rev = valid_idx[::-1][rev_pos_in_compact] + + # Same "count real values, use the count as a lookup index" + # trick as the prefix case, mirrored: count real values from s + # to the end (a running count from the right, with a trailing 0 + # so index n means "none"), then look that many entries into the + # *reversed* running-position array. + # dtype=np.int64: see the matching comment in _prefix_argext. + valid_count_suffix = np.concatenate( + (np.cumsum((~extended_null)[::-1], dtype=np.int64)[::-1], [0]) + ) + has_any = valid_count_suffix[:n] > 0 + k = np.clip(valid_count_suffix[:n] - 1, 0, max(running_pos_rev.size - 1, 0)) + clean_suffix = np.where(has_any, running_pos_rev[k], -1) + else: + clean_suffix = np.full(n, -1, dtype=np.int64) + + # ELI5: now decide, per row, which answer actually applies. + # `frozen_mask[s]` = "does [s, n) contain any non-null value at + # all?" (next_valid[s] < n). If so, check whether THAT row's own + # anchor (its first non-null value) happens to be NaN -- if it is, + # the frozen-NaN answer wins; otherwise fall back to the shared + # "clean" scan computed above. + frozen_mask = next_valid < n + anchor_isnan = np.zeros(n, dtype=bool) + if np.issubdtype(arr.dtype, np.floating): + valid_next = next_valid[frozen_mask] + anchor_isnan[frozen_mask] = np.isnan(arr[valid_next]) + + result[:n] = np.where( + ~frozen_mask, -1, np.where(anchor_isnan, next_valid, clean_suffix) + ) + return result + + +def _argext_dispatch( + arr: np.ndarray, + booleans: np.ndarray, + indexer: np.ndarray, + is_starts: bool, + is_max: bool, + mapping: dict, +) -> np.ndarray | None: + """ + Try the O(n) NumPy prefix/suffix path shared by _min_starts/_min_ends/ + _max_starts/_max_ends; return None to signal "fall through to the + Rust `mapping` dict instead". + + ELI5: one shared decision point instead of four near-identical + copies. A dtype is eligible exactly when it's a key in `mapping` -- + so there's no second, separately-maintained list of supported + dtypes that could drift out of sync with it -- and the NumPy + precompute only pays off once there's enough total query width (see + `_use_argext`). `is_starts=True` means `indexer` is `starts` + (suffix scans via `_suffix_argext`); `is_starts=False` means + `indexer` is `ends` (prefix scans via `_prefix_argext`). + """ + dtype_name = arr.dtype.name + if dtype_name not in mapping or indexer.size <= _ARGEXT_WORK_FACTOR: + return None + if is_starts: + # sum(n - start) == n*len(indexer) - sum(indexer) + total_width = (arr.size * indexer.size) - indexer.sum(dtype=np.int64) + else: + # each row's Rust scan covers `end` elements (range [0, end)) + total_width = int(indexer.sum(dtype=np.int64)) + if not _use_argext(arr_size=arr.size, total_width=total_width): + return None + argext = _suffix_argext if is_starts else _prefix_argext + return argext(arr=arr, booleans=booleans, is_max=is_max)[indexer] + def _sum_starts( arr: np.ndarray, @@ -162,7 +442,12 @@ def _min_starts( booleans: np.ndarray, ) -> tuple: """ - Compute min + Compute min. + + ELI5: for a handful of narrow ranges, just ask Rust to scan each one + directly. For lots of wide/overlapping ranges, it's cheaper to walk + the whole array once (`_suffix_argext`) and look up every row's answer. + See `_argext_dispatch` for the shared decision logic. """ mapping = { "int64": janitor_rs.compute_min_start_int64, @@ -176,6 +461,16 @@ def _min_starts( "float64": janitor_rs.compute_min_start_f64, "float32": janitor_rs.compute_min_start_f32, } + result = _argext_dispatch( + arr=arr, + booleans=booleans, + indexer=starts, + is_starts=True, + is_max=False, + mapping=mapping, + ) + if result is not None: + return result dtype_name = arr.dtype.name try: func = mapping[dtype_name] @@ -190,7 +485,12 @@ def _min_ends( booleans: np.ndarray, ) -> tuple: """ - Compute min + Compute min. + + ELI5: same trade-off as `_min_starts`, mirrored for prefixes -- Rust + for a few narrow ranges, `_prefix_argext`'s one-pass precompute once + there are enough wide/overlapping ones to make it worth it. See + `_argext_dispatch` for the shared decision logic. """ mapping = { "int64": janitor_rs.compute_min_end_int64, @@ -204,6 +504,16 @@ def _min_ends( "float64": janitor_rs.compute_min_end_f64, "float32": janitor_rs.compute_min_end_f32, } + result = _argext_dispatch( + arr=arr, + booleans=booleans, + indexer=ends, + is_starts=False, + is_max=False, + mapping=mapping, + ) + if result is not None: + return result dtype_name = arr.dtype.name try: func = mapping[dtype_name] @@ -218,7 +528,11 @@ def _max_starts( booleans: np.ndarray, ) -> tuple: """ - Compute max + Compute max. + + ELI5: same trade-off as `_min_starts`, just tracking the running + maximum instead of the minimum. See `_argext_dispatch` for the + shared decision logic. """ mapping = { "int64": janitor_rs.compute_max_start_int64, @@ -232,6 +546,16 @@ def _max_starts( "float64": janitor_rs.compute_max_start_f64, "float32": janitor_rs.compute_max_start_f32, } + result = _argext_dispatch( + arr=arr, + booleans=booleans, + indexer=starts, + is_starts=True, + is_max=True, + mapping=mapping, + ) + if result is not None: + return result dtype_name = arr.dtype.name try: func = mapping[dtype_name] @@ -246,7 +570,11 @@ def _max_ends( booleans: np.ndarray, ) -> tuple: """ - Compute max + Compute max. + + ELI5: same trade-off as `_min_ends`, just tracking the running + maximum instead of the minimum. See `_argext_dispatch` for the + shared decision logic. """ mapping = { "int64": janitor_rs.compute_max_end_int64, @@ -260,6 +588,16 @@ def _max_ends( "float64": janitor_rs.compute_max_end_f64, "float32": janitor_rs.compute_max_end_f32, } + result = _argext_dispatch( + arr=arr, + booleans=booleans, + indexer=ends, + is_starts=False, + is_max=True, + mapping=mapping, + ) + if result is not None: + return result dtype_name = arr.dtype.name try: func = mapping[dtype_name] diff --git a/pixi.lock b/pixi.lock index 8b8f1a44c..9c2eca750 100644 --- a/pixi.lock +++ b/pixi.lock @@ -22775,8 +22775,8 @@ packages: timestamp: 1774796815820 - pypi: ./ name: pyjanitor - version: 0.32.23 - sha256: 2c551db6e9b84114e74ba359a19fafa89130d1d5ce2b322f906d8b934e72926a + version: 0.32.24 + sha256: 048f055a9e41a71f61809b58114fe6dd79913a4232e76f8a5cf2e4b661ff05c5 requires_dist: - pandas>=3.0.0 - natsort>=8.4.0,<9 diff --git a/tests/functions/test_conditional_join_agg_min_max.py b/tests/functions/test_conditional_join_agg_min_max.py new file mode 100644 index 000000000..60bcda851 --- /dev/null +++ b/tests/functions/test_conditional_join_agg_min_max.py @@ -0,0 +1,402 @@ +"""Focused unit tests for the prefix/suffix min/max kernels backing +`join_agg`. + +Covers `_min_starts`, `_min_ends`, `_max_starts`, `_max_ends` in +`janitor.functions._conditional_join._agg_functions` (Issue #1653) -- the +O(n + m) NumPy replacements for the Rust `compute_min_start*`, +`compute_min_end*`, `compute_max_start*`, and `compute_max_end*` kernels, +used once the number/width of requested ranges makes precomputing worth it +(see `_use_argext`); sparse range requests continue to use the Rust +kernels directly. +""" + +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from janitor.functions._conditional_join import _agg_functions + +ALL_DTYPES = [ + "int64", + "int32", + "int16", + "int8", + "uint64", + "uint32", + "uint16", + "uint8", + "float64", + "float32", +] +FLOAT_DTYPES = ["float64", "float32"] + + +def naive_suffix(arr, booleans, start, is_max): + """Direct transliteration of the Rust min_starts/max_starts loop.""" + n = len(arr) + if start >= n: + return -1 + base = -1 + base_val = arr[start] + for nn in range(start, n): + if booleans[nn]: + continue + current = arr[nn] + if base == -1 or (current > base_val if is_max else current < base_val): + base_val = current + base = nn + return base + + +def naive_prefix(arr, booleans, end, is_max): + """Direct transliteration of the Rust min_ends/max_ends loop.""" + n = len(arr) + if n == 0: + return -1 + base = -1 + base_val = arr[0] + for nn in range(0, end): + if booleans[nn]: + continue + current = arr[nn] + if base == -1 or (current > base_val if is_max else current < base_val): + base_val = current + base = nn + return base + + +def _dense_query(n, size=200): + """Enough queries/width to clear `_use_argext`'s work-factor gate.""" + rng = np.random.default_rng(0) + return rng.integers(0, n if n else 1, size=size).astype("int64") + + +@pytest.mark.parametrize("dtype", ALL_DTYPES) +@pytest.mark.parametrize("is_max", [True, False]) +def test_starts_matches_naive_dense(dtype, is_max): + """Dense queries (forces the NumPy path) match the naive Rust-loop + reference, including ties and nulls.""" + rng = np.random.default_rng(hash((dtype, is_max)) % (2**32)) + n = 40 + if dtype.startswith("float"): + arr = rng.integers(-5, 6, size=n).astype(dtype) + arr[rng.random(n) < 0.15] = np.nan + else: + info = np.iinfo(dtype) + choices = np.array( + [info.min, info.min + 1, info.max - 1, info.max, 0, 1], dtype=dtype + ) + arr = choices[rng.integers(0, len(choices), size=n)] + booleans = rng.random(n) < 0.2 + starts = _dense_query(n) + func = _agg_functions._max_starts if is_max else _agg_functions._min_starts + + expected = np.array( + [naive_suffix(arr, booleans, int(s), is_max) for s in starts], dtype=np.int64 + ) + actual = func(arr=arr, starts=starts, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", ALL_DTYPES) +@pytest.mark.parametrize("is_max", [True, False]) +def test_ends_matches_naive_dense(dtype, is_max): + rng = np.random.default_rng(hash((dtype, is_max, "ends")) % (2**32)) + n = 40 + if dtype.startswith("float"): + arr = rng.integers(-5, 6, size=n).astype(dtype) + arr[rng.random(n) < 0.15] = np.nan + else: + info = np.iinfo(dtype) + choices = np.array( + [info.min, info.min + 1, info.max - 1, info.max, 0, 1], dtype=dtype + ) + arr = choices[rng.integers(0, len(choices), size=n)] + booleans = rng.random(n) < 0.2 + ends = rng.integers(0, n + 1, size=200).astype("int64") + func = _agg_functions._max_ends if is_max else _agg_functions._min_ends + + expected = np.array( + [naive_prefix(arr, booleans, int(e), is_max) for e in ends], dtype=np.int64 + ) + actual = func(arr=arr, ends=ends, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", ALL_DTYPES) +def test_all_null_range_is_minus_one(dtype): + arr = np.array([1, 2, 3, 4] * 15, dtype=dtype) + booleans = np.ones(60, dtype=bool) + starts = _dense_query(60) + ends = np.arange(60, dtype="int64") + + assert ( + _agg_functions._min_starts(arr=arr, starts=starts, booleans=booleans) == -1 + ).all() + assert (_agg_functions._min_ends(arr=arr, ends=ends, booleans=booleans) == -1).all() + assert ( + _agg_functions._max_starts(arr=arr, starts=starts, booleans=booleans) == -1 + ).all() + assert (_agg_functions._max_ends(arr=arr, ends=ends, booleans=booleans) == -1).all() + + +@pytest.mark.parametrize("dtype", ALL_DTYPES) +def test_empty_range_is_minus_one(dtype): + # Mix a few empty-range queries into a dense batch of real ones, so the + # whole call routes through the NumPy path (which handles start == n / + # end == 0 gracefully). An all-empty batch would instead route to the + # Rust fallback (total_width is always 0 there), and the existing Rust + # kernel panics on `arr[start_]` when start_ == n -- a pre-existing + # issue in the Rust kernel itself, unrelated to this change, and not a + # shape of input this suite needs to force through that path. + n = 60 + arr = np.array(list(range(60)), dtype=dtype) + booleans = np.zeros(n, dtype=bool) + starts = _dense_query(n) + starts[:5] = n # empty suffix + ends = _dense_query(n) + ends[:5] = 0 # empty prefix + + assert ( + _agg_functions._min_starts(arr=arr, starts=starts, booleans=booleans)[:5] == -1 + ).all() + assert ( + _agg_functions._min_ends(arr=arr, ends=ends, booleans=booleans)[:5] == -1 + ).all() + assert ( + _agg_functions._max_starts(arr=arr, starts=starts, booleans=booleans)[:5] == -1 + ).all() + assert ( + _agg_functions._max_ends(arr=arr, ends=ends, booleans=booleans)[:5] == -1 + ).all() + + +@pytest.mark.parametrize("dtype", FLOAT_DTYPES) +def test_ties_keep_first_occurrence(dtype): + """Duplicate extrema resolve to the smaller/earlier position.""" + arr = np.array([5, 1, 3, 1, 1, 9, 1], dtype=dtype) + booleans = np.zeros(7, dtype=bool) + starts = _dense_query(7) + ends = np.arange(8, dtype="int64").repeat(30)[:200] + + expected_starts = np.array( + [naive_suffix(arr, booleans, int(s), False) for s in starts], dtype=np.int64 + ) + expected_ends = np.array( + [naive_prefix(arr, booleans, int(e), False) for e in ends], dtype=np.int64 + ) + np.testing.assert_array_equal( + _agg_functions._min_starts(arr=arr, starts=starts, booleans=booleans), + expected_starts, + ) + np.testing.assert_array_equal( + _agg_functions._min_ends(arr=arr, ends=ends, booleans=booleans), expected_ends + ) + + +@pytest.mark.parametrize("dtype", FLOAT_DTYPES) +def test_nan_freezes_only_when_first_in_range(dtype): + """A NaN that is the first non-null value of a *specific row's* range + poisons that row's result (Rust: nothing compares less/greater than + NaN); a NaN found later in the same row's range is just skipped.""" + arr = np.array([3.0, np.nan, 1.0], dtype=dtype) + booleans = np.zeros(3, dtype=bool) + # pad with dense filler queries so this exercises the NumPy path too + filler = _dense_query(3) + starts = np.concatenate([[0, 1, 2], filler]) + + expected = np.array( + [naive_suffix(arr, booleans, int(s), False) for s in starts], dtype=np.int64 + ) + actual = _agg_functions._min_starts(arr=arr, starts=starts, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + # hand-verified contract for the first three (non-filler) queries: + # start=0 -> real anchor (3.0), skips the NaN, finds 1.0 at index 2 + # start=1 -> anchor IS the NaN itself -> frozen there, index 1 + # start=2 -> real anchor (1.0), index 2 + assert actual[0] == 2 + assert actual[1] == 1 + assert actual[2] == 2 + + +@pytest.mark.parametrize("dtype", FLOAT_DTYPES) +def test_nan_first_in_prefix_poisons_rest(dtype): + """Prefix scans always start at index 0, so if the array's first + non-null value is NaN, every end past that point is frozen there.""" + arr = np.array([np.nan, 1.0, 2.0, 3.0], dtype=dtype) + booleans = np.zeros(4, dtype=bool) + ends = np.array([0, 1, 2, 3, 4], dtype="int64") + filler = np.full(200, 4, dtype="int64") + ends = np.concatenate([ends, filler]) + + actual = _agg_functions._min_ends(arr=arr, ends=ends, booleans=booleans) + assert actual[0] == -1 # end=0 -> empty prefix + assert (actual[1:5] == 0).all() # frozen at the NaN's own position + assert (actual[5:] == 0).all() + + +@pytest.mark.parametrize( + "func_name,rust_name,indexers", + [ + ("_min_starts", "compute_min_start_int64", {"starts": [5]}), + ("_min_ends", "compute_min_end_int64", {"ends": [1]}), + ("_max_starts", "compute_max_start_int64", {"starts": [5]}), + ("_max_ends", "compute_max_end_int64", {"ends": [1]}), + ], +) +def test_sparse_ranges_use_rust(monkeypatch, func_name, rust_name, indexers): + """A handful of narrow queries should not pay for the full precompute.""" + expected = np.array([7], dtype=np.int64) + + def fake_rust(**kwargs): + return expected + + monkeypatch.setattr(_agg_functions.janitor_rs, rust_name, fake_rust) + monkeypatch.setattr( + _agg_functions, + "_suffix_argext" if "start" in func_name else "_prefix_argext", + lambda **kwargs: pytest.fail("sparse ranges should stay in Rust"), + ) + actual = getattr(_agg_functions, func_name)( + arr=np.arange(100, dtype=np.int64), + booleans=np.zeros(100, dtype=bool), + **{name: np.array(values, dtype=np.int64) for name, values in indexers.items()}, + ) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + "func_name,rust_name,indexers,expected", + [ + ("_min_starts", "compute_min_start_int64", {"starts": [0] * 200}, 0), + ("_min_ends", "compute_min_end_int64", {"ends": [100] * 200}, 0), + ("_max_starts", "compute_max_start_int64", {"starts": [0] * 200}, 99), + ("_max_ends", "compute_max_end_int64", {"ends": [100] * 200}, 99), + ], +) +def test_dense_ranges_use_numpy(monkeypatch, func_name, rust_name, indexers, expected): + """Heavily overlapping/wide ranges should use the O(n) precompute. + + `arr = arange(100)`, so over the full range [0, 100) the min is 0 (at + index 0) and the max is 99 (at index 99). + """ + + def fail_rust(**kwargs): + pytest.fail("dense ranges should use the NumPy prefix/suffix path") + + monkeypatch.setattr(_agg_functions.janitor_rs, rust_name, fail_rust) + actual = getattr(_agg_functions, func_name)( + arr=np.arange(100, dtype=np.int64), + booleans=np.zeros(100, dtype=bool), + **{name: np.array(values, dtype=np.int64) for name, values in indexers.items()}, + ) + assert (actual == expected).all() + + +@pytest.mark.parametrize( + "func_name,indexer_name,indexer_value", + [ + ("_min_starts", "starts", 0), + ("_min_ends", "ends", 100), + ("_max_starts", "starts", 0), + ("_max_ends", "ends", 100), + ], +) +@pytest.mark.parametrize("query_count", [1, 200]) +def test_unsupported_dtype_error_does_not_depend_on_density( + func_name, indexer_name, indexer_value, query_count +): + """The performance gate must not silently expand supported dtypes.""" + arr = np.ones(100, dtype=bool) + booleans = np.zeros(100, dtype=bool) + indexers = np.full(query_count, indexer_value, dtype=np.int64) + + with pytest.raises(KeyError, match="Unsupported data type -> bool"): + getattr(_agg_functions, func_name)( + arr=arr, + booleans=booleans, + **{indexer_name: indexers}, + ) + + +@pytest.mark.parametrize("is_max", [True, False]) +def test_suffix_uses_min_or_max_accumulate_semantics(is_max): + """Sanity check that _max_* actually picks the maximum, not the + minimum (i.e. the is_max plumbing is wired correctly end to end).""" + arr = np.array([1, 9, 2, 9, 3], dtype="int64") + booleans = np.zeros(5, dtype=bool) + starts = _dense_query(5) + func = _agg_functions._max_starts if is_max else _agg_functions._min_starts + expected = np.array( + [naive_suffix(arr, booleans, int(s), is_max) for s in starts], dtype=np.int64 + ) + actual = func(arr=arr, starts=starts, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +# --------------------------------------------------------------------------- +# Property-based tests: the real correctness gate for the NaN/tie/null +# semantics above, which are easy to get subtly wrong by hand. +# --------------------------------------------------------------------------- + +_HYP_DTYPES = ["int64", "int8", "uint64", "uint8", "float64", "float32"] + + +@st.composite +def _array_and_booleans(draw, dtype): + n = draw(st.integers(min_value=0, max_value=25)) + if dtype.startswith("float"): + values = draw( + st.lists( + st.one_of( + st.floats( + min_value=-5, max_value=5, allow_nan=False, allow_infinity=False + ), + st.just(float("nan")), + ), + min_size=n, + max_size=n, + ) + ) + arr = np.array(values, dtype=dtype) + else: + info = np.iinfo(dtype) + values = draw( + st.lists( + st.sampled_from( + [info.min, info.min + 1, -2, -1, 0, 1, 2, info.max - 1, info.max] + ), + min_size=n, + max_size=n, + ) + ) + arr = np.array([v for v in values if info.min <= v <= info.max], dtype=dtype) + n = arr.size + booleans = np.array(draw(st.lists(st.booleans(), min_size=n, max_size=n))) + return arr, booleans + + +@given(dtype=st.sampled_from(_HYP_DTYPES), is_max=st.booleans(), data=st.data()) +@settings(max_examples=300, deadline=None) +def test_prefix_suffix_property_against_naive(dtype, is_max, data): + arr, booleans = data.draw(_array_and_booleans(dtype)) + n = arr.size + + pre = _agg_functions._prefix_argext(arr=arr, booleans=booleans, is_max=is_max) + suf = _agg_functions._suffix_argext(arr=arr, booleans=booleans, is_max=is_max) + + for e in range(n + 1): + assert pre[e] == naive_prefix(arr, booleans, e, is_max), ( + arr, + booleans, + "end", + e, + ) + for s in range(n + 1): + assert suf[s] == naive_suffix(arr, booleans, s, is_max), ( + arr, + booleans, + "start", + s, + )