simplify test

Author: benhurdelhey

python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py

@@ -61,7 +61,12 @@ def conf(cls):
         cfg.set("spark.sql.shuffle.partitions", "5")
         return cfg
 
-    def _test_apply_in_pandas_with_state_basic(self, func, check_results):
+    def _test_apply_in_pandas_with_state_basic(self, func, check_results, output_type=None):
+        if output_type is None:
+            output_type = StructType(
+                [StructField("key", StringType()), StructField("countAsString", StringType())]
+            )
+
         input_path = tempfile.mkdtemp()
 
         def prepare_test_resource():
@@ -77,9 +82,6 @@ def prepare_test_resource():
             q.stop()
         self.assertTrue(df.isStreaming)
 
-        output_type = StructType(
-            [StructField("key", StringType()), StructField("countAsString", StringType())]
-        )
         state_type = StructType([StructField("c", LongType())])
 
         q = (
@@ -316,90 +318,25 @@ def assert_test():
         finally:
             q.stop()
 
-    def _test_apply_in_pandas_with_state_decimal_coercion(self, coercion_enabled, should_succeed):
-        input_path = tempfile.mkdtemp()
-
-        with open(input_path + "/numeric-test.txt", "w") as fw:
-            fw.write("group1,123\ngroup2,456\ngroup1,789\n")
-
-        df = (
-            self.spark.readStream.format("csv")
-            .option("header", "false")
-            .schema("key string, value int")
-            .load(input_path)
-        )
+    def test_apply_in_pandas_with_state_int_to_decimal_coercion(self):
+        def func(key, pdf_iter, state):
+            assert isinstance(state, GroupState)
+            yield pd.DataFrame({"key": [key[0]], "decimal_sum": [1]})
 
-        for q in self.spark.streams.active:
-            q.stop()
-        self.assertTrue(df.isStreaming)
+        def check_results(batch_df, _):
+            assert set(batch_df.sort("key").collect()) == {
+                Row(key="hello", decimal_sum=Decimal("1.00")),
+                Row(key="this", decimal_sum=Decimal("1.00")),
+            }, "Decimal coercion failed: " + str(batch_df.sort("key").collect())
 
         output_type = StructType(
-            [
-                StructField("key", StringType()),
-                StructField("decimal_sum", DecimalType(10, 2)),
-                StructField("count", LongType()),
-            ]
+            [StructField("key", StringType()), StructField("decimal_sum", DecimalType(10, 2))]
         )
-        state_type = StructType([StructField("sum", LongType()), StructField("count", LongType())])
-
-        def stateful_func(key, pdf_iter, state):
-            current_sum = state.get[0] if state.exists else 0
-            current_count = state.get[1] if state.exists else 0
-
-            for pdf in pdf_iter:
-                current_sum += pdf["value"].sum()
-                current_count += len(pdf)
-
-            state.update((current_sum, current_count))
-            yield pd.DataFrame(
-                {"key": [key[0]], "decimal_sum": [current_sum], "count": [current_count]}
-            )
-
-        def check_results(batch_df, _):
-            if should_succeed:
-                results = batch_df.sort("key").collect()
-                for row in results:
-                    assert isinstance(row["decimal_sum"], Decimal)
-                    if row["key"] == "group1":
-                        assert row["decimal_sum"] == Decimal("912.00")
-                    elif row["key"] == "group2":
-                        assert row["decimal_sum"] == Decimal("456.00")
 
         with self.sql_conf(
-            {"spark.sql.execution.pythonUDF.pandas.intToDecimalCoercionEnabled": coercion_enabled}
+            {"spark.sql.execution.pythonUDF.pandas.intToDecimalCoercionEnabled": True}
         ):
-            q = (
-                df.groupBy(df["key"])
-                .applyInPandasWithState(
-                    stateful_func, output_type, state_type, "Update", GroupStateTimeout.NoTimeout
-                )
-                .writeStream.queryName(f"test_coercion_{coercion_enabled}")
-                .foreachBatch(check_results)
-                .outputMode("update")
-                .start()
-            )
-
-            self.assertTrue(q.isActive)
-
-            if should_succeed:
-                q.processAllAvailable()
-                self.assertTrue(q.exception() is None)
-            else:
-                with self.assertRaises(Exception) as context:
-                    q.processAllAvailable()
-                self.assertIn("STREAM_FAILED", str(context.exception))
-
-            q.stop()
-
-    def test_apply_in_pandas_with_state_int_to_decimal_coercion(self):
-        self._test_apply_in_pandas_with_state_decimal_coercion(
-            coercion_enabled=True, should_succeed=True
-        )
-
-        self._test_apply_in_pandas_with_state_decimal_coercion(
-            coercion_enabled=False, should_succeed=False
-        )
-
+            self._test_apply_in_pandas_with_state_basic(func, check_results, output_type)
 
 class GroupedApplyInPandasWithStateTests(
     GroupedApplyInPandasWithStateTestsMixin, ReusedSQLTestCase