[SPARK-46517][PS][TESTS][FOLLOWUPS] Reorganize IndexingTest: factor out test_loc2d*

dev/sparktestsupport/modules.py

@@ -819,6 +819,7 @@ def __hash__(self):
         "pyspark.pandas.tests.indexes.test_indexing_basic",
         "pyspark.pandas.tests.indexes.test_indexing_iloc",
         "pyspark.pandas.tests.indexes.test_indexing_loc",
+        "pyspark.pandas.tests.indexes.test_indexing_loc_2d",
         "pyspark.pandas.tests.indexes.test_indexing_loc_multi_idx",
         "pyspark.pandas.tests.indexes.test_reindex",
         "pyspark.pandas.tests.indexes.test_rename",
@@ -1095,6 +1096,7 @@ def __hash__(self):
         "pyspark.pandas.tests.connect.indexes.test_parity_indexing_basic",
         "pyspark.pandas.tests.connect.indexes.test_parity_indexing_iloc",
         "pyspark.pandas.tests.connect.indexes.test_parity_indexing_loc",
+        "pyspark.pandas.tests.connect.indexes.test_parity_indexing_loc_2d",
         "pyspark.pandas.tests.connect.indexes.test_parity_indexing_loc_multi_idx",
         "pyspark.pandas.tests.connect.indexes.test_parity_reindex",
         "pyspark.pandas.tests.connect.indexes.test_parity_rename",

python/pyspark/pandas/tests/connect/indexes/test_parity_indexing_loc_2d.py

@@ -0,0 +1,41 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import unittest
+
+from pyspark.pandas.tests.indexes.test_indexing_loc_2d import IndexingLoc2DMixin
+from pyspark.testing.connectutils import ReusedConnectTestCase
+from pyspark.testing.pandasutils import PandasOnSparkTestUtils
+
+
+class IndexingLoc2DParityTests(
+    IndexingLoc2DMixin,
+    PandasOnSparkTestUtils,
+    ReusedConnectTestCase,
+):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.connect.indexes.test_parity_indexing_loc_2d import *  # noqa
+
+    try:
+        import xmlrunner  # type: ignore[import]
+
+        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+    except ImportError:
+        testRunner = None
+    unittest.main(testRunner=testRunner, verbosity=2)

python/pyspark/pandas/tests/indexes/test_indexing_loc_2d.py

@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class IndexingLoc2DMixin:
+    @property
+    def pdf(self):
+        return pd.DataFrame(
+            {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]},
+            index=[0, 1, 3, 5, 6, 8, 9, 9, 9],
+        )
+
+    @property
+    def pdf2(self):
+        return pd.DataFrame(
+            {0: [1, 2, 3, 4, 5, 6, 7, 8, 9], 1: [4, 5, 6, 3, 2, 1, 0, 0, 0]},
+            index=[0, 1, 3, 5, 6, 8, 9, 9, 9],
+        )
+
+    @property
+    def psdf(self):
+        return ps.from_pandas(self.pdf)
+
+    @property
+    def psdf2(self):
+        return ps.from_pandas(self.pdf2)
+
+    def test_loc2d_multiindex(self):
+        psdf = self.psdf
+        psdf = psdf.set_index("b", append=True)
+        pdf = self.pdf
+        pdf = pdf.set_index("b", append=True)
+
+        self.assert_eq(psdf.loc[:, :], pdf.loc[:, :])
+        self.assert_eq(psdf.loc[:, "a"], pdf.loc[:, "a"])
+        self.assert_eq(psdf.loc[5:5, "a"], pdf.loc[5:5, "a"])
+
+        self.assert_eq(psdf.loc[:, "a":"a"], pdf.loc[:, "a":"a"])
+        self.assert_eq(psdf.loc[:, "a":"c"], pdf.loc[:, "a":"c"])
+        self.assert_eq(psdf.loc[:, "b":"c"], pdf.loc[:, "b":"c"])
+
+    def test_loc2d(self):
+        psdf = self.psdf
+        pdf = self.pdf
+
+        # index indexer is always regarded as slice for duplicated values
+        self.assert_eq(psdf.loc[5:5, "a"], pdf.loc[5:5, "a"])
+        self.assert_eq(psdf.loc[[5], "a"], pdf.loc[[5], "a"])
+        self.assert_eq(psdf.loc[5:5, ["a"]], pdf.loc[5:5, ["a"]])
+        self.assert_eq(psdf.loc[[5], ["a"]], pdf.loc[[5], ["a"]])
+        self.assert_eq(psdf.loc[:, :], pdf.loc[:, :])
+
+        self.assert_eq(psdf.loc[3:8, "a"], pdf.loc[3:8, "a"])
+        self.assert_eq(psdf.loc[:8, "a"], pdf.loc[:8, "a"])
+        self.assert_eq(psdf.loc[3:, "a"], pdf.loc[3:, "a"])
+        self.assert_eq(psdf.loc[[8], "a"], pdf.loc[[8], "a"])
+
+        self.assert_eq(psdf.loc[3:8, ["a"]], pdf.loc[3:8, ["a"]])
+        self.assert_eq(psdf.loc[:8, ["a"]], pdf.loc[:8, ["a"]])
+        self.assert_eq(psdf.loc[3:, ["a"]], pdf.loc[3:, ["a"]])
+        # TODO?: self.assert_eq(psdf.loc[[3, 4, 3], ['a']], pdf.loc[[3, 4, 3], ['a']])
+
+        self.assertRaises(SparkPandasIndexingError, lambda: psdf.loc[3, 3, 3])
+        self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[3, 3])
+        self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[3:, 3])
+        self.assertRaises(SparkPandasIndexingError, lambda: psdf.a.loc[psdf.a % 2 == 0, 3])
+
+        self.assert_eq(psdf.loc[5, "a"], pdf.loc[5, "a"])
+        self.assert_eq(psdf.loc[9, "a"], pdf.loc[9, "a"])
+        self.assert_eq(psdf.loc[5, ["a"]], pdf.loc[5, ["a"]])
+        self.assert_eq(psdf.loc[9, ["a"]], pdf.loc[9, ["a"]])
+
+        self.assert_eq(psdf.loc[:, "a":"a"], pdf.loc[:, "a":"a"])
+        self.assert_eq(psdf.loc[:, "a":"d"], pdf.loc[:, "a":"d"])
+        self.assert_eq(psdf.loc[:, "c":"d"], pdf.loc[:, "c":"d"])
+
+        # bool list-like column select
+        bool_list = [True, False]
+        self.assert_eq(psdf.loc[:, bool_list], pdf.loc[:, bool_list])
+        self.assert_eq(psdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)])
+
+        pser = pd.Series(bool_list, index=pdf.columns)
+        self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser])
+        pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns)))
+        self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser])
+
+        self.assertRaises(IndexError, lambda: psdf.loc[:, bool_list[:-1]])
+        self.assertRaises(IndexError, lambda: psdf.loc[:, np.array(bool_list + [True])])
+        self.assertRaises(SparkPandasIndexingError, lambda: psdf.loc[:, pd.Series(bool_list)])
+
+        # non-string column names
+        psdf = self.psdf2
+        pdf = self.pdf2
+
+        self.assert_eq(psdf.loc[5:5, 0], pdf.loc[5:5, 0])
+        self.assert_eq(psdf.loc[5:5, [0]], pdf.loc[5:5, [0]])
+        self.assert_eq(psdf.loc[3:8, 0], pdf.loc[3:8, 0])
+        self.assert_eq(psdf.loc[3:8, [0]], pdf.loc[3:8, [0]])
+
+        self.assert_eq(psdf.loc[:, 0:0], pdf.loc[:, 0:0])
+        self.assert_eq(psdf.loc[:, 0:3], pdf.loc[:, 0:3])
+        self.assert_eq(psdf.loc[:, 2:3], pdf.loc[:, 2:3])
+
+    def test_loc2d_multiindex_columns(self):
+        arrays = [np.array(["bar", "bar", "baz", "baz"]), np.array(["one", "two", "one", "two"])]
+
+        pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays)
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf.loc["B":"B", "bar"], pdf.loc["B":"B", "bar"])
+        self.assert_eq(psdf.loc["B":"B", ["bar"]], pdf.loc["B":"B", ["bar"]])
+
+        self.assert_eq(psdf.loc[:, "bar":"bar"], pdf.loc[:, "bar":"bar"])
+        self.assert_eq(psdf.loc[:, "bar":("baz", "one")], pdf.loc[:, "bar":("baz", "one")])
+        self.assert_eq(
+            psdf.loc[:, ("bar", "two"):("baz", "one")], pdf.loc[:, ("bar", "two"):("baz", "one")]
+        )
+        self.assert_eq(psdf.loc[:, ("bar", "two"):"bar"], pdf.loc[:, ("bar", "two"):"bar"])
+        self.assert_eq(psdf.loc[:, "a":"bax"], pdf.loc[:, "a":"bax"])
+        self.assert_eq(
+            psdf.loc[:, ("bar", "x"):("baz", "a")],
+            pdf.loc[:, ("bar", "x"):("baz", "a")],
+            almost=True,
+        )
+
+        pdf = pd.DataFrame(
+            np.random.randn(3, 4),
+            index=["A", "B", "C"],
+            columns=pd.MultiIndex.from_tuples(
+                [("bar", "two"), ("bar", "one"), ("baz", "one"), ("baz", "two")]
+            ),
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf.loc[:, "bar":"baz"], pdf.loc[:, "bar":"baz"])
+
+        self.assertRaises(KeyError, lambda: psdf.loc[:, "bar":("baz", "one")])
+        self.assertRaises(KeyError, lambda: psdf.loc[:, ("bar", "two"):"bar"])
+
+        # bool list-like column select
+        bool_list = [True, False, True, False]
+        self.assert_eq(psdf.loc[:, bool_list], pdf.loc[:, bool_list])
+        self.assert_eq(psdf.loc[:, np.array(bool_list)], pdf.loc[:, np.array(bool_list)])
+
+        pser = pd.Series(bool_list, index=pdf.columns)
+        self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser])
+
+        pser = pd.Series(list(reversed(bool_list)), index=list(reversed(pdf.columns)))
+        self.assert_eq(psdf.loc[:, pser], pdf.loc[:, pser])
+
+        # non-string column names
+        arrays = [np.array([0, 0, 1, 1]), np.array([1, 2, 1, 2])]
+
+        pdf = pd.DataFrame(np.random.randn(3, 4), index=["A", "B", "C"], columns=arrays)
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf.loc["B":"B", 0], pdf.loc["B":"B", 0])
+        self.assert_eq(psdf.loc["B":"B", [0]], pdf.loc["B":"B", [0]])
+        self.assert_eq(psdf.loc[:, 0:0], pdf.loc[:, 0:0])
+        self.assert_eq(psdf.loc[:, 0:(1, 1)], pdf.loc[:, 0:(1, 1)])
+        self.assert_eq(psdf.loc[:, (0, 2):(1, 1)], pdf.loc[:, (0, 2):(1, 1)])
+        self.assert_eq(psdf.loc[:, (0, 2):0], pdf.loc[:, (0, 2):0])
+        self.assert_eq(psdf.loc[:, -1:2], pdf.loc[:, -1:2])
+
+    def test_loc2d_with_known_divisions(self):
+        pdf = pd.DataFrame(
+            np.random.randn(20, 5), index=list("abcdefghijklmnopqrst"), columns=list("ABCDE")
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf.loc[["a"], "A"], pdf.loc[["a"], "A"])
+        self.assert_eq(psdf.loc[["a"], ["A"]], pdf.loc[["a"], ["A"]])
+        self.assert_eq(psdf.loc["a":"o", "A"], pdf.loc["a":"o", "A"])
+        self.assert_eq(psdf.loc["a":"o", ["A"]], pdf.loc["a":"o", ["A"]])
+        self.assert_eq(psdf.loc[["n"], ["A"]], pdf.loc[["n"], ["A"]])
+        self.assert_eq(psdf.loc[["a", "c", "n"], ["A"]], pdf.loc[["a", "c", "n"], ["A"]])
+        # TODO?: self.assert_eq(psdf.loc[['t', 'b'], ['A']], pdf.loc[['t', 'b'], ['A']])
+        # TODO?: self.assert_eq(psdf.loc[['r', 'r', 'c', 'g', 'h'], ['A']],
+        # TODO?:                pdf.loc[['r', 'r', 'c', 'g', 'h'], ['A']])
+
+    @unittest.skip("TODO: should handle duplicated columns properly")
+    def test_loc2d_duplicated_columns(self):
+        pdf = pd.DataFrame(
+            np.random.randn(20, 5), index=list("abcdefghijklmnopqrst"), columns=list("AABCD")
+        )
+        psdf = ps.from_pandas(pdf)
+
+        # TODO?: self.assert_eq(psdf.loc[['a'], 'A'], pdf.loc[['a'], 'A'])
+        # TODO?: self.assert_eq(psdf.loc[['a'], ['A']], pdf.loc[['a'], ['A']])
+        self.assert_eq(psdf.loc[["j"], "B"], pdf.loc[["j"], "B"])
+        self.assert_eq(psdf.loc[["j"], ["B"]], pdf.loc[["j"], ["B"]])
+
+        # TODO?: self.assert_eq(psdf.loc['a':'o', 'A'], pdf.loc['a':'o', 'A'])
+        # TODO?: self.assert_eq(psdf.loc['a':'o', ['A']], pdf.loc['a':'o', ['A']])
+        self.assert_eq(psdf.loc["j":"q", "B"], pdf.loc["j":"q", "B"])
+        self.assert_eq(psdf.loc["j":"q", ["B"]], pdf.loc["j":"q", ["B"]])
+
+        # TODO?: self.assert_eq(psdf.loc['a':'o', 'B':'D'], pdf.loc['a':'o', 'B':'D'])
+        # TODO?: self.assert_eq(psdf.loc['a':'o', 'B':'D'], pdf.loc['a':'o', 'B':'D'])
+        # TODO?: self.assert_eq(psdf.loc['j':'q', 'B':'A'], pdf.loc['j':'q', 'B':'A'])
+        # TODO?: self.assert_eq(psdf.loc['j':'q', 'B':'A'], pdf.loc['j':'q', 'B':'A'])
+
+        self.assert_eq(psdf.loc[psdf.B > 0, "B"], pdf.loc[pdf.B > 0, "B"])
+        # TODO?: self.assert_eq(psdf.loc[psdf.B > 0, ['A', 'C']], pdf.loc[pdf.B > 0, ['A', 'C']])
+
+
+class IndexingLoc2DTests(
+    IndexingLoc2DMixin,
+    PandasOnSparkTestCase,
+    SQLTestUtils,
+):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.indexes.test_indexing_loc_2d import *  # noqa: F401
+
+    try:
+        import xmlrunner
+
+        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+    except ImportError:
+        testRunner = None
+    unittest.main(testRunner=testRunner, verbosity=2)