kafka custom partitioner with pyspark structured streaming - apache-spark

Question

I want to use kafka custom partitioner for my pyspark application reading from kafka pushing to another kafka topic. Data is transformed from source to sink with pyspark processing. I want to have control over partition to which data should be pushed based on some key in data/message. In spark structured streaming documentation I couldn't find any reference or example for such use case. I am using python processing along with pyspark and confluent-kafka-python is used as kafka client there, but it also lacks documentation/example for custom partitioner.
is solution available to achieve this?
Below spark code is tried with partition column and is not pushing data as per partition column.
df = spark.range(5)
df = (df
.withColumn("topic", F.lit("test_temp"))
.withColumn("partition", (F.col("id")%2).cast("int"))
.withColumn("key", F.lit("test"))
.withColumn("value", F.lit("test_data"))
).select(["topic", "key", "value", "partition"])
df.printSchema()
(df.write.format("kafka").partitionBy("partition")
.option("kafka.bootstrap.servers", kafka_endpoint)
#.option("topic", "test_temp")
.save())
Output:
+---------+----+---------+---------+
| topic| key| value|partition|
+---------+----+---------+---------+
|test_temp|test|test_data| 0|
|test_temp|test|test_data| 1|
|test_temp|test|test_data| 0|
|test_temp|test|test_data| 1|
|test_temp|test|test_data| 0|
+---------+----+---------+---------+
root
|-- topic: string (nullable = false)
|-- key: string (nullable = false)
|-- value: string (nullable = false)
|-- partition: integer (nullable = true)
Kafka console consumer output:
./kafka-console-consumer --bootstrap-server <broker>:9092 --topic test_temp --partition 1

Answer 1

As written in the Structured Streaming documentation, you'd either add a column to the dataframe being written called partition of an int type, and this controls which partition that Spark will write to. Or you can add a JVM paritioner to the classpath.
If a “partition” column is not specified (or its value is null) then the partition is calculated by the Kafka producer. A Kafka partitioner can be specified in Spark by setting the kafka.partitioner.class option. If not present, Kafka default partitioner will be used.
Regarding the Confluent Python client - https://github.com/confluentinc/confluent-kafka-python/issues/1107
But kafka-python module does support custom partitioners