Applying Pandas UDF without shuffling - apache-spark

Question

I am trying to apply a pandas UDF to each partition of a Spark (3.3.0) DataFrame separately so as to avoid any shuffling requirements. However, when I run the query below, a lot of data is getting shuffled around. The execution plan contains a SORT stage; this might be the culprit.
from pyspark.sql.functions import spark_partition_id
query = df.groupBy(spark_partition_id())\
.applyInPandas(lambda x: pd.DataFrame([x.shape]), "n_rows long, n_cols long")
query.explain()
Output:
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- FlatMapGroupsInPandas [SPARK_PARTITION_ID()#1562], <lambda>(id#0L, date#1L, feature#2, partition_id#926)#1561, [nr#1563L, nc#1564L]
+- Sort [SPARK_PARTITION_ID()#1562 ASC NULLS FIRST], false, 0
+- Exchange hashpartitioning(SPARK_PARTITION_ID()#1562, 200), ENSURE_REQUIREMENTS, [id=#748]
+- Project [SPARK_PARTITION_ID() AS SPARK_PARTITION_ID()#1562, id#0L, date#1L, feature#2, partition_id#926]
+- Scan ExistingRDD[id#0L,date#1L,feature#2,partition_id#926]
In contrast, if I request the execution plan for a very similar query below, the SORT stage is not there and I detect no shuffling upon execution.
df.groupBy(spark_partition_id()).count().explain()
Output:
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[_nondeterministic#1532], functions=[count(1)])
+- Exchange hashpartitioning(_nondeterministic#1532, 200), ENSURE_REQUIREMENTS, [id=#704]
+- HashAggregate(keys=[_nondeterministic#1532], functions=[partial_count(1)])
+- Project [SPARK_PARTITION_ID() AS _nondeterministic#1532]
+- Scan ExistingRDD[id#0L,date#1L,feature#2,partition_id#926]
What is happening here and how do I achieve the goal I had stated? Thank you!

Answer 1

After some tinkering, it seems I am able to do what I want as follows, although probably it is not ideal.
spark_session.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch",'0')
def get_shape(iterator):
for pdf in iterator:
yield pd.DataFrame([pdf.shape])
df.mapInPandas(get_shape, "nr long, nc long").toPandas()