[PERF] Speed up complete for unsorted domains and many groups

[samukweku]

Summary

complete is already fast for ordinary ungrouped, sorted, two-key inputs because pandas' outer merge dominates and is well optimized. Three narrower changes are worth pursuing:

1. Pre-sort unsorted completion domains before constructing a large Cartesian grid.
2. Vectorize plain-column grouped expansion when the group count is large, while retaining the current loop for small groups and dynamic specifications.
3. Restrict dictionary fill_value work to the requested columns; this also fixes a current KeyError.

The shared np.indices Cartesian-product allocation should be evaluated separately in expand/cartesian_product, rather than counting it as a complete-specific win.

stack / unstack was tested and rejected as the general engine

For unique two-key inputs, set_index(...).unstack().stack(future_stack=True) produced the same values and order, but it was much slower:

1M-row completed output	Outer merge	Stack/unstack
1 payload column	17.6 ms	112.9 ms
32 payload columns	45.7 ms	339.6 ms
unsorted domains, 1 payload column	102.2 ms	117.7 ms

It also raises ValueError: Index contains duplicate entries for duplicate completion keys, which complete currently supports. With three or more independent completion keys, unstacking the remaining levels only completes observed tuples, not the full cross-product; explicitly reindexing to that product was also slower than the merge. Stack/unstack is useful as a correctness oracle in tests, but not as the production fast path.

1. Pre-sort large unsorted domains

_computations_complete passes sort into expand, then always uses an outer merge. With pandas 3, an outer merge sorts the union keys lexicographically even when sort=False. If the domains arrive unsorted, pandas ends up sorting the large completed grid during the merge. Sorting the small unique domains before constructing that grid is substantially cheaper.

Median results on pandas 3.0.3:

Input and completed output	Current sort=False	Pre-sorted domains	Change
1,000 input rows, 1M-row output, 1 payload column	101.4 ms	55.5 ms	45% faster
same, 32 payload columns	130.3 ms	85.1 ms	35% faster
1,000 input rows, 1M-row output, 3 completion keys	94.5 ms	64.1 ms	32% faster
10,000 input rows, 10M-row output	1,502 ms	747 ms	2.0x faster
10M-row tall input, 100 x 100 domains	323 ms	322 ms	neutral
1,000-column wide payload, 10.9K-row output	19.57 ms	19.57 ms	neutral

Peak RSS for the 10M-row output fell from about 718 MiB to 573 MiB. Outputs were element-for-element identical in 432 comparisons spanning tiny through 1,000-row inputs; plain, nested, external, and grouped specifications; integer, float/null, string, categorical, nullable integer, and datetime keys; and both fill modes.

Suggested implementation: inside complete, sort only non-monotonic domain objects before Cartesian expansion when the projected product is large. Preserve the public sort contract and verify against pandas 3 outer-merge ordering. Avoid sorting the fully materialized grid.

2. Vectorize grouped expansion for many groups

_expand_groupby currently loops over every group in Python, builds a Cartesian product for each frame, stores arrays in lists, then concatenates them. Runtime grows sharply with group count:

100K-row input	Current	Vectorized plain-column prototype	Change
10 groups	6.6 ms	7.6 ms	16% slower
100 groups	13.0 ms	8.6 ms	34% faster
1,000 groups	70.3 ms	15.1 ms	4.7x faster
10,000 groups, 1M-row output	590 ms	83.8 ms	7.0x faster
10,000 groups, 20 x 20 domains	1,549 ms	91 ms	17x faster
1,000 groups, 32 payload columns	245 ms	32.6 ms	7.5x faster
1,000 groups, 100 x 100 domains, 10M-row output	926 ms	935 ms	neutral

The prototype builds each observed (group keys + completion object) table with drop_duplicates, then performs sequential many-to-many merges on the group keys. It matched the current output exactly in 864 comparisons covering 1–1,000 rows, 1–7 groups, scalar and nested specifications, both sort modes, both group-order modes, and NumPy/string/categorical/nullable/datetime dtypes.

This needs a guarded fast path. Small group counts can regress by 15–25%, and dictionaries/callables must retain the current per-group evaluation semantics. A structural cutoff based on group count and projected product should be benchmarked rather than applying the vectorized path universally.

3. Fix and narrow dictionary fills

The current code first scans every non-key output column with hasnans, then indexes the user dictionary for every null-bearing column. This fails when the dictionary intentionally names only a subset:

df = pd.DataFrame({"a": [1], "b": [1], "v0": [1.0], "v1": [2.0]})
df.complete(
    {"a": [1, 2]},
    {"b": [1, 2]},
    fill_value={"v0": 0},
)
# KeyError: 'v1'

For a 100K x 512 payload frame, scanning all columns before filling one requested column took 13.2 ms; selecting only the dictionary key and calling fillna took 0.23 ms with the same result. The same unnecessary scan costs 13.4 ms even when the frame has no nulls.

Suggested fix:

• For dictionary fills, intersect its keys with non-merge columns first and inspect/fill only those columns.
• Keep the current no-null short circuit for scalar fills; directly filling all columns is 10–14% faster when nulls exist but slower when none exist.

Also replace indicator = "".join(columns) with a collision-safe label generator. explicit=False currently raises TypeError when any source column label is non-string.

Additional cleanup

Calling df.groupby(...).complete(...) currently emits the warning that the deprecated by= argument was used. The warning condition should distinguish an actual by argument from an already-grouped input.

Verification performed

• tests/functions/test_complete.py: 31 passed, 1 expected failure.
• Current behavior profiled from 1-row/tiny cases through 10M-row tall inputs, 10M-row completed outputs, 1,000-column payloads, 10 completion dimensions, and up to 10,000 groups.
• Time measurements are medians of repeated isolated runs; peak memory used sampled process RSS and, where useful, tracemalloc.