Skip to content

Improve performance and typing in _protect_dataset_variables_inplace - #9069

Draft
Illviljan wants to merge 2 commits into
pydata:mainfrom
Illviljan:type_protect_dataset_variables_inplace
Draft

Illviljan wants to merge 2 commits into
pydata:mainfrom
Illviljan:type_protect_dataset_variables_inplace

Conversation

@Illviljan

@Illviljan Illviljan commented Jun 4, 2024 •

Copy link
Copy Markdown
Contributor

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
Comment thread xarray/backends/api.py Resolved
@Illviljan Illviljan added the run-benchmark Run the ASV benchmark workflow label Jun 4, 2024
@Illviljan

Illviljan commented Jun 4, 2024 •

Copy link
Copy Markdown
Contributor Author
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

This branch has not been deployed

No deployments
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