Repository: dask/dask

Issue:
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.



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)
Thx, for your quick response.
I'm gonna try to explain why I used rechunk with this example, maybe "we" can figure out how to handle the chunks.

Let's assume that this example is representative of the general use case. 

You have at first a big chunked array of *m* elements and a few *n* choices (booleans expression over the same array) and choices (expressions over same array).

My idea was to keep the chunks of the big array over the conditions and the choices *n X m* "matrices", just apply `np.select` over every chunk and join them in a new big array keeping the same chunks.

But first, by default, `da.array(condlist)` creates chunks using the first dimension so generally I get *n* chunks of *m*-sized arrays, instead of creating *n X chunk_size* matrices.
`rechunk` was the only way that I found to "avoid" chunking in the first dimension ( I failed transposing but I think that it won't work for multi dimensional array )

And I'm still assuming that both parameters are expressions from the same array, that it couldn't be the case. 

So I think that by default I should take a chunk parameter for the result similar to `map_blocks` and if it's not provided then rechunk.

I'm gonna sleep on it and try again tomorrow :) I'll review `broadcast_arrays` and look into more code.
Ok, I looked into `blockwise` and its tests and it looks like it's the answer.

```python
da.blockwise(np.select,'m', condlist, 'nm', choicelist, 'nm', concatenate=True).compute()
```

Works as expected for one-dimensional arrays, so the PR should be coming soon.
Yeah `blockwise` may be a good choice. Would see if `concatenate=True` could be avoided (as that can be kind of expensive)
I tried without `concatenate=True` but I got an `ValueError: could not broadcast input array from shape (10,) into shape (1,)`

I'm not sure if that was expected, but with `concatenate` I was already happy :) , considering what you said I'll take a look to add a benchmark in the PR. 
It might be a matter of adding singleton dimensions to the input arrays. Alternatively using only `'n'` or `'m'` might make sense for one or both of these arrays.

Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- 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
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1843,6 +1843,51 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+def _select(*args, **kwargs):
+    """
+    This is a version of :func:`numpy.select` that acceptes an arbitrary number of arguments and
+    splits them in half to create ``condlist`` and ``choicelist`` params.
+    """
+    split_at = len(args) // 2
+    condlist = args[:split_at]
+    choicelist = args[split_at:]
+    return np.select(condlist, choicelist, **kwargs)
+
+
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    # Making the same checks that np.select
+    # Check the size of condlist and choicelist are the same, or abort.
+    if len(condlist) != len(choicelist):
+        raise ValueError("list of cases must be same length as list of conditions")
+
+    if len(condlist) == 0:
+        raise ValueError("select with an empty condition list is not possible")
+
+    choicelist = [asarray(choice) for choice in choicelist]
+
+    try:
+        intermediate_dtype = result_type(*choicelist)
+    except TypeError as e:
+        msg = "Choicelist elements do not have a common dtype."
+        raise TypeError(msg) from e
+
+    blockwise_shape = tuple(range(choicelist[0].ndim))
+
+    condargs = [arg for elem in condlist for arg in (elem, blockwise_shape)]
+    choiceargs = [arg for elem in choicelist for arg in (elem, blockwise_shape)]
+
+    return blockwise(
+        _select,
+        blockwise_shape,
+        *condargs,
+        *choiceargs,
+        dtype=intermediate_dtype,
+        name="select",
+        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


Test patch:
diff --git a/dask/array/tests/test_array_function.py b/dask/array/tests/test_array_function.py
--- a/dask/array/tests/test_array_function.py
+++ b/dask/array/tests/test_array_function.py
@@ -34,6 +34,9 @@
         lambda x: np.round(x),
         lambda x: np.insert(x, 0, 3, axis=0),
         lambda x: np.delete(x, 0, axis=0),
+        lambda x: np.select(
+            [x < 0.3, x < 0.6, x > 0.7], [x * 2, x, x / 2], default=0.65
+        ),
     ],
 )
 def test_array_function_dask(func):
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1491,6 +1491,50 @@ 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_select_multidimension():
+    x = np.random.random((100, 50, 2))
+    y = da.from_array(x, chunks=(50, 50, 1))
+    res_x = np.select([x < 0, x > 2, x > 1], [x, x * 2, x * 3], default=1)
+    res_y = da.select([y < 0, y > 2, y > 1], [y, y * 2, y * 3], default=1)
+    assert isinstance(res_y, da.Array)
+    assert_eq(res_y, res_x)
+
+
+def test_select_return_dtype():
+    d = np.array([1, 2, 3, np.nan, 5, 7])
+    m = np.isnan(d)
+    d_d = da.from_array(d)
+    d_m = da.isnan(d_d)
+    assert_eq(np.select([m], [d]), da.select([d_m], [d_d]), equal_nan=True)
+
+
+@pytest.mark.xfail(reason="broadcasting in da.select() not implemented yet")
+def test_select_broadcasting():
+    conditions = [np.array(True), np.array([False, True, False])]
+    choices = [1, np.arange(12).reshape(4, 3)]
+    d_conditions = da.from_array(conditions)
+    d_choices = da.from_array(choices)
+    assert_eq(np.select(conditions, choices), da.select(d_conditions, d_choices))
+    # default can broadcast too:
+    assert_eq(np.select([True], [0], default=[0]), da.select([True], [0], default=[0]))
+
+
 def test_argwhere():
     for shape, chunks in [(0, ()), ((0, 0), (0, 0)), ((15, 16), (4, 5))]:
         x = np.random.randint(10, size=shape)