Skip to content

[PERF] Add guarded fast paths for cartesian_product and expand #1670

Description

@samukweku

ELI5 summary

cartesian_product currently builds a full grid of row numbers first, then uses that grid to copy the real values. For a large NumPy-only result, we can copy the values directly. It is like laying out the books themselves instead of first laying out millions of sticky notes saying which book goes where.

expand has two additional costs:

  1. It uses drop_duplicates() when it only needs the unique values from each scalar column.
  2. Grouped expansion runs one Python loop per group. Thousands of tiny groups spend much more time entering/leaving Python than building results.

The useful changes should be separate, guarded PRs. None should replace the general dtype-sensitive path.

Environment

  • pyjanitor 8a3b83e
  • Python 3.14.5, pandas 3.0.3, NumPy 2.3.5
  • macOS arm64
  • Medians of 3-7 warm runs unless noted; peak RSS uses a fresh process/one run

1. Direct NumPy materialization in cartesian_product

For NumPy-backed columns, replace the shared np.indices(lengths) matrix with per-column values built from np.repeat/np.tile.

Example for x=[1,2], y=[3,4,5]:

x_out = np.repeat(x, 3)  # [1,1,1,2,2,2]
y_out = np.tile(y, 2)    # [3,4,5,3,4,5]

This preserves the current element-for-element order without retaining a (number_of_inputs, output_rows) integer matrix.

Shape Current Prototype Result
2 x 3 (6 rows) 0.155 ms 0.156 ms neutral
100 x 100 (10k rows) 0.170 ms 0.183 ms 8% slower
~100k rows, 2 cols 0.354 ms 0.221 ms 37% faster
1M rows, 2 cols 1.72 ms 1.47 ms 15% faster
10M rows, 2 cols 25.9 ms 15.5 ms 40% faster
ten factors, 4^10 rows, 10 cols 13.47 ms 5.05 ms 2.7x faster
1M rows, 64 cols 49.1 ms 36.8 ms 25% faster

Fresh-process 10M-row peak RSS fell from about 306 MiB to 154 MiB. End-to-end expand producing 10.08M x 3 fell from 39.1 to 18.9 ms and about 454 to 152 MiB peak RSS.

Guardrails:

  • Use only when every materialized column is NumPy-backed and the projected product is large enough (the measured crossover was around 100k output rows).
  • Keep the current path for extension arrays, MultiIndex, and mixed inputs. Direct generation regressed string extension input 5.27 -> 6.80 ms and categorical input 1.32 -> 2.46 ms in the prototype.
  • Handle zero-length inputs without dividing by a factor length.
  • Avoid redundant np.tile(..., 1) / np.repeat(..., 1) copies.

2. Use unique() for safe scalar-column expand specs

For a scalar spec such as df.expand("x", "y"), the original row index discarded by drop_duplicates() is never used. A named Series around df[column].unique() is faster and retains order/dtype.

On 1M input rows with ten unique values per column:

Dtype drop_duplicates() path unique() prototype
int64 5.14 ms 3.21 ms
float64 6.17 ms 4.26 ms
pandas string 18.13 ms 10.20 ms
category 5.42 ms 3.16 ms
nullable Int64 5.75 ms 3.31 ms

Fresh-process peak RSS for int fell about 3.0 -> 1.0 MiB; pandas string fell about 21.9 -> 0.9 MiB. Combined with direct NumPy Cartesian generation, a 1M-input/1M-output numeric case improved 9.12 -> 6.16 ms.

Guardrail: retain drop_duplicates() for object dtype. drop_duplicates() accepts unhashable objects such as lists, while Series.unique() raises TypeError.

3. Vectorize supported grouped expand

For plain scalar/list column specs and column-based groups, build each per-group distinct table once and join the tables on the group keys. This removes one Python call to _compute_cartesian_product per group.

Shape Current loop Join prototype
10 groups, 90 output rows 0.88 ms 2.10 ms (regression)
50 groups, 5k rows 2.78 ms 2.20 ms
1k groups, 1M x 3 67.3 ms 25.6 ms
10k groups, 40k x 2, int 484 ms 4.56 ms
same, pandas string 1098 ms 4.59 ms
same, category 615 ms 4.14 ms
10 groups, 100k x 64 32.2 ms 9.75 ms

The join path is not universally better. On the narrow 1M x 3 case, clean peak RSS rose from roughly 39 to 61 MiB. On the wide 100k x 64 case it fell from about 96 to 45 MiB. A guarded heuristic should therefore include group count, output width/cardinality, and memory, rather than dispatching every grouped call to joins.

Required fallbacks:

  • callable, dictionary, explicit Series/Index/DataFrame, or expression specs;
  • grouping by an index level or external grouper instead of columns;
  • specs overlapping group-key columns;
  • unobserved categorical groups (observed=False), unless zero-sized groups are explicitly removed while preserving current ordering.

Compatibility evidence

  • Direct NumPy Cartesian prototype: 222 exact assert_frame_equal comparisons across NumPy dtypes, object arrays, datetime/timedelta, sort modes, DataFrames, tuple columns, empty inputs, and extension fallbacks.
  • Scalar unique() prototype: 3,380 exact comparisons across empty/small/tall inputs, sort modes, scalar/tuple labels, NumPy, timezone datetime, period, categorical, nullable, and string dtypes.
  • Grouped join prototype: 2,880 exact comparisons across one/two group keys, NA groups, both group/expand sort modes, scalar/nested specs, and six payload dtype families.
  • Existing suites: test_cartesian_product.py: 9 passed, 1 xpassed; test_expand.py: 6 passed.

Rejected experiments

  • A smaller per-factor integer indexer reduced memory only for high-dimensional products, but was 16-73% slower on common two-factor 1M-10M shapes.
  • pandas cross merges were about 9-10x slower and used more memory.
  • MultiIndex.from_product().to_frame() was about 2x slower at 1M x 2 and did not help the ten-factor case.
  • Batched grouped joins lost much of the speedup and did not reliably reduce peak memory.

Suggested PR order

  1. Scalar non-object unique() fast path: smallest change and broad tall-frame win.
  2. Large all-NumPy direct Cartesian path with an explicit crossover threshold.
  3. Grouped join path only after a conservative dispatch heuristic and tests for every fallback above.

Each PR should preserve output values, ordering, labels, index, dtypes, categorical metadata, and current exceptions exactly.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions