snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/test_dataframe_spark_io.py

@@ -0,0 +1,487 @@
+#
+# 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 glob
+import os
+from distutils.version import LooseVersion
+
+import numpy as np
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils
+
+
+class DataFrameSparkIOTestsMixin:
+    """Test cases for big data I/O using Spark."""
+
+    @property
+    def test_column_order(self):
+        return ["i32", "i64", "f", "bhello"]
+
+    @property
+    def test_pdf(self):
+        pdf = pd.DataFrame(
+            {
+                "i32": np.arange(20, dtype=np.int32) % 3,
+                "i64": np.arange(20, dtype=np.int64) % 5,
+                "f": np.arange(20, dtype=np.float64),
+                "bhello": np.random.choice(["hello", "yo", "people"], size=20).astype("O"),
+            },
+            columns=self.test_column_order,
+            index=np.random.rand(20),
+        )
+        return pdf
+
+    def test_parquet_read(self):
+        with self.temp_dir() as tmp:
+            data = self.test_pdf
+            self.spark.createDataFrame(data, "i32 int, i64 long, f double, bhello string").coalesce(
+                1
+            ).write.parquet(tmp, mode="overwrite")
+
+            def check(columns):
+                expected = pd.read_parquet(tmp, columns=columns)
+                actual = ps.read_parquet(tmp, columns=columns)
+                self.assertPandasEqual(expected, actual._to_pandas())
+
+            check(None)
+            check(["i32", "i64"])
+            check(["i64", "i32"])
+
+            # check with pyspark patch.
+            expected = pd.read_parquet(tmp)
+            actual = ps.read_parquet(tmp)
+            self.assertPandasEqual(expected, actual._to_pandas())
+
+            # When index columns are known
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+
+            expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]]
+            actual_idx = ps.read_parquet(tmp, index_col="bhello")[["f", "i32", "i64"]]
+            self.assert_eq(
+                actual_idx.sort_values(by="f").to_spark().toPandas(),
+                expected_idx.sort_values(by="f").to_spark().toPandas(),
+            )
+
+    def test_parquet_read_with_pandas_metadata(self):
+        with self.temp_dir() as tmp:
+            expected1 = self.test_pdf
+
+            path1 = "{}/file1.parquet".format(tmp)
+            expected1.to_parquet(path1)
+
+            self.assert_eq(ps.read_parquet(path1, pandas_metadata=True), expected1)
+
+            expected2 = expected1.reset_index()
+
+            path2 = "{}/file2.parquet".format(tmp)
+            expected2.to_parquet(path2)
+
+            self.assert_eq(ps.read_parquet(path2, pandas_metadata=True), expected2)
+
+            expected3 = expected2.set_index("index", append=True)
+            # There is a bug in `to_parquet` from pandas 1.5.0 when writing MultiIndex.
+            # See https://github.com/pandas-dev/pandas/issues/48848 for the reported issue.
+            if LooseVersion(pd.__version__) > LooseVersion("1.5.0"):
+                expected_psdf = ps.read_parquet(path2, pandas_metadata=True).set_index(
+                    "index", append=True
+                )
+            else:
+                path3 = "{}/file3.parquet".format(tmp)
+                expected3.to_parquet(path3)
+                expected_psdf = ps.read_parquet(path3, pandas_metadata=True)
+
+            self.assert_eq(expected_psdf, expected3)
+
+    def test_parquet_write(self):
+        with self.temp_dir() as tmp:
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+
+            # Write out partitioned by one column
+            expected.to_parquet(tmp, mode="overwrite", partition_cols="i32")
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Write out partitioned by two columns
+            expected.to_parquet(tmp, mode="overwrite", partition_cols=["i32", "bhello"])
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Set `compression` with string
+            expected.to_parquet(tmp, mode="overwrite", partition_cols="i32", compression="none")
+            actual = ps.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Test `options` parameter
+            expected.to_parquet(
+                tmp, mode="overwrite", partition_cols="i32", options={"compression": "none"}
+            )
+            actual = ps.read_parquet(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+    def test_table(self):
+        with self.table("test_table"):
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+
+            # Write out partitioned by one column
+            expected.spark.to_table("test_table", mode="overwrite", partition_cols="i32")
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_table("test_table")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Write out partitioned by two columns
+            expected.to_table("test_table", mode="overwrite", partition_cols=["i32", "bhello"])
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_table("test_table")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # When index columns are known
+            expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]]
+            actual_idx = ps.read_table("test_table", index_col="bhello")[["f", "i32", "i64"]]
+            self.assert_eq(
+                actual_idx.sort_values(by="f").to_spark().toPandas(),
+                expected_idx.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            expected_idx = expected.set_index(["bhello"])[["f", "i32", "i64"]]
+            actual_idx = ps.read_table("test_table", index_col=["bhello"])[["f", "i32", "i64"]]
+            self.assert_eq(
+                actual_idx.sort_values(by="f").to_spark().toPandas(),
+                expected_idx.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            expected_idx = expected.set_index(["i32", "bhello"])[["f", "i64"]]
+            actual_idx = ps.read_table("test_table", index_col=["i32", "bhello"])[["f", "i64"]]
+            self.assert_eq(
+                actual_idx.sort_values(by="f").to_spark().toPandas(),
+                expected_idx.sort_values(by="f").to_spark().toPandas(),
+            )
+
+    def test_spark_io(self):
+        with self.temp_dir() as tmp:
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+
+            # Write out partitioned by one column
+            expected.to_spark_io(tmp, format="json", mode="overwrite", partition_cols="i32")
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_spark_io(tmp, format="json")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Write out partitioned by two columns
+            expected.to_spark_io(
+                tmp, format="json", mode="overwrite", partition_cols=["i32", "bhello"]
+            )
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_spark_io(path=tmp, format="json")
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # When index columns are known
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+            col_order = ["f", "i32", "i64"]
+
+            expected_idx = expected.set_index("bhello")[col_order]
+            actual_idx = ps.read_spark_io(tmp, format="json", index_col="bhello")[col_order]
+            self.assert_eq(
+                actual_idx.sort_values(by="f").to_spark().toPandas(),
+                expected_idx.sort_values(by="f").to_spark().toPandas(),
+            )
+
+    # TODO(SPARK-40353): re-enabling the `test_read_excel`.
+    @unittest.skip("openpyxl")
+    def test_read_excel(self):
+        with self.temp_dir() as tmp:
+
+            path1 = "{}/file1.xlsx".format(tmp)
+            self.test_pdf[["i32"]].to_excel(path1)
+
+            self.assert_eq(ps.read_excel(open(path1, "rb")), pd.read_excel(open(path1, "rb")))
+            self.assert_eq(
+                ps.read_excel(open(path1, "rb"), index_col=0),
+                pd.read_excel(open(path1, "rb"), index_col=0),
+            )
+            self.assert_eq(
+                ps.read_excel(open(path1, "rb"), index_col=0, squeeze=True),
+                pd.read_excel(open(path1, "rb"), index_col=0, squeeze=True),
+            )
+
+            self.assert_eq(ps.read_excel(path1), pd.read_excel(path1))
+            self.assert_eq(ps.read_excel(path1, index_col=0), pd.read_excel(path1, index_col=0))
+            self.assert_eq(
+                ps.read_excel(path1, index_col=0, squeeze=True),
+                pd.read_excel(path1, index_col=0, squeeze=True),
+            )
+
+            self.assert_eq(ps.read_excel(tmp), pd.read_excel(path1))
+
+            path2 = "{}/file2.xlsx".format(tmp)
+            self.test_pdf[["i32"]].to_excel(path2)
+            self.assert_eq(
+                ps.read_excel(tmp, index_col=0).sort_index(),
+                pd.concat(
+                    [pd.read_excel(path1, index_col=0), pd.read_excel(path2, index_col=0)]
+                ).sort_index(),
+            )
+            self.assert_eq(
+                ps.read_excel(tmp, index_col=0, squeeze=True).sort_index(),
+                pd.concat(
+                    [
+                        pd.read_excel(path1, index_col=0, squeeze=True),
+                        pd.read_excel(path2, index_col=0, squeeze=True),
+                    ]
+                ).sort_index(),
+            )
+
+        with self.temp_dir() as tmp:
+            path1 = "{}/file1.xlsx".format(tmp)
+            with pd.ExcelWriter(path1) as writer:
+                self.test_pdf.to_excel(writer, sheet_name="Sheet_name_1")
+                self.test_pdf[["i32"]].to_excel(writer, sheet_name="Sheet_name_2")
+
+            sheet_names = [["Sheet_name_1", "Sheet_name_2"], None]
+
+            pdfs1 = pd.read_excel(open(path1, "rb"), sheet_name=None, index_col=0)
+            pdfs1_squeezed = pd.read_excel(
+                open(path1, "rb"), sheet_name=None, index_col=0, squeeze=True
+            )
+
+            for sheet_name in sheet_names:
+                psdfs = ps.read_excel(open(path1, "rb"), sheet_name=sheet_name, index_col=0)
+                self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"])
+                self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"])
+
+                psdfs = ps.read_excel(
+                    open(path1, "rb"), sheet_name=sheet_name, index_col=0, squeeze=True
+                )
+                self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"])
+                self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"])
+
+            self.assert_eq(
+                ps.read_excel(tmp, index_col=0, sheet_name="Sheet_name_2"),
+                pdfs1["Sheet_name_2"],
+            )
+
+            for sheet_name in sheet_names:
+                psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0)
+                self.assert_eq(psdfs["Sheet_name_1"], pdfs1["Sheet_name_1"])
+                self.assert_eq(psdfs["Sheet_name_2"], pdfs1["Sheet_name_2"])
+
+                psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True)
+                self.assert_eq(psdfs["Sheet_name_1"], pdfs1_squeezed["Sheet_name_1"])
+                self.assert_eq(psdfs["Sheet_name_2"], pdfs1_squeezed["Sheet_name_2"])
+
+            path2 = "{}/file2.xlsx".format(tmp)
+            with pd.ExcelWriter(path2) as writer:
+                self.test_pdf.to_excel(writer, sheet_name="Sheet_name_1")
+                self.test_pdf[["i32"]].to_excel(writer, sheet_name="Sheet_name_2")
+
+            pdfs2 = pd.read_excel(path2, sheet_name=None, index_col=0)
+            pdfs2_squeezed = pd.read_excel(path2, sheet_name=None, index_col=0, squeeze=True)
+
+            self.assert_eq(
+                ps.read_excel(tmp, sheet_name="Sheet_name_2", index_col=0).sort_index(),
+                pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(),
+            )
+            self.assert_eq(
+                ps.read_excel(
+                    tmp, sheet_name="Sheet_name_2", index_col=0, squeeze=True
+                ).sort_index(),
+                pd.concat(
+                    [pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]]
+                ).sort_index(),
+            )
+
+            for sheet_name in sheet_names:
+                psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0)
+                self.assert_eq(
+                    psdfs["Sheet_name_1"].sort_index(),
+                    pd.concat([pdfs1["Sheet_name_1"], pdfs2["Sheet_name_1"]]).sort_index(),
+                )
+                self.assert_eq(
+                    psdfs["Sheet_name_2"].sort_index(),
+                    pd.concat([pdfs1["Sheet_name_2"], pdfs2["Sheet_name_2"]]).sort_index(),
+                )
+
+                psdfs = ps.read_excel(tmp, sheet_name=sheet_name, index_col=0, squeeze=True)
+                self.assert_eq(
+                    psdfs["Sheet_name_1"].sort_index(),
+                    pd.concat(
+                        [pdfs1_squeezed["Sheet_name_1"], pdfs2_squeezed["Sheet_name_1"]]
+                    ).sort_index(),
+                )
+                self.assert_eq(
+                    psdfs["Sheet_name_2"].sort_index(),
+                    pd.concat(
+                        [pdfs1_squeezed["Sheet_name_2"], pdfs2_squeezed["Sheet_name_2"]]
+                    ).sort_index(),
+                )
+
+    def test_read_orc(self):
+        with self.temp_dir() as tmp:
+            path = "{}/file1.orc".format(tmp)
+            data = self.test_pdf
+            self.spark.createDataFrame(data, "i32 int, i64 long, f double, bhello string").coalesce(
+                1
+            ).write.orc(path, mode="overwrite")
+
+            expected = data.reset_index()[data.columns]
+            actual = ps.read_orc(path)
+            self.assertPandasEqual(expected, actual._to_pandas())
+
+            # columns
+            columns = ["i32", "i64"]
+            expected = data.reset_index()[columns]
+            actual = ps.read_orc(path, columns=columns)
+            self.assertPandasEqual(expected, actual._to_pandas())
+
+            # index_col
+            expected = data.set_index("i32")
+            actual = ps.read_orc(path, index_col="i32")
+            self.assert_eq(actual, expected)
+
+            expected = data.set_index(["i32", "f"])
+            actual = ps.read_orc(path, index_col=["i32", "f"])
+            self.assert_eq(actual, expected)
+
+            # index_col with columns
+            expected = data.set_index("i32")[["i64", "bhello"]]
+            actual = ps.read_orc(path, index_col=["i32"], columns=["i64", "bhello"])
+            self.assert_eq(actual, expected)
+
+            expected = data.set_index(["i32", "f"])[["bhello", "i64"]]
+            actual = ps.read_orc(path, index_col=["i32", "f"], columns=["bhello", "i64"])
+            self.assert_eq(actual, expected)
+
+            msg = "Unknown column name 'i'"
+            with self.assertRaises(ValueError, msg=msg):
+                ps.read_orc(path, columns="i32")
+            msg = "Unknown column name 'i34'"
+            with self.assertRaises(ValueError, msg=msg):
+                ps.read_orc(path, columns=["i34", "i64"])
+
+    def test_orc_write(self):
+        with self.temp_dir() as tmp:
+            pdf = self.test_pdf
+            expected = ps.DataFrame(pdf)
+
+            # Write out partitioned by one column
+            expected.to_orc(tmp, mode="overwrite", partition_cols="i32")
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_orc(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Write out partitioned by two columns
+            expected.to_orc(tmp, mode="overwrite", partition_cols=["i32", "bhello"])
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_orc(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+            # Test `options` parameter
+            expected.to_orc(
+                tmp, mode="overwrite", partition_cols="i32", options={"compression": "none"}
+            )
+            # Reset column order, as once the data is written out, Spark rearranges partition
+            # columns to appear first.
+            actual = ps.read_orc(tmp)
+            self.assertFalse((actual.columns == self.test_column_order).all())
+            actual = actual[self.test_column_order]
+            self.assert_eq(
+                actual.sort_values(by="f").to_spark().toPandas(),
+                expected.sort_values(by="f").to_spark().toPandas(),
+            )
+
+
+class DataFrameSparkIOTests(DataFrameSparkIOTestsMixin, PandasOnSparkTestCase, TestUtils):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.test_dataframe_spark_io 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)

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/test_default_index.py

@@ -0,0 +1,109 @@
+#
+# 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 pandas as pd
+
+from pyspark.sql import functions as F
+from pyspark import pandas as ps
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+
+
+class DefaultIndexTestsMixin:
+    def test_default_index_sequence(self):
+        with ps.option_context("compute.default_index_type", "sequence"):
+            sdf = self.spark.range(1000)
+            self.assert_eq(ps.DataFrame(sdf), pd.DataFrame({"id": list(range(1000))}))
+
+    def test_default_index_distributed_sequence(self):
+        with ps.option_context("compute.default_index_type", "distributed-sequence"):
+            sdf = self.spark.range(1000)
+            self.assert_eq(ps.DataFrame(sdf), pd.DataFrame({"id": list(range(1000))}))
+
+    def test_default_index_distributed(self):
+        with ps.option_context("compute.default_index_type", "distributed"):
+            sdf = self.spark.range(1000)
+            pdf = ps.DataFrame(sdf)._to_pandas()
+            self.assertEqual(len(set(pdf.index)), len(pdf))
+
+    def test_index_distributed_sequence_cleanup(self):
+        with ps.option_context(
+            "compute.default_index_type", "distributed-sequence"
+        ), ps.option_context("compute.ops_on_diff_frames", True):
+
+            with ps.option_context("compute.default_index_cache", "LOCAL_CHECKPOINT"):
+                cached_rdd_ids = [rdd_id for rdd_id in self.spark._jsc.getPersistentRDDs()]
+
+                psdf1 = (
+                    self.spark.range(0, 100, 1, 10).withColumn("Key", F.col("id") % 33).pandas_api()
+                )
+
+                psdf2 = psdf1["Key"].reset_index()
+                psdf2["index"] = (psdf2.groupby(["Key"]).cumcount() == 0).astype(int)
+                psdf2["index"] = psdf2["index"].cumsum()
+
+                psdf3 = ps.merge(psdf1, psdf2, how="inner", left_on=["Key"], right_on=["Key"])
+                _ = len(psdf3)
+
+                # newly cached rdd
+                self.assertTrue(
+                    any(
+                        rdd_id not in cached_rdd_ids
+                        for rdd_id in self.spark._jsc.getPersistentRDDs()
+                    )
+                )
+
+            for storage_level in ["NONE", "DISK_ONLY_2", "MEMORY_AND_DISK_SER"]:
+                with ps.option_context("compute.default_index_cache", storage_level):
+                    cached_rdd_ids = [rdd_id for rdd_id in self.spark._jsc.getPersistentRDDs()]
+
+                    psdf1 = (
+                        self.spark.range(0, 100, 1, 10)
+                        .withColumn("Key", F.col("id") % 33)
+                        .pandas_api()
+                    )
+
+                    psdf2 = psdf1["Key"].reset_index()
+                    psdf2["index"] = (psdf2.groupby(["Key"]).cumcount() == 0).astype(int)
+                    psdf2["index"] = psdf2["index"].cumsum()
+
+                    psdf3 = ps.merge(psdf1, psdf2, how="inner", left_on=["Key"], right_on=["Key"])
+                    _ = len(psdf3)
+
+                    # no newly cached rdd
+                    self.assertTrue(
+                        all(
+                            rdd_id in cached_rdd_ids
+                            for rdd_id in self.spark._jsc.getPersistentRDDs()
+                        )
+                    )
+
+
+class DefaultIndexTests(DefaultIndexTestsMixin, PandasOnSparkTestCase):
+    pass
+
+
+if __name__ == "__main__":
+    import unittest
+    from pyspark.pandas.tests.test_default_index 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)