How to make Spark write by partition and group by a column

Question

I have roughly 100Gb of data that I'm trying to process. The data has the form:

| timestamp | social_profile_id | json_payload     |
|-----------|-------------------|------------------|
| 123       | 1                 | {"json":"stuff"} |
| 124       | 2                 | {"json":"stuff"} |
| 125       | 3                 | {"json":"stuff"} |

I'm trying to split this data frame into folders in S3 by social_profile_id. There are roughly 430,000 social_profile_ids.

I've loaded the data no problem into a Dataset. However when I'm writing it out, and trying to partition it, it takes forever! Here's what I've tried:

messagesDS
      .write
      .partitionBy("socialProfileId")
      .mode(sparkSaveMode)

I don't really care how many files are in each folder at the end of the job. My theory is that each node can group by the social_profile_id, then write out to its respective folder without having to do a shuffle or communicate with other nodes. But this isn't happening as evidenced by the long job time. Ideally the end result would look a little something like this:

├── social_id_1 (only two partitions had id_1 data)
|   ├── partition1_data.parquet
|   └── partition3_data.parquet
├── social_id_2 (more partitions have this data in it)
|   ├── partition3_data.parquet
|   └── partition4_data.parquet
|   ├── etc.
├── social_id_3
|   ├── partition2_data.parquet
|   └── partition4_data.parquet
|   ├── etc.
├── etc.

I've tried increasing the compute resources a few times, both increasing instances sizes and # of instances. What I've been able to see form the spark UI that is the majority of time time is being taken by the write operation. It seems that all of the executors are being used, but they take an absurdly long time to execute (like taking 3-5 hours to write ~150Mb) Any help would be appreciated! Sorry if I mixed up some of the spark terminology.