Propogating NaN values when using str.split (pandas-dev#18450)

Author: WillAyd

pandas/core/strings.py

@@ -1407,7 +1407,6 @@ def _wrap_result(self, result, use_codes=True,
         if not hasattr(result, 'ndim') or not hasattr(result, 'dtype'):
             return result
         assert result.ndim < 3
-
         if expand is None:
             # infer from ndim if expand is not specified
             expand = False if result.ndim == 1 else True
@@ -1423,6 +1422,10 @@ def cons_row(x):
                     return [x]
 
             result = [cons_row(x) for x in result]
+            if result:
+                # propogate nan values to match longest sequence (GH 18450)
+                max_len = max(len(x) for x in result)
+                result = [x * max_len if x[0] is np.nan else x for x in result]
 
         if not isinstance(expand, bool):
             raise ValueError("expand must be True or False")

pandas/tests/test_strings.py

@@ -2002,6 +2002,12 @@ def test_split_to_dataframe(self):
                          5: [NA, 'not']})
         tm.assert_frame_equal(result, exp)
 
+        # make sure we propogate NaN values across all columns
+        s = Series(["foo,bar,baz", NA])
+        result = s.str.split(",", expand=True)
+        exp = DataFrame([["foo", "bar", "baz"], [NA, NA, NA]])
+        tm.assert_frame_equal(result, exp)
+
         s = Series(['some_splits', 'with_index'], index=['preserve', 'me'])
         result = s.str.split('_', expand=True)
         exp = DataFrame({0: ['some', 'with'], 1: ['splits', 'index']},