[SPARK-51983][PS] Prepare the test environment for pandas API on Spark with ANSI mode enabled

### What changes were proposed in this pull request?

Prepares the test environment for pandas API on Spark with ANSI mode enabled.

- Remove forcibly disabling ANSI mode in tests
- Add a new option `compute.ansi_mode_support` to keep the current behavior (default `False`)
    - eventually it should be `True` by default
- Skip the failed tests affected by ANSI mode
- Make `pyspark-pandas` tests run in the nightly Non-ANSI test to also run skipped tests with Non-ANSI mode

### Why are the changes needed?

Currently pandas API on Spark doesn't support ANSI mode and show warnings if it's enabled.

```py
>>> import pyspark.pandas as ps
>>> ps.range(10)
...: PandasAPIOnSparkAdviceWarning: The config 'spark.sql.ansi.enabled' is set to True. This can cause unexpected behavior from pandas API on Spark since pandas API on Spark follows the behavior of pandas, not SQL.
  warnings.warn(message, PandasAPIOnSparkAdviceWarning)

...
```

Now ANSI mode is enabled by default, pandas API on Spark should also support it.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

The existing tests should pass.

### Was this patch authored or co-authored using generative AI tooling?

Closes #50779 from ueshin/issues/SPARK-51983/pandas_ansi_test.

Authored-by: Takuya Ueshin <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: ueshin

.github/workflows/build_non_ansi.yml

@@ -46,6 +46,7 @@ jobs:
           "build": "true",
           "docs": "true",
           "pyspark": "true",
+          "pyspark-pandas": "true",
           "sparkr": "true",
           "tpcds-1g": "true",
           "docker-integration-tests": "true",

python/docs/source/tutorial/pandas_on_spark/options.rst

@@ -319,6 +319,10 @@ compute.isin_limit              80                      'compute.isin_limit' set
                                                         better performance.
 compute.pandas_fallback         False                   'compute.pandas_fallback' sets whether or not to
                                                         fallback automatically to Pandas' implementation.
+compute.ansi_mode_support       False                   'compute.ansi_mode_support' sets whether or not to
+                                                        support the ANSI mode of the underlying Spark. If
+                                                        False, pandas API on Spark may hit unexpected results
+                                                        or errors. The default is False.
 plotting.max_rows               1000                    'plotting.max_rows' sets the visual limit on top-n-
                                                         based plots such as `plot.bar` and `plot.pie`. If it
                                                         is set to 1000, the first 1000 data points will be

python/pyspark/pandas/config.py

@@ -20,9 +20,10 @@
 """
 from contextlib import contextmanager
 import json
-from typing import Any, Callable, Dict, Iterator, List, Tuple, Union
+from typing import Any, Callable, Dict, Iterator, List, Tuple, Union, Optional
 
 from pyspark._globals import _NoValue, _NoValueType
+from pyspark.sql.session import SparkSession
 from pyspark.pandas.utils import default_session
 
 
@@ -268,6 +269,17 @@ def validate(self, v: Any) -> None:
         default=False,
         types=bool,
     ),
+    Option(
+        key="compute.ansi_mode_support",
+        doc=(
+            "'compute.ansi_mode_support' sets whether or not to support the ANSI mode of "
+            "the underlying Spark. "
+            "If False, pandas API on Spark may hit unexpected results or errors. "
+            "The default is False."
+        ),
+        default=False,
+        types=bool,
+    ),
     Option(
         key="plotting.max_rows",
         doc=(
@@ -351,7 +363,12 @@ def show_options() -> None:
     print(row_format.format("=" * 31, "=" * 23, "=" * 53))
 
 
-def get_option(key: str, default: Union[Any, _NoValueType] = _NoValue) -> Any:
+def get_option(
+    key: str,
+    default: Union[Any, _NoValueType] = _NoValue,
+    *,
+    spark_session: Optional[SparkSession] = None,
+) -> Any:
     """
     Retrieves the value of the specified option.
 
@@ -361,6 +378,9 @@ def get_option(key: str, default: Union[Any, _NoValueType] = _NoValue) -> Any:
         The key which should match a single option.
     default : object
         The default value if the option is not set yet. The value should be JSON serializable.
+    spark_session : :class:`SparkSession`, optional
+        The explicit :class:`SparkSession` object to get the option.
+        If not specified, the default session will be used.
 
     Returns
     -------
@@ -374,12 +394,12 @@ def get_option(key: str, default: Union[Any, _NoValueType] = _NoValue) -> Any:
     if default is _NoValue:
         default = _options_dict[key].default
     _options_dict[key].validate(default)
-    spark_session = default_session()
+    spark_session = spark_session or default_session()
 
     return json.loads(spark_session.conf.get(_key_format(key), default=json.dumps(default)))
 
 
-def set_option(key: str, value: Any) -> None:
+def set_option(key: str, value: Any, *, spark_session: Optional[SparkSession] = None) -> None:
     """
     Sets the value of the specified option.
 
@@ -389,19 +409,22 @@ def set_option(key: str, value: Any) -> None:
         The key which should match a single option.
     value : object
         New value of option. The value should be JSON serializable.
+    spark_session : :class:`SparkSession`, optional
+        The explicit :class:`SparkSession` object to set the option.
+        If not specified, the default session will be used.
 
     Returns
     -------
     None
     """
     _check_option(key)
     _options_dict[key].validate(value)
-    spark_session = default_session()
+    spark_session = spark_session or default_session()
 
     spark_session.conf.set(_key_format(key), json.dumps(value))
 
 
-def reset_option(key: str) -> None:
+def reset_option(key: str, *, spark_session: Optional[SparkSession] = None) -> None:
     """
     Reset one option to their default value.
 
@@ -411,13 +434,17 @@ def reset_option(key: str) -> None:
     ----------
     key : str
         If specified only option will be reset.
+    spark_session : :class:`SparkSession`, optional
+        The explicit :class:`SparkSession` object to reset the option.
+        If not specified, the default session will be used.
 
     Returns
     -------
     None
     """
     _check_option(key)
-    default_session().conf.unset(_key_format(key))
+    spark_session = spark_session or default_session()
+    spark_session.conf.unset(_key_format(key))
 
 
 @contextmanager

python/pyspark/pandas/frame.py

@@ -13778,11 +13778,32 @@ def _test() -> None:
     import uuid
     from pyspark.sql import SparkSession
     import pyspark.pandas.frame
+    from pyspark.testing.utils import is_ansi_mode_test
 
     os.chdir(os.environ["SPARK_HOME"])
 
     globs = pyspark.pandas.frame.__dict__.copy()
     globs["ps"] = pyspark.pandas
+
+    if is_ansi_mode_test:
+        del pyspark.pandas.frame.DataFrame.add.__doc__
+        del pyspark.pandas.frame.DataFrame.div.__doc__
+        del pyspark.pandas.frame.DataFrame.floordiv.__doc__
+        del pyspark.pandas.frame.DataFrame.melt.__doc__
+        del pyspark.pandas.frame.DataFrame.mod.__doc__
+        del pyspark.pandas.frame.DataFrame.mul.__doc__
+        del pyspark.pandas.frame.DataFrame.pow.__doc__
+        del pyspark.pandas.frame.DataFrame.radd.__doc__
+        del pyspark.pandas.frame.DataFrame.rdiv.__doc__
+        del pyspark.pandas.frame.DataFrame.rfloordiv.__doc__
+        del pyspark.pandas.frame.DataFrame.rmod.__doc__
+        del pyspark.pandas.frame.DataFrame.rmul.__doc__
+        del pyspark.pandas.frame.DataFrame.rpow.__doc__
+        del pyspark.pandas.frame.DataFrame.rsub.__doc__
+        del pyspark.pandas.frame.DataFrame.rtruediv.__doc__
+        del pyspark.pandas.frame.DataFrame.sub.__doc__
+        del pyspark.pandas.frame.DataFrame.truediv.__doc__
+
     spark = (
         SparkSession.builder.master("local[4]").appName("pyspark.pandas.frame tests").getOrCreate()
     )

python/pyspark/pandas/groupby.py

@@ -4595,12 +4595,17 @@ def _test() -> None:
     import numpy
     from pyspark.sql import SparkSession
     import pyspark.pandas.groupby
+    from pyspark.testing.utils import is_ansi_mode_test
 
     os.chdir(os.environ["SPARK_HOME"])
 
     globs = pyspark.pandas.groupby.__dict__.copy()
     globs["np"] = numpy
     globs["ps"] = pyspark.pandas
+
+    if is_ansi_mode_test:
+        del pyspark.pandas.groupby.DataFrameGroupBy.corr.__doc__
+
     spark = (
         SparkSession.builder.master("local[4]")
         .appName("pyspark.pandas.groupby tests")

python/pyspark/pandas/namespace.py

@@ -3874,6 +3874,7 @@ def _test() -> None:
     from pyspark.sql import SparkSession
     import pyspark.pandas.namespace
     from pandas.util.version import Version
+    from pyspark.testing.utils import is_ansi_mode_test
 
     os.chdir(os.environ["SPARK_HOME"])
 
@@ -3886,6 +3887,11 @@ def _test() -> None:
     globs = pyspark.pandas.namespace.__dict__.copy()
     globs["ps"] = pyspark.pandas
     globs["sf"] = F
+
+    if is_ansi_mode_test:
+        del pyspark.pandas.namespace.melt.__doc__
+        del pyspark.pandas.namespace.to_numeric.__doc__
+
     spark = (
         SparkSession.builder.master("local[4]")
         .appName("pyspark.pandas.namespace tests")

python/pyspark/pandas/series.py

@@ -7336,11 +7336,16 @@ def _test() -> None:
     import sys
     from pyspark.sql import SparkSession
     import pyspark.pandas.series
+    from pyspark.testing.utils import is_ansi_mode_test
 
     os.chdir(os.environ["SPARK_HOME"])
 
     globs = pyspark.pandas.series.__dict__.copy()
     globs["ps"] = pyspark.pandas
+
+    if is_ansi_mode_test:
+        del pyspark.pandas.series.Series.autocorr.__doc__
+
     spark = (
         SparkSession.builder.master("local[4]").appName("pyspark.pandas.series tests").getOrCreate()
     )

python/pyspark/pandas/tests/computation/test_corr.py

@@ -22,9 +22,11 @@
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase, SPARK_CONF_ARROW_ENABLED
 from pyspark.testing.sqlutils import SQLTestUtils
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 
 
 class FrameCorrMixin:
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_dataframe_corr(self):
         pdf = pd.DataFrame(
             index=[

python/pyspark/pandas/tests/computation/test_missing_data.py

@@ -23,6 +23,7 @@
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
 from pyspark.testing.sqlutils import SQLTestUtils
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 
 
 # This file contains test cases for 'Missing data handling'
@@ -273,6 +274,7 @@ def test_fillna(self):
             pdf.fillna({("x", "a"): -1, ("x", "b"): -2, ("y", "c"): -5}),
         )
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_replace(self):
         pdf = pd.DataFrame(
             {

python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py

@@ -25,6 +25,7 @@
 from pyspark import pandas as ps
 from pyspark.pandas import option_context
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
 from pyspark.pandas.typedef.typehints import (
     extension_float_dtypes_available,
@@ -99,6 +100,7 @@ def test_mul(self):
             else:
                 self.assertRaises(TypeError, lambda: b_psser * psser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_truediv(self):
         pdf, psdf = self.pdf, self.psdf
 
@@ -114,6 +116,7 @@ def test_truediv(self):
         for col in self.non_numeric_df_cols:
             self.assertRaises(TypeError, lambda: b_psser / psdf[col])
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_floordiv(self):
         pdf, psdf = self.pdf, self.psdf
 
@@ -134,6 +137,7 @@ def test_floordiv(self):
         for col in self.non_numeric_df_cols:
             self.assertRaises(TypeError, lambda: b_psser // psdf[col])
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_mod(self):
         pdf, psdf = self.pdf, self.psdf
 
@@ -233,6 +237,7 @@ def test_rpow(self):
         self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) ** b_psser)
         self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) ** b_psser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_rmod(self):
         psdf = self.psdf
 

python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py

@@ -21,6 +21,7 @@
 
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
 
 
@@ -40,6 +41,7 @@ def float_pser(self):
     def float_psser(self):
         return ps.from_pandas(self.float_pser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_add(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:
@@ -57,6 +59,7 @@ def test_add(self):
                 else:
                     self.assertRaises(TypeError, lambda: psser + psdf[n_col])
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_sub(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:

python/pyspark/pandas/tests/data_type_ops/test_num_mod.py

@@ -22,6 +22,7 @@
 
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
 
 
@@ -34,6 +35,7 @@ def float_pser(self):
     def float_psser(self):
         return ps.from_pandas(self.float_pser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_mod(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:

python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py

@@ -22,6 +22,7 @@
 
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
 
 
@@ -34,6 +35,7 @@ def float_pser(self):
     def float_psser(self):
         return ps.from_pandas(self.float_pser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_mul(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:

python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py

@@ -22,6 +22,7 @@
 
 from pyspark import pandas as ps
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.utils import is_ansi_mode_test, ansi_mode_not_supported_message
 from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
 
 
@@ -41,6 +42,7 @@ def float_pser(self):
     def float_psser(self):
         return ps.from_pandas(self.float_pser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_radd(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:
@@ -53,6 +55,7 @@ def test_radd(self):
             self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) + psser)
             self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) + psser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_rsub(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:
@@ -65,6 +68,7 @@ def test_rsub(self):
             self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) - psser)
             self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) - psser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_rmul(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:
@@ -113,6 +117,7 @@ def test_rpow(self):
             self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) ** psser)
             self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) ** psser)
 
+    @unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
     def test_rmod(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols: