| # |
| # 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. |
| # |
| |
| from contextlib import contextmanager |
| import inspect |
| import tempfile |
| import unittest |
| import os |
| import sys |
| import warnings |
| from io import StringIO |
| from typing import Iterator, cast |
| |
| from pyspark import SparkConf |
| from pyspark.errors import PySparkValueError |
| from pyspark.sql import SparkSession |
| from pyspark.sql.functions import col, arrow_udf, pandas_udf, udf |
| from pyspark.sql.window import Window |
| from pyspark.profiler import UDFBasicProfiler |
| from pyspark.testing.sqlutils import ReusedSQLTestCase |
| from pyspark.testing.utils import ( |
| have_pandas, |
| have_pyarrow, |
| have_flameprof, |
| pandas_requirement_message, |
| pyarrow_requirement_message, |
| ) |
| |
| |
| def _do_computation(spark, *, action=lambda df: df.collect(), use_arrow=False): |
| @udf("long", useArrow=use_arrow) |
| def add1(x): |
| return x + 1 |
| |
| @udf("long", useArrow=use_arrow) |
| def add2(x): |
| return x + 2 |
| |
| df = spark.range(10).select(add1("id"), add2("id"), add1("id"), add2(col("id") + 1)) |
| action(df) |
| |
| |
| class UDFProfilerTests(unittest.TestCase): |
| def setUp(self): |
| self._old_sys_path = list(sys.path) |
| class_name = self.__class__.__name__ |
| conf = SparkConf().set("spark.python.profile", "true") |
| self.spark = ( |
| SparkSession.builder.master("local[4]") |
| .config(conf=conf) |
| .appName(class_name) |
| .getOrCreate() |
| ) |
| self.sc = self.spark.sparkContext |
| |
| def tearDown(self): |
| self.spark.stop() |
| sys.path = self._old_sys_path |
| |
| def test_udf_profiler(self): |
| _do_computation(self.spark) |
| |
| profilers = self.sc.profiler_collector.profilers |
| self.assertEqual(4, len(profilers)) |
| |
| old_stdout = sys.stdout |
| try: |
| sys.stdout = io = StringIO() |
| self.sc.show_profiles() |
| finally: |
| sys.stdout = old_stdout |
| |
| with tempfile.TemporaryDirectory(prefix="test_udf_profiler") as d: |
| self.sc.dump_profiles(d) |
| |
| for i, udf_name in enumerate(["add1", "add2", "add1", "add2"]): |
| id, profiler, _ = profilers[i] |
| with self.subTest(id=id, udf_name=udf_name): |
| stats = profiler.stats() |
| self.assertTrue(stats is not None) |
| width, stat_list = stats.get_print_list([]) |
| func_names = [func_name for fname, n, func_name in stat_list] |
| self.assertTrue(udf_name in func_names) |
| |
| self.assertTrue(udf_name in io.getvalue()) |
| self.assertTrue("udf_%d.pstats" % id in os.listdir(d)) |
| |
| def test_custom_udf_profiler(self): |
| class TestCustomProfiler(UDFBasicProfiler): |
| def show(self, id): |
| self.result = "Custom formatting" |
| |
| self.sc.profiler_collector.udf_profiler_cls = TestCustomProfiler |
| |
| _do_computation(self.spark) |
| |
| profilers = self.sc.profiler_collector.profilers |
| self.assertEqual(4, len(profilers)) |
| _, profiler, _ = profilers[0] |
| self.assertTrue(isinstance(profiler, TestCustomProfiler)) |
| |
| self.sc.show_profiles() |
| self.assertEqual("Custom formatting", profiler.result) |
| |
| # Unsupported |
| def exec_pandas_udf_iter_to_iter(self): |
| import pandas as pd |
| |
| @pandas_udf("int") |
| def iter_to_iter(batch_ser: Iterator[pd.Series]) -> Iterator[pd.Series]: |
| for ser in batch_ser: |
| yield ser + 1 |
| |
| self.spark.range(10).select(iter_to_iter("id")).collect() |
| |
| def exec_arrow_udf_iter_to_iter(self): |
| import pyarrow as pa |
| |
| @arrow_udf("int") |
| def iter_to_iter(iter: Iterator[pa.Array]) -> Iterator[pa.Array]: |
| for s in iter: |
| yield pa.compute.add(s, 1) |
| |
| self.spark.range(10).select(iter_to_iter("id")).collect() |
| |
| # Unsupported |
| def exec_map(self): |
| import pandas as pd |
| |
| def map(pdfs: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: |
| for pdf in pdfs: |
| yield pdf[pdf.id == 1] |
| |
| df = self.spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0)], ("id", "v")) |
| df.mapInPandas(map, schema=df.schema).collect() |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| def test_unsupported(self): |
| with warnings.catch_warnings(record=True) as warns: |
| warnings.simplefilter("always") |
| self.exec_pandas_udf_iter_to_iter() |
| user_warns = [warn.message for warn in warns if isinstance(warn.message, UserWarning)] |
| self.assertTrue(len(user_warns) > 0) |
| self.assertTrue( |
| "Profiling UDFs with iterators input/output is not supported" in str(user_warns[0]) |
| ) |
| |
| with warnings.catch_warnings(record=True) as warns: |
| warnings.simplefilter("always") |
| self.exec_arrow_udf_iter_to_iter() |
| user_warns = [warn.message for warn in warns if isinstance(warn.message, UserWarning)] |
| self.assertTrue(len(user_warns) > 0) |
| self.assertTrue( |
| "Profiling UDFs with iterators input/output is not supported" in str(user_warns[0]) |
| ) |
| |
| with warnings.catch_warnings(record=True) as warns: |
| warnings.simplefilter("always") |
| self.exec_map() |
| user_warns = [warn.message for warn in warns if isinstance(warn.message, UserWarning)] |
| self.assertTrue(len(user_warns) > 0) |
| self.assertTrue( |
| "Profiling UDFs with iterators input/output is not supported" in str(user_warns[0]) |
| ) |
| |
| |
| class UDFProfiler2TestsMixin: |
| @contextmanager |
| def trap_stdout(self): |
| old_stdout = sys.stdout |
| sys.stdout = io = StringIO() |
| try: |
| yield io |
| finally: |
| sys.stdout = old_stdout |
| |
| @property |
| def profile_results(self): |
| return self.spark._profiler_collector._perf_profile_results |
| |
| def assert_udf_profile_present(self, udf_id, expected_line_count_prefix): |
| """ |
| Assert that the performance profile for a given UDF ID is present and correctly formatted. |
| |
| This checks the output of `spark.profile.show()` for the specified UDF ID, ensures that |
| it includes a line matching the expected line count prefix and file name, and optionally |
| verifies that `spark.profile.render()` produces SVG output when flameprof is available. |
| """ |
| with self.trap_stdout() as io: |
| self.spark.profile.show(udf_id, type="perf") |
| self.assertIn(f"Profile of UDF<id={udf_id}>", io.getvalue()) |
| self.assertRegex( |
| io.getvalue(), |
| f"{expected_line_count_prefix}.*{os.path.basename(inspect.getfile(_do_computation))}", |
| ) |
| |
| if have_flameprof: |
| self.assertIn("svg", self.spark.profile.render(udf_id)) |
| |
| def test_perf_profiler_udf(self): |
| _do_computation(self.spark) |
| |
| # Without the conf enabled, no profile results are collected. |
| self.assertEqual(0, len(self.profile_results), str(list(self.profile_results))) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark) |
| |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| with self.trap_stdout(): |
| self.spark.profile.show(type="perf") |
| |
| with tempfile.TemporaryDirectory(prefix="test_perf_profiler_udf") as d: |
| self.spark.profile.dump(d, type="perf") |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=10) |
| self.assertTrue(f"udf_{id}_perf.pstats" in os.listdir(d)) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_udf_with_arrow(self): |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark, use_arrow=True) |
| |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=10) |
| |
| def test_perf_profiler_udf_multiple_actions(self): |
| def action(df): |
| df.collect() |
| df.show() |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark, action=action) |
| |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=20) |
| |
| def test_perf_profiler_udf_registered(self): |
| @udf("long") |
| def add1(x): |
| return x + 1 |
| |
| self.spark.udf.register("add1", add1) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| self.spark.sql("SELECT id, add1(id) add1 FROM range(10)").collect() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=10) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_pandas_udf(self): |
| @pandas_udf("long") |
| def add1(x): |
| return x + 1 |
| |
| @pandas_udf("long") |
| def add2(x): |
| return x + 2 |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.range(10, numPartitions=2).select( |
| add1("id"), add2("id"), add1("id"), add2(col("id") + 1) |
| ) |
| df.collect() |
| |
| self.assertEqual(3, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| def test_perf_profiler_arrow_udf(self): |
| import pyarrow as pa |
| |
| @arrow_udf("long") |
| def add1(x): |
| return pa.compute.add(x, 1) |
| |
| @arrow_udf("long") |
| def add2(x): |
| return pa.compute.add(x, 2) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.range(10, numPartitions=2).select( |
| add1("id"), add2("id"), add1("id"), add2(col("id") + 1) |
| ) |
| df.collect() |
| |
| self.assertEqual(3, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_pandas_udf_iterator_not_supported(self): |
| import pandas as pd |
| |
| @pandas_udf("long") |
| def add1(x): |
| return x + 1 |
| |
| @pandas_udf("long") |
| def add2(iter: Iterator[pd.Series]) -> Iterator[pd.Series]: |
| for s in iter: |
| yield s + 2 |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.range(10, numPartitions=2).select( |
| add1("id"), add2("id"), add1("id"), add2(col("id") + 1) |
| ) |
| df.collect() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| def test_perf_profiler_arrow_udf_iterator_not_supported(self): |
| import pyarrow as pa |
| |
| @arrow_udf("long") |
| def add1(x): |
| return pa.compute.add(x, 1) |
| |
| @arrow_udf("long") |
| def add2(iter: Iterator[pa.Array]) -> Iterator[pa.Array]: |
| for s in iter: |
| yield pa.compute.add(s, 2) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.range(10, numPartitions=2).select( |
| add1("id"), add2("id"), add1("id"), add2(col("id") + 1) |
| ) |
| df.collect() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_map_in_pandas_not_supported(self): |
| df = self.spark.createDataFrame([(1, 21), (2, 30)], ("id", "age")) |
| |
| def filter_func(iterator): |
| for pdf in iterator: |
| yield pdf[pdf.id == 1] |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df.mapInPandas(filter_func, df.schema).show() |
| |
| self.assertEqual(0, len(self.profile_results), str(self.profile_results.keys())) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_pandas_udf_window(self): |
| # WindowInPandasExec |
| import pandas as pd |
| |
| @pandas_udf("double") |
| def mean_udf(v: pd.Series) -> float: |
| return v.mean() |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| w = Window.partitionBy("id").orderBy("v").rowsBetween(-1, 0) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df.withColumn("mean_v", mean_udf("v").over(w)).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=5) |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| def test_perf_profiler_arrow_udf_window(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double") |
| def mean_udf(v: pa.Array) -> float: |
| return pa.compute.mean(v) |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| w = Window.partitionBy("id").orderBy("v").rowsBetween(-1, 0) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df.withColumn("mean_v", mean_udf("v").over(w)).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=5) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_aggregate_in_pandas(self): |
| # AggregateInPandasExec |
| import pandas as pd |
| |
| @pandas_udf("double") |
| def min_udf(v: pd.Series) -> float: |
| return v.min() |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.createDataFrame( |
| [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"] |
| ) |
| df.groupBy(df.name).agg(min_udf(df.age)).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| def test_perf_profiler_arrow_udf_agg(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double") |
| def min_udf(v: pa.Array) -> float: |
| return pa.compute.min(v) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df = self.spark.createDataFrame( |
| [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"] |
| ) |
| df.groupBy(df.name).agg(min_udf(df.age)).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_group_apply_in_pandas(self): |
| # FlatMapGroupsInBatchExec |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| |
| def normalize(pdf): |
| v = pdf.v |
| return pdf.assign(v=(v - v.mean()) / v.std()) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df.groupby("id").applyInPandas(normalize, schema="id long, v double").show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_cogroup_apply_in_pandas(self): |
| # FlatMapCoGroupsInBatchExec |
| import pandas as pd |
| |
| df1 = self.spark.createDataFrame( |
| [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)], |
| ("time", "id", "v1"), |
| ) |
| df2 = self.spark.createDataFrame( |
| [(20000101, 1, "x"), (20000101, 2, "y")], ("time", "id", "v2") |
| ) |
| |
| def asof_join(left, right): |
| return pd.merge_asof(left, right, on="time", by="id") |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas( |
| asof_join, schema="time int, id int, v1 double, v2 string" |
| ).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_group_apply_in_arrow(self): |
| # FlatMapGroupsInBatchExec |
| import pyarrow.compute as pc |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| |
| def normalize(table): |
| v = table.column("v") |
| norm = pc.divide(pc.subtract(v, pc.mean(v)), pc.stddev(v, ddof=1)) |
| return table.set_column(1, "v", norm) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df.groupby("id").applyInArrow(normalize, schema="id long, v double").show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| @unittest.skipIf( |
| not have_pandas or not have_pyarrow, |
| cast(str, pandas_requirement_message or pyarrow_requirement_message), |
| ) |
| def test_perf_profiler_cogroup_apply_in_arrow(self): |
| import pyarrow as pa |
| |
| df1 = self.spark.createDataFrame([(1, 1.0), (2, 2.0), (1, 3.0), (2, 4.0)], ("id", "v1")) |
| df2 = self.spark.createDataFrame([(1, "x"), (2, "y")], ("id", "v2")) |
| |
| def summarize(left, right): |
| return pa.Table.from_pydict({"left": [left.num_rows], "right": [right.num_rows]}) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| df1.groupby("id").cogroup(df2.groupby("id")).applyInArrow( |
| summarize, schema="left long, right long" |
| ).show() |
| |
| self.assertEqual(1, len(self.profile_results), str(self.profile_results.keys())) |
| |
| for id in self.profile_results: |
| self.assert_udf_profile_present(udf_id=id, expected_line_count_prefix=2) |
| |
| def test_perf_profiler_render(self): |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark) |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| id = list(self.profile_results.keys())[0] |
| |
| if have_flameprof: |
| self.assertIn("svg", self.spark.profile.render(id)) |
| self.assertIn("svg", self.spark.profile.render(id, type="perf")) |
| self.assertIn("svg", self.spark.profile.render(id, renderer="flameprof")) |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.spark.profile.render(id, type="unknown") |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="VALUE_NOT_ALLOWED", |
| messageParameters={ |
| "arg_name": "type", |
| "allowed_values": "['perf', 'memory']", |
| }, |
| ) |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.spark.profile.render(id, type="memory") |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="VALUE_NOT_ALLOWED", |
| messageParameters={ |
| "arg_name": "(type, renderer)", |
| "allowed_values": "[('perf', None), ('perf', 'flameprof')]", |
| }, |
| ) |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.spark.profile.render(id, renderer="unknown") |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="VALUE_NOT_ALLOWED", |
| messageParameters={ |
| "arg_name": "(type, renderer)", |
| "allowed_values": "[('perf', None), ('perf', 'flameprof')]", |
| }, |
| ) |
| |
| with self.trap_stdout() as io: |
| self.spark.profile.show(id, type="perf") |
| show_value = io.getvalue() |
| |
| with self.trap_stdout() as io: |
| self.spark.profile.render( |
| id, renderer=lambda s: s.sort_stats("time", "cumulative").print_stats() |
| ) |
| render_value = io.getvalue() |
| |
| self.assertIn(render_value, show_value) |
| |
| def test_perf_profiler_clear(self): |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark) |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| for id in self.profile_results: |
| self.spark.profile.clear(id) |
| self.assertNotIn(id, self.profile_results) |
| self.assertEqual(0, len(self.profile_results), str(list(self.profile_results))) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark) |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| self.spark.profile.clear(type="memory") |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| self.spark.profile.clear(type="perf") |
| self.assertEqual(0, len(self.profile_results), str(list(self.profile_results))) |
| |
| with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}): |
| _do_computation(self.spark) |
| self.assertEqual(3, len(self.profile_results), str(list(self.profile_results))) |
| |
| self.spark.profile.clear() |
| self.assertEqual(0, len(self.profile_results), str(list(self.profile_results))) |
| |
| |
| class UDFProfiler2Tests(UDFProfiler2TestsMixin, ReusedSQLTestCase): |
| def setUp(self) -> None: |
| super().setUp() |
| self.spark._profiler_collector._accumulator._value = None |
| |
| |
| if __name__ == "__main__": |
| from pyspark.sql.tests.test_udf_profiler import * # noqa: F401 |
| |
| try: |
| import xmlrunner # type: ignore |
| |
| testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) |
| except ImportError: |
| testRunner = None |
| unittest.main(testRunner=testRunner, verbosity=2) |
