[SPARK-46687][PYTHON][CONNECT] Basic support of SparkSession-based memory profiler

[xinrong-meng]

What changes were proposed in this pull request?

Basic support of SparkSession-based memory profiler in both Spark Connect and non-Spark-Connect.

Why are the changes needed?

We need to make the memory profiler SparkSession-based to support memory profiling in Spark Connect.

Does this PR introduce any user-facing change?

Yes, the SparkSession-based memory profiler is available.

An example is as shown below

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.taskcontext import TaskContext

spark.conf.set("spark.sql.pyspark.udf.profiler", "memory")

@udf("string")
def f(x):
  if TaskContext.get().partitionId() % 2 == 0:
    return str(x)
  else:
    return None

spark.range(10).select(f(col("id"))).show()

spark.showMemoryProfiles()

shows profile result:

============================================================
Profile of UDF<id=2>
============================================================
Filename: /var/folders/h_/60n1p_5s7751jx1st4_sk0780000gp/T/ipykernel_72839/2848225169.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     7    113.2 MiB    113.2 MiB          10   @udf("string")
     8                                         def f(x):
     9    114.4 MiB      1.3 MiB          10     if TaskContext.get().partitionId() % 2 == 0:
    10     31.8 MiB      0.1 MiB           4       return str(x)
    11                                           else:
    12     82.8 MiB      0.1 MiB           6       return None

How was this patch tested?

New and existing unit tests:

• pyspark.tests.test_memory_profiler
• pyspark.sql.tests.connect.test_parity_memory_profiler

And manual tests on Jupyter notebook.

Was this patch authored or co-authored using generative AI tooling?

No.

[xinrong-meng]

https://github.com/xinrong-meng/spark/actions/runs/7648782322/job/20842144027 failure is irrelevant to the PR changes. I will rebase master.
@ueshin would you please review when you are free?

[ueshin]

We may want to delay to generate the full linenos until showing the results to reduce the intermediate data?

[xinrong-meng]

Good idea! Updated.

[xinrong-meng]

============================================================
Profile of UDF<id=2>
============================================================
Filename: /var/folders/h_/60n1p_5s7751jx1st4_sk0780000gp/T/ipykernel_69451/109011680.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     8    147.7 MiB    147.7 MiB          20   @udf("string")
     9                                         def a(x):
    10    149.6 MiB      1.8 MiB          20     if TaskContext.get().partitionId() % 2 == 0:
    11     59.9 MiB      0.1 MiB           8       return str(x)
    12                                           else:
    13     89.9 MiB      0.1 MiB          12       return None

tested on Jupyter.

[ueshin]

LGTM, pending tests.

[ueshin]

Thanks! merging to master.

[xinrong-meng]

Thank you @ueshin !