Skip to content

Implement preferred_chunks for netcdf 4 backends#7948

Merged
Illviljan merged 21 commits intopydata:mainfrom
mraspaud:feature-nc-preferred-chunks
Sep 11, 2023
Merged

Implement preferred_chunks for netcdf 4 backends#7948
Illviljan merged 21 commits intopydata:mainfrom
mraspaud:feature-nc-preferred-chunks

Conversation

@mraspaud
Copy link
Contributor

@mraspaud mraspaud commented Jun 28, 2023 •
edited by dcherian
Loading

According to the open_dataset documentation, using chunks="auto" or chunks={} should yield datasets with variables chunked depending on the preferred chunks of the backend. However neither the netcdf4 nor the h5netcdf backend seem to implement the preferred_chunks encoding attribute needed for this to work.

This PR adds this attribute to the encoding upon data reading. This results in chunks="auto" in open_dataset returning variables with chunk sizes multiples of the chunks in the nc file, and for chunks={}, returning the variables with then exact nc chunk sizes.

@mraspaud
Copy link
Contributor Author

mraspaud commented Jun 28, 2023

Hope this is up to the standards! Please tell me if there is anything I can do to improve this PR.

@mraspaud mraspaud changed the title Implement peferred_chunks for netcdf 4 backend Implement preferred_chunks for netcdf 4 backend Jun 28, 2023
@mraspaud mraspaud changed the title Implement preferred_chunks for netcdf 4 backend Implement preferred_chunks for netcdf 4 backends Jun 28, 2023
@ghiggi ghiggi mentioned this pull request Jun 28, 2023
Closed
@djhoese
Copy link
Contributor

djhoese commented Jun 28, 2023

Is the resulting chunk size a multiple of the on-disk file chunk size or is it exactly the file chunk size? So if I open a file with a file chunk size of 256, do I get a dask array with chunk size 256 or some larger chunk size that is a multiple of 256. By itself, 256 would perform worse than one large chunk for a lot of arrays sizes just from the overhead.

@mraspaud
Copy link
Contributor Author

mraspaud commented Jun 28, 2023

@djhoese if chunks={}, the same chunks size as the on-disk file will be used. if chunks="auto", a multiple of the on-disk chunk size will be used to accommodate the dask chunk size preferences. This is in line with what is said in the open_dataset docstring as I understand it.
For our case, using chunks="auto" with this patch resulted in a significant performance boost in processing the data as compared to the main branch.

jhamman
jhamman reviewed Jun 28, 2023
View reviewed changes
Copy link
Member

@jhamman jhamman left a comment

Choose a reason for hiding this comment

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

Thanks for taking this on @mraspaud. This has been something folks have wanted for some time now!

@Illviljan Illviljan added the run-benchmark Run the ASV benchmark workflow label Jun 28, 2023
jhamman
jhamman reviewed Jun 30, 2023
View reviewed changes
@dcherian dcherian requested a review from jhamman July 3, 2023 15:40
@dcherian
Copy link
Contributor

dcherian commented Jul 26, 2023 

Found 5 errors in 1 file (checked 142 source files)
xarray/tests/test_backends.py:1568: error: Need type annotation for "tmp_file"  [var-annotated]
xarray/tests/test_backends.py:1577: error: Argument 1 to "contextmanager" has incompatible type "Callable[[NetCDF4Base, Any, Any], None]"; expected "Callable[[NetCDF4Base, Any, Any], Iterator[<nothing>]]"  [arg-type]
xarray/tests/test_backends.py:1578: error: The return type of a generator function should be "Generator" or one of its supertypes  [misc]
xarray/tests/test_backends.py:1599: error: Need type annotation for "tmp_file"  [var-annotated]

@Illviljan @headtr1ck help!

headtr1ck
headtr1ck reviewed Jul 26, 2023
View reviewed changes
xarray/tests/test_backends.py Outdated
assert all(np.asanyarray(chunksizes) == expected)

@contextlib.contextmanager
def create_chunked_file(self, array_shape, chunk_sizes) -> None:
Copy link
Collaborator

@headtr1ck headtr1ck Jul 26, 2023

Choose a reason for hiding this comment

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

Suggested change
def create_chunked_file(self, array_shape, chunk_sizes) -> None:
def create_chunked_file(self, array_shape: tuple[int, int, int], chunk_sizes: tuple[int, int, int]) -> Generator[str, None, None]:

I haven't checked it though, but it should bring you in the right direction.

@jhamman jhamman mentioned this pull request Jul 26, 2023
Open
@mraspaud
Copy link
Contributor Author

mraspaud commented Aug 14, 2023

I attempted to fix the type annotation problem, please tell me if there is more.

@mraspaud
Copy link
Contributor Author

mraspaud commented Aug 21, 2023

@dcherian @jhamman anything more I can do on this?

dcherian
dcherian approved these changes Aug 31, 2023
View reviewed changes
Copy link
Contributor

@dcherian dcherian left a comment

Choose a reason for hiding this comment

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

LGTM! Thanks for your patience here @mraspaud

Can you add a note to whats-new to advertise this great improvement please?

dcherian
dcherian reviewed Aug 31, 2023
View reviewed changes
@dcherian dcherian added the plan to merge Final call for comments label Aug 31, 2023
@mraspaud
Copy link
Contributor Author

mraspaud commented Aug 31, 2023

What's new now updated

@dcherian
Copy link
Contributor

dcherian commented Sep 8, 2023

The mypy failures are. related to these changes I think:

xarray/tests/test_backends.py:1559: error: "str" has no attribute "data"  [attr-defined]
xarray/tests/test_backends.py:1559: error: Invalid index type "str" for "str"; expected type "SupportsIndex | slice"  [index]
xarray/tests/test_backends.py:1576: error: "str" has no attribute "data"  [attr-defined]
xarray/tests/test_backends.py:1576: error: Invalid index type "str" for "str"; expected type "SupportsIndex | slice"  [index]
xarray/tests/test_backends.py:1610: error: "str" has no attribute "encoding"  [attr-defined]
xarray/tests/test_backends.py:1610: error: Invalid index type "str" for "str"; expected type "SupportsIndex | slice"  [index]

@Illviljan can you take a look please?

@dcherian dcherian removed the plan to merge Final call for comments label Sep 9, 2023
Illviljan
Illviljan reviewed Sep 11, 2023
View reviewed changes
Illviljan
Illviljan reviewed Sep 11, 2023
View reviewed changes
@Illviljan Illviljan enabled auto-merge (squash) September 11, 2023 18:30
@Illviljan Illviljan disabled auto-merge September 11, 2023 23:05
@Illviljan Illviljan merged commit de66dae into pydata:main Sep 11, 2023
@Illviljan
Copy link
Contributor

Illviljan commented Sep 11, 2023

Thanks, @mraspaud !

@mraspaud mraspaud deleted the feature-nc-preferred-chunks branch September 12, 2023 09:00
@mraspaud
Copy link
Contributor Author

mraspaud commented Sep 12, 2023

My pleasure! thanks for merging!

This was referenced Sep 12, 2023
Closed
Closed
@dougiesquire dougiesquire mentioned this pull request Oct 12, 2023
Merged
@dave-andersen dave-andersen mentioned this pull request Nov 29, 2023
Closed
5 tasks
@weiji14 weiji14 mentioned this pull request Jan 31, 2024
Closed
5 tasks
@charles-turner-1 charles-turner-1 mentioned this pull request Apr 8, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@dcherian dcherian dcherian approved these changes

@Illviljan Illviljan Illviljan left review comments

@headtr1ck headtr1ck headtr1ck left review comments

@jhamman jhamman Awaiting requested review from jhamman

Assignees

No one assigned

Labels

io run-benchmark Run the ASV benchmark workflow topic-backends

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

If a NetCDF file is chunked on disk, open it with compatible dask chunks

7 participants

@mraspaud @djhoese @dcherian @Illviljan @jhamman @headtr1ck @max-sixty