[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,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.adapter;
+
+import org.apache.fluss.flink.FlinkConnectorOptions;
+import org.apache.fluss.flink.source.FlinkLookupShuffleTableSource;
+import org.apache.fluss.flink.source.FlinkTableSource;
+
+import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.table.api.ValidationException;
+
+/**
+ * Flink 2.x implementation: wraps the source so it implements {@code SupportsLookupCustomShuffle},
+ * but only when the table exposes the metadata required to reproduce Fluss bucket routing. Shadows
+ * the common (no-op) class of the same fully-qualified name at package time.
+ *
+ * <p>Eligibility is decided here, at planning time, rather than inside {@code getPartitioner()}:
+ *
+ * <ul>
+ *   <li>the table is not lookup-shuffle eligible, or {@code bucket.num} is not configured -&gt;
+ *       leave the source unwrapped so that an explicitly requested lookup shuffle falls back to
+ *       Flink's default hash distribution;
+ *   <li>{@code bucket.num} configured but not a positive integer -&gt; fail fast with a {@link
+ *       ValidationException} (a real misconfiguration);
+ *   <li>otherwise wrap with the validated bucket count, so the wrapped source can always return a
+ *       partitioner.
+ * </ul>
+ *
+ * <p>The unwrapped (hash-fallback) branch is a defensive safety net rather than a reachable state
+ * for a Fluss table: a Fluss table can only be created through the {@code FlussCatalog}, which
+ * always defaults {@code bucket.num} and derives the bucket key from the primary key, while Flink
+ * rejects creating a Fluss source as a temporary/{@code connector}-based table. It is therefore
+ * covered by adapter unit tests rather than an end-to-end planner test.
+ */
+public class LookupShuffleSourceAdapter {
+
+    private LookupShuffleSourceAdapter() {}
+
+    /**
+     * Wraps an eligible source with the Flink 2.x custom lookup shuffle ability.
+     *
+     * <p>An explicitly configured invalid bucket number is rejected regardless of eligibility.
+     * Missing custom-shuffle metadata leaves the source unwrapped so Flink can apply its normal
+     * hash fallback when lookup shuffle is requested.
+     *
+     * @param source source to wrap
+     * @param tableOptions resolved table options
+     * @param lookupShuffleEligible whether the table has a primary key and a non-empty bucket key
+     * @return the wrapped source when custom bucket routing is available, otherwise the original
+     *     source
+     * @throws ValidationException if {@code bucket.num} is configured but is not a positive integer
+     */
+    public static FlinkTableSource maybeWithCustomShuffle(
+            FlinkTableSource source, ReadableConfig tableOptions, boolean lookupShuffleEligible) {
+        final Integer numBuckets;
+        try {
+            numBuckets = tableOptions.get(FlinkConnectorOptions.BUCKET_NUMBER);
+        } catch (IllegalArgumentException e) {
+            // bucket.num is present but not a valid integer.
+            throw new ValidationException(
+                    String.format(
+                            "Invalid value for '%s': it must be a positive integer to enable the "
+                                    + "Fluss custom lookup shuffle.",
+                            FlinkConnectorOptions.BUCKET_NUMBER.key()),
+                    e);
+        }
+        if (numBuckets != null && numBuckets <= 0) {
+            throw new ValidationException(
+                    String.format(
+                            "Invalid value '%d' for '%s': it must be a positive integer to enable "
+                                    + "the Fluss custom lookup shuffle.",
+                            numBuckets, FlinkConnectorOptions.BUCKET_NUMBER.key()));
+        }
+        if (!lookupShuffleEligible || numBuckets == null) {
+            // Custom-shuffle metadata is incomplete; do not advertise the ability so that Flink
+            // falls back to a hash lookup shuffle instead of no shuffle at all.
+            return source;
+        }
+        return new FlinkLookupShuffleTableSource(source, numBuckets);
+    }
+}

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

@@ -0,0 +1,120 @@
+/*
+ * 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.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;
+
+import static org.apache.fluss.utils.Preconditions.checkArgument;
+import static org.apache.fluss.utils.Preconditions.checkState;
+
+/**
+ * A {@link FlinkTableSource} variant for Flink 2.x that implements {@link
+ * SupportsLookupCustomShuffle}, shuffling the lookup-join probe stream by Fluss bucket so that rows
+ * of the same bucket are co-located on the same lookup subtask (better cache locality and lower RPC
+ * fan-out).
+ *
+ * <p>This subtype is only created (by {@code LookupShuffleSourceAdapter}) when the table exposes
+ * the metadata required to reproduce Fluss bucket routing: a non-empty bucket key and a positive
+ * {@code bucket.num}. When that metadata is incomplete the source is left unwrapped so Flink falls
+ * back to its default (hash) lookup shuffle; when it is present but illegal the adapter fails fast.
+ * As a result the invariant of this class is that {@link #getPartitioner()} always returns a
+ * partitioner ({@link Optional#empty()} would suppress the shuffle entirely rather than fall back
+ * to a hash shuffle, because the planner stops inserting its own shuffle once this ability is
+ * advertised).
+ */
+public class FlinkLookupShuffleTableSource extends FlinkTableSource
+        implements SupportsLookupCustomShuffle {
+
+    /** Number of buckets of the Fluss table; validated positive by the wrapping adapter. */
+    private final int numBuckets;
+
+    /**
+     * Creates a table source with the Flink 2.x custom lookup shuffle ability.
+     *
+     * @param base base Fluss table source
+     * @param numBuckets positive number of buckets in the Fluss table
+     */
+    public FlinkLookupShuffleTableSource(FlinkTableSource base, int numBuckets) {
+        super(base);
+        checkArgument(numBuckets > 0, "numBuckets must be positive, but was %s.", numBuckets);
+        this.numBuckets = numBuckets;
+    }
+
+    @Override
+    public DynamicTableSource copy() {
+        return new FlinkLookupShuffleTableSource(this, numBuckets);
+    }
+
+    @Override
+    public Optional<InputDataPartitioner> getPartitioner() {
+        LookupNormalizer normalizer = lastLookupNormalizer();
+        // The planner always calls getLookupRuntimeProvider() (which stashes the normalizer) before
+        // getPartitioner(); a null normalizer would mean the call order this source relies on was
+        // violated. Likewise this source is only ever created for a table with a bucket key. Both
+        // are internal invariants: surface them instead of silently returning empty (which the
+        // planner would treat as "keep the arbitrary distribution", i.e. no shuffle at all).
+        checkState(
+                normalizer != null,
+                "getPartitioner() was invoked before getLookupRuntimeProvider(); "
+                        + "the Fluss custom lookup shuffle cannot determine the lookup keys.");
+        int[] bucketKeyIndexes = bucketKeyIndexes();
+        checkState(
+                bucketKeyIndexes.length > 0,
+                "Fluss custom lookup shuffle was enabled for a table without a bucket key.");
+
+        RowType keyFlinkRowType =
+                FlinkUtils.projectRowType(tableOutputType(), normalizer.getLookupKeyIndexes());
+
+        List<String> allNames = tableOutputType().getFieldNames();
+        List<String> bucketKeyNames = new ArrayList<>(bucketKeyIndexes.length);
+        for (int idx : bucketKeyIndexes) {
+            bucketKeyNames.add(allNames.get(idx));
+        }
+
+        // Partitioned tables route by (partitionId, bucketId): rows with the same bucket id but
+        // different partitions target different tablets, so the partition keys must join the bucket
+        // key in the channel computation to avoid collapsing them onto one subtask.
+        int[] partitionKeyIndexes = partitionKeyIndexes();
+        List<String> partitionKeyNames = new ArrayList<>(partitionKeyIndexes.length);
+        for (int idx : partitionKeyIndexes) {
+            partitionKeyNames.add(allNames.get(idx));
+        }
+
+        DataLakeFormat lakeFormat = tableConfigInternal().getDataLakeFormat().orElse(null);
+
+        return Optional.of(
+                new FlussLookupInputPartitioner(
+                        normalizer,
+                        keyFlinkRowType,
+                        bucketKeyNames,
+                        partitionKeyNames,
+                        lakeFormat,
+                        numBuckets));
+    }
+}

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

@@ -0,0 +1,148 @@
+/*
+ * 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.fluss.utils.MathUtils;
+import org.apache.fluss.utils.MurmurHashUtils;
+
+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;
+
+import static org.apache.fluss.utils.Preconditions.checkArgument;
+import static org.apache.fluss.utils.UnsafeUtils.BYTE_ARRAY_BASE_OFFSET;
+
+/**
+ * Partitions the lookup-join probe stream by Fluss bucket, consistent with the client-side
+ * bucketing used by {@code PrimaryKeyLookuper}/{@code PrefixKeyLookuper} (bucket key encoding +
+ * {@link BucketingFunction}). Rows mapping to the same Fluss bucket are always routed to the same
+ * partition.
+ *
+ * <p>For partitioned tables the client routes by {@code (partitionId, bucketId)}: the same bucket
+ * id in different partitions is a different tablet. The partition keys therefore join the bucket
+ * key in the channel computation, so rows targeting the same tablet stay co-located while different
+ * partitions are spread across subtasks instead of collapsing onto one.
+ */
+public class FlussLookupInputPartitioner implements InputDataPartitioner {
+
+    private static final long serialVersionUID = 1L;
+
+    private final LookupNormalizer normalizer;
+    // Flink row type of the normalized lookup key row (the expected lookup keys in Fluss key order,
+    // i.e. the full primary key for a primary-key lookup, or the bucket keys + partition keys for a
+    // prefix lookup).
+    private final RowType keyFlinkRowType;
+    private final List<String> bucketKeyNames;
+    // Partition-key field names within the normalized lookup key; empty for non-partitioned tables.
+    private final List<String> partitionKeyNames;
+    @Nullable private final DataLakeFormat lakeFormat;
+    private final int numBuckets;
+
+    private transient KeyEncoder bucketKeyEncoder;
+    // null when the table is not partitioned.
+    @Nullable private transient KeyEncoder partitionKeyEncoder;
+    private transient BucketingFunction bucketingFunction;
+    private transient FlinkAsFlussRow reuseRow;
+
+    /**
+     * Creates a partitioner consistent with Fluss client-side bucket routing.
+     *
+     * @param normalizer normalizes Flink lookup keys into Fluss lookup-key order
+     * @param keyFlinkRowType row type of the normalized lookup key
+     * @param bucketKeyNames bucket-key field names within the normalized lookup key
+     * @param partitionKeyNames partition-key field names within the normalized lookup key; empty
+     *     for non-partitioned tables
+     * @param lakeFormat optional lake format that defines key encoding and bucketing behavior
+     * @param numBuckets positive number of buckets in the Fluss table
+     */
+    public FlussLookupInputPartitioner(
+            LookupNormalizer normalizer,
+            RowType keyFlinkRowType,
+            List<String> bucketKeyNames,
+            List<String> partitionKeyNames,
+            @Nullable DataLakeFormat lakeFormat,
+            int numBuckets) {
+        this.normalizer = normalizer;
+        this.keyFlinkRowType = keyFlinkRowType;
+        this.bucketKeyNames = bucketKeyNames;
+        this.partitionKeyNames = partitionKeyNames;
+        this.lakeFormat = lakeFormat;
+        checkArgument(numBuckets > 0, "numBuckets must be positive, but was %s.", numBuckets);
+        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);
+            if (!partitionKeyNames.isEmpty()) {
+                partitionKeyEncoder =
+                        KeyEncoder.ofBucketKeyEncoder(flussKeyType, partitionKeyNames, lakeFormat);
+            }
+            bucketingFunction = BucketingFunction.of(lakeFormat);
+            reuseRow = new FlinkAsFlussRow();
+        }
+    }
+
+    @Override
+    public int partition(RowData joinKeys, int numPartitions) {
+        // Null lookup keys cannot match, but LEFT lookup joins still need to reach the operator.
+        for (int i = 0; i < joinKeys.getArity(); i++) {
+            if (joinKeys.isNullAt(i)) {
+                return 0;
+            }
+        }
+        ensureInitialized();
+        // normalize the projected join keys into the Fluss key order
+        RowData normalizedKey = normalizer.normalizeLookupKey(joinKeys);
+        InternalRow flussKeyRow = reuseRow.replace(normalizedKey);
+        byte[] bucketKeyBytes = bucketKeyEncoder.encodeKey(flussKeyRow);
+        // BucketingFunction always returns a non-negative bucket id.
+        int bucketId = bucketingFunction.bucketing(bucketKeyBytes, numBuckets);
+        if (partitionKeyEncoder == null) {
+            return bucketId % numPartitions;
+        }
+        // Route by (partition, bucket) so different partitions of the same bucket id are spread
+        // across subtasks while rows targeting the same tablet stay co-located. Mix in the bucket
+        // id using the same Murmur hash family the client bucketing uses (FlussBucketingFunction).
+        byte[] partitionKeyBytes = partitionKeyEncoder.encodeKey(flussKeyRow);
+        int partitionHash =
+                MurmurHashUtils.hashUnsafeBytes(
+                        partitionKeyBytes, BYTE_ARRAY_BASE_OFFSET, partitionKeyBytes.length);
+        // murmurHash returns a non-negative int, so the modulo is always a valid channel.
+        return MathUtils.murmurHash(partitionHash * 31 + bucketId) % numPartitions;
+    }
+
+    @Override
+    public boolean isDeterministic() {
+        return true;
+    }
+}