Repository: databricks/koalas

Issue:
str.split(expand = True)
in dataframe string split into multiple columns is not available. what is alternative to that


Discussion:
Would you mind if you show me a little bit of your code examples?

I'm not sure what you exactly want to, but maybe the below way will work for you??

```python
>>> import databricks.koalas as ks
>>> from databricks.koalas.config import option_context

>>> kdf = ks.DataFrame({'col1': ['A_1', 'A_2', 'A_3', 'A_4', 'A_5']})
>>> kdf
  col1
0  A_1
1  A_2
2  A_3
3  A_4
4  A_5

>>> with option_context('compute.ops_on_diff_frames', True):
...     kdf['first_split'] = kdf.map_in_pandas(lambda pdf: pdf.col1.str.split('_').str[0].to_frame()).col1
...     kdf['second_split'] = kdf.map_in_pandas(lambda pdf: pdf.col1.str.split('_').str[1].to_frame()).col1
...
>>> kdf
  col1 first_split second_split
0  A_1           A            1
1  A_2           A            2
3  A_4           A            4
2  A_3           A            3
4  A_5           A            5
```

Thanks. It was super helpful.

I believe its possible to use most of the pandas function using map_in_pandas ?

More specifically I was looking for this 

`Temp = kdf.map_in_pandas(lambda pdf:pdf.col1.str.split("_", expand = True))`

`Temp.columns  = ['S1', 'S2']`

```
Temp.head()

        S1     S2
0	A	1
1	A	2
3	A	4
2	A	3
4	A	5
```
@mks2192 

> I believe its possible to use most of the pandas function using map_in_pandas ?

It doesn't support for aggregation-like methods which needs operation through the whole dataset.

In this case, It's possible to use `map_in_pandas` because the `str.split` is applied to row by row.

If you wanna know about more detail, please check [this documents](https://koalas.readthedocs.io/en/latest/reference/api/databricks.koalas.DataFrame.map_in_pandas.html).

It's really happy that It helped you :D

Resolving patch:
diff --git a/databricks/koalas/strings.py b/databricks/koalas/strings.py
--- a/databricks/koalas/strings.py
+++ b/databricks/koalas/strings.py
@@ -17,7 +17,7 @@
 """
 String functions on Koalas Series
 """
-from typing import TYPE_CHECKING
+from typing import Union, TYPE_CHECKING
 
 import numpy as np
 
@@ -1156,9 +1156,16 @@ def join(self, sep) -> "ks.Series":
             Series with list entries concatenated by intervening occurrences of
             the delimiter.
 
+        See Also
+        --------
+        str.split : Split strings around given separator/delimiter.
+        str.rsplit : Splits string around given separator/delimiter,
+            starting from the right.
+
         Examples
         --------
         Example with a list that contains a None element.
+
         >>> s = ks.Series([['lion', 'elephant', 'zebra'],
         ...                ['cat', None, 'dog']])
         >>> s
@@ -1167,6 +1174,7 @@ def join(self, sep) -> "ks.Series":
         Name: 0, dtype: object
 
         Join all lists using a ‘-‘. The list containing None will produce None.
+
         >>> s.str.join('-')
         0    lion-elephant-zebra
         1                   None
@@ -1776,7 +1784,7 @@ def slice_replace(self, start=None, stop=None, repl=None) -> "ks.Series":
             self, lambda x: x.str.slice_replace(start, stop, repl), StringType()
         ).alias(self.name)
 
-    def split(self, pat=None, n=-1, expand=False) -> "ks.Series":
+    def split(self, pat=None, n=-1, expand=False) -> Union["ks.Series", "ks.DataFrame"]:
         """
         Split strings around given separator/delimiter.
 
@@ -1791,13 +1799,38 @@ def split(self, pat=None, n=-1, expand=False) -> "ks.Series":
         n : int, default -1 (all)
             Limit number of splits in output. None, 0 and -1 will be
             interpreted as return all splits.
-        expand : bool, currently only False supported
+        expand : bool, default False
             Expand the splitted strings into separate columns.
 
+            * If ``True``, `n` must be a positive integer, and return DataFrame expanding
+              dimensionality.
+            * If ``False``, return Series, containing lists of strings.
+
         Returns
         -------
-        Series of object
-            Series with split strings.
+        Series, DataFrame
+            Type matches caller unless `expand=True` (see Notes).
+
+        See Also
+        --------
+        str.rsplit : Splits string around given separator/delimiter,
+            starting from the right.
+        str.join : Join lists contained as elements in the Series/Index
+            with passed delimiter.
+
+        Notes
+        -----
+        The handling of the `n` keyword depends on the number of found splits:
+
+        - If found splits > `n`,  make first `n` splits only
+        - If found splits <= `n`, make all splits
+        - If for a certain row the number of found splits < `n`,
+          append `None` for padding up to `n` if ``expand=True``
+
+        If using ``expand=True``, Series callers return DataFrame objects with `n + 1` columns.
+
+        .. note:: Even if `n` is much larger than found splits, the number of columns does NOT
+            shrink unlike pandas.
 
         Examples
         --------
@@ -1843,15 +1876,54 @@ def split(self, pat=None, n=-1, expand=False) -> "ks.Series":
         1    [https:, , docs.python.org, 3, tutorial, index...
         2                                                 None
         Name: 0, dtype: object
+
+        When using ``expand=True``, the split elements will expand out into
+        separate columns. If NaN is present, it is propagated throughout
+        the columns during the split.
+
+        >>> s.str.split(n=4, expand=True)
+                                                       0     1     2        3         4
+        0                                           this    is     a  regular  sentence
+        1  https://docs.python.org/3/tutorial/index.html  None  None     None      None
+        2                                           None  None  None     None      None
+
+        For slightly more complex use cases like splitting the html document name
+        from a url, a combination of parameter settings can be used.
+
+        >>> s.str.rsplit("/", n=1, expand=True)
+                                            0           1
+        0          this is a regular sentence        None
+        1  https://docs.python.org/3/tutorial  index.html
+        2                                None        None
+
+        Remember to escape special characters when explicitly using regular
+        expressions.
+
+        >>> s = pd.Series(["1+1=2"])
+        >>> s.str.split(r"\\+|=", n=2, expand=True)
+           0  1  2
+        0  1  1  2
         """
-        if expand:
-            raise NotImplementedError("expand=True is currently not supported.")
+        from databricks.koalas.frame import DataFrame
 
-        return _wrap_accessor_pandas(
-            self, lambda x: x.str.split(pat, n, expand), ArrayType(StringType(), containsNull=True)
+        if expand and n <= 0:
+            raise NotImplementedError("expand=True is currently only supported with n > 0.")
+
+        kser = _wrap_accessor_pandas(
+            self, lambda x: x.str.split(pat, n), ArrayType(StringType(), containsNull=True)
         ).alias(self.name)
 
-    def rsplit(self, pat=None, n=-1, expand=False) -> "ks.Series":
+        if expand:
+            kdf = kser.to_frame()
+            scol = kdf._internal.data_spark_columns[0]
+            spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)]
+            column_labels = [(str(i),) for i in range(n + 1)]
+            internal = kdf._internal.with_new_columns(spark_columns, column_labels)
+            return DataFrame(internal)
+        else:
+            return kser
+
+    def rsplit(self, pat=None, n=-1, expand=False) -> Union["ks.Series", "ks.DataFrame"]:
         """
         Split strings around given separator/delimiter.
 
@@ -1866,13 +1938,37 @@ def rsplit(self, pat=None, n=-1, expand=False) -> "ks.Series":
         n : int, default -1 (all)
             Limit number of splits in output. None, 0 and -1 will be
             interpreted as return all splits.
-        expand : bool, currently only False supported
+        expand : bool, default False
             Expand the splitted strings into separate columns.
 
+            * If ``True``, `n` must be a positive integer, and return DataFrame expanding
+              dimensionality.
+            * If ``False``, return Series, containing lists of strings.
+
         Returns
         -------
-        Series of object
-            Series with split strings.
+        Series, DataFrame
+            Type matches caller unless `expand=True` (see Notes).
+
+        See Also
+        --------
+        str.split : Split strings around given separator/delimiter.
+        str.join : Join lists contained as elements in the Series/Index
+            with passed delimiter.
+
+        Notes
+        -----
+        The handling of the `n` keyword depends on the number of found splits:
+
+        - If found splits > `n`,  make first `n` splits only
+        - If found splits <= `n`, make all splits
+        - If for a certain row the number of found splits < `n`,
+          append `None` for padding up to `n` if ``expand=True``
+
+        If using ``expand=True``, Series callers return DataFrame objects with `n + 1` columns.
+
+        .. note:: Even if `n` is much larger than found splits, the number of columns does NOT
+            shrink unlike pandas.
 
         Examples
         --------
@@ -1910,14 +2006,53 @@ def rsplit(self, pat=None, n=-1, expand=False) -> "ks.Series":
         1    [https://docs.python.org/3/tutorial/index.html]
         2                                               None
         Name: 0, dtype: object
+
+        When using ``expand=True``, the split elements will expand out into
+        separate columns. If NaN is present, it is propagated throughout
+        the columns during the split.
+
+        >>> s.str.split(n=4, expand=True)
+                                                       0     1     2        3         4
+        0                                           this    is     a  regular  sentence
+        1  https://docs.python.org/3/tutorial/index.html  None  None     None      None
+        2                                           None  None  None     None      None
+
+        For slightly more complex use cases like splitting the html document name
+        from a url, a combination of parameter settings can be used.
+
+        >>> s.str.rsplit("/", n=1, expand=True)
+                                            0           1
+        0          this is a regular sentence        None
+        1  https://docs.python.org/3/tutorial  index.html
+        2                                None        None
+
+        Remember to escape special characters when explicitly using regular
+        expressions.
+
+        >>> s = pd.Series(["1+1=2"])
+        >>> s.str.split(r"\\+|=", n=2, expand=True)
+           0  1  2
+        0  1  1  2
         """
-        if expand:
-            raise NotImplementedError("expand=True is currently not supported.")
+        from databricks.koalas.frame import DataFrame
 
-        return _wrap_accessor_pandas(
-            self, lambda x: x.str.rsplit(pat, n, expand), ArrayType(StringType(), containsNull=True)
+        if expand and n <= 0:
+            raise NotImplementedError("expand=True is currently only supported with n > 0.")
+
+        kser = _wrap_accessor_pandas(
+            self, lambda x: x.str.rsplit(pat, n), ArrayType(StringType(), containsNull=True)
         ).alias(self.name)
 
+        if expand:
+            kdf = kser.to_frame()
+            scol = kdf._internal.data_spark_columns[0]
+            spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)]
+            column_labels = [(str(i),) for i in range(n + 1)]
+            internal = kdf._internal.with_new_columns(spark_columns, column_labels)
+            return DataFrame(internal)
+        else:
+            return kser
+
     def translate(self, table) -> "ks.Series":
         """
         Map all characters in the string through the given mapping table.


Test patch:
diff --git a/databricks/koalas/tests/test_series_string.py b/databricks/koalas/tests/test_series_string.py
--- a/databricks/koalas/tests/test_series_string.py
+++ b/databricks/koalas/tests/test_series_string.py
@@ -45,9 +45,14 @@ def pser(self):
     def check_func(self, func):
         self.check_func_on_series(func, self.pser)
 
-    def check_func_on_series(self, func, pser):
+    def check_func_on_series(self, func, pser, result_is_dataframe=False):
         kser = ks.from_pandas(pser)
-        mt.assert_series_equal(func(kser).toPandas(), func(pser), check_names=False)
+        if result_is_dataframe:
+            mt.assert_frame_equal(
+                func(kser).to_pandas(), func(pser).rename(columns=str), check_names=False
+            )
+        else:
+            mt.assert_series_equal(func(kser).to_pandas(), func(pser), check_names=False)
 
     def test_string_add_str_num(self):
         pdf = pd.DataFrame(dict(col1=["a"], col2=[1]))
@@ -291,6 +296,9 @@ def test_string_split(self):
         pser = pd.Series(["This is a sentence.", "This-is-a-long-word."])
         self.check_func_on_series(lambda x: x.str.split(n=2), pser)
         self.check_func_on_series(lambda x: x.str.split(pat="-", n=2), pser)
+        self.check_func_on_series(
+            lambda x: x.str.split(n=2, expand=True), pser, result_is_dataframe=True,
+        )
         with self.assertRaises(NotImplementedError):
             self.check_func(lambda x: x.str.split(expand=True))
 
@@ -300,6 +308,9 @@ def test_string_rsplit(self):
         pser = pd.Series(["This is a sentence.", "This-is-a-long-word."])
         self.check_func_on_series(lambda x: x.str.rsplit(n=2), pser)
         self.check_func_on_series(lambda x: x.str.rsplit(pat="-", n=2), pser)
+        self.check_func_on_series(
+            lambda x: x.str.rsplit(n=2, expand=True), pser, result_is_dataframe=True
+        )
         with self.assertRaises(NotImplementedError):
             self.check_func(lambda x: x.str.rsplit(expand=True))