fix + test

Author: xinrong-meng

python/pyspark/pandas/strings.py

@@ -2031,7 +2031,13 @@ def pudf(s: pd.Series) -> pd.Series:
         if expand:
             psdf = psser.to_frame()
             scol = psdf._internal.data_spark_columns[0]
-            spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)]
+
+            if ps.get_option("compute.ansi_mode_support"):
+                spark_columns = [
+                    F.try_element_at(scol, F.lit(i + 1)).alias(str(i)) for i in range(n + 1)
+                ]
+            else:
+                spark_columns = [scol[i].alias(str(i)) for i in range(n + 1)]
             column_labels = [(i,) for i in range(n + 1)]
             internal = psdf._internal.with_new_columns(
                 spark_columns,

python/pyspark/pandas/tests/series/test_string_ops_adv.py

@@ -185,6 +185,8 @@ def test_string_split(self):
         with self.assertRaises(NotImplementedError):
             self.check_func(lambda x: x.str.split(expand=True))
 
+        self.check_func_on_series(lambda x: repr(x.str.split("-", n=1, expand=True)), pser)
+
     @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_string_rsplit(self):
         self.check_func_on_series(lambda x: repr(x.str.rsplit()), self.pser[:-1])