test partition function

Author: dapeng

kafka-base/kafka-base-sink/src/main/java/com/dtstack/flink/sql/sink/kafka/CustomerFlinkPartition.java

@@ -3,15 +3,36 @@
 
 import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkFixedPartitioner;
 
+import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
 import org.apache.flink.util.Preconditions;
 
-public class CustomerFlinkPartition<T> extends FlinkFixedPartitioner<T> {
-    @Override
+public class CustomerFlinkPartition<T> extends FlinkKafkaPartitioner<T> {
+    private static final long serialVersionUID = 1L;
+    private int parallelInstanceId;
+
+    public CustomerFlinkPartition() {
+    
+    }
+
+    public void open(int parallelInstanceId, int parallelInstances) {
+        Preconditions.checkArgument(parallelInstanceId >= 0, "Id of this subtask cannot be negative.");
+        Preconditions.checkArgument(parallelInstances > 0, "Number of subtasks must be larger than 0.");
+        this.parallelInstanceId = parallelInstanceId;
+    }
+
     public int partition(T record, byte[] key, byte[] value, String targetTopic, int[] partitions) {
         Preconditions.checkArgument(partitions != null && partitions.length > 0, "Partitions of the target topic is empty.");
         if(key == null){
-            return super.partition(record, key, value, targetTopic, partitions);
+            return partitions[this.parallelInstanceId % partitions.length];
         }
         return partitions[key.hashCode() % partitions.length];
     }
+
+    public boolean equals(Object o) {
+        return this == o || o instanceof CustomerFlinkPartition;
+    }
+
+    public int hashCode() {
+        return CustomerFlinkPartition.class.hashCode();
+    }
 }