content: #
  # 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.sql.functions import array, explode, col, lit, udf, pandas_udf
  from pyspark.sql.types import DoubleType, StructType, StructField, Row
  from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
      pandas_requirement_message, pyarrow_requirement_message
  from pyspark.testing.utils import QuietTest
  
  if have_pandas:
      import pandas as pd
      from pandas.testing import assert_frame_equal
  
  if have_pyarrow:
      import pyarrow as pa  # noqa: F401
  
  
  @unittest.skipIf(
      not have_pandas or not have_pyarrow,
      pandas_requirement_message or pyarrow_requirement_message)  # type: ignore[arg-type]
  class CogroupedMapInPandasTests(ReusedSQLTestCase):
  
      @property
      def data1(self):
          return self.spark.range(10).toDF('id') \
              .withColumn("ks", array([lit(i) for i in range(20, 30)])) \
              .withColumn("k", explode(col('ks')))\
              .withColumn("v", col('k') * 10)\
              .drop('ks')
  
      @property
      def data2(self):
          return self.spark.range(10).toDF('id') \
              .withColumn("ks", array([lit(i) for i in range(20, 30)])) \
              .withColumn("k", explode(col('ks'))) \
              .withColumn("v2", col('k') * 100) \
              .drop('ks')
  
      def test_simple(self):
          self._test_merge(self.data1, self.data2)
  
      def test_left_group_empty(self):
          left = self.data1.where(col("id") % 2 == 0)
          self._test_merge(left, self.data2)
  
      def test_right_group_empty(self):
          right = self.data2.where(col("id") % 2 == 0)
          self._test_merge(self.data1, right)
  
      def test_different_schemas(self):
          right = self.data2.withColumn('v3', lit('a'))
          self._test_merge(self.data1, right, 'id long, k int, v int, v2 int, v3 string')
  
      def test_complex_group_by(self):
          left = pd.DataFrame.from_dict({
              'id': [1, 2, 3],
              'k':  [5, 6, 7],
              'v': [9, 10, 11]
          })
  
          right = pd.DataFrame.from_dict({
              'id': [11, 12, 13],
              'k': [5, 6, 7],
              'v2': [90, 100, 110]
          })
  
          left_gdf = self.spark\
              .createDataFrame(left)\
              .groupby(col('id') % 2 == 0)
  
          right_gdf = self.spark \
              .createDataFrame(right) \
              .groupby(col('id') % 2 == 0)
  
          def merge_pandas(l, r):
              return pd.merge(l[['k', 'v']], r[['k', 'v2']], on=['k'])
  
          result = left_gdf \
              .cogroup(right_gdf) \
              .applyInPandas(merge_pandas, 'k long, v long, v2 long') \
              .sort(['k']) \
              .toPandas()
  
          expected = pd.DataFrame.from_dict({
              'k': [5, 6, 7],
              'v': [9, 10, 11],
              'v2': [90, 100, 110]
          })
  
          assert_frame_equal(expected, result)
  
      def test_empty_group_by(self):
          left = self.data1
          right = self.data2
  
          def merge_pandas(l, r):
              return pd.merge(l, r, on=['id', 'k'])
  
          result = left.groupby().cogroup(right.groupby())\
              .applyInPandas(merge_pandas, 'id long, k int, v int, v2 int') \
              .sort(['id', 'k']) \
              .toPandas()
  
          left = left.toPandas()
          right = right.toPandas()
  
          expected = pd \
              .merge(left, right, on=['id', 'k']) \
              .sort_values(by=['id', 'k'])
  
          assert_frame_equal(expected, result)
  
      def test_mixed_scalar_udfs_followed_by_cogrouby_apply(self):
          df = self.spark.range(0, 10).toDF('v1')
          df = df.withColumn('v2', udf(lambda x: x + 1, 'int')(df['v1'])) \
              .withColumn('v3', pandas_udf(lambda x: x + 2, 'int')(df['v1']))
  
          result = df.groupby().cogroup(df.groupby()) \
              .applyInPandas(lambda x, y: pd.DataFrame([(x.sum().sum(), y.sum().sum())]),
                             'sum1 int, sum2 int').collect()
  
          self.assertEqual(result[0]['sum1'], 165)
          self.assertEqual(result[0]['sum2'], 165)
  
      def test_with_key_left(self):
          self._test_with_key(self.data1, self.data1, isLeft=True)
  
      def test_with_key_right(self):
          self._test_with_key(self.data1, self.data1, isLeft=False)
  
      def test_with_key_left_group_empty(self):
          left = self.data1.where(col("id") % 2 == 0)
          self._test_with_key(left, self.data1, isLeft=True)
  
      def test_with_key_right_group_empty(self):
          right = self.data1.where(col("id") % 2 == 0)
          self._test_with_key(self.data1, right, isLeft=False)
  
      def test_with_key_complex(self):
  
          def left_assign_key(key, l, _):
              return l.assign(key=key[0])
  
          result = self.data1 \
              .groupby(col('id') % 2 == 0)\
              .cogroup(self.data2.groupby(col('id') % 2 == 0)) \
              .applyInPandas(left_assign_key, 'id long, k int, v int, key boolean') \
              .sort(['id', 'k']) \
              .toPandas()
  
          expected = self.data1.toPandas()
          expected = expected.assign(key=expected.id % 2 == 0)
  
          assert_frame_equal(expected, result)
  
      def test_wrong_return_type(self):
          # Test that we get a sensible exception invalid values passed to apply
          left = self.data1
          right = self.data2
          with QuietTest(self.sc):
              with self.assertRaisesRegex(
                      NotImplementedError,
                      'Invalid return type.*ArrayType.*TimestampType'):
                  left.groupby('id').cogroup(right.groupby('id')).applyInPandas(
                      lambda l, r: l, 'id long, v array<timestamp>')
  
      def test_wrong_args(self):
          left = self.data1
          right = self.data2
          with self.assertRaisesRegex(ValueError, 'Invalid function'):
              left.groupby('id').cogroup(right.groupby('id')) \
                  .applyInPandas(lambda: 1, StructType([StructField("d", DoubleType())]))
  
      def test_case_insensitive_grouping_column(self):
          # SPARK-31915: case-insensitive grouping column should work.
          df1 = self.spark.createDataFrame([(1, 1)], ("column", "value"))
  
          row = df1.groupby("ColUmn").cogroup(
              df1.groupby("COLUMN")
          ).applyInPandas(lambda r, l: r + l, "column long, value long").first()
          self.assertEqual(row.asDict(), Row(column=2, value=2).asDict())
  
          df2 = self.spark.createDataFrame([(1, 1)], ("column", "value"))
  
          row = df1.groupby("ColUmn").cogroup(
              df2.groupby("COLUMN")
          ).applyInPandas(lambda r, l: r + l, "column long, value long").first()
          self.assertEqual(row.asDict(), Row(column=2, value=2).asDict())
  
      def test_self_join(self):
          # SPARK-34319: self-join with FlatMapCoGroupsInPandas
          df = self.spark.createDataFrame([(1, 1)], ("column", "value"))
  
          row = df.groupby("ColUmn").cogroup(
              df.groupby("COLUMN")
          ).applyInPandas(lambda r, l: r + l, "column long, value long")
  
          row = row.join(row).first()
  
          self.assertEqual(row.asDict(), Row(column=2, value=2).asDict())
  
      @staticmethod
      def _test_with_key(left, right, isLeft):
  
          def right_assign_key(key, l, r):
              return l.assign(key=key[0]) if isLeft else r.assign(key=key[0])
  
          result = left \
              .groupby('id') \
              .cogroup(right.groupby('id')) \
              .applyInPandas(right_assign_key, 'id long, k int, v int, key long') \
              .toPandas()
  
          expected = left.toPandas() if isLeft else right.toPandas()
          expected = expected.assign(key=expected.id)
  
          assert_frame_equal(expected, result)
  
      @staticmethod
      def _test_merge(left, right, output_schema='id long, k int, v int, v2 int'):
  
          def merge_pandas(l, r):
              return pd.merge(l, r, on=['id', 'k'])
  
          result = left \
              .groupby('id') \
              .cogroup(right.groupby('id')) \
              .applyInPandas(merge_pandas, output_schema)\
              .sort(['id', 'k']) \
              .toPandas()
  
          left = left.toPandas()
          right = right.toPandas()
  
          expected = pd \
              .merge(left, right, on=['id', 'k']) \
              .sort_values(by=['id', 'k'])
  
          assert_frame_equal(expected, result)
  
  
  if __name__ == "__main__":
      from pyspark.sql.tests.test_pandas_cogrouped_map import *  # noqa: F401
  
      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)
  