Repository: dask/dask
Pull request: Add index to meta after .str.split with expand
This PR proposes a fix for https://github.com/dask/dask/issues/7021. 

Without the proposed fix, `dask/dataframe/tests/test_accessors.py::test_str_accessor_expand` fails if the index is of any other type than `int`. 

- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`


Problem statement:
`index.dtype` changes from date to int after `.str.split`
When I apply `.str.split` to a column with datetime index, the dtype of the index changes to int64.

Example:

```
import pandas as pd
import dask.dataframe as dd

series_with_strings = dd.from_pandas(
    pd.Series(
        data=["some_text", "split_it"],
        index=[pd.Timestamp(2020, 1, 1), pd.Timestamp(2020, 1, 1)]
    ),
    npartitions=2
)

df_with_splitted_strings = series_with_strings.str.split(pat="_", n=1, expand=True)

print(series_with_strings.index.dtype)
# datetime64[ns]

print(df_with_splitted_strings.index.dtype)
# int64
```

The problem is also that this cannot be fixed by:

```df_with_splitted_strings.index = dd.to_datetime(df_with_splitted_strings.index)```

since `.set_index(..)` needs to be used for that.

I'm happy to (try to) find a fix for this one, but I'll need some pointers to begin (the dask code base is still largely unfamiliar to me). 

- Dask version: 2020.12.0
- Python version: 3.8.6
- Operating System: MacOS Big Sure, Ubuntu
- Install method (conda, pip, source): pip with pyenv

Hints from issue discussion:

Commit message: iloc[:0] more expressive than .drop(0)

Patch to review:
diff --git a/dask/dataframe/accessor.py b/dask/dataframe/accessor.py
index 28573b1e5a2..0c8a3ade3a1 100644
--- a/dask/dataframe/accessor.py
+++ b/dask/dataframe/accessor.py
@@ -129,6 +129,7 @@ def split(self, pat=None, n=-1, expand=False):
                 delimiter = " " if pat is None else pat
                 meta = type(self._series._meta)([delimiter.join(["a"] * (n + 1))])
                 meta = meta.str.split(n=n, expand=expand, pat=pat)
+                meta = meta.iloc[:0].set_index(self._series._meta.index)
         else:
             meta = (self._series.name, object)
         return self._function_map("split", pat=pat, n=n, expand=expand, meta=meta)
diff --git a/dask/dataframe/tests/test_accessors.py b/dask/dataframe/tests/test_accessors.py
index 59dbc2c2c2a..e477589b271 100644
--- a/dask/dataframe/tests/test_accessors.py
+++ b/dask/dataframe/tests/test_accessors.py
@@ -233,7 +233,9 @@ def test_str_accessor_noexpand():
 
 
 def test_str_accessor_expand():
-    s = pd.Series(["a b c d", "aa bb cc dd", "aaa bbb ccc dddd"])
+    s = pd.Series(
+        ["a b c d", "aa bb cc dd", "aaa bbb ccc dddd"], index=["row1", "row2", "row3"]
+    )
     ds = dd.from_pandas(s, npartitions=2)
 
     for n in [1, 2, 3]:


Code review comments:
diff_hunk: @@ -129,6 +129,7 @@ def split(self, pat=None, n=-1, expand=False):
                 delimiter = " " if pat is None else pat
                 meta = type(self._series._meta)([delimiter.join(["a"] * (n + 1))])
path: dask/dataframe/accessor.py
text: @user1:
Small suggestion, to make this more consistent with the rest of dask.

```suggestion
                meta = self._series._constructor([delimiter.join(["a"] * (n + 1))], index=self._series._meta_nonempty[:1].index)
```

That may not be properly linted.

@author:
The `self._series._meta_nonempty[:1].index` is certainly nice. But if I'm not mistaken, the `self._series._constructor` needs to be called with `dsk, name, meta, divisions` (cf. `new_dd_object(dsk, name, meta, divisions)`). Hence, your suggestion results in `TypeError: new_dd_object() got an unexpected keyword argument 'index'`.

Maybe I miss the point. 

@user1:
Whoops, sorry. Should be `self._series._meta._constructor`. That's similar to `type(self._series._meta)(...)`, but a bit more correct (newer parts of dask are using the `._constructor` approach).

diff_hunk: @@ -129,6 +129,7 @@ def split(self, pat=None, n=-1, expand=False):
                 delimiter = " " if pat is None else pat
                 meta = type(self._series._meta)([delimiter.join(["a"] * (n + 1))])
                 meta = meta.str.split(n=n, expand=expand, pat=pat)
+                meta = meta.iloc[:0].set_index(self._series._meta.index)
path: dask/dataframe/accessor.py
text: @user1:
This can be removed if you apply my suggestion from above. It's not needed now that we create a `meta` with the proper index.