Spark: Weird partitioning on join - apache-spark

Question

I have a Spark SQL query that works somewhat like this (actual fields have been omitted):
SELECT
a1.fieldA,
a1.fieldB,
a1.fieldC,
a1.fieldD,
a1.joinType
FROM sample_a a1
WHERE a1.joinType != "test"
UNION
SELECT
a2.fieldA,
a2.fieldB,
a2.fieldC,
b.fieldD,
a2.joinType
FROM sample_a a2
INNER JOIN sample_b b ON b.joinField = a2.joinField
WHERE a2.joinType = "test"
This is working perfectly fine but Spark will read sample_a twice. (From cache or disk)
I'm trying to get rid of the union and came up with the following solution:
SELECT
a.fieldA,
a.fieldB,
a.fieldC,
a.joinType,
CASE WHEN a.joinType = "test" THEN b.fieldD ELSE a.fieldD END as fieldD
FROM sample_a a
LEFT JOIN sample_b b ON a.joinType = "test" AND a.joinField = b.joinField
WHERE a.joinType != "test" OR (a.joinType = "test" AND b.joinField IS NOT NULL)
This should basically do the same thing but Spark is being very weird about it. While the first one keeps the partitions the same as sample_a (~1200) the second one will go down to 200 partitions, which is what sample_b has. It will also put a lot of data into a single partition. (Around 90% of data is in one of the 200 partitions)
The input data is stored in parquet files and not partitioned in any way. While sample_a has a much bigger file size, the joinField values for our joinType = "test" part are a subset of the joinField values in sample_b.
Edit: The physical plans look like this.
First Query:
Union
:- *(1) Project [fieldA#0, fieldD#1 joinType#2, joinField#3]
: +- *(1) Filter (isnotnull(joinType#2) && NOT (joinType#2 = test))
: +- *(1) FileScan parquet [fieldA#0, fieldD#1 joinType#2, joinField#3] Batched: true, Format: Parquet, Location: InMemoryFileIndex[...], PartitionFilters: [], PushedFilters: [IsNotNull(joinType), Not(EqualTo(joinType,test))], ReadSchema: struct<fieldA:string,fieldD:string,joinType:string,joinField:string>
+- *(6) Project [fieldA#0, fieldD#4 joinType#2, joinField#3]
+- *(6) SortMergeJoin [joinField#3], [joinField#5], Inner
:- *(3) Sort [joinField#3 ASC NULLS FIRST], false, 0
: +- Exchange hashpartitioning(joinField#3, 200)
: +- *(2) Project [fieldA#0, fieldD#1 joinType#2, joinField#3]
: +- *(2) Filter ((isnotnull(joinType#2) && (joinType#2 = test)) && isnotnull(joinField#3))
: +- *(2) FileScan parquet [fieldA#0, fieldD#1 joinType#2, joinField#3] Batched: true, Format: Parquet, Location: InMemoryFileIndex[...], PartitionFilters: [], PushedFilters: [IsNotNull(joinType), EqualTo(joinType,test), IsNotNull(joinField)], ReadSchema: struct<fieldA:string,fieldD:string,joinType:string,joinField:string>
+- *(5) Sort [joinField#5 ASC NULLS FIRST], false, 0
+- Exchange hashpartitioning(joinField#5, 200)
+- *(4) FileScan parquet [fieldD#4, joinField#5] Batched: true, Format: Parquet, Location: InMemoryFileIndex[...], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<fieldD:string,joinField:string>
Second Query:
*(5) Project [fieldA#0, CASE WHEN (joinType#2 = test) THEN fieldD#4 ELSE fieldD#1 END AS fieldD#6, joinType#2, joinField#3]
+- *(5) Filter (NOT (joinType#2 = test) || ((joinType#2 = test) && isnotnull(joinField#5)))
+- SortMergeJoin [joinField#3], [joinField#5], LeftOuter, (joinType#2 = test)
:- *(2) Sort [joinField#3 ASC NULLS FIRST], false, 0
: +- Exchange hashpartitioning(joinField#3, 200)
: +- *(1) FileScan parquet [fieldA#0, fieldD#1 joinType#2, joinField#3] Batched: true, Format: Parquet, Location: InMemoryFileIndex[...], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<fieldA:string,fieldD:string,joinType:string,joinField:string>
+- *(4) Sort [joinField#5 ASC NULLS FIRST], false, 0
+- Exchange hashpartitioning(joinField#5, 200)
+- *(3) FileScan parquet [fieldD#4, joinField#5] Batched: true, Format: Parquet, Location: InMemoryFileIndex[...], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<fieldD:string,joinField:string>