Skip to content

A4: Implement reshape fast path for _unstack_once #9

Description

@FBumann

Finding

Investigation under #3 isolated where the 48% of multiindex_stack workload time inside _unstack_once actually goes. Result:

variable shape: (20, 40000) dtype: float64
multiindex levels: [200, 200] is_monotonic: True is_unique: True
new_shape: (20, 200, 200) is_missing: False
Time per step (over 20 iterations):
transpose only : 0.002ms
allocate template only : 0.001ms
allocate + scatter write : 0.920ms ← all of it
(alternative) reshape : 0.000ms ← essentially free
full _unstack_once : 0.912ms
Reshape == scatter result? True

100% of the time is in data[(..., *indexer)] = reordered — numpy advanced-indexing scatter. But for the common case (lex-sorted MultiIndex, no missing combinations — i.e. a normal .stack(...).unstack(...) round trip) the codes form a row-major sequence and a plain reshape produces the same array.

Proposed fix

Two-line guard in xarray/core/variable.py::_unstack_once:

ifnotsparseandnotis_missing_valuesandindex.is_monotonic_increasing:
# Complete lex-sorted Cartesian product — codes are row-major,# so reshape produces the same result as the scatter write.data=reordered.data.reshape(new_shape)
else:
data=create_template(self.data, shape=new_shape, dtype=dtype)
data[(..., *indexer)] =reordered.data

Existing scatter path stays as the fallback for sparse output, missing combinations, or a user-constructed non-lex-sorted MultiIndex.

Expected impact

  • _unstack_once: ~900× speedup on stack→unstack round trips (the common case)
  • multiindex_stack workload: ~48% of total time eliminated
  • Zero behavior change on the fallback paths

Acceptance

  • Existing variable + multiindex tests pass
  • New test covers both fast path (lex-sorted, complete) and fallback (missing or non-lex-sorted) producing identical results
  • asv benchmark for stack/unstack shows the speedup

Closes #3.


[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