Repository: dask/dask
Pull request: Initial da.select() implementation and test (#6067)
- [X] Closes #6067 
- [X] Tests added / passed
- [X] Passes `black dask` / `flake8 dask` / `isort dask`

̉First contribution! :muscle: 
I just implemented it as much as `da.piecewise()`, there are potential improvements as benchmarks, tests for broadcasting, but I think that it works for a first implementation.

Problem statement:
Please implement numpy.select()
4.  **Feature Request:** then please do a quick search of existing issues to
    make sure that this has not been asked before.
This is a feature request. I have been convert my pandas code for dask and couldnt find a substitution for numpy.select(). My current workaround is to use .apply() with a custom function, but it would be nice to have native dask.select() method.

        cond = [(self._sell_shift == 1), (self._buy_shift == 1)]
        out = ["Sell", "Buy"]
        df = np.select(cond, out, default=0) # now 2 arrays become 1


Thank you.

Hints from issue discussion:
Thanks for raising an issue @Biarys. Adding a `da.select` method would be a welcome contribution, is this something you're interesting in working on (no obligation though)? There's already a `da.where` method, as `da.select` provides similar functionality, looking at how things are implemented in `da.where` would be a good starting point 
@jrbourbeau I will try to work on it once I learn about your library a bit better :) 
We already implement [`np.choose`]( https://numpy.org/doc/1.18/reference/generated/numpy.choose.html ) for Dask Array's [here]( https://github.com/dask/dask/blob/149abe507a3988073a9acc9d5b8483f44c747f6b/dask/array/routines.py#L1243 ), which is pretty similar. Probably a simple matter to implement `np.select` on top of that.

Alternatively this could be implemented directly in NumPy and all array-like libraries could benefit 🙂
I'd like to pick this up if @Biarys is not planning to go ahead with this. @Biarys? 
Would say go for it @Madhu94! 😄
@Madhu94 please go ahead :) 
Thanks @jakirkham I'll wait till after gsoc applications close in case a student wants to pick this up.
@Madhu94 I'd like to take this up. could you guide me to the correct documentation for this piece so that I can get a quick hang of the use-case for this functionality.
I was trying to make my first contribution to dask, and first I wanted to understand better numpy and dask. 

So I migrated a "bigger" version of the numpy example to dask:
```python
import dask.array as da
import numpy as np

def da_select(condlist, choicelist, chunk_size, default=0):
    def np_select(condlist_da, choicelist_da):
        return np.select(np.array(condlist_da), np.array(choicelist_da), default)

    assert len(condlist) == len(choicelist)
    n_choices = len(condlist)
    result_type = da.array(choicelist)[0].dtype
    
    return da.map_blocks(
        np_select, 
        da.array(condlist).rechunk({0: n_choices, 1: chunk_size}), 
        da.array(choicelist).rechunk({0: n_choices, 1: chunk_size}), 
        meta=result_type,
        drop_axis=[1],
        chunks=(1,chunk_size))
        
chunk_size = 10
array_size = 10 * chunk_size
y = da.arange(array_size, chunks=(chunk_size))
condlist = [y<10, y>20]
choicelist = [y, y%10]
da_select(condlist, choicelist, chunk_size).compute()
```

There is a lot missing to make a PR, like to test it with multiple dimensional arrays, param checks, write doc and tests, etc. and I would love to do it, but I wanted to check a few things before continuing:

* is the `map_blocks` approach ok?  @jakirkham @jrbourbeau what do you think? ( or should I ask about it in slack? I didn't think is relevant for `dev-scheduler` channel )
* how can I ask `chunk_size` dynamically? Are there are some technical docs about how to handle chunks in the APIs? I read the code for `da.where` and `da.choose` and the docs for [Array Chunks](https://docs.dask.org/en/latest/array-chunks.html) but I think I'm still not fully getting it. (The hacking dask tutorial in the summit helped too)

`map_blocks` sounds good. Using GH is preferred for technical discussions

Ideally we wouldn't rechunk, but instead use the chunking of the existing arrays.

[`da.broadcast_arrays`]( https://docs.dask.org/en/latest/array-api.html#dask.array.broadcast_arrays ) may be helpful in broadcasting arrays to match (though it's possible this code can be written without that)

Commit message: Initial da.select() implementation and test (#6067)

Patch to review:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 5996e09243a..a4fb819509f 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -134,6 +134,7 @@
         rot90,
         round,
         searchsorted,
+        select,
         shape,
         squeeze,
         swapaxes,
diff --git a/dask/array/routines.py b/dask/array/routines.py
index ba1b22bb735..d79e6d20dd4 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1843,6 +1843,22 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    return blockwise(
+        np.select,
+        "m",
+        condlist,
+        "nm",
+        choicelist,
+        "nm",
+        dtype=result_type(*choicelist),
+        name="select",
+        concatenate=True,
+        default=default,
+    )
+
+
 def _partition(total: int, divisor: int) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
     """
     Given a total and a divisor, return two tuples: A tuple containing `divisor` repeated
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 60ded385408..0e9f1d86952 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1491,6 +1491,22 @@ def test_piecewise_otherwise():
     )
 
 
+def test_select():
+    conditions = [
+        np.array([False, False, False, False]),
+        np.array([False, True, False, True]),
+        np.array([False, False, True, True]),
+    ]
+    choices = [
+        np.array([1, 2, 3, 4]),
+        np.array([5, 6, 7, 8]),
+        np.array([9, 10, 11, 12]),
+    ]
+    d_conditions = da.from_array(conditions, chunks=(3, 2))
+    d_choices = da.from_array(choices)
+    assert_eq(np.select(conditions, choices), da.select(d_conditions, d_choices))
+
+
 def test_argwhere():
     for shape, chunks in [(0, ()), ((0, 0), (0, 0)), ((15, 16), (4, 5))]:
         x = np.random.randint(10, size=shape)


Code review comments:
diff_hunk: @@ -1843,6 +1843,22 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    return blockwise(
+        np.select,
+        "m",
+        condlist,
+        "nm",
path: dask/array/routines.py
text: @user1:
See comment below
```suggestion
        tuple(range(choicelist[0].ndim + 1)),
```

diff_hunk: @@ -1843,6 +1843,22 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    return blockwise(
+        np.select,
+        "m",
+        condlist,
+        "nm",
+        choicelist,
+        "nm",
path: dask/array/routines.py
text: @user1:
To support arrays with ndims > 1, you can use a tuple that depends on `ndim` something like:
```suggestion
        tuple(range(choicelist[0].ndim + 1)),
```

diff_hunk: @@ -1843,6 +1843,22 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    return blockwise(
+        np.select,
+        "m",
path: dask/array/routines.py
text: @user1:
Given the suggestion below, this would then become:
```suggestion
        tuple(range(1, choicelist[0].ndim + 1)),
```