dask / dask/distributed

[Tasks from tasks] Create collections with number of chunks not known at graph definition time

Open
#6,204 1 comment 3 reactions 0 assignees View on GitHub
discussion feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.