Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve performance and typing in _protect_dataset_variables_inplace #9069

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

Illviljan
Copy link
Contributor

@Illviljan Illviljan commented Jun 4, 2024

Improve performance by removing type and shape checking in _protect_dataset_variables_inplace by setting directly to self._data.
The function already has valid variables, there should be no need to recheck when adding small wrapper classes like CopyOnWriteArray and MemoryCachedArray.

variable.data = data does this, and as_compatible_data is a known performance bottleneck:

@data.setter
def data(self, data: T_DuckArray | ArrayLike) -> None:
data = as_compatible_data(data)
self._check_shape(data)
self._data = data

Add some typing to make sure strange types are not input to these functions.

@Illviljan Illviljan marked this pull request as draft June 4, 2024 19:33
@Illviljan Illviljan added the run-benchmark Run the ASV benchmark workflow label Jun 4, 2024
@Illviljan
Copy link
Contributor Author

Illviljan commented Jun 4, 2024

Found 2 errors in 1 file (checked 171 source files)
xarray/backends/api.py:241: error: Incompatible types in assignment (expression has type "MemoryCachedArray", variable has type "_arrayfunction[Any, Any] | _arrayapi[Any, Any]")  [assignment]
xarray/backends/api.py:243: error: Incompatible types in assignment (expression has type "CopyOnWriteArray", variable has type "_arrayfunction[Any, Any] | _arrayapi[Any, Any]")  [assignment]

And here's the rabbit hole... :(
I suspect it probably is not as easy to just add some methods to these classes.

@Illviljan Illviljan changed the title Type protect dataset variables inplace Improve performance and typing in protect_ dataset_variables_inplace Jun 4, 2024
@Illviljan Illviljan changed the title Improve performance and typing in protect_ dataset_variables_inplace Improve performance and typing in _protect_dataset_variables_inplace Jun 4, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
run-benchmark Run the ASV benchmark workflow
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants