[Tasks from tasks] Create collections with number of chunks not known at graph definition time
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
A common issue is that a user does not know in advance how many chunks a collection needs to have. The typical workaround to this is to generate as many chunks as it could *possibly* have and then leave some (most) of them either empty or oversplit.
Same problem is for
- da.Arrays with an unknown number of dimensions
- df.DataFrame with unknown columns
An elegant solution, which would require #5671 as a prerequisite, is the following:
## Proposed core design
Enhance ``Client.publish_dataset`` to accept nameless collection(s). If you do, the function returns a single future and the collection is not listed by ``Client.list_datasets``.
Calling result() on the future returns the persisted collection or tuple of collections.
When the future is descoped, the published collection is automatically forgotten.
#### Sample usage
```python
def dynamic_ones(shape, chunks):
a = da.ones(shape, chunks)
return get_client().publish_dataset(a)
c = Client()
a = c.submit(
dynamic_ones,
c.submit(dynamic_shape),
c.submit(dynamic_chunks)
).result() # Return persisted collection
```
## Proposed extension: compatibility with pure dask
In dask/dask, add ``dask.publish(*collections)`` and the shorthand ``collection.publish()``.
If scheduler==distributed, call ``Client.publish_dataset`` under the hood and return a delayed wrapping its output distributed.Future.
If scheduler in ("threads", "synchronous" "processes"), call persist() and then return a single dummy delayed object to the persisted collection(s).
## Proposed extension: rechunk() nan chunks
``Array.rechunk`` can't be applied to an array where one or more chunks have size nan, since it would create an unknown amount of output chunks. To fix this, it could gain an optional bool parameter ``delay``:
```python
@overload
def rechunk(self: Array, new_chunks: int | tuple, delay: Literal[False] = False) -> Array:
...
@overload
def rechunk(self: Array, new_chunks: int | tuple, delay: Literal[True]) -> Delayed[Array]:
...
def rechunk(self, new_chunks, delay=False):
if not delay:
return ... # current implementation
if not any(math.isnan(j) for i in self.chunks for j in i):
return delayed(self.rechunk)(new_chunks)
@delayed
def dshape(chunk: np.ndarray) -> tuple[int, ...]:
return chunk.shape
@delayed
def drechunk(arr: Array, chunk_shapes: list[tuple[int, ...]], new_chunks: int | tuple) -> Delayed[Array]:
# Replace nan chunks with known chunks
arr = Array(arr.dask, arr.name, to_chunks(chunk_shapes, arr), meta=arr)
arr = arr.rechunk(new_chunks)
return arr.publish()
persisted = self.persist()
return drechunk(
persisted.publish(),
[dshape(c) for c in flatten(persisted.to_delayed())],
new_chunks,
)
```
Similar treatment could be done for all functions in dask.array and dask.dataframe that currently don't work with nan chunks.
Contributor guide
Assessment
This issue has not been assessed yet.