This is an automated email from the ASF dual-hosted git repository.
dockerzhang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/inlong.git
The following commit(s) were added to refs/heads/master by this push:
new 4da7fc3230 [INLONG-7900][Sort] Support partition by custom fields when
upsert single table of Kafka (#8209)
4da7fc3230 is described below
commit 4da7fc32309491041c6ce2bf044f212a8be7e901
Author: Yizhou Yang <[email protected]>
AuthorDate: Wed Aug 2 18:39:51 2023 +0800
[INLONG-7900][Sort] Support partition by custom fields when upsert single
table of Kafka (#8209)
Co-authored-by: Yizhou Yang <[email protected]>
---
.../kafka/DynamicKafkaSerializationSchema.java | 16 ++-
.../SingleTableCustomFieldsPartitioner.java | 150 +++++++++++++++++++++
.../sort/kafka/table/KafkaDynamicTableFactory.java | 33 ++++-
.../table/UpsertKafkaDynamicTableFactory.java | 42 +++++-
.../SingleTableCustomFieldsPartitionerTest.java | 80 +++++++++++
5 files changed, 308 insertions(+), 13 deletions(-)
diff --git
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/DynamicKafkaSerializationSchema.java
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/DynamicKafkaSerializationSchema.java
index d80cc9d047..0413bb4cbb 100644
---
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/DynamicKafkaSerializationSchema.java
+++
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/DynamicKafkaSerializationSchema.java
@@ -93,10 +93,10 @@ class DynamicKafkaSerializationSchema implements
KafkaSerializationSchema<RowDat
*/
private final int[] metadataPositions;
private final String sinkMultipleFormat;
- private boolean multipleSink;
- private JsonDynamicSchemaFormat jsonDynamicSchemaFormat;
private final DirtyOptions dirtyOptions;
private final @Nullable DirtySink<Object> dirtySink;
+ private boolean multipleSink;
+ private JsonDynamicSchemaFormat jsonDynamicSchemaFormat;
private int[] partitions;
private int parallelInstanceId;
@@ -141,10 +141,6 @@ class DynamicKafkaSerializationSchema implements
KafkaSerializationSchema<RowDat
this.policyMap = policyMap;
}
- public void setMetricData(SinkTopicMetricData metricData) {
- this.metricData = metricData;
- }
-
static RowData createProjectedRow(
RowData consumedRow, RowKind kind, RowData.FieldGetter[]
fieldGetters) {
final int arity = fieldGetters.length;
@@ -155,6 +151,10 @@ class DynamicKafkaSerializationSchema implements
KafkaSerializationSchema<RowDat
return genericRowData;
}
+ public void setMetricData(SinkTopicMetricData metricData) {
+ this.metricData = metricData;
+ }
+
@Override
public void open(SerializationSchema.InitializationContext context) throws
Exception {
if (keySerialization != null) {
@@ -174,6 +174,10 @@ class DynamicKafkaSerializationSchema implements
KafkaSerializationSchema<RowDat
jsonDynamicSchemaFormat =
(JsonDynamicSchemaFormat)
DynamicSchemaFormatFactory.getFormat(sinkMultipleFormat);
}
+
+ if (partitioner instanceof SingleTableCustomFieldsPartitioner) {
+ ((SingleTableCustomFieldsPartitioner<?>)
partitioner).setValueFieldGetters(valueFieldGetters);
+ }
}
@Override
diff --git
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitioner.java
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitioner.java
new file mode 100644
index 0000000000..19c5ae9382
--- /dev/null
+++
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitioner.java
@@ -0,0 +1,150 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.inlong.sort.kafka;
+
+import
org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
+import org.apache.flink.table.api.TableSchema;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.RowData.FieldGetter;
+import org.apache.flink.util.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+
+/**
+ * The PrimaryKey Partitioner is used to extract primary key from single table
messages
+ *
+ * @param <T>
+ */
+public class SingleTableCustomFieldsPartitioner<T> extends
FlinkKafkaPartitioner<T> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(SingleTableCustomFieldsPartitioner.class);
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * schema entry names which will be used to extract fields from schema for
hashing.
+ */
+ private String partitionKey;
+
+ /**
+ * If a partitionNumber is specified, then use this number directly
instead.
+ */
+ private int partitionNumber = -1;
+
+ /**
+ * the fieldnames of the schema
+ */
+ private String[] fieldNames;
+
+ /**
+ * the value field getters used to extract rowdata
+ */
+ private FieldGetter[] valueFieldGetters;
+
+ @Override
+ public void open(int parallelInstanceId, int parallelInstances) {
+ super.open(parallelInstanceId, parallelInstances);
+ }
+
+ /**
+ * precondition: record is of rowData Type, schema and partitionKey are
not null
+ */
+ @Override
+ 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.");
+ // parse the partition key fields
+ return getPartition((RowData) record, partitions);
+ }
+
+ public void setPartitionKey(String partitionKey) {
+ this.partitionKey = partitionKey;
+ }
+
+ /**
+ * input: rowdata, partitions
+ * output: the hashed partition number
+ */
+ private int getPartition(RowData data, int[] partitions) {
+ if (partitionNumber >= 0) {
+ return partitionNumber;
+ }
+ List<Integer> pos = getFieldPos();
+ // parse out the List<String> from partitionkey and then get list of
positions in schema.
+ long hashCode = 0;
+ for (int i : pos) {
+ Object fieldValue = valueFieldGetters[i].getFieldOrNull(data);
+ if (fieldValue != null) {
+ hashCode += fieldValue.hashCode();
+ }
+ }
+ return partitions[((int) ((hashCode & Integer.MAX_VALUE) %
partitions.length))];
+ }
+
+ /**
+ * output: integer list of partitionKeys' corresponding positions within
schema
+ */
+ private List<Integer> getFieldPos() {
+ // the positions of the partition keys.
+ List<Integer> positions = new ArrayList<>();
+ if (partitionKey == null) {
+ LOG.error("primaryKeyPartioner:failed to fetch partitionKey");
+ return positions;
+ }
+ // a map storing field names and their position in the schema
+ HashMap<String, Integer> map = new HashMap<>();
+ String[] keys = partitionKey.split(",");
+ for (int i = 0; i < fieldNames.length; i++) {
+ map.put(fieldNames[i], i);
+ }
+ for (String key : keys) {
+ positions.add(map.get(key));
+ }
+ return positions;
+ }
+
+ @SuppressWarnings("deprecation")
+ public void setSchema(TableSchema schema) {
+ this.fieldNames = schema.getFieldNames();
+ }
+
+ public void setPartitionNumber(String partitionNumber) {
+ // partition Number is optional parameter and not always needed, so it
can be null.
+ if (partitionNumber != null) {
+ this.partitionNumber = Integer.parseInt(partitionNumber);
+ }
+ }
+
+ public void setValueFieldGetters(FieldGetter[] fieldGetters) {
+ this.valueFieldGetters = fieldGetters;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ return o instanceof SingleTableCustomFieldsPartitioner;
+ }
+
+ @Override
+ public int hashCode() {
+ return SingleTableCustomFieldsPartitioner.class.hashCode();
+ }
+}
diff --git
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/KafkaDynamicTableFactory.java
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/KafkaDynamicTableFactory.java
index bf8b8225ed..3340298305 100644
---
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/KafkaDynamicTableFactory.java
+++
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/KafkaDynamicTableFactory.java
@@ -22,6 +22,7 @@ import org.apache.inlong.sort.base.dirty.sink.DirtySink;
import org.apache.inlong.sort.base.dirty.utils.DirtySinkFactoryUtils;
import org.apache.inlong.sort.base.format.DynamicSchemaFormatFactory;
import org.apache.inlong.sort.kafka.KafkaDynamicSink;
+import org.apache.inlong.sort.kafka.SingleTableCustomFieldsPartitioner;
import
org.apache.inlong.sort.kafka.partitioner.InLongFixedPartitionPartitioner;
import org.apache.inlong.sort.kafka.partitioner.RawDataHashPartitioner;
import org.apache.inlong.sort.protocol.enums.SchemaChangePolicy;
@@ -43,6 +44,7 @@ import
org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartiti
import org.apache.flink.streaming.connectors.kafka.table.KafkaOptions;
import org.apache.flink.streaming.connectors.kafka.table.KafkaSinkSemantic;
import org.apache.flink.streaming.connectors.kafka.table.SinkBufferFlushMode;
+import org.apache.flink.table.api.TableSchema;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.catalog.CatalogTable;
import org.apache.flink.table.catalog.ObjectIdentifier;
@@ -108,6 +110,8 @@ import static
org.apache.flink.table.factories.FactoryUtil.FORMAT;
import static org.apache.flink.table.factories.FactoryUtil.SINK_PARALLELISM;
import static org.apache.inlong.sort.base.Constants.*;
import static org.apache.inlong.sort.base.Constants.DATASOURCE_PARTITION_MAP;
+import static org.apache.inlong.sort.base.Constants.PATTERN_PARTITION_MAP;
+import static org.apache.inlong.sort.base.Constants.SINK_MULTIPLE_FORMAT;
import static org.apache.inlong.sort.base.Constants.SINK_SCHEMA_CHANGE_ENABLE;
import static
org.apache.inlong.sort.base.Constants.SINK_SCHEMA_CHANGE_POLICIES;
import static
org.apache.inlong.sort.kafka.table.KafkaOptions.KAFKA_IGNORE_ALL_CHANGELOG;
@@ -124,14 +128,20 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
public static final String IDENTIFIER = "kafka-inlong";
public static final String SINK_PARTITIONER_VALUE_RAW_HASH = "raw-hash";
-
+ public static final String SINK_PARTITIONER_VALUE_PRIMARY_KEY =
"primaryKey";
public static final String SINK_PARTITIONER_VALUE_INLONG_FIXED_PARTITION =
"inlong-fixed-partition";
public static final ConfigOption<String> SINK_MULTIPLE_PARTITION_PATTERN =
ConfigOptions.key("sink.multiple.partition-pattern")
.stringType()
.noDefaultValue()
- .withDescription("option 'sink.multiple.partition-pattern'
used when the partitioner is raw-hash.");
+ .withDescription(
+ "option 'sink.multiple.partition-pattern' used
either when the partitioner is raw-hash, or when passing in designated
partition field names for custom field partitions");
+
+ public static final ConfigOption<String> SINK_FIXED_IDENTIFIER =
+ ConfigOptions.key("sink.fixed.identifier")
+ .stringType()
+ .defaultValue("-1");
private static final Set<String> SINK_SEMANTIC_ENUMS =
new HashSet<>(
@@ -249,7 +259,7 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
}
private Optional<FlinkKafkaPartitioner<RowData>> getFlinkKafkaPartitioner(
- ReadableConfig tableOptions, ClassLoader classLoader) {
+ ReadableConfig tableOptions, ClassLoader classLoader, TableSchema
schema) {
if (tableOptions.getOptional(SINK_PARTITIONER).isPresent() &&
SINK_PARTITIONER_VALUE_INLONG_FIXED_PARTITION
.equals(tableOptions.getOptional(SINK_PARTITIONER).get())) {
InLongFixedPartitionPartitioner<RowData>
inLongFixedPartitionPartitioner =
@@ -260,8 +270,7 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
inLongFixedPartitionPartitioner.setSinkMultipleFormat(tableOptions.getOptional(SINK_MULTIPLE_FORMAT)
.orElse(null));
return Optional.of(inLongFixedPartitionPartitioner);
- }
- if (tableOptions.getOptional(SINK_PARTITIONER).isPresent()
+ } else if (tableOptions.getOptional(SINK_PARTITIONER).isPresent()
&&
SINK_PARTITIONER_VALUE_RAW_HASH.equals(tableOptions.getOptional(SINK_PARTITIONER).get()))
{
RawDataHashPartitioner<RowData> rawHashPartitioner = new
RawDataHashPartitioner<>(tableOptions
.getOptional(DATASOURCE_PARTITION_MAP).orElse(Collections.emptyMap()));
@@ -269,6 +278,16 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
rawHashPartitioner.setPartitionPattern(tableOptions.getOptional(SINK_MULTIPLE_PARTITION_PATTERN)
.orElse(null));
return Optional.of(rawHashPartitioner);
+ } else if (tableOptions.getOptional(SINK_PARTITIONER).isPresent()
+ &&
SINK_PARTITIONER_VALUE_PRIMARY_KEY.equals(tableOptions.getOptional(SINK_PARTITIONER).get()))
{
+ SingleTableCustomFieldsPartitioner<RowData>
customFieldsPartitioner =
+ new SingleTableCustomFieldsPartitioner<>();
+ // the pattern is different from the ${} pattern, it is a
truncated string of schema fields.
+
customFieldsPartitioner.setPartitionNumber(tableOptions.getOptional(SINK_FIXED_IDENTIFIER).orElse(null));
+
customFieldsPartitioner.setPartitionKey(tableOptions.getOptional(SINK_MULTIPLE_PARTITION_PATTERN)
+ .orElse(null));
+ customFieldsPartitioner.setSchema(schema);
+ return Optional.of(customFieldsPartitioner);
}
Optional<FlinkKafkaPartitioner<RowData>> partitioner = KafkaOptions
.getFlinkKafkaPartitioner(tableOptions, classLoader);
@@ -338,6 +357,7 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
options.add(DATASOURCE_PARTITION_MAP);
options.add(SINK_SCHEMA_CHANGE_ENABLE);
options.add(SINK_SCHEMA_CHANGE_POLICIES);
+ options.add(SINK_FIXED_IDENTIFIER);
return options;
}
@@ -472,7 +492,8 @@ public class KafkaDynamicTableFactory implements
DynamicTableSourceFactory, Dyna
tableOptions.get(TOPIC).get(0),
getKafkaProperties(context.getCatalogTable().getOptions()),
context.getCatalogTable(),
- getFlinkKafkaPartitioner(tableOptions,
context.getClassLoader()).orElse(null),
+ getFlinkKafkaPartitioner(tableOptions,
context.getClassLoader(),
+ context.getCatalogTable().getSchema()).orElse(null),
getSinkSemantic(tableOptions),
parallelism,
inlongMetric,
diff --git
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/UpsertKafkaDynamicTableFactory.java
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/UpsertKafkaDynamicTableFactory.java
index 9f527e3a29..cb5b054867 100644
---
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/UpsertKafkaDynamicTableFactory.java
+++
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/main/java/org/apache/inlong/sort/kafka/table/UpsertKafkaDynamicTableFactory.java
@@ -21,6 +21,7 @@ import org.apache.inlong.sort.base.dirty.DirtyOptions;
import org.apache.inlong.sort.base.dirty.sink.DirtySink;
import org.apache.inlong.sort.base.dirty.utils.DirtySinkFactoryUtils;
import org.apache.inlong.sort.kafka.KafkaDynamicSink;
+import org.apache.inlong.sort.kafka.SingleTableCustomFieldsPartitioner;
import org.apache.inlong.sort.protocol.enums.SchemaChangePolicy;
import org.apache.inlong.sort.protocol.enums.SchemaChangeType;
import org.apache.inlong.sort.util.SchemaChangeUtils;
@@ -29,9 +30,11 @@ import
org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.streaming.connectors.kafka.config.StartupMode;
+import
org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.streaming.connectors.kafka.table.KafkaOptions;
import org.apache.flink.streaming.connectors.kafka.table.KafkaSinkSemantic;
import org.apache.flink.streaming.connectors.kafka.table.SinkBufferFlushMode;
@@ -59,6 +62,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
+import java.util.Optional;
import java.util.Properties;
import java.util.Set;
@@ -68,6 +72,7 @@ import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.KEY
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.PROPS_BOOTSTRAP_SERVERS;
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.SINK_BUFFER_FLUSH_INTERVAL;
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.SINK_BUFFER_FLUSH_MAX_ROWS;
+import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.SINK_PARTITIONER;
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.TOPIC;
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.VALUE_FIELDS_INCLUDE;
import static
org.apache.flink.streaming.connectors.kafka.table.KafkaOptions.VALUE_FORMAT;
@@ -82,6 +87,7 @@ import static
org.apache.inlong.sort.base.Constants.INLONG_METRIC;
import static org.apache.inlong.sort.base.Constants.SINK_MULTIPLE_FORMAT;
import static org.apache.inlong.sort.base.Constants.SINK_SCHEMA_CHANGE_ENABLE;
import static
org.apache.inlong.sort.base.Constants.SINK_SCHEMA_CHANGE_POLICIES;
+import static
org.apache.inlong.sort.kafka.table.KafkaDynamicTableFactory.SINK_PARTITIONER_VALUE_PRIMARY_KEY;
import static
org.apache.inlong.sort.kafka.table.KafkaOptions.KAFKA_IGNORE_ALL_CHANGELOG;
/**
@@ -96,6 +102,17 @@ public class UpsertKafkaDynamicTableFactory
DynamicTableSinkFactory {
public static final String IDENTIFIER = "upsert-kafka-inlong";
+ public static final ConfigOption<String> SINK_MULTIPLE_PARTITION_PATTERN =
+ ConfigOptions.key("sink.multiple.partition-pattern")
+ .stringType()
+ .noDefaultValue()
+ .withDescription(
+ "option 'sink.multiple.partition-pattern' used
either when the partitioner is raw-hash, or when passing in designated
partition field names for custom field partitions");
+
+ public static final ConfigOption<String> SINK_FIXED_IDENTIFIER =
+ ConfigOptions.key("sink.fixed.identifier")
+ .stringType()
+ .defaultValue("-1");
private static void validateSource(
ReadableConfig tableOptions, Format keyFormat, Format valueFormat,
TableSchema schema) {
@@ -199,6 +216,9 @@ public class UpsertKafkaDynamicTableFactory
options.add(SINK_BUFFER_FLUSH_MAX_ROWS);
options.add(KAFKA_IGNORE_ALL_CHANGELOG);
options.add(INLONG_METRIC);
+ options.add(SINK_PARTITIONER);
+ options.add(SINK_MULTIPLE_PARTITION_PATTERN);
+ options.add(SINK_FIXED_IDENTIFIER);
return options;
}
@@ -250,6 +270,23 @@ public class UpsertKafkaDynamicTableFactory
auditKeys);
}
+ private Optional<FlinkKafkaPartitioner<RowData>> getFlinkKafkaPartitioner(
+ ReadableConfig tableOptions, ClassLoader classLoader, TableSchema
schema) {
+ if (tableOptions.getOptional(SINK_PARTITIONER).isPresent()
+ &&
SINK_PARTITIONER_VALUE_PRIMARY_KEY.equals(tableOptions.getOptional(SINK_PARTITIONER).get()))
{
+ SingleTableCustomFieldsPartitioner<RowData>
customFieldsPartitioner =
+ new SingleTableCustomFieldsPartitioner<>();
+
customFieldsPartitioner.setPartitionNumber(tableOptions.getOptional(SINK_FIXED_IDENTIFIER).orElse(null));
+
customFieldsPartitioner.setPartitionKey(tableOptions.getOptional(SINK_MULTIPLE_PARTITION_PATTERN)
+ .orElse(null));
+ customFieldsPartitioner.setSchema(schema);
+ return Optional.of(customFieldsPartitioner);
+ }
+ Optional<FlinkKafkaPartitioner<RowData>> partitioner = KafkaOptions
+ .getFlinkKafkaPartitioner(tableOptions, classLoader);
+ return partitioner;
+ }
+
@Override
public DynamicTableSink createDynamicTableSink(Context context) {
FactoryUtil.TableFactoryHelper helper =
@@ -289,6 +326,9 @@ public class UpsertKafkaDynamicTableFactory
final boolean enableSchemaChange =
tableOptions.get(SINK_SCHEMA_CHANGE_ENABLE);
final Map<SchemaChangeType, SchemaChangePolicy> policyMap =
enableSchemaChange ?
SchemaChangeUtils.deserialize(schemaChangePolicies) : Collections.emptyMap();
+ final FlinkKafkaPartitioner<RowData> partitioner =
+ getFlinkKafkaPartitioner(tableOptions,
context.getClassLoader(),
+ context.getCatalogTable().getSchema()).orElse(null);
// use {@link
org.apache.kafka.clients.producer.internals.DefaultPartitioner}.
// it will use hash partition if key is set else in round-robin
behaviour.
@@ -303,7 +343,7 @@ public class UpsertKafkaDynamicTableFactory
tableOptions.get(TOPIC).get(0),
properties,
context.getCatalogTable(),
- null,
+ partitioner,
KafkaSinkSemantic.AT_LEAST_ONCE,
true,
flushMode,
diff --git
a/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/test/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitionerTest.java
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/test/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitionerTest.java
new file mode 100644
index 0000000000..1df59dfe5b
--- /dev/null
+++
b/inlong-sort/sort-flink/sort-flink-v1.13/sort-connectors/kafka/src/test/java/org/apache/inlong/sort/kafka/partitioner/SingleTableCustomFieldsPartitionerTest.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.inlong.sort.kafka.partitioner;
+
+import org.apache.inlong.sort.kafka.SingleTableCustomFieldsPartitioner;
+
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.api.TableSchema;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.RowData.FieldGetter;
+import org.apache.flink.table.data.binary.BinaryRowData;
+import org.apache.flink.table.data.writer.BinaryRowWriter;
+import org.apache.flink.table.types.logical.IntType;
+import org.junit.Assert;
+import org.junit.Test;
+
+/**
+ * The unit tests for {@link
+ *
org.apache.inlong.sort.kafka.partitioner.SingleTableCustomFieldsPartitioner}.
+ */
+public class SingleTableCustomFieldsPartitionerTest {
+
+ @Test
+ public void testCustomFieldsPartitioner() {
+ SingleTableCustomFieldsPartitioner singleTableCustomFieldsPartitioner =
+ new SingleTableCustomFieldsPartitioner();
+
+ TableSchema schema = TableSchema.builder()
+ .field("id", DataTypes.INT())
+ .field("age", DataTypes.INT())
+ .build();
+
+ singleTableCustomFieldsPartitioner.setSchema(schema);
+ singleTableCustomFieldsPartitioner.setPartitionKey("age");
+
+ FieldGetter getter0 = RowData.createFieldGetter(new IntType(), 0);
+ FieldGetter getter1 = RowData.createFieldGetter(new IntType(), 1);
+ FieldGetter[] valuefieldgetters = {getter0, getter1};
+
+
singleTableCustomFieldsPartitioner.setValueFieldGetters(valuefieldgetters);
+
+ BinaryRowData rowData1 = new BinaryRowData(2);
+ BinaryRowWriter writer1 = new BinaryRowWriter(rowData1);
+ BinaryRowData rowData2 = new BinaryRowData(2);
+ BinaryRowWriter writer2 = new BinaryRowWriter(rowData2);
+
+ writer1.writeInt(0, 1);
+ writer1.writeInt(1, 786819156);
+
+ writer2.writeInt(0, 2);
+ writer2.writeInt(1, 786819156);
+
+ // key is null since the actual key is being calculated from
deserialized json
+ int partition1 =
singleTableCustomFieldsPartitioner.partition(rowData1, null,
+ null, null, new int[]{0, 1, 2, 3, 4});
+
+ int partition2 =
singleTableCustomFieldsPartitioner.partition(rowData2, null,
+ null, null, new int[]{0, 1, 2, 3, 4});
+
+ writer1.complete();
+ writer2.complete();
+
+ Assert.assertEquals(partition1, partition2);
+ }
+}
