Apache Spark dataframe repartition with UDF functions

I have a misunderstanding in the proper usage of the Apache Spark in case of the repartition together with the UDF calls

Let's say there is a simple registered UDF function:

public static DataFrame compute(SQLContext sqlContext, DataFrame df) {
    sqlContext.udf().register("compute", new MyUdf(), DataTypes.StringType);
    return df.withColumn("CLIENT", functions.callUDF("compute", new Column("CLIENT")));
}

class MyUdf implements UDF1<String, String> {
Set<String> cache = new HashSet<>();
@Override
public String call(String client) {
    cache.add(client);
    System.out.println("cache size: " + cache.size() + ", obj: " + this.toString() + ", client: " + client);
    return client;
}

}

As you can see I am going to pass a client field, add it into a local cache of the the function to track it's execution.

Now I am going to create a dataframe with a repartition by client:

df = df.repartition(col("CLIENT"));
df = compute(sqlContext, df);
df.show(false);

I am calling the repartition without the first numPartitions arg, and by default it is 200:

System.out.println(df.toJavaRDD().getNumPartitions()); // will show 200

The log of the function seems fine (each function runs with it's own client, right?):

cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@11633b65, client: cl2
cache size: 1, obj: demo.MyUdf@11633b65, client: cl2
cache size: 1, obj: demo.MyUdf@11633b65, client: cl2

But if I try to set a numPartitions to:

df = df.repartition(2, col("CLIENT"));

I see the "unexpected" result for me (single function for the both partitions):

cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 1, obj: demo.MyUdf@3a9de0bf, client: cl1
cache size: 2, obj: demo.MyUdf@3a9de0bf, client: cl2
cache size: 2, obj: demo.MyUdf@3a9de0bf, client: cl2
cache size: 2, obj: demo.MyUdf@3a9de0bf, client: cl2

Why does it happen?