Skip to content

ENH: Support DataTree in apply_ufunc - #11552

Merged
shoyer merged 2 commits into
pydata:mainfrom
QinXi-ai:feat/datatree-apply-ufunc
Sep 9, 2026
Merged

shoyer merged 2 commits into
pydata:mainfrom
QinXi-ai:feat/datatree-apply-ufunc

Conversation

@QinXi-ai

Copy link
Copy Markdown
Contributor

Closes #9789

Problem

apply_ufunc currently treats DataTree as a dict-like object. Unary calls silently return only the root Dataset, while core-dimension and multi-tree calls can fail when child nodes reach the variable dispatch path.

Solution

  • Dispatch DataTree inputs before the dict-like Dataset path.
  • Reuse map_over_datasets to apply the existing apply_ufunc machinery at corresponding nodes, preserving tree structure, names, inherited coordinates, scalar arguments, kwargs, core dimensions, and tuple outputs.
  • Enable the public xarray.ufuncs wrappers to operate over complete trees.
  • Add regressions for identity, inherited coordinates, core dimensions plus kwargs/scalars, two isomorphic trees, non-isomorphic errors, multiple outputs, and the ufunc wrapper.

Tests

  • python -m pytest -o addopts='' --basetemp D:/CodexCaches/pytest-xarray-9789-full xarray/tests/test_computation.py xarray/tests/test_ufuncs.py xarray/tests/test_datatree_mapping.py -q (323 passed, 131 skipped, 1 xfailed, 2 xpassed)
  • ruff check xarray/computation/apply_ufunc.py xarray/ufuncs.py xarray/tests/test_computation.py xarray/tests/test_ufuncs.py
  • ruff format --check xarray/computation/apply_ufunc.py xarray/ufuncs.py xarray/tests/test_computation.py xarray/tests/test_ufuncs.py
  • git diff --check

@shoyer shoyer left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm sure there are some edge cases to figure out, but this looks good to me for now!

@shoyer
shoyer enabled auto-merge (squash) September 9, 2026 15:50
@shoyer
shoyer disabled auto-merge September 9, 2026 16:25
@shoyer
shoyer merged commit 503b262 into pydata:main Sep 9, 2026
43 of 44 checks passed
@welcome

welcome Bot commented Sep 9, 2026

Copy link
Copy Markdown

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

TomNicholas added a commit that referenced this pull request Sep 29, 2026
* Update whats-new for v2026.09.0 release

- Rename the unreleased section to v2026.09.0 and add the release summary
  and contributor list
- Move entries for #11348, #11292, #11290 and #11239 that were merged into
  older release sections
- Add missing entries for #11552, #11521, #11486 and #11513
- Add missing PR links, fix the #11547 link and other markup
- Fix the documented default of the use_bottleneck option
- Add contributor names to the typos allow-list

* Mention plot.lines and last Python 3.11 release in whats-new
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support DataTree in apply_ufunc

2 participants