[SPARK-35537][PYTHON] Introduce a util function spark_column_equals

python/pyspark/pandas/indexing.py

@@ -51,6 +51,7 @@
     name_like_string,
     same_anchor,
     scol_for,
+    spark_column_equals,
     verify_temp_column_name,
 )
 
@@ -710,7 +711,7 @@ def __setitem__(self, key, value):
                 self._internal.data_dtypes,
             ):
                 for scol in data_spark_columns:
-                    if new_scol._jc.equals(scol._jc):
+                    if spark_column_equals(new_scol, scol):
                         new_scol = F.when(cond, value).otherwise(scol).alias(spark_column_name)
                         new_dtype = spark_type_to_pandas_dtype(
                             self._internal.spark_frame.select(new_scol).schema[0].dataType,

python/pyspark/pandas/internal.py

@@ -54,6 +54,7 @@
     lazy_property,
     name_like_string,
     scol_for,
+    spark_column_equals,
     verify_temp_column_name,
 )
 
@@ -547,7 +548,7 @@ def __init__(
                 scol_for(spark_frame, col)
                 for col in spark_frame.columns
                 if all(
-                    not scol_for(spark_frame, col)._jc.equals(index_scol._jc)  # type: ignore
+                    not spark_column_equals(scol_for(spark_frame, col), index_scol)
                     for index_scol in index_spark_columns
                 )
                 and col not in HIDDEN_COLUMNS
@@ -879,7 +880,7 @@ def spark_columns(self) -> List[spark.Column]:
             spark_column
             for spark_column in self.data_spark_columns
             if all(
-                not spark_column._jc.equals(scol._jc)  # type: ignore
+                not spark_column_equals(spark_column, scol)
                 for scol in index_spark_columns
             )
         ]
@@ -929,7 +930,7 @@ def to_internal_spark_frame(self) -> spark.DataFrame:
         data_columns = []
         for spark_column in self.data_spark_columns:
             if all(
-                not spark_column._jc.equals(scol._jc)  # type: ignore
+                not spark_column_equals(spark_column, scol)
                 for scol in index_spark_columns
             ):
                 data_columns.append(spark_column)
@@ -967,7 +968,7 @@ def arguments_for_restore_index(self) -> Dict:
             for index_spark_column_name, index_spark_column in zip(
                 self.index_spark_column_names, self.index_spark_columns
             ):
-                if spark_column._jc.equals(index_spark_column._jc):  # type: ignore
+                if spark_column_equals(spark_column, index_spark_column):
                     column_names.append(index_spark_column_name)
                     break
             else:

python/pyspark/pandas/tests/test_internal.py

@@ -22,6 +22,7 @@
     SPARK_DEFAULT_INDEX_NAME,
     SPARK_INDEX_NAME_FORMAT,
 )
+from pyspark.pandas.utils import spark_column_equals
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
 from pyspark.testing.sqlutils import SQLTestUtils
 
@@ -37,8 +38,8 @@ def test_from_pandas(self):
         self.assert_eq(internal.index_names, [None])
         self.assert_eq(internal.column_labels, [("a",), ("b",)])
         self.assert_eq(internal.data_spark_column_names, ["a", "b"])
-        self.assertTrue(internal.spark_column_for(("a",))._jc.equals(sdf["a"]._jc))
-        self.assertTrue(internal.spark_column_for(("b",))._jc.equals(sdf["b"]._jc))
+        self.assertTrue(spark_column_equals(internal.spark_column_for(("a",)), sdf["a"]))
+        self.assertTrue(spark_column_equals(internal.spark_column_for(("b",)), sdf["b"]))
 
         self.assert_eq(internal.to_pandas_frame, pdf)
 
@@ -52,8 +53,8 @@ def test_from_pandas(self):
         self.assert_eq(internal.index_names, [None])
         self.assert_eq(internal.column_labels, [(0,), (1,)])
         self.assert_eq(internal.data_spark_column_names, ["0", "1"])
-        self.assertTrue(internal.spark_column_for((0,))._jc.equals(sdf["0"]._jc))
-        self.assertTrue(internal.spark_column_for((1,))._jc.equals(sdf["1"]._jc))
+        self.assertTrue(spark_column_equals(internal.spark_column_for((0,)), sdf["0"]))
+        self.assertTrue(spark_column_equals(internal.spark_column_for((1,)), sdf["1"]))
 
         self.assert_eq(internal.to_pandas_frame, pdf1)
 
@@ -70,7 +71,7 @@ def test_from_pandas(self):
         self.assert_eq(internal.index_names, [None, ("a",)])
         self.assert_eq(internal.column_labels, [("b",)])
         self.assert_eq(internal.data_spark_column_names, ["b"])
-        self.assertTrue(internal.spark_column_for(("b",))._jc.equals(sdf["b"]._jc))
+        self.assertTrue(spark_column_equals(internal.spark_column_for(("b",)), sdf["b"]))
 
         self.assert_eq(internal.to_pandas_frame, pdf)
 
@@ -87,7 +88,7 @@ def test_from_pandas(self):
         self.assert_eq(internal.index_names, [None, ("a",)])
         self.assert_eq(internal.column_labels, [("x", "b")])
         self.assert_eq(internal.data_spark_column_names, ["(x, b)"])
-        self.assertTrue(internal.spark_column_for(("x", "b"))._jc.equals(sdf["(x, b)"]._jc))
+        self.assertTrue(spark_column_equals(internal.spark_column_for(("x", "b")), sdf["(x, b)"]))
 
         self.assert_eq(internal.to_pandas_frame, pdf)
 

python/pyspark/pandas/tests/test_namespace.py

@@ -21,6 +21,7 @@
 
 from pyspark import pandas as ps
 from pyspark.pandas.namespace import _get_index_map
+from pyspark.pandas.utils import spark_column_equals
 from pyspark.testing.pandasutils import PandasOnSparkTestCase
 from pyspark.testing.sqlutils import SQLTestUtils
 
@@ -317,7 +318,7 @@ def check(actual, expected):
             self.assertEqual(len(actual_scols), len(expected_column_names))
             for actual_scol, expected_column_name in zip(actual_scols, expected_column_names):
                 expected_scol = sdf[expected_column_name]
-                self.assertTrue(actual_scol._jc.equals(expected_scol._jc))
+                self.assertTrue(spark_column_equals(actual_scol, expected_scol))
             self.assertEqual(actual_labels, expected_labels)
 
         check(_get_index_map(sdf, "year"), (["year"], [("year",)]))

python/pyspark/pandas/utils.py

@@ -96,7 +96,7 @@ def same_anchor(
         this_internal.spark_frame is that_internal.spark_frame
         and this_internal.index_level == that_internal.index_level
         and all(
-            this_scol._jc.equals(that_scol._jc)  # type: ignore
+            spark_column_equals(this_scol, that_scol)
             for this_scol, that_scol in zip(
                 this_internal.index_spark_columns, that_internal.index_spark_columns
             )
@@ -861,6 +861,26 @@ def verify_temp_column_name(
     return column_name_or_label
 
 
+def spark_column_equals(left: spark.Column, right: spark.Column) -> bool:
+    """
+    Check both `left` and `right` have the same expressions.
+
+    >>> spark_column_equals(F.lit(0), F.lit(0))
+    True
+    >>> spark_column_equals(F.lit(0) + 1, F.lit(0) + 1)
+    True
+    >>> spark_column_equals(F.lit(0) + 1, F.lit(0) + 2)
+    False
+    >>> sdf1 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark()
+    >>> spark_column_equals(sdf1["x"] + 1, sdf1["x"] + 1)
+    True
+    >>> sdf2 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark()
+    >>> spark_column_equals(sdf1["x"] + 1, sdf2["x"] + 1)
+    False
+    """
+    return left._jc.equals(right._jc)  # type: ignore
+
+
 def compare_null_first(
     left: spark.Column,
     right: spark.Column,