[flink] Support custom lookup shuffle by bucket for Flink 2.x

fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java

@@ -0,0 +1,34 @@
+/*
+ * 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.fluss.flink.adapter;
+
+import org.apache.fluss.flink.source.FlinkLookupShuffleTableSource;
+import org.apache.fluss.flink.source.FlinkTableSource;
+
+/**
+ * Flink 2.x implementation: wraps the source so it implements {@code SupportsLookupCustomShuffle}.
+ * Shadows the common (no-op) class of the same fully-qualified name at package time.
+ */
+public class LookupShuffleSourceAdapter {
+
+    private LookupShuffleSourceAdapter() {}
+
+    public static FlinkTableSource maybeWithCustomShuffle(FlinkTableSource source) {
+        return new FlinkLookupShuffleTableSource(source);
+    }
+}

fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSource.java

@@ -0,0 +1,100 @@
+/*
+ * 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.fluss.flink.source;
+
+import org.apache.fluss.client.lookup.LookupType;
+import org.apache.fluss.flink.FlinkConnectorOptions;
+import org.apache.fluss.flink.source.lookup.FlussLookupInputPartitioner;
+import org.apache.fluss.flink.source.lookup.LookupNormalizer;
+import org.apache.fluss.flink.utils.FlinkUtils;
+import org.apache.fluss.metadata.DataLakeFormat;
+
+import org.apache.flink.table.connector.source.DynamicTableSource;
+import org.apache.flink.table.connector.source.abilities.SupportsLookupCustomShuffle;
+import org.apache.flink.table.types.logical.RowType;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * A {@link FlinkTableSource} variant for Flink 2.x that implements {@link
+ * SupportsLookupCustomShuffle}, shuffling the lookup-join probe stream by Fluss bucket key so that
+ * rows of the same bucket are co-located on the same lookup subtask (better cache locality and
+ * lower RPC fan-out).
+ */
+public class FlinkLookupShuffleTableSource extends FlinkTableSource
+        implements SupportsLookupCustomShuffle {
+
+    public FlinkLookupShuffleTableSource(FlinkTableSource base) {
+        super(base);
+    }
+
+    @Override
+    public DynamicTableSource copy() {
+        // Copy once via the copy constructor; this preserves the subtype (and thus the ability).
+        return new FlinkLookupShuffleTableSource(this);
+    }
+
+    @Override
+    public Optional<InputDataPartitioner> getPartitioner() {
+        LookupNormalizer normalizer = lastLookupNormalizer();
+        int[] bucketKeyIndexes = bucketKeyIndexes();
+        // Not ready or no bucket key -> keep Flink's default distribution.
+        if (normalizer == null || bucketKeyIndexes.length == 0) {
+            return Optional.empty();
+        }
+        // Only full primary-key lookups are shuffled here. A prefix lookup's normalized key does
+        // contain all bucket-key fields, so it could be shuffled too; supporting it would mean
+        // keying off the normalizer's expected-key row type instead of the full PK projection.
+        // TODO: extend custom shuffle to prefix lookups on the bucket key.
+        if (normalizer.getLookupType() != LookupType.LOOKUP) {
+            return Optional.empty();
+        }
+
+        // Number of buckets of the Fluss table. Catalog tables always have this injected; a
+        // temporary table created without an explicit (or with a malformed) 'bucket.num' silently
+        // keeps Flink's default distribution.
+        String bucketNumStr = tableOptions().get(FlinkConnectorOptions.BUCKET_NUMBER.key());
+        if (bucketNumStr == null) {
+            return Optional.empty();
+        }
+        final int numBuckets;
+        try {
+            numBuckets = Integer.parseInt(bucketNumStr);
+        } catch (NumberFormatException e) {
+            return Optional.empty();
+        }
+
+        // The normalized lookup key row is the primary key in Fluss key order.
+        RowType keyFlinkRowType = FlinkUtils.projectRowType(tableOutputType(), primaryKeyIndexes());
+
+        // bucket key field names (bucket key is a subset of the primary key)
+        List<String> allNames = tableOutputType().getFieldNames();
+        List<String> bucketKeyNames = new ArrayList<>(bucketKeyIndexes.length);
+        for (int idx : bucketKeyIndexes) {
+            bucketKeyNames.add(allNames.get(idx));
+        }
+
+        DataLakeFormat lakeFormat = tableConfigInternal().getDataLakeFormat().orElse(null);
+
+        return Optional.of(
+                new FlussLookupInputPartitioner(
+                        normalizer, keyFlinkRowType, bucketKeyNames, lakeFormat, numBuckets));
+    }
+}

fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java

@@ -0,0 +1,96 @@
+/*
+ * 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.fluss.flink.source.lookup;
+
+import org.apache.fluss.bucketing.BucketingFunction;
+import org.apache.fluss.flink.row.FlinkAsFlussRow;
+import org.apache.fluss.flink.utils.FlinkConversions;
+import org.apache.fluss.metadata.DataLakeFormat;
+import org.apache.fluss.row.InternalRow;
+import org.apache.fluss.row.encode.KeyEncoder;
+
+import org.apache.flink.table.connector.source.abilities.SupportsLookupCustomShuffle.InputDataPartitioner;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.types.logical.RowType;
+
+import javax.annotation.Nullable;
+
+import java.util.List;
+
+/**
+ * Partitions the lookup-join probe stream by Fluss bucket key, consistent with the client-side
+ * bucketing used by {@code PrimaryKeyLookuper} (bucket key encoding + {@link BucketingFunction}).
+ * Rows mapping to the same Fluss bucket are always routed to the same partition.
+ */
+public class FlussLookupInputPartitioner implements InputDataPartitioner {
+
+    private static final long serialVersionUID = 1L;
+
+    private final LookupNormalizer normalizer;
+    // Flink row type of the normalized lookup key row (primary key in Fluss key order).
+    private final RowType keyFlinkRowType;
+    private final List<String> bucketKeyNames;
+    @Nullable private final DataLakeFormat lakeFormat;
+    private final int numBuckets;
+
+    private transient KeyEncoder bucketKeyEncoder;
+    private transient BucketingFunction bucketingFunction;
+    private transient FlinkAsFlussRow reuseRow;
+
+    public FlussLookupInputPartitioner(
+            LookupNormalizer normalizer,
+            RowType keyFlinkRowType,
+            List<String> bucketKeyNames,
+            @Nullable DataLakeFormat lakeFormat,
+            int numBuckets) {
+        this.normalizer = normalizer;
+        this.keyFlinkRowType = keyFlinkRowType;
+        this.bucketKeyNames = bucketKeyNames;
+        this.lakeFormat = lakeFormat;
+        this.numBuckets = numBuckets;
+    }
+
+    private void ensureInitialized() {
+        if (bucketKeyEncoder == null) {
+            org.apache.fluss.types.RowType flussKeyType =
+                    FlinkConversions.toFlussRowType(keyFlinkRowType);
+            // bucketing uses the bucket-key encoder consistent with the client's bucket routing
+            bucketKeyEncoder =
+                    KeyEncoder.ofBucketKeyEncoder(flussKeyType, bucketKeyNames, lakeFormat);
+            bucketingFunction = BucketingFunction.of(lakeFormat);
+            reuseRow = new FlinkAsFlussRow();
+        }
+    }
+
+    @Override
+    public int partition(RowData joinKeys, int numPartitions) {
+        ensureInitialized();
+        // normalize the projected join keys into Fluss primary key order
+        RowData normalizedKey = normalizer.normalizeLookupKey(joinKeys);
+        InternalRow flussKeyRow = reuseRow.replace(normalizedKey);
+        byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(flussKeyRow);
+        int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets);
+        int p = bucketId % numPartitions;
+        return p < 0 ? p + numPartitions : p;
+    }
+
+    @Override
+    public boolean isDeterministic() {
+        return true;
+    }
+}