Generate multiple columns using partition-batch process without shuffles or custom join logic

The Setup

I have several use cases in my pipelines where I need to add a couple of columns to a very wide DataFrame using the output of an external tool that has to operate on large batches.

The external tool takes as input a CSV file with several required columns and an arbitrary ID column and spits out a CSV file with the output columns and the same ID I passed in. Not every input row will make it into the output file.

The Concept

Conceptually this is pretty simple:

1. Select only the columns I need from the source DataFrame
2. Obtain a derived DataFrame by calling mapPartitions on the projected DataFrame
   1. Use the provided iterator to produce the input file
   2. Invoke the external binary
   3. Stream the output CSV file in as an iterator for the output partition
3. Left join the derived DataFrame onto source DataFrame using my id column

The Problem

Actually performing a left join will cause a shuffle of the source DataFrame to repartition it by the ID column. Given that the external tool is basically a flatMap that emits one or 0 items for each input row, the shuffle is not semantically necessary. I can guarantee that this "join" only needs data from the partitions as they currently exist. In the general case, I cannot assume anything about the distribution of the input DataFrame

Current Workaround

The obvious solution here is to implement the partition join manually.

Simplest example would be to materialize the entire (wide) DataFrame partition (inside mapPartition) as a List of entire rows as values (either as Row objects or case classes). Use that collection to generate the input CSV, call the binary, materialize the output CSV as a HashMap keyed by the input ID, and manually do a left join to the list, creating new Rows or output case class objects.

This works but has several bad characteristics. Having to materialize the entire partition is prohibitively expensive since the input DataFrame has potentially 300+ columns. Furthermore, and this is true of any solution that uses mapPartitions, if the input DataFrame did happen to be partitioned by certain key columns, mapPartitions discards that information even though in our case we're not changing the key. That means that any downstream joins will have to reshuffle the DataFrame to restore the proper distribution (which will likely be exactly the same).

The memory problem can be worked around by using sortWithinPartition to efficiently sort the 2 DataFrames' partitions by our join key (this will fallback to external sorting when needed), then dropping down into the RDD API to use zipPartitions, and finally manually implement whatever join mechanism is needed using the supplied iterators. This has the benefit that we only materialize a few rows at a time. This solution is what we're currently using and while it works, it requires a lot of custom code that I would rather not have to maintain. I'm also not a fan of jumping through the RDD layer since I would like to express as much of my pipeline in catalyst as possible. And as mentioned before, this solution also suffers from the lost distribution problem since it's based on mapPartitions.

The Question

Is there a well established idiom of doing something like this without a shuffle? Ideally while maintaining the distribution knowledge of the input DataFrame?

My research

As far as I can tell, the only way to maintain the distribution knowledge is to use withColumn. This makes sense since it's obvious to Spark that we're only adding columns, not changing any that might be part of the partition expression.

withColumn, however, doesn't work nicely out of the box here since we need to handle a whole batch of records at once.

In cases where the input DataFrame is partitioned by a low cardinality expression (think group/custom partition ID), I can aggregate the groups by id into a list, use a UDF to invoke the binary, explode the result back out, and the final DataFrame does exactly what I want (no shuffle and it's still distributed by the same expression).

Unfortunately, I don't have the guarantee that the input DataFrame is partitioned by any such column. I was hopeful that I could use a spark_partition_id column which serves as the perfect groupId for the above situation, and groupBy that to use my UDF approach but after some experiments it seems that Spark doesn't understand that grouping by partition id semantically could be shuffle-free since all of the rows are already co-located.

At one point I even attempted to make the source DataFrame look like it was partitioned correctly by forcing a HashClusteredDistribution onto it but that got so nasty I gave up.