Skip to content

A1: Vectorize _codes_to_group_indices #1

Description

@FBumann

Problem

xarray/core/groupby.py:91 _codes_to_group_indices is a Python-level loop that builds per-group index lists via list.append. Profile (200k rows, 1k groups) shows it consumes 89% of groupby() setup time:

89.00% 1739 _codes_to_group_indices xarray/core/groupby.py
3.48% 68 factorize_array pandas/core/algorithms.py

200,001 list.append calls visible in the profile.

Proposed fix

Replace the loop with vectorized NumPy:

  • order = np.argsort(codes, kind="stable")
  • split_points = np.searchsorted(codes[order], np.arange(n_groups))
  • groups = np.split(order, split_points[1:])

Or use np.unique(codes, return_inverse=True) + a single bincount + argpartition pass.

Expected outcome

~10× faster groupby construction for moderate-to-large group counts. Most user-visible in time_init benchmarks and as setup cost before .mean()/.sum() agg calls.

Benchmark

Existing `asv_bench/benchmarks/groupby.py::GroupBy.time_init` covers low-group-count case. Add a high-cardinality variant (200k rows, 1k groups) as the realistic regression target.

Acceptance

  • All existing groupby tests pass
  • New asv benchmark shows the speedup
  • No behavioral change for empty groups, NaN codes, or unordered codes

[This is Claude Code on behalf of Felix Bumann]

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions