Repository: dask/dask

Issue:
BUG: Dask dataframe cannot handle string index
I write the following code,but it doesn't work.
```python
import dask.bag as db
import dask.dataframe as dd
import dask.array as da
import pandas as pd
if __name__ == '__main__':
	#Just for test
	pWord = dd.from_pandas(pd.DataFrame({'word':["a","b","c","bc","abc"],'probability':[0.1,0.2,0.3,0.4,0.5]}),npartitions=3)
	wordInfoTable = dd.from_pandas(pd.DataFrame({'word':["bc","abc"]}),npartitions=3)
	pWord = pWord.set_index(pWord.word)
	pIfRandom = pWord.loc[wordInfoTable.word.str.slice(stop=1)].probability.reset_index(drop=True) * pWord.loc[wordInfoTable.word.str.slice(start=1)].probability.reset_index(drop=True)
	wordInfoTable["adhesiveness"] = pWord.loc[wordInfoTable.word].probability.reset_index(drop=True) / pIfRandom
	print(wordInfoTable.compute())
```

The following code works well, but very slow
```python
import dask.bag as db
import dask.dataframe as dd
import dask.array as da
import pandas as pd
if __name__ == '__main__':
	#Just for test
	pWord = dd.from_pandas(pd.DataFrame({'word':["a","b","c","bc","abc"],'probability':[0.1,0.2,0.3,0.4,0.5]}),npartitions=3)
	wordInfoTable = dd.from_pandas(pd.DataFrame({'word':["bc","abc"]}),npartitions=3)
	pWord = pWord.set_index(pWord.word)
	def getP(t):
		return next(pWord.loc[t].itertuples()).probability
	def applyAdhesiveness(row):
		word = row["word"]
		pIfRandom = getP(word[0]) * getP(word[1:])
		adhesiveness = getP(word) / pIfRandom
		row["adhesiveness"] = adhesiveness
		return row
	wordInfoTable["adhesiveness"] = None
	wordInfoTable = wordInfoTable.apply(applyAdhesiveness,axis=1,meta=wordInfoTable._meta)
	print(wordInfoTable.compute())
```

I know it may be not a suitable place to ask for help here, but I think it's the only place which I can get the answer.



Discussion:
For faster feedback, I recommend constructing a minimal example.  

-  https://stackoverflow.com/help/mcve
-  http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
When I try to use `pWord.loc[wordInfoTable.word.str.slice(stop=1)]`
The code doesn't work.

```pWord``` is indexed over ```word``` column. While ```wordInfoTable``` is indexed over some random integer index
I think the problem is related to strings, I believe it's a bug in dask handling of string indices.



```
import dask.dataframe as dd
import pandas as pd

a = dd.from_pandas(pd.DataFrame({'x':[0.1,0.2,0.3,0.4,0.5]}, index=["a","b","c","bc","abc"]), npartitions=3)
b = dd.from_pandas(pd.DataFrame({'word':["bc","abc"]}), npartitions=3)
a.loc[b.word]

TypeError: '<' not supported between instances of 'int' and 'str'
```

Another issue happens even with standard pd.Series

```
import dask.dataframe as dd
import pandas as pd

a = dd.from_pandas(pd.DataFrame({'x':[0.1,0.2,0.3,0.4,0.5]}, index=["a","b","c","bc","abc"]), npartitions=3)
b = pd.DataFrame({'word':["bc","abc"]})
a.loc[b.word]
ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

```


@gabrielelanaro 
When both `a` and `b` are the type from pandas
```python
import pandas as pd

a = pd.DataFrame({'x':[0.1,0.2,0.3,0.4,0.5]}, index=["a","b","c","bc","abc"])
b = pd.DataFrame({'word':["bc","abc"]})
print(a.loc[b.word])
```
It can work well

Thanks for the minimal example in https://github.com/dask/dask/issues/3269#issuecomment-372082058 @gabrielelanaro !

I agree that this looks like a bug in dask dataframe.  When we call loc on a dask.dataframe.Series we are currently trying to align the two dataframe/series objects by divisions, which is where the exception is occurring.  Instead we actually need to do a more complex operation where every partition of the indexer talks to every partition of the indexed.  

This should be fixed if anyone has interest. 

It's worth pointing out though that this will never be nearly as fast as in the Pandas case.
I've been experimenting with this issue on my local copy. I would say the main issue is that ```DataFrame.map_partitions``` does not actually just map a functions to all partitions in the dataframe, but it also performs some alignment under the hood. That's because it dispatches to ```core.map_partitions``` which perform alignment automatically (a more specific name would be ```map_aligned_partitions```) 

I feel the best way to do this is to actually implement different versions of 
```map_partitions```:

Something on these lines:
1. "align" version, will take the _Frame arguments and align their index
2. "forward" version, which will forward the arguments unchanged (warning: this could blow up)
3. "broadcast" version, which will do the complex operation you were referring to. So the function will be applied to "paired partitions"l 

With regard to the present issue, if the Series is of type bool, you need to align partitions, otherwise, you need to broadcast  and concat the results 
Sorry for the necro, just wanted to put my 2 cents in.
I ran into this problem when I wanted to select the groupwise minima over a certain column like this:
```
indexes = df.groupby("grp")["val"].idxmin()
df.loc[indexes]
```

I implemented the suggested fix of having a `broadcast_map_partitions()` function which maps the given function over all combinations of partitions. I implemented this naively by copying `map_partitions`,  removing the alignment that happens, and changing `partitionwise_graph` to create a partition-to-partition graph. This worked great, but memory usage was insane. My adaptive cluster peaked at 11 terabyte of memory usage on a calculation that previously spiked at ~100 gigs. I guess each partition is copied to the worker where each other partition lives? Might be an oversight in my implementation too, but at least it worked.

Anyway, I found that just relying on Dask's divisions and alignment works way better.
```
def _loc(_part, _selector_part):
    try:
        return _part.loc[_part.index.intersection(_selector_part)]
    except KeyError:
         return []

indexes = df.groupby("grp")["val"].idxmin()
indexes_partitioned = indexes.to_frame().set_index("val", drop=False)["val"]
df.map_partitions(_loc, indexes_partitioned)
```
On my big dataset, this takes the same time but uses ~100x less memory. And in my specific case, I could even cut out the extra steps by combining `idxmin` and `loc` in a single pass using `groupby()._aca_agg()` directly.
@dask/maintenance can someone edit the title of this issue to better reflect the contents? I'd suggest something like *"BUG: Dask dataframe cannot handle string index"*

A minimal, reproducible example for this issue is given here: https://github.com/dask/dask/issues/3269#issuecomment-372082058
(title updated)
This is an issue that I think could be pretty decently handled with a good explanation in an error message. I am opening a PR with one option. 

Resolving patch:
diff --git a/dask/dataframe/indexing.py b/dask/dataframe/indexing.py
--- a/dask/dataframe/indexing.py
+++ b/dask/dataframe/indexing.py
@@ -11,7 +11,7 @@
 from . import methods
 from ._compat import PANDAS_GT_130
 from .core import Series, new_dd_object
-from .utils import is_index_like, meta_nonempty
+from .utils import is_index_like, is_series_like, meta_nonempty
 
 
 class _IndexerBase:
@@ -114,14 +114,21 @@ def _loc(self, iindexer, cindexer):
                 return self._loc_slice(iindexer, cindexer)
             elif isinstance(iindexer, (list, np.ndarray)):
                 return self._loc_list(iindexer, cindexer)
+            elif is_series_like(iindexer) and iindexer.dtype != bool:
+                return self._loc_list(iindexer.values, cindexer)
             else:
                 # element should raise KeyError
                 return self._loc_element(iindexer, cindexer)
         else:
-            if isinstance(iindexer, (list, np.ndarray)):
+            if isinstance(iindexer, (list, np.ndarray)) or (
+                is_series_like(iindexer) and iindexer.dtype != bool
+            ):
                 # applying map_partitions to each partition
                 # results in duplicated NaN rows
-                msg = "Cannot index with list against unknown division"
+                msg = (
+                    "Cannot index with list against unknown division. "
+                    "Try setting divisions using ``ddf.set_index``"
+                )
                 raise KeyError(msg)
             elif not isinstance(iindexer, slice):
                 iindexer = slice(iindexer, iindexer)
@@ -141,6 +148,11 @@ def _maybe_partial_time_string(self, iindexer):
         return iindexer
 
     def _loc_series(self, iindexer, cindexer):
+        if iindexer.dtype != bool:
+            raise KeyError(
+                "Cannot index with non-boolean dask Series. Try passing computed "
+                "values instead (e.g. ``ddf.loc[iindexer.compute()]``)"
+            )
         meta = self._make_meta(iindexer, cindexer)
         return self.obj.map_partitions(
             methods.loc, iindexer, cindexer, token="loc-series", meta=meta


Test patch:
diff --git a/dask/dataframe/tests/test_indexing.py b/dask/dataframe/tests/test_indexing.py
--- a/dask/dataframe/tests/test_indexing.py
+++ b/dask/dataframe/tests/test_indexing.py
@@ -84,15 +84,19 @@ def test_loc_with_text_dates():
 def test_loc_with_series():
     assert_eq(d.loc[d.a % 2 == 0], full.loc[full.a % 2 == 0])
 
-    assert sorted(d.loc[d.a % 2].dask) == sorted(d.loc[d.a % 2].dask)
-    assert sorted(d.loc[d.a % 2].dask) != sorted(d.loc[d.a % 3].dask)
+    assert sorted(d.loc[d.a % 2 == 0].dask) == sorted(d.loc[d.a % 2 == 0].dask)
+    assert sorted(d.loc[d.a % 2 == 0].dask) != sorted(d.loc[d.a % 3 == 0].dask)
 
 
 def test_loc_with_array():
     assert_eq(d.loc[(d.a % 2 == 0).values], full.loc[(full.a % 2 == 0).values])
 
-    assert sorted(d.loc[(d.a % 2).values].dask) == sorted(d.loc[(d.a % 2).values].dask)
-    assert sorted(d.loc[(d.a % 2).values].dask) != sorted(d.loc[(d.a % 3).values].dask)
+    assert sorted(d.loc[(d.a % 2 == 0).values].dask) == sorted(
+        d.loc[(d.a % 2 == 0).values].dask
+    )
+    assert sorted(d.loc[(d.a % 2 == 0).values].dask) != sorted(
+        d.loc[(d.a % 3 == 0).values].dask
+    )
 
 
 def test_loc_with_function():
@@ -131,6 +135,39 @@ def test_loc_with_series_different_partition():
     )
 
 
+def test_loc_with_non_boolean_series():
+    df = pd.Series(
+        np.random.randn(20),
+        index=list("abcdefghijklmnopqrst"),
+    )
+    ddf = dd.from_pandas(df, 3)
+
+    s = pd.Series(list("bdmnat"))
+    ds = dd.from_pandas(s, npartitions=3)
+
+    msg = (
+        "Cannot index with non-boolean dask Series. Try passing computed values instead"
+    )
+    with pytest.raises(KeyError, match=msg):
+        ddf.loc[ds]
+
+    assert_eq(ddf.loc[s], df.loc[s])
+
+    with pytest.raises(KeyError, match=msg):
+        ddf.loc[ds.values]
+
+    assert_eq(ddf.loc[s.values], df.loc[s])
+
+    ddf = ddf.clear_divisions()
+    with pytest.raises(KeyError, match=msg):
+        ddf.loc[ds]
+
+    with pytest.raises(
+        KeyError, match="Cannot index with list against unknown division"
+    ):
+        ddf.loc[s]
+
+
 def test_loc2d():
     # index indexer is always regarded as slice for duplicated values
     assert_eq(d.loc[5, "a"], full.loc[5:5, "a"])