[connector] Flink source support detect bucket changes

Author: LiebingYu

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/FlinkConnectorOptions.java

@@ -100,6 +100,15 @@ public class FlinkConnectorOptions {
                                     + "the new partitions for partitioned table while scanning."
                                     + " A non-positive value disables the partition discovery.");
 
+    public static final ConfigOption<Duration> SCAN_BUCKET_DISCOVERY_INTERVAL =
+            ConfigOptions.key("scan.bucket.discovery.interval")
+                    .durationType()
+                    .defaultValue(Duration.ofSeconds(10))
+                    .withDescription(
+                            "The time interval for the Fluss source to discover "
+                                    + "the new buckets for table while scanning."
+                                    + " A non-positive value disables the bucket discovery.");
+
     public static final ConfigOption<Boolean> SINK_IGNORE_DELETE =
             ConfigOptions.key("sink.ignore-delete")
                     .booleanType()

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/catalog/FlinkTableFactory.java

@@ -127,6 +127,8 @@ public DynamicTableSource createDynamicTableSource(Context context) {
                 tableOptions
                         .get(FlinkConnectorOptions.SCAN_PARTITION_DISCOVERY_INTERVAL)
                         .toMillis();
+        long bucketDiscoveryIntervalMs =
+                tableOptions.get(FlinkConnectorOptions.SCAN_BUCKET_DISCOVERY_INTERVAL).toMillis();
 
         return new FlinkTableSource(
                 toFlussTablePath(context.getObjectIdentifier()),
@@ -142,6 +144,7 @@ public DynamicTableSource createDynamicTableSource(Context context) {
                 tableOptions.get(FlinkConnectorOptions.LOOKUP_ASYNC),
                 cache,
                 partitionDiscoveryIntervalMs,
+                bucketDiscoveryIntervalMs,
                 tableOptions.get(toFlinkOption(ConfigOptions.TABLE_DATALAKE_ENABLED)),
                 tableOptions.get(toFlinkOption(ConfigOptions.TABLE_MERGE_ENGINE)),
                 context.getCatalogTable().getOptions());

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/source/FlinkSource.java

@@ -67,6 +67,7 @@ public class FlinkSource<OUT>
     @Nullable private final int[] projectedFields;
     protected final OffsetsInitializer offsetsInitializer;
     protected final long scanPartitionDiscoveryIntervalMs;
+    protected final long scanBucketDiscoveryIntervalMs;
     private final boolean streaming;
     private final FlussDeserializationSchema<OUT> deserializationSchema;
 
@@ -83,6 +84,7 @@ public FlinkSource(
             @Nullable int[] projectedFields,
             OffsetsInitializer offsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             FlussDeserializationSchema<OUT> deserializationSchema,
             boolean streaming,
             List<FieldEqual> partitionFilters) {
@@ -95,6 +97,7 @@ public FlinkSource(
                 projectedFields,
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 deserializationSchema,
                 streaming,
                 partitionFilters,
@@ -110,6 +113,7 @@ public FlinkSource(
             @Nullable int[] projectedFields,
             OffsetsInitializer offsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             FlussDeserializationSchema<OUT> deserializationSchema,
             boolean streaming,
             List<FieldEqual> partitionFilters,
@@ -122,6 +126,7 @@ public FlinkSource(
         this.projectedFields = projectedFields;
         this.offsetsInitializer = offsetsInitializer;
         this.scanPartitionDiscoveryIntervalMs = scanPartitionDiscoveryIntervalMs;
+        this.scanBucketDiscoveryIntervalMs = scanBucketDiscoveryIntervalMs;
         this.deserializationSchema = deserializationSchema;
         this.streaming = streaming;
         this.partitionFilters = checkNotNull(partitionFilters);
@@ -144,6 +149,7 @@ public SplitEnumerator<SourceSplitBase, SourceEnumeratorState> createEnumerator(
                 splitEnumeratorContext,
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 streaming,
                 partitionFilters,
                 lakeSource);
@@ -163,6 +169,7 @@ public SplitEnumerator<SourceSplitBase, SourceEnumeratorState> restoreEnumerator
                 sourceEnumeratorState.getAssignedPartitions(),
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 streaming,
                 partitionFilters,
                 lakeSource);

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/source/FlinkTableSource.java

@@ -112,6 +112,7 @@ public class FlinkTableSource
     @Nullable private final LookupCache cache;
 
     private final long scanPartitionDiscoveryIntervalMs;
+    private final long scanBucketDiscoveryIntervalMs;
     private final boolean isDataLakeEnabled;
     @Nullable private final MergeEngineType mergeEngineType;
 
@@ -150,6 +151,7 @@ public FlinkTableSource(
             boolean lookupAsync,
             @Nullable LookupCache cache,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             boolean isDataLakeEnabled,
             @Nullable MergeEngineType mergeEngineType,
             Map<String, String> tableOptions) {
@@ -168,6 +170,7 @@ public FlinkTableSource(
         this.cache = cache;
 
         this.scanPartitionDiscoveryIntervalMs = scanPartitionDiscoveryIntervalMs;
+        this.scanBucketDiscoveryIntervalMs = scanBucketDiscoveryIntervalMs;
         this.isDataLakeEnabled = isDataLakeEnabled;
         this.mergeEngineType = mergeEngineType;
         this.tableOptions = tableOptions;
@@ -280,6 +283,7 @@ public boolean isBounded() {
                         projectedFields,
                         offsetsInitializer,
                         scanPartitionDiscoveryIntervalMs,
+                        scanBucketDiscoveryIntervalMs,
                         new RowDataDeserializationSchema(),
                         streaming,
                         partitionFilters,
@@ -371,6 +375,7 @@ public DynamicTableSource copy() {
                         lookupAsync,
                         cache,
                         scanPartitionDiscoveryIntervalMs,
+                        scanBucketDiscoveryIntervalMs,
                         isDataLakeEnabled,
                         mergeEngineType,
                         tableOptions);

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/source/FlussSource.java

@@ -44,6 +44,7 @@
  *     .setProjectedFields("orderId", "amount")
  *     .setStartingOffsets(OffsetsInitializer.earliest())
  *     .setScanPartitionDiscoveryIntervalMs(1000L)
+ *     .setScanBucketDiscoveryIntervalMs(1000L)
  *     .setDeserializationSchema(new OrderDeserializationSchema())
  *     .build();
  *
@@ -68,6 +69,7 @@ public class FlussSource<OUT> extends FlinkSource<OUT> {
             @Nullable int[] projectedFields,
             OffsetsInitializer offsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             FlussDeserializationSchema<OUT> deserializationSchema,
             boolean streaming) {
         // TODO: Support partition pushDown in datastream
@@ -80,6 +82,7 @@ public class FlussSource<OUT> extends FlinkSource<OUT> {
                 projectedFields,
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 deserializationSchema,
                 streaming,
                 Collections.emptyList());

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/source/FlussSourceBuilder.java

@@ -54,6 +54,7 @@
  *     .setTable("orders")
  *     .setProjectedFields("orderId", "amount")
  *     .setScanPartitionDiscoveryIntervalMs(1000L)
+ *     .setScanBucketDiscoveryIntervalMs(1000L)
  *     .setStartingOffsets(OffsetsInitializer.earliest())
  *     .setDeserializationSchema(new OrderDeserializationSchema())
  *     .build();
@@ -69,6 +70,7 @@ public class FlussSourceBuilder<OUT> {
     private int[] projectedFields;
     private String[] projectedFieldNames;
     private Long scanPartitionDiscoveryIntervalMs;
+    private Long scanBucketDiscoveryIntervalMs;
     private OffsetsInitializer offsetsInitializer;
     private FlussDeserializationSchema<OUT> deserializationSchema;
 
@@ -131,6 +133,21 @@ public FlussSourceBuilder<OUT> setScanPartitionDiscoveryIntervalMs(
         return this;
     }
 
+    /**
+     * Sets the scan bucket discovery interval in milliseconds.
+     *
+     * <p>If not specified, the default value from {@link
+     * FlinkConnectorOptions#SCAN_BUCKET_DISCOVERY_INTERVAL} is used.
+     *
+     * @param scanBucketDiscoveryIntervalMs interval in milliseconds
+     * @return this builder
+     */
+    public FlussSourceBuilder<OUT> setScanBucketDiscoveryIntervalMs(
+            long scanBucketDiscoveryIntervalMs) {
+        this.scanBucketDiscoveryIntervalMs = scanBucketDiscoveryIntervalMs;
+        return this;
+    }
+
     /**
      * Sets the starting offsets strategy for the Fluss source.
      *
@@ -225,6 +242,12 @@ public FlussSource<OUT> build() {
                             .toMillis();
         }
 
+        // if null use the default value:
+        if (scanBucketDiscoveryIntervalMs == null) {
+            scanBucketDiscoveryIntervalMs =
+                    FlinkConnectorOptions.SCAN_BUCKET_DISCOVERY_INTERVAL.defaultValue().toMillis();
+        }
+
         if (this.flussConf == null) {
             this.flussConf = new Configuration();
         }
@@ -299,6 +322,7 @@ public FlussSource<OUT> build() {
                 projectedFields,
                 offsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 deserializationSchema,
                 true);
     }

fluss-flink/fluss-flink-common/src/main/java/com/alibaba/fluss/flink/source/enumerator/FlinkSourceEnumerator.java

@@ -111,6 +111,7 @@ public class FlinkSourceEnumerator
     private final Set<TableBucket> assignedTableBuckets;
 
     private final long scanPartitionDiscoveryIntervalMs;
+    private final long scanBucketDiscoveryIntervalMs;
 
     private final boolean streaming;
     private final OffsetsInitializer startingOffsetsInitializer;
@@ -142,6 +143,7 @@ public FlinkSourceEnumerator(
             SplitEnumeratorContext<SourceSplitBase> context,
             OffsetsInitializer startingOffsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             boolean streaming,
             List<FieldEqual> partitionFilters) {
         this(
@@ -152,6 +154,7 @@ public FlinkSourceEnumerator(
                 context,
                 startingOffsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 streaming,
                 partitionFilters,
                 null);
@@ -165,6 +168,7 @@ public FlinkSourceEnumerator(
             SplitEnumeratorContext<SourceSplitBase> context,
             OffsetsInitializer startingOffsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             boolean streaming,
             List<FieldEqual> partitionFilters,
             LakeSource<LakeSplit> lakeSource) {
@@ -178,6 +182,7 @@ public FlinkSourceEnumerator(
                 Collections.emptyMap(),
                 startingOffsetsInitializer,
                 scanPartitionDiscoveryIntervalMs,
+                scanBucketDiscoveryIntervalMs,
                 streaming,
                 partitionFilters,
                 lakeSource);
@@ -193,6 +198,7 @@ public FlinkSourceEnumerator(
             Map<Long, String> assignedPartitions,
             OffsetsInitializer startingOffsetsInitializer,
             long scanPartitionDiscoveryIntervalMs,
+            long scanBucketDiscoveryIntervalMs,
             boolean streaming,
             List<FieldEqual> partitionFilters,
             @Nullable LakeSource<LakeSplit> lakeSource) {
@@ -206,6 +212,7 @@ public FlinkSourceEnumerator(
         this.startingOffsetsInitializer = startingOffsetsInitializer;
         this.assignedPartitions = new HashMap<>(assignedPartitions);
         this.scanPartitionDiscoveryIntervalMs = scanPartitionDiscoveryIntervalMs;
+        this.scanBucketDiscoveryIntervalMs = scanBucketDiscoveryIntervalMs;
         this.streaming = streaming;
         this.partitionFilters = checkNotNull(partitionFilters);
         this.stoppingOffsetsInitializer =
@@ -258,8 +265,26 @@ public void start() {
             if (!streaming) {
                 startInBatchMode();
             } else {
-                // init bucket splits and assign
-                context.callAsync(this::initNonPartitionedSplits, this::handleSplitsAdd);
+                // Currently, only log table support alter bucket.
+                // So we only enable bucket discovery for log table.
+                if (scanBucketDiscoveryIntervalMs > 0 && !hasPrimaryKey) {
+                    LOG.info(
+                            "Starting the FlussSourceEnumerator for table {} "
+                                    + "with new bucket discovery interval of {} ms.",
+                            tablePath,
+                            scanBucketDiscoveryIntervalMs);
+                    context.callAsync(
+                            this::listBuckets,
+                            this::checkBucketChanges,
+                            0,
+                            scanBucketDiscoveryIntervalMs);
+                } else {
+                    LOG.info(
+                            "Starting the FlussSourceEnumerator for table {} without bucket discovery.",
+                            tablePath);
+                    // init bucket splits and assign
+                    context.callAsync(this::initNonPartitionedSplits, this::handleSplitsAdd);
+                }
             }
         }
     }
@@ -292,6 +317,66 @@ private List<SourceSplitBase> initNonPartitionedSplits() {
         }
     }
 
+    private Set<TableBucket> listBuckets() {
+        try {
+            // Always refresh tableInfo to discovery bucket num change
+            tableInfo = flussAdmin.getTableInfo(tablePath).get();
+            int bucketNum = tableInfo.getNumBuckets();
+            Set<TableBucket> tableBuckets = new HashSet<>(bucketNum);
+            for (int bucketId = 0; bucketId < bucketNum; bucketId++) {
+                TableBucket tableBucket = new TableBucket(tableInfo.getTableId(), null, bucketId);
+                tableBuckets.add(tableBucket);
+            }
+            return tableBuckets;
+        } catch (Exception e) {
+            throw new FlinkRuntimeException(
+                    String.format("Failed to list buckets for %s", tablePath),
+                    ExceptionUtils.stripCompletionException(e));
+        }
+    }
+
+    private void checkBucketChanges(Set<TableBucket> tableBuckets, Throwable t) {
+        if (closed) {
+            // skip if the enumerator is closed to avoid unnecessary error logs
+            return;
+        }
+        if (t != null) {
+            LOG.error("Failed to list buckets for {}", tablePath, t);
+            return;
+        }
+        final BucketChange bucketChange = getBucketChange(tableBuckets);
+        if (bucketChange.isEmpty()) {
+            return;
+        }
+
+        // handle new buckets
+        context.callAsync(() -> initBucketSplits(bucketChange.newBuckets), this::handleSplitsAdd);
+    }
+
+    private BucketChange getBucketChange(Set<TableBucket> tableBuckets) {
+        return null;
+    }
+
+    private List<SourceSplitBase> initBucketSplits(Collection<TableBucket> newBuckets) {
+        List<Integer> bucketsNeedInitOffset =
+                newBuckets.stream().map(TableBucket::getBucket).collect(Collectors.toList());
+
+        List<SourceSplitBase> splits = new ArrayList<>(bucketsNeedInitOffset.size());
+        if (!bucketsNeedInitOffset.isEmpty()) {
+            startingOffsetsInitializer
+                    .getBucketOffsets(null, bucketsNeedInitOffset, bucketOffsetsRetriever)
+                    .forEach(
+                            (bucketId, startingOffset) ->
+                                    splits.add(
+                                            new LogSplit(
+                                                    new TableBucket(
+                                                            tableInfo.getTableId(), null, bucketId),
+                                                    null,
+                                                    startingOffset)));
+        }
+        return splits;
+    }
+
     private Set<PartitionInfo> listPartitions() {
         try {
             List<PartitionInfo> partitionInfos = flussAdmin.listPartitionInfos(tablePath).get();
@@ -775,6 +860,22 @@ public void close() throws IOException {
     }
 
     // --------------- private class ---------------
+
+    /** A container class to hold the newly added buckets and removed buckets. */
+    private static class BucketChange {
+        private final Collection<TableBucket> newBuckets;
+        private final Collection<TableBucket> removedBuckets;
+
+        BucketChange(Collection<TableBucket> newBuckets, Collection<TableBucket> removedBuckets) {
+            this.newBuckets = newBuckets;
+            this.removedBuckets = removedBuckets;
+        }
+
+        public boolean isEmpty() {
+            return newBuckets.isEmpty() && removedBuckets.isEmpty();
+        }
+    }
+
     /** A container class to hold the newly added partitions and removed partitions. */
     private static class PartitionChange {
         private final Collection<Partition> newPartitions;