[flink] Flink sink supports dynamic Fluss sink supports dynamic shuffle.

Author: loserwang1024

fluss-flink/fluss-flink-1.18/src/main/java/org/apache/fluss/flink/adapter/FlinkCompatibilityUtil.java

@@ -18,6 +18,7 @@
 package org.apache.fluss.flink.adapter;
 
 import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 
 /**
  * An adapter for Flink {@link RuntimeContext} class. The {@link RuntimeContext} class added the
@@ -31,4 +32,8 @@ public class RuntimeContextAdapter {
     public static int getAttemptNumber(RuntimeContext runtimeContext) {
         return runtimeContext.getAttemptNumber();
     }
+
+    public static int getIndexOfThisSubtask(StreamingRuntimeContext runtimeContext) {
+        return runtimeContext.getIndexOfThisSubtask();
+    }
 }

fluss-flink/fluss-flink-1.18/src/main/java/org/apache/fluss/flink/adapter/RuntimeContextAdapter.java

@@ -18,6 +18,7 @@
 package org.apache.fluss.flink.adapter;
 
 import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 
 /**
  * An adapter for Flink {@link RuntimeContext} class. The {@link RuntimeContext} class added the
@@ -31,4 +32,8 @@ public class RuntimeContextAdapter {
     public static int getAttemptNumber(RuntimeContext runtimeContext) {
         return runtimeContext.getTaskInfo().getAttemptNumber();
     }
+
+    public static int getIndexOfThisSubtask(StreamingRuntimeContext runtimeContext) {
+        return runtimeContext.getTaskInfo().getIndexOfThisSubtask();
+    }
 }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/adapter/FlinkCompatibilityUtil.java

@@ -19,6 +19,7 @@
 
 import org.apache.fluss.annotation.Internal;
 import org.apache.fluss.config.Configuration;
+import org.apache.fluss.flink.adapter.SinkAdapter;
 import org.apache.fluss.flink.sink.serializer.FlussSerializationSchema;
 import org.apache.fluss.flink.sink.shuffle.DataStatisticsOperatorFactory;
 import org.apache.fluss.flink.sink.shuffle.DistributionMode;
@@ -34,7 +35,6 @@
 import org.apache.flink.api.common.functions.FlatMapFunction;
 import org.apache.flink.api.common.operators.MailboxExecutor;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
-import org.apache.flink.api.connector.sink2.Sink;
 import org.apache.flink.api.connector.sink2.SinkWriter;
 import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
 import org.apache.flink.runtime.metrics.groups.InternalSinkWriterMetricGroup;
@@ -52,7 +52,7 @@
 import static org.apache.fluss.flink.utils.FlinkConversions.toFlussRowType;
 
 /** Flink sink for Fluss. */
-class FlinkSink<InputT> implements Sink<InputT>, SupportsPreWriteTopology<InputT> {
+class FlinkSink<InputT> extends SinkAdapter<InputT> implements SupportsPreWriteTopology<InputT> {
 
     private static final long serialVersionUID = 1L;
 
@@ -165,13 +165,13 @@ public DataStream<InputT> addPreWriteTopology(DataStream<InputT> input) {
                             new StatisticsOrRecordTypeInformation<>(rowTypeInformation);
                     SingleOutputStreamOperator<StatisticsOrRecord<InputT>> shuffleStream =
                             input.transform(
-                                            "Range shuffle Collector",
+                                            "Dynamic shuffle data statistics",
                                             statisticsOrRecordTypeInformation,
                                             new DataStatisticsOperatorFactory<>(
                                                     toFlussRowType(tableRowType),
                                                     partitionKeys,
                                                     flussSerializationSchema))
-                                    .uid("Range shuffle Collector" + tablePath)
+                                    .uid("Dynamic shuffle data statistics" + tablePath)
                                     // Set the parallelism same as input operator to encourage
                                     // chaining
                                     .setParallelism(input.getParallelism());

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/adapter/RuntimeContextAdapter.java

@@ -199,14 +199,13 @@ public SinkRuntimeProvider getSinkRuntimeProvider(Context context) {
                 @Override
                 public DataStreamSink<?> consumeDataStream(
                         ProviderContext providerContext, DataStream<RowData> dataStream) {
-                    // todo: add prefix to the sink to avoid conflict in same flink job.
-                    String defaultSuffix =
+                    String sinkName =
                             "sink"
                                     + "-"
                                     + tablePath.getTableName()
                                     + "-"
                                     + tablePath.getDatabaseName();
-                    return dataStream.sinkTo(flinkSink).uid(defaultSuffix).name(defaultSuffix);
+                    return dataStream.sinkTo(flinkSink).uid(sinkName).name(sinkName);
                 }
             };
         }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkSink.java

@@ -74,7 +74,7 @@ public class FlussSinkBuilder<InputT> {
     private String tableName;
     private final Map<String, String> configOptions = new HashMap<>();
     private FlussSerializationSchema<InputT> serializationSchema;
-    private DistributionMode shuffleMode = DistributionMode.BUCKET_SHUFFLE;
+    private DistributionMode distributionMode = DistributionMode.BUCKET_SHUFFLE;
     private TypeInformation<InputT> rowTypeInformation;
 
     /** Set the bootstrap server for the sink. */
@@ -97,17 +97,17 @@ public FlussSinkBuilder<InputT> setTable(String table) {
 
     /**
      * Set shuffle by bucket id. Deprecated use {@link
-     * FlussSinkBuilder#setShuffleMode(DistributionMode) } instead.
+     * FlussSinkBuilder#setDistributionMode(DistributionMode) } instead.
      */
     @Deprecated
     public FlussSinkBuilder<InputT> setShuffleByBucketId(boolean shuffleByBucketId) {
-        this.shuffleMode =
+        this.distributionMode =
                 shuffleByBucketId ? DistributionMode.BUCKET_SHUFFLE : DistributionMode.NONE;
         return this;
     }
 
-    public FlussSinkBuilder<InputT> setShuffleMode(DistributionMode shuffleByBucketId) {
-        this.shuffleMode = shuffleByBucketId;
+    public FlussSinkBuilder<InputT> setDistributionMode(DistributionMode distributionMode) {
+        this.distributionMode = distributionMode;
         return this;
     }
 
@@ -145,9 +145,7 @@ public FlussSink<InputT> build() {
         FlinkSink.SinkWriterBuilder<? extends FlinkSinkWriter<InputT>, InputT> writerBuilder;
 
         TablePath tablePath = new TablePath(database, tableName);
-        if (bootstrapServers != null) {
-            flussConfig.setString(ConfigOptions.BOOTSTRAP_SERVERS.key(), bootstrapServers);
-        }
+        flussConfig.setString(ConfigOptions.BOOTSTRAP_SERVERS.key(), bootstrapServers);
 
         TableInfo tableInfo;
         try (Connection connection = ConnectionFactory.createConnection(flussConfig);
@@ -184,7 +182,7 @@ public FlussSink<InputT> build() {
                             bucketKeys,
                             partitionKeys,
                             lakeFormat,
-                            shuffleMode,
+                            distributionMode,
                             serializationSchema);
         } else {
             LOG.info("Initializing Fluss append sink writer ...");
@@ -197,7 +195,7 @@ public FlussSink<InputT> build() {
                             bucketKeys,
                             partitionKeys,
                             lakeFormat,
-                            shuffleMode,
+                            distributionMode,
                             serializationSchema,
                             rowTypeInformation);
         }

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlinkTableSink.java

@@ -210,12 +210,12 @@ public long size(RowData value, RowType rowType) {
                     size += 2;
                     break;
                 case INTEGER:
+                case FLOAT:
                 case DATE:
                 case TIME_WITHOUT_TIME_ZONE:
                     size += 4;
                     break;
                 case BIGINT:
-                case FLOAT:
                 case DOUBLE:
                 case TIMESTAMP_WITHOUT_TIME_ZONE:
                 case TIMESTAMP_WITH_LOCAL_TIME_ZONE:

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/FlussSinkBuilder.java

@@ -40,7 +40,7 @@
 
 import java.util.List;
 
-import static org.apache.fluss.flink.adapter.FlinkCompatibilityUtil.getIndexOfThisSubtask;
+import static org.apache.fluss.flink.adapter.RuntimeContextAdapter.getIndexOfThisSubtask;
 import static org.apache.fluss.utils.Preconditions.checkArgument;
 
 /**

fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/sink/serializer/RowDataSerializationSchema.java

@@ -108,6 +108,10 @@ public void setup(int numChannels) {
 
     @Override
     public int channel(StatisticsOrRecord<InputT> wrapper) {
+        if (wrapper == null) {
+            throw new FlussRuntimeException("StatisticsOrRecord wrapper must not be null");
+        }
+
         try {
             if (wrapper.hasStatistics()) {
                 this.delegatePartitioner = delegatePartitioner(wrapper.statistics());