diff --git a/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java b/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java new file mode 100644 index 0000000000..3321b62253 --- /dev/null +++ b/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. + * + *

Eligibility is decided here, at planning time, rather than inside {@code getPartitioner()}: + * + *

+ * + *

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. + * + *

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); + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSource.java b/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSource.java new file mode 100644 index 0000000000..f17f405587 --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSource.java @@ -0,0 +1,113 @@ +/* + * 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). + * + *

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() { + // Copy once via the copy constructor; this preserves the subtype (and thus the ability) and + // the stashed lookup normalizer carried over by the base copy constructor. + return new FlinkLookupShuffleTableSource(this, numBuckets); + } + + @Override + public Optional 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."); + + // The normalized lookup key row is the expected lookup keys in Fluss key order: the full + // primary key for a primary-key lookup, or the bucket keys + partition keys for a prefix + // lookup on the bucket key. Both are shuffled by the same bucket-key routing, so build the + // key row type from the normalizer's expected-key indexes (this is what makes prefix + // lookups shuffle rather than being excluded). + RowType keyFlinkRowType = + FlinkUtils.projectRowType(tableOutputType(), normalizer.getLookupKeyIndexes()); + + List allNames = tableOutputType().getFieldNames(); + List 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)); + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java b/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java new file mode 100644 index 0000000000..8c9c6a5a29 --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/main/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitioner.java @@ -0,0 +1,111 @@ +/* + * 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; + +import static org.apache.fluss.utils.Preconditions.checkArgument; + +/** + * 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. + */ +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 bucketKeyNames; + @Nullable private final DataLakeFormat lakeFormat; + private final int numBuckets; + + private transient KeyEncoder bucketKeyEncoder; + 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 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 bucketKeyNames, + @Nullable DataLakeFormat lakeFormat, + int numBuckets) { + this.normalizer = normalizer; + this.keyFlinkRowType = keyFlinkRowType; + this.bucketKeyNames = bucketKeyNames; + 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); + bucketingFunction = BucketingFunction.of(lakeFormat); + reuseRow = new FlinkAsFlussRow(); + } + } + + @Override + public int partition(RowData joinKeys, int numPartitions) { + 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); + return bucketId % numPartitions; + } + + @Override + public boolean isDeterministic() { + return true; + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapterTest.java b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapterTest.java new file mode 100644 index 0000000000..2eb0977d84 --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapterTest.java @@ -0,0 +1,170 @@ +/* + * 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.config.Configuration; +import org.apache.fluss.config.TableConfig; +import org.apache.fluss.flink.FlinkConnectorOptions; +import org.apache.fluss.flink.source.FlinkLookupShuffleTableSource; +import org.apache.fluss.flink.source.FlinkTableSource; +import org.apache.fluss.flink.utils.FlinkConnectorOptionsUtils; +import org.apache.fluss.metadata.TablePath; + +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.ValidationException; +import org.apache.flink.table.connector.source.abilities.SupportsLookupCustomShuffle; +import org.apache.flink.table.types.logical.RowType; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for the Flink 2.x {@link LookupShuffleSourceAdapter}, which decides at planning time + * whether to advertise the custom lookup shuffle ability. + * + *

The behavior matrix under test: + * + *

+ */ +class LookupShuffleSourceAdapterTest { + + private static FlinkTableSource tableSource() { + RowType tableOutputType = + (RowType) + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.INT()), + DataTypes.FIELD("name", DataTypes.STRING())) + .getLogicalType(); + Configuration flussConfig = new Configuration(); + flussConfig.setString(FlinkConnectorOptions.BOOTSTRAP_SERVERS.key(), "localhost:9092"); + FlinkConnectorOptionsUtils.StartupOptions startupOptions = + new FlinkConnectorOptionsUtils.StartupOptions(); + startupOptions.startupMode = FlinkConnectorOptions.ScanStartupMode.EARLIEST; + return new FlinkTableSource( + TablePath.of("test_db", "dim"), + flussConfig, + new TableConfig(new Configuration()), + tableOutputType, + new int[] {0}, // primary key indexes (id) + new int[] {0}, // bucket key indexes (id) + new int[] {}, // partition key indexes + true, // streaming + startupOptions, + false, // lookup async + false, // insert if not exists + null, // cache + 1000L, // scan partition discovery interval + false, // is data lake enabled + null, // merge engine type + new HashMap<>(), + null); // lease context + } + + private static org.apache.flink.configuration.Configuration options(String bucketNum) { + org.apache.flink.configuration.Configuration options = + new org.apache.flink.configuration.Configuration(); + if (bucketNum != null) { + options.setString(FlinkConnectorOptions.BUCKET_NUMBER.key(), bucketNum); + } + return options; + } + + @Test + void testWrapsWhenEligibleWithValidBucketNumber() { + FlinkTableSource source = tableSource(); + FlinkTableSource result = + LookupShuffleSourceAdapter.maybeWithCustomShuffle(source, options("3"), true); + assertThat(result) + .isInstanceOf(SupportsLookupCustomShuffle.class) + .isInstanceOf(FlinkLookupShuffleTableSource.class); + } + + @Test + void testNotWrappedWhenBucketNumberMissing() { + // Eligible table, but the optional bucket.num is not configured: leave the source unwrapped + // so Flink applies its default (hash) lookup shuffle rather than no shuffle at all. + FlinkTableSource source = tableSource(); + FlinkTableSource result = + LookupShuffleSourceAdapter.maybeWithCustomShuffle(source, options(null), true); + assertThat(result).isSameAs(source); + assertThat(result).isNotInstanceOf(SupportsLookupCustomShuffle.class); + } + + @Test + void testNotWrappedWhenNotEligible() { + // The factory reports the table has no primary key and/or no bucket key, so bucket routing + // cannot be reproduced even if bucket.num is present. + FlinkTableSource source = tableSource(); + FlinkTableSource result = + LookupShuffleSourceAdapter.maybeWithCustomShuffle(source, options("3"), false); + assertThat(result).isSameAs(source); + assertThat(result).isNotInstanceOf(SupportsLookupCustomShuffle.class); + } + + @Test + void testFailFastOnNonPositiveBucketNumber() { + FlinkTableSource source = tableSource(); + for (String invalid : new String[] {"0", "-1"}) { + assertThatThrownBy( + () -> + LookupShuffleSourceAdapter.maybeWithCustomShuffle( + source, options(invalid), true)) + .as("bucket.num=%s must fail fast", invalid) + .isInstanceOf(ValidationException.class) + .hasMessageContaining(FlinkConnectorOptions.BUCKET_NUMBER.key()); + } + } + + @Test + void testFailFastOnNonNumericBucketNumber() { + FlinkTableSource source = tableSource(); + assertThatThrownBy( + () -> + LookupShuffleSourceAdapter.maybeWithCustomShuffle( + source, options("abc"), true)) + .isInstanceOf(ValidationException.class) + .hasMessageContaining(FlinkConnectorOptions.BUCKET_NUMBER.key()); + } + + @Test + void testFailFastOnInvalidBucketNumberEvenWhenNotEligible() { + // An explicitly configured invalid bucket.num is a real misconfiguration and must be + // rejected before the eligibility check, so it fails even for a non-eligible table. + FlinkTableSource source = tableSource(); + for (String invalid : new String[] {"abc", "0"}) { + assertThatThrownBy( + () -> + LookupShuffleSourceAdapter.maybeWithCustomShuffle( + source, options(invalid), false)) + .as("bucket.num=%s must fail fast even when not eligible", invalid) + .isInstanceOf(ValidationException.class) + .hasMessageContaining(FlinkConnectorOptions.BUCKET_NUMBER.key()); + } + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/Flink22LookupShuffleITCase.java b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/Flink22LookupShuffleITCase.java new file mode 100644 index 0000000000..10134cc45f --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/Flink22LookupShuffleITCase.java @@ -0,0 +1,555 @@ +/* + * 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.bucketing.BucketingFunction; +import org.apache.fluss.client.table.Table; +import org.apache.fluss.client.table.writer.UpsertWriter; +import org.apache.fluss.config.ConfigOptions; +import org.apache.fluss.flink.row.FlinkAsFlussRow; +import org.apache.fluss.flink.source.lookup.FlussLookupInputPartitioner; +import org.apache.fluss.flink.source.lookup.LookupNormalizer; +import org.apache.fluss.flink.utils.FlinkConversions; +import org.apache.fluss.flink.utils.FlinkTestBase; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.encode.KeyEncoder; + +import org.apache.flink.api.common.functions.Partitioner; +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.RowTypeInfo; +import org.apache.flink.streaming.api.datastream.DataStream; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.graph.StreamEdge; +import org.apache.flink.streaming.api.graph.StreamGraph; +import org.apache.flink.streaming.api.graph.StreamNode; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.EnvironmentSettings; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.table.api.config.ExecutionConfigOptions; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import static org.apache.fluss.flink.FlinkConnectorOptions.BOOTSTRAP_SERVERS; +import static org.apache.fluss.flink.source.testutils.FlinkRowAssertionsUtils.assertResultsIgnoreOrder; +import static org.apache.fluss.server.testutils.FlussClusterExtension.BUILTIN_DATABASE; +import static org.apache.fluss.testutils.DataTestUtils.row; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * IT case for {@code SupportsLookupCustomShuffle} (custom lookup shuffle) in Flink 2.2. + * + *

Runs lookup joins on a real Fluss + Flink mini-cluster with the {@code LOOKUP(..,'shuffle' = + * 'true')} hint enabled (and lookup cache disabled) to verify that: + * + *

+ */ +public class Flink22LookupShuffleITCase extends FlinkTestBase { + + private static final String CATALOG_NAME = "test_catalog"; + + private StreamExecutionEnvironment execEnv; + private StreamTableEnvironment tEnv; + + @BeforeEach + public void beforeEach() { + bootstrapServers = conn.getConfiguration().get(ConfigOptions.BOOTSTRAP_SERVERS).get(0); + execEnv = StreamExecutionEnvironment.getExecutionEnvironment(); + tEnv = StreamTableEnvironment.create(execEnv, EnvironmentSettings.inStreamingMode()); + tEnv.executeSql( + String.format( + "create catalog %s with ('type' = 'fluss', '%s' = '%s')", + CATALOG_NAME, BOOTSTRAP_SERVERS.key(), bootstrapServers)); + tEnv.useCatalog(CATALOG_NAME); + tEnv.executeSql(String.format("create database if not exists `%s`", DEFAULT_DB)); + tEnv.useDatabase(DEFAULT_DB); + // parallelism > 1 so that the custom shuffle actually redistributes the probe stream + tEnv.getConfig().set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4); + } + + @AfterEach + void after() { + tEnv.useDatabase(BUILTIN_DATABASE); + tEnv.executeSql(String.format("drop database `%s` cascade", DEFAULT_DB)); + } + + @Test + void testLookupShuffleOnNonPartitionedPkTable() throws Exception { + String dim = "dim_pk"; + // pk == bucket key == id, 3 buckets, lookup cache disabled (not set) + tEnv.executeSql( + String.format( + "create table %s (" + + " id int not null," + + " address varchar," + + " name varchar," + + " primary key (id) NOT ENFORCED" + + ") with ('bucket.num' = '3', 'bucket.key' = 'id'," + + " 'lookup.async' = 'false')", + dim)); + try (Table dimTable = conn.getTable(TablePath.of(DEFAULT_DB, dim))) { + UpsertWriter writer = dimTable.newUpsert().createWriter(); + for (int i = 1; i <= 5; i++) { + writer.upsert(row(i, "address" + i, "name" + (i % 4))); + } + writer.flush(); + } + + registerNonPartitionedSrc(); + + List expected = + Arrays.asList("+I[1, 11, name1]", "+I[2, 2, name2]", "+I[3, 33, name3]"); + String columns = "src.a, src.c, %s.name"; + String from = "FROM src JOIN %s FOR SYSTEM_TIME AS OF src.proc ON src.a = %s.id"; + + // with the custom shuffle enabled + String withShuffle = + String.format( + "SELECT /*+ LOOKUP('table' = '%s', 'shuffle' = 'true') */ " + + columns + + " " + + from, + dim, + dim, + dim, + dim); + // baseline without shuffle + String withoutShuffle = String.format("SELECT " + columns + " " + from, dim, dim, dim); + + // The custom lookup shuffle must actually be applied: Flink wraps our InputDataPartitioner + // in a RowDataCustomStreamPartitioner on the probe-side edge feeding the lookup join. + // Checking the plan text is not enough here, because the 'shuffle=[true]' digest is driven + // by the hint alone and Flink would fall back to a hash shuffle (still shown as + // shuffle=[true]) when the source provides no custom partitioner. + assertThat(usesCustomShufflePartitioner(withShuffle)) + .as("probe stream should be repartitioned by the Fluss custom partitioner") + .isTrue(); + assertThat(usesCustomShufflePartitioner(withoutShuffle)) + .as("no custom partitioner without the shuffle hint") + .isFalse(); + + assertResultsIgnoreOrder(tEnv.executeSql(withShuffle).collect(), expected, true); + // parity: same results without the shuffle + assertResultsIgnoreOrder(tEnv.executeSql(withoutShuffle).collect(), expected, true); + } + + @Test + void testLookupShuffleOnPartitionedPkTable() throws Exception { + String dim = "dim_pk_part"; + tEnv.executeSql( + String.format( + "create table %s (" + + " id int not null," + + " address varchar," + + " name varchar," + + " p_date varchar," + + " primary key (id, p_date) NOT ENFORCED" + + ") partitioned by (p_date) with (" + + " 'bucket.num' = '3', 'bucket.key' = 'id', 'lookup.async' = 'false'," + + " 'table.auto-partition.enabled' = 'true'," + + " 'table.auto-partition.time-unit' = 'year')", + dim)); + + TablePath dimPath = TablePath.of(DEFAULT_DB, dim); + Map partitionNameById = + waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), dimPath); + Iterator partitionIterator = partitionNameById.values().iterator(); + String partition1 = partitionIterator.next(); + String partition2 = partitionIterator.next(); + + // dim data only lives in partition1 + try (Table dimTable = conn.getTable(dimPath)) { + UpsertWriter writer = dimTable.newUpsert().createWriter(); + for (int i = 1; i <= 5; i++) { + writer.upsert(row(i, "address" + i, "name" + (i % 4), partition1)); + } + writer.flush(); + } + + registerPartitionedSrc(partition1, partition2); + + // only src rows in partition1 match dim data + List expected = Arrays.asList("+I[1, 11, name1]", "+I[2, 2, name2]"); + String columns = "src.a, src.c, %s.name"; + String from = + "FROM src JOIN %s FOR SYSTEM_TIME AS OF src.proc ON src.a = %s.id" + + " AND src.p_date = %s.p_date"; + + String withShuffle = + String.format( + "SELECT /*+ LOOKUP('table' = '%s', 'shuffle' = 'true') */ " + + columns + + " " + + from, + dim, + dim, + dim, + dim, + dim); + String withoutShuffle = String.format("SELECT " + columns + " " + from, dim, dim, dim, dim); + + assertThat(usesCustomShufflePartitioner(withShuffle)) + .as("probe stream should be repartitioned by the Fluss custom partitioner") + .isTrue(); + assertThat(usesCustomShufflePartitioner(withoutShuffle)) + .as("no custom partitioner without the shuffle hint") + .isFalse(); + + assertResultsIgnoreOrder(tEnv.executeSql(withShuffle).collect(), expected, true); + assertResultsIgnoreOrder(tEnv.executeSql(withoutShuffle).collect(), expected, true); + } + + @Test + void testLookupShuffleOnPrefixKeyLookup() throws Exception { + String dim = "dim_prefix"; + // primary key (name, id), bucket key = name (a strict prefix of the PK). A lookup on the + // bucket key alone is a prefix lookup, which must still be shuffled by the bucket key. + tEnv.executeSql( + String.format( + "create table %s (" + + " name varchar not null," + + " id int not null," + + " address varchar," + + " primary key (name, id) NOT ENFORCED" + + ") with ('bucket.num' = '3', 'bucket.key' = 'name'," + + " 'lookup.async' = 'false')", + dim)); + try (Table dimTable = conn.getTable(TablePath.of(DEFAULT_DB, dim))) { + UpsertWriter writer = dimTable.newUpsert().createWriter(); + writer.upsert(row("name1", 1, "address1")); + writer.upsert(row("name1", 5, "address5")); + writer.upsert(row("name2", 2, "address2")); + writer.upsert(row("name0", 10, "address4")); + writer.flush(); + } + + registerNonPartitionedSrc(); + + // prefix lookup on the bucket key (name) returns every dim row sharing that name + List expected = + Arrays.asList( + "+I[1, name1, address1]", + "+I[1, name1, address5]", + "+I[2, name2, address2]", + "+I[10, name0, address4]"); + String columns = "src.a, src.b, %s.address"; + String from = "FROM src JOIN %s FOR SYSTEM_TIME AS OF src.proc ON src.b = %s.name"; + + String withShuffle = + String.format( + "SELECT /*+ LOOKUP('table' = '%s', 'shuffle' = 'true') */ " + + columns + + " " + + from, + dim, + dim, + dim, + dim); + String withoutShuffle = String.format("SELECT " + columns + " " + from, dim, dim, dim); + + // The prefix lookup must actually be shuffled by the custom partitioner (before this + // change prefix lookups were excluded and produced no shuffle at all). + assertThat(usesCustomShufflePartitioner(withShuffle)) + .as("prefix lookup probe stream should be repartitioned by the Fluss partitioner") + .isTrue(); + assertThat(usesCustomShufflePartitioner(withoutShuffle)) + .as("no custom partitioner without the shuffle hint") + .isFalse(); + + assertResultsIgnoreOrder(tEnv.executeSql(withShuffle).collect(), expected, true); + assertResultsIgnoreOrder(tEnv.executeSql(withoutShuffle).collect(), expected, true); + } + + @Test + void testLookupShuffleOnPartitionedPrefixKeyLookup() throws Exception { + String dim = "dim_prefix_part"; + // partitioned table; primary key (name, id, p_date), bucket key = name (a strict prefix of + // the PK). A lookup on the bucket key + partition key (name, p_date) is a partitioned + // prefix + // lookup, which must still be shuffled by the bucket key. + tEnv.executeSql( + String.format( + "create table %s (" + + " name varchar not null," + + " id int not null," + + " address varchar," + + " p_date varchar," + + " primary key (name, id, p_date) NOT ENFORCED" + + ") partitioned by (p_date) with (" + + " 'bucket.num' = '3', 'bucket.key' = 'name', 'lookup.async' = 'false'," + + " 'table.auto-partition.enabled' = 'true'," + + " 'table.auto-partition.time-unit' = 'year')", + dim)); + + TablePath dimPath = TablePath.of(DEFAULT_DB, dim); + Map partitionNameById = + waitUntilPartitions(FLUSS_CLUSTER_EXTENSION.getZooKeeperClient(), dimPath); + Iterator partitionIterator = partitionNameById.values().iterator(); + String partition1 = partitionIterator.next(); + String partition2 = partitionIterator.next(); + + try (Table dimTable = conn.getTable(dimPath)) { + UpsertWriter writer = dimTable.newUpsert().createWriter(); + writer.upsert(row("name1", 1, "address1", partition1)); + writer.upsert(row("name1", 5, "address5", partition1)); + writer.upsert(row("name2", 2, "address2", partition1)); + writer.upsert(row("name0", 10, "address4", partition2)); + writer.flush(); + } + + registerPartitionedSrc(partition1, partition2); + + // prefix lookup on bucket key + partition key (name, p_date) + List expected = + Arrays.asList( + "+I[1, name1, address1]", + "+I[1, name1, address5]", + "+I[2, name2, address2]", + "+I[10, name0, address4]"); + String columns = "src.a, src.b, %s.address"; + String from = + "FROM src JOIN %s FOR SYSTEM_TIME AS OF src.proc ON src.b = %s.name" + + " AND src.p_date = %s.p_date"; + + String withShuffle = + String.format( + "SELECT /*+ LOOKUP('table' = '%s', 'shuffle' = 'true') */ " + + columns + + " " + + from, + dim, + dim, + dim, + dim, + dim); + String withoutShuffle = String.format("SELECT " + columns + " " + from, dim, dim, dim, dim); + + assertThat(usesCustomShufflePartitioner(withShuffle)) + .as("partitioned prefix lookup should use the Fluss custom partitioner") + .isTrue(); + assertThat(usesCustomShufflePartitioner(withoutShuffle)) + .as("no custom partitioner without the shuffle hint") + .isFalse(); + + assertResultsIgnoreOrder(tEnv.executeSql(withShuffle).collect(), expected, true); + assertResultsIgnoreOrder(tEnv.executeSql(withoutShuffle).collect(), expected, true); + } + + /** + * Builds the {@link StreamGraph} for the given query and reports whether the probe stream + * feeding the lookup join is repartitioned by Fluss' custom {@code InputDataPartitioner} (which + * Flink wraps in a {@code RowDataCustomStreamPartitioner}). Returns {@code false} when the edge + * uses a forward/hash partitioner instead, i.e. when the custom shuffle was not applied. + */ + private boolean usesCustomShufflePartitioner(String query) { + tEnv.toChangelogStream(tEnv.sqlQuery(query)); + // clearTransformations=true so the next inspection starts from a clean environment + StreamGraph streamGraph = execEnv.getStreamGraph(true); + for (StreamNode node : streamGraph.getStreamNodes()) { + for (StreamEdge edge : node.getInEdges()) { + if ("RowDataCustomStreamPartitioner" + .equals(edge.getPartitioner().getClass().getSimpleName())) { + return true; + } + } + } + return false; + } + + @Test + void testSameBucketKeysAreRoutedToSameSubtask() throws Exception { + int numBuckets = 4; + int parallelism = 4; + int numKeys = 12; + execEnv.setParallelism(parallelism); + + // The production partitioner, wired exactly as FlinkLookupShuffleTableSource#getPartitioner + // builds it for a table whose primary key == bucket key (identity normalizer). + RowType keyRowType = + RowType.of(new LogicalType[] {new IntType(false)}, new String[] {"id"}); + LookupNormalizer normalizer = + LookupNormalizer.createPrimaryKeyLookupNormalizer(new int[] {0}, keyRowType); + FlussLookupInputPartitioner flussPartitioner = + new FlussLookupInputPartitioner( + normalizer, + keyRowType, + Collections.singletonList("id"), + /* lakeFormat */ null, + numBuckets); + + List ids = new ArrayList<>(); + for (int i = 1; i <= numKeys; i++) { + ids.add(i); + } + + // Route the keys through a real Flink job using the production partitioner. This mirrors + // RowDataCustomStreamPartitioner#selectChannel (partition(key, numberOfChannels)); the + // downstream map is a FORWARD chain, so the subtask it observes IS the channel the key was + // routed to. + DataStream tagged = + execEnv.fromCollection(ids) + .partitionCustom( + new DelegatingPartitioner(flussPartitioner), new IdKeySelector()) + .map(new SubtaskTagger()) + .returns(Types.ROW(Types.INT, Types.INT)); + + Map subtaskById = new HashMap<>(); + try (CloseableIterator it = tagged.executeAndCollect()) { + while (it.hasNext()) { + Row r = it.next(); + subtaskById.put((Integer) r.getField(1), (Integer) r.getField(0)); + } + } + assertThat(subtaskById).as("every key should be observed once").hasSize(numKeys); + + // Core property: same Fluss bucket -> same subtask (subtask == bucketId % parallelism). + Map subtaskByBucket = new HashMap<>(); + for (Map.Entry e : subtaskById.entrySet()) { + int id = e.getKey(); + int subtask = e.getValue(); + int bucket = flussBucketOf(id, numBuckets); + assertThat(subtask) + .as("id=%d (bucket=%d) should be routed to bucketId %% parallelism", id, bucket) + .isEqualTo(Math.floorMod(bucket, parallelism)); + Integer prev = subtaskByBucket.putIfAbsent(bucket, subtask); + if (prev != null) { + assertThat(subtask) + .as("all keys of bucket %d must share one subtask", bucket) + .isEqualTo(prev); + } + } + + // Sanity: keys actually spread across multiple subtasks (not trivially co-located). + assertThat(new HashSet<>(subtaskById.values()).size()) + .as("keys should be spread across more than one subtask") + .isGreaterThan(1); + } + + /** Independently computes the Fluss bucket id for an int key, mirroring the client routing. */ + private static int flussBucketOf(int id, int numBuckets) { + RowType keyRowType = + RowType.of(new LogicalType[] {new IntType(false)}, new String[] {"id"}); + org.apache.fluss.types.RowType flussKeyType = FlinkConversions.toFlussRowType(keyRowType); + KeyEncoder encoder = + KeyEncoder.ofBucketKeyEncoder(flussKeyType, Collections.singletonList("id"), null); + byte[] bytes = encoder.encodeKey(new FlinkAsFlussRow().replace(GenericRowData.of(id))); + return BucketingFunction.of(null).bucketing(bytes, numBuckets); + } + + /** Flink partitioner that delegates channel selection to Fluss' custom lookup partitioner. */ + private static class DelegatingPartitioner implements Partitioner { + private final FlussLookupInputPartitioner partitioner; + + private DelegatingPartitioner(FlussLookupInputPartitioner partitioner) { + this.partitioner = partitioner; + } + + @Override + public int partition(RowData key, int numPartitions) { + return partitioner.partition(key, numPartitions); + } + } + + /** Builds the single-column int lookup key row from an id. */ + private static class IdKeySelector implements KeySelector { + @Override + public RowData getKey(Integer id) { + return GenericRowData.of(id); + } + } + + /** Tags each element with the subtask index that processed it. */ + private static class SubtaskTagger extends RichMapFunction { + @Override + public Row map(Integer id) { + return Row.of(getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(), id); + } + } + + private void registerNonPartitionedSrc() { + List testData = + Arrays.asList( + Row.of(1, "name1", 11), + Row.of(2, "name2", 2), + Row.of(3, "name33", 33), + Row.of(10, "name0", 44)); + RowTypeInfo typeInfo = + new RowTypeInfo( + new TypeInformation[] {Types.INT, Types.STRING, Types.INT}, + new String[] {"a", "b", "c"}); + Schema schema = + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.INT()) + .columnByExpression("proc", "PROCTIME()") + .build(); + DataStream srcDs = execEnv.fromCollection(testData).returns(typeInfo); + tEnv.createTemporaryView("src", tEnv.fromDataStream(srcDs, schema)); + } + + private void registerPartitionedSrc(String partition1, String partition2) { + List testData = + Arrays.asList( + Row.of(1, "name1", 11, partition1), + Row.of(2, "name2", 2, partition1), + Row.of(3, "name33", 33, partition2), + Row.of(10, "name0", 44, partition2)); + RowTypeInfo typeInfo = + new RowTypeInfo( + new TypeInformation[] {Types.INT, Types.STRING, Types.INT, Types.STRING}, + new String[] {"a", "b", "c", "p_date"}); + Schema schema = + Schema.newBuilder() + .column("a", DataTypes.INT()) + .column("b", DataTypes.STRING()) + .column("c", DataTypes.INT()) + .column("p_date", DataTypes.STRING()) + .columnByExpression("proc", "PROCTIME()") + .build(); + DataStream srcDs = execEnv.fromCollection(testData).returns(typeInfo); + tEnv.createTemporaryView("src", tEnv.fromDataStream(srcDs, schema)); + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSourceTest.java b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSourceTest.java new file mode 100644 index 0000000000..35e59ec26d --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/FlinkLookupShuffleTableSourceTest.java @@ -0,0 +1,193 @@ +/* + * 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.bucketing.BucketingFunction; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.config.TableConfig; +import org.apache.fluss.flink.FlinkConnectorOptions; +import org.apache.fluss.flink.row.FlinkAsFlussRow; +import org.apache.fluss.flink.utils.FlinkConnectorOptionsUtils; +import org.apache.fluss.flink.utils.FlinkConversions; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.encode.KeyEncoder; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.connector.source.DynamicTableSource; +import org.apache.flink.table.connector.source.LookupTableSource; +import org.apache.flink.table.connector.source.abilities.SupportsLookupCustomShuffle.InputDataPartitioner; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link FlinkLookupShuffleTableSource}, focusing on the invariant that a wrapped + * source (which has already advertised {@code SupportsLookupCustomShuffle}, so Flink suppresses its + * own hash shuffle) must always return a partitioner, and that a copy made between {@code + * getLookupRuntimeProvider()} and {@code getPartitioner()} keeps the stashed lookup normalizer. + */ +class FlinkLookupShuffleTableSourceTest { + + private static FlinkTableSource baseSource(int[] primaryKeys, int[] bucketKeys) { + RowType tableOutputType = + (RowType) + DataTypes.ROW( + DataTypes.FIELD("id", DataTypes.INT()), + DataTypes.FIELD("name", DataTypes.STRING()), + DataTypes.FIELD("address", DataTypes.STRING())) + .getLogicalType(); + Configuration flussConfig = new Configuration(); + flussConfig.setString(FlinkConnectorOptions.BOOTSTRAP_SERVERS.key(), "localhost:9092"); + FlinkConnectorOptionsUtils.StartupOptions startupOptions = + new FlinkConnectorOptionsUtils.StartupOptions(); + startupOptions.startupMode = FlinkConnectorOptions.ScanStartupMode.EARLIEST; + return new FlinkTableSource( + TablePath.of("test_db", "dim"), + flussConfig, + new TableConfig(new Configuration()), + tableOutputType, + primaryKeys, + bucketKeys, + new int[] {}, // partition key indexes + true, // streaming + startupOptions, + false, // lookup async (sync avoids any async runtime setup) + false, // insert if not exists + null, // cache + 1000L, // scan partition discovery interval + false, // is data lake enabled + null, // merge engine type + new HashMap<>(), + null); // lease context + } + + @Test + void testGetPartitionerPresentForFullPkLookup() { + // primary key == bucket key == id; lookup on the full primary key. + FlinkLookupShuffleTableSource source = + new FlinkLookupShuffleTableSource(baseSource(new int[] {0}, new int[] {0}), 3); + source.getLookupRuntimeProvider(new TestLookupContext(new int[][] {{0}})); + assertThat(source.getPartitioner()).isPresent(); + } + + @Test + void testGetPartitionerPresentForPrefixLookup() { + // primary key (id, name), bucket key = id; a lookup on id alone is a prefix lookup. + FlinkLookupShuffleTableSource source = + new FlinkLookupShuffleTableSource(baseSource(new int[] {0, 1}, new int[] {0}), 3); + source.getLookupRuntimeProvider(new TestLookupContext(new int[][] {{0}})); + assertThat(source.getPartitioner()).isPresent(); + } + + @Test + void testGetPartitionerFailsBeforeRuntimeProvider() { + // getPartitioner() before getLookupRuntimeProvider() violates the call order the source + // relies on: it must surface an internal-state error, not silently return empty (which the + // planner treats as "no shuffle"). + FlinkLookupShuffleTableSource source = + new FlinkLookupShuffleTableSource(baseSource(new int[] {0}, new int[] {0}), 3); + assertThatThrownBy(source::getPartitioner) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("getLookupRuntimeProvider"); + } + + @Test + void testCopyPreservesAbilityAndStashedNormalizer() { + int numBuckets = 3; + FlinkLookupShuffleTableSource source = + new FlinkLookupShuffleTableSource( + baseSource(new int[] {0}, new int[] {0}), numBuckets); + // stash the normalizer, then copy (as the planner may do) before getPartitioner(). + source.getLookupRuntimeProvider(new TestLookupContext(new int[][] {{0}})); + + DynamicTableSource copy = source.copy(); + assertThat(copy).isInstanceOf(FlinkLookupShuffleTableSource.class); + + Optional partitioner = + ((FlinkLookupShuffleTableSource) copy).getPartitioner(); + assertThat(partitioner) + .as("the copy must keep the stashed normalizer so the shuffle stays enabled") + .isPresent(); + + // The copy must also carry numBuckets over: route several keys and confirm they land on the + // bucket computed with numBuckets == 3 (a dropped or defaulted numBuckets would misroute). + InputDataPartitioner copiedPartitioner = partitioner.get(); + int numPartitions = 2; + for (int id = 0; id < 30; id++) { + assertThat(copiedPartitioner.partition(GenericRowData.of(id), numPartitions)) + .as("id=%d must route by numBuckets=%d", id, numBuckets) + .isEqualTo(Math.floorMod(flussBucketOf(id, numBuckets), numPartitions)); + } + } + + /** Independently computes the Fluss bucket id for an int key (matching the client routing). */ + private static int flussBucketOf(int id, int numBuckets) { + RowType keyRowType = RowType.of(new LogicalType[] {new IntType()}, new String[] {"id"}); + org.apache.fluss.types.RowType flussKeyType = FlinkConversions.toFlussRowType(keyRowType); + KeyEncoder encoder = + KeyEncoder.ofBucketKeyEncoder(flussKeyType, Collections.singletonList("id"), null); + byte[] bytes = encoder.encodeKey(new FlinkAsFlussRow().replace(GenericRowData.of(id))); + return BucketingFunction.of(null).bucketing(bytes, numBuckets); + } + + /** Minimal {@link LookupTableSource.LookupContext} exposing only the lookup key indexes. */ + private static class TestLookupContext implements LookupTableSource.LookupContext { + private final int[][] keys; + + private TestLookupContext(int[][] keys) { + this.keys = keys; + } + + @Override + public int[][] getKeys() { + return keys; + } + + @Override + public boolean preferCustomShuffle() { + return true; + } + + @Override + public TypeInformation createTypeInformation(DataType producedDataType) { + throw new UnsupportedOperationException(); + } + + @Override + public TypeInformation createTypeInformation(LogicalType producedLogicalType) { + throw new UnsupportedOperationException(); + } + + @Override + public DynamicTableSource.DataStructureConverter createDataStructureConverter( + DataType producedDataType) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitionerTest.java b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitionerTest.java new file mode 100644 index 0000000000..daa0bb048d --- /dev/null +++ b/fluss-flink/fluss-flink-2.2/src/test/java/org/apache/fluss/flink/source/lookup/FlussLookupInputPartitionerTest.java @@ -0,0 +1,364 @@ +/* + * 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.flink.utils.FlinkUtils; +import org.apache.fluss.metadata.DataLakeFormat; +import org.apache.fluss.row.encode.KeyEncoder; + +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.data.StringData; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; +import org.apache.flink.table.types.logical.VarCharType; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link FlussLookupInputPartitioner}, verifying that the lookup-join probe stream + * is partitioned consistently with the client-side Fluss bucket assignment. + */ +class FlussLookupInputPartitionerTest { + + private static final int GOLDEN_NUM_BUCKETS = 16; + + /** + * Frozen golden vectors: for a single {@code INT} bucket key and {@code numBuckets = 16}, these + * are the exact bucket ids produced by the client-compatible routing for each lake format that + * has its own bucketing implementation ({@code null}/default Fluss, {@code PAIMON}, {@code + * ICEBERG} and {@code HUDI}). {@code LANCE} shares the default Fluss implementation, so it is + * not covered separately. Baking these values in as literals makes them a regression guard that + * is independent of re-deriving the value with the same encoder/bucketing code (addressing the + * "same-implementation oracle" concern). + * + *

This test pins only the exact bucket id ({@code numPartitions == numBuckets}); the {@code + * bucket % numPartitions} channel selection is already covered by {@link + * #testPartitionEqualsFlussBucketModNumPartitions()}. + */ + @Test + void testGoldenBucketVectors() { + RowType keyRowType = RowType.of(new LogicalType[] {new IntType()}, new String[] {"id"}); + List bucketKeyNames = Collections.singletonList("id"); + + Map flussGolden = new LinkedHashMap<>(); + flussGolden.put(1, 12); + flussGolden.put(2, 10); + flussGolden.put(42, 10); + flussGolden.put(100, 8); + flussGolden.put(12345, 5); + assertGoldenRouting(keyRowType, bucketKeyNames, null, flussGolden); + + Map paimonGolden = new LinkedHashMap<>(); + paimonGolden.put(1, 14); + paimonGolden.put(2, 0); + paimonGolden.put(42, 5); + paimonGolden.put(100, 5); + paimonGolden.put(12345, 14); + assertGoldenRouting(keyRowType, bucketKeyNames, DataLakeFormat.PAIMON, paimonGolden); + + Map icebergGolden = new LinkedHashMap<>(); + icebergGolden.put(1, 4); + icebergGolden.put(2, 4); + icebergGolden.put(42, 14); + icebergGolden.put(100, 0); + icebergGolden.put(12345, 1); + assertGoldenRouting(keyRowType, bucketKeyNames, DataLakeFormat.ICEBERG, icebergGolden); + + Map hudiGolden = new LinkedHashMap<>(); + hudiGolden.put(1, 0); + hudiGolden.put(2, 1); + hudiGolden.put(42, 13); + hudiGolden.put(100, 0); + hudiGolden.put(12345, 2); + assertGoldenRouting(keyRowType, bucketKeyNames, DataLakeFormat.HUDI, hudiGolden); + } + + private static void assertGoldenRouting( + RowType keyRowType, + List bucketKeyNames, + DataLakeFormat lakeFormat, + Map goldenBucketById) { + LookupNormalizer normalizer = + LookupNormalizer.createPrimaryKeyLookupNormalizer(new int[] {0}, keyRowType); + FlussLookupInputPartitioner partitioner = + new FlussLookupInputPartitioner( + normalizer, keyRowType, bucketKeyNames, lakeFormat, GOLDEN_NUM_BUCKETS); + for (Map.Entry entry : goldenBucketById.entrySet()) { + int id = entry.getKey(); + int goldenBucket = entry.getValue(); + // numPartitions == numBuckets pins the exact bucket id. + assertThat(partitioner.partition(GenericRowData.of(id), GOLDEN_NUM_BUCKETS)) + .as("id=%d fmt=%s", id, lakeFormat) + .isEqualTo(goldenBucket); + } + } + + /** Independently computes the Fluss bucket id for the given key row. */ + private static int flussBucketOf( + RowType keyRowType, List bucketKeyNames, RowData key, int numBuckets) { + org.apache.fluss.types.RowType flussKeyType = FlinkConversions.toFlussRowType(keyRowType); + KeyEncoder encoder = KeyEncoder.ofBucketKeyEncoder(flussKeyType, bucketKeyNames, null); + byte[] bytes = encoder.encodeKey(new FlinkAsFlussRow().replace(key)); + return BucketingFunction.of(null).bucketing(bytes, numBuckets); + } + + private static FlussLookupInputPartitioner partitionerFor( + RowType keyRowType, List bucketKeyNames, int numBuckets) { + // primary key == bucket key -> identity normalizer (no reordering) + int[] pkIndexes = new int[bucketKeyNames.size()]; + for (int i = 0; i < pkIndexes.length; i++) { + pkIndexes[i] = i; + } + LookupNormalizer normalizer = + LookupNormalizer.createPrimaryKeyLookupNormalizer(pkIndexes, keyRowType); + return new FlussLookupInputPartitioner( + normalizer, keyRowType, bucketKeyNames, /* lakeFormat */ null, numBuckets); + } + + @Test + void testPartitionEqualsFlussBucketModNumPartitions() { + RowType keyRowType = RowType.of(new LogicalType[] {new IntType()}, new String[] {"id"}); + List bucketKeyNames = Collections.singletonList("id"); + int numBuckets = 8; + FlussLookupInputPartitioner partitioner = + partitionerFor(keyRowType, bucketKeyNames, numBuckets); + + for (int numPartitions : new int[] {1, 2, 3, 5, 8}) { + for (int id = 0; id < 100; id++) { + RowData key = GenericRowData.of(id); + int actual = partitioner.partition(key, numPartitions); + int expected = + Math.floorMod( + flussBucketOf(keyRowType, bucketKeyNames, key, numBuckets), + numPartitions); + assertThat(actual) + .as("id=%d, numPartitions=%d", id, numPartitions) + .isEqualTo(expected) + .isBetween(0, numPartitions - 1); + } + } + } + + @Test + void testDeterministicAndStable() { + RowType keyRowType = RowType.of(new LogicalType[] {new IntType()}, new String[] {"id"}); + FlussLookupInputPartitioner partitioner = + partitionerFor(keyRowType, Collections.singletonList("id"), 8); + + assertThat(partitioner.isDeterministic()).isTrue(); + // same key -> same partition across repeated calls + for (int id = 0; id < 50; id++) { + int first = partitioner.partition(GenericRowData.of(id), 4); + for (int i = 0; i < 5; i++) { + assertThat(partitioner.partition(GenericRowData.of(id), 4)).isEqualTo(first); + } + } + } + + @Test + void testRowsInSameBucketGoToSamePartition() { + RowType keyRowType = RowType.of(new LogicalType[] {new IntType()}, new String[] {"id"}); + List bucketKeyNames = Collections.singletonList("id"); + int numBuckets = 4; + int numPartitions = 2; + FlussLookupInputPartitioner partitioner = + partitionerFor(keyRowType, bucketKeyNames, numBuckets); + + // group ids by their fluss bucket, then assert every id in the same bucket maps to the + // same partition (co-partitioning by bucket key). + Map bucketToPartition = new HashMap<>(); + for (int id = 0; id < 200; id++) { + RowData key = GenericRowData.of(id); + int bucket = flussBucketOf(keyRowType, bucketKeyNames, key, numBuckets); + int partition = partitioner.partition(key, numPartitions); + Integer previous = bucketToPartition.putIfAbsent(bucket, partition); + if (previous != null) { + assertThat(partition) + .as("all ids in bucket %d must share a partition", bucket) + .isEqualTo(previous); + } + } + } + + @Test + void testCompositeBucketKey() { + RowType keyRowType = + RowType.of( + new LogicalType[] {new IntType(), new VarCharType(VarCharType.MAX_LENGTH)}, + new String[] {"id", "region"}); + List bucketKeyNames = Arrays.asList("id", "region"); + int numBuckets = 6; + FlussLookupInputPartitioner partitioner = + partitionerFor(keyRowType, bucketKeyNames, numBuckets); + + int numPartitions = 3; + for (int id = 0; id < 50; id++) { + RowData key = GenericRowData.of(id, StringData.fromString("region-" + (id % 4))); + int actual = partitioner.partition(key, numPartitions); + org.apache.fluss.types.RowType flussKeyType = + FlinkConversions.toFlussRowType(keyRowType); + KeyEncoder encoder = KeyEncoder.ofBucketKeyEncoder(flussKeyType, bucketKeyNames, null); + byte[] bytes = encoder.encodeKey(new FlinkAsFlussRow().replace(key)); + int expected = + Math.floorMod( + BucketingFunction.of(null).bucketing(bytes, numBuckets), numPartitions); + assertThat(actual).as("id=%d", id).isEqualTo(expected).isBetween(0, numPartitions - 1); + } + } + + /** + * A prefix lookup (bucket key is a strict prefix of the primary key) must be shuffled too: the + * normalized key row is the bucket keys (+ partition keys), and routing must match the Fluss + * bucket assignment for that prefix. + */ + @Test + void testPrefixLookupPartitionMatchesBucket() { + // schema: region (bucket key), id, name; primary key (region, id), bucket key (region) + RowType tableSchema = + RowType.of( + new LogicalType[] { + new VarCharType(VarCharType.MAX_LENGTH), + new IntType(), + new VarCharType(VarCharType.MAX_LENGTH) + }, + new String[] {"region", "id", "name"}); + int[] pkIndexes = {0, 1}; + int[] bucketKeyIndexes = {0}; + // prefix lookup: only the bucket key (region) is used as the lookup key + int[][] lookupKeyIndexes = {{0}}; + int numBuckets = 5; + + FlussLookupInputPartitioner partitioner = + buildPartitioner( + tableSchema, lookupKeyIndexes, pkIndexes, bucketKeyIndexes, numBuckets); + + RowType keyRowType = + RowType.of( + new LogicalType[] {new VarCharType(VarCharType.MAX_LENGTH)}, + new String[] {"region"}); + List bucketKeyNames = Collections.singletonList("region"); + int numPartitions = 3; + for (int i = 0; i < 40; i++) { + RowData joinKeys = GenericRowData.of(StringData.fromString("region-" + i)); + int actual = partitioner.partition(joinKeys, numPartitions); + int expected = + Math.floorMod( + flussBucketOf(keyRowType, bucketKeyNames, joinKeys, numBuckets), + numPartitions); + assertThat(actual) + .as("region-%d", i) + .isEqualTo(expected) + .isBetween(0, numPartitions - 1); + } + } + + /** + * Flink 2.2 delivers lookup keys in ascending field-index order, so a key reorder is driven by + * the primary key being declared in a different order than the table fields (here PK + * {@code (region, id)} over fields {@code (id, region)}). The normalizer must reorder the probe + * key into Fluss primary-key order before bucketing, and routing must be computed on that + * reordered (normalized) key rather than on the raw ascending join-key order. + */ + @Test + void testKeyReorderMatchesNormalizedBucket() { + // schema field order: id(0), region(1), name(2) + RowType tableSchema = + RowType.of( + new LogicalType[] { + new IntType(), + new VarCharType(VarCharType.MAX_LENGTH), + new VarCharType(VarCharType.MAX_LENGTH) + }, + new String[] {"id", "region", "name"}); + // primary key & bucket key declared as (region, id) -> differs from schema field order + int[] pkIndexes = {1, 0}; + int[] bucketKeyIndexes = {1, 0}; + // Flink sorts lookup keys ascending by field index, so the probe side delivers (id, region) + int[][] lookupKeyIndexes = {{0}, {1}}; + int numBuckets = 7; + + FlussLookupInputPartitioner partitioner = + buildPartitioner( + tableSchema, lookupKeyIndexes, pkIndexes, bucketKeyIndexes, numBuckets); + + // the normalized key row type is (region, id) in Fluss primary-key order + RowType normalizedKeyRowType = + RowType.of( + new LogicalType[] {new VarCharType(VarCharType.MAX_LENGTH), new IntType()}, + new String[] {"region", "id"}); + List bucketKeyNames = Arrays.asList("region", "id"); + int numPartitions = 4; + for (int id = 0; id < 40; id++) { + StringData region = StringData.fromString("region-" + (id % 5)); + // probe delivers the join key row in ascending field order (id, region) + RowData joinKeys = GenericRowData.of(id, region); + int actual = partitioner.partition(joinKeys, numPartitions); + // expected: bucket computed on the normalized (region, id) row + RowData normalizedKey = GenericRowData.of(region, id); + int expected = + Math.floorMod( + flussBucketOf( + normalizedKeyRowType, + bucketKeyNames, + normalizedKey, + numBuckets), + numPartitions); + assertThat(actual).as("id=%d", id).isEqualTo(expected).isBetween(0, numPartitions - 1); + } + } + + /** Builds a partitioner exactly the way {@code FlinkLookupShuffleTableSource} does. */ + private static FlussLookupInputPartitioner buildPartitioner( + RowType tableSchema, + int[][] lookupKeyIndexes, + int[] pkIndexes, + int[] bucketKeyIndexes, + int numBuckets) { + LookupNormalizer normalizer = + LookupNormalizer.validateAndCreateLookupNormalizer( + lookupKeyIndexes, + pkIndexes, + bucketKeyIndexes, + new int[0], + tableSchema, + null); + RowType keyRowType = + FlinkUtils.projectRowType(tableSchema, normalizer.getLookupKeyIndexes()); + List allNames = tableSchema.getFieldNames(); + List bucketKeyNames = new ArrayList<>(bucketKeyIndexes.length); + for (int idx : bucketKeyIndexes) { + bucketKeyNames.add(allNames.get(idx)); + } + return new FlussLookupInputPartitioner( + normalizer, keyRowType, bucketKeyNames, /* lakeFormat */ null, numBuckets); + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java new file mode 100644 index 0000000000..be65680a1a --- /dev/null +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/adapter/LookupShuffleSourceAdapter.java @@ -0,0 +1,51 @@ +/* + * 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.FlinkTableSource; + +import org.apache.flink.configuration.ReadableConfig; + +/** + * Adapter that optionally wraps a {@link FlinkTableSource} with a version-specific ability. + * + *

This is the default (Flink 1.x) no-op implementation: {@code SupportsLookupCustomShuffle} does + * not exist before Flink 2.0, so the source is returned unchanged. The {@code fluss-flink-2.2} + * module ships a class with the same fully-qualified name that wraps the source with the custom + * lookup shuffle ability; that class shadows this one at shade/package time. + * + *

TODO: remove this class when no longer supporting Flink 1.x. + */ +public class LookupShuffleSourceAdapter { + + private LookupShuffleSourceAdapter() {} + + /** + * Optionally wraps the source with a version-specific custom lookup shuffle ability. + * + * @param source source to wrap + * @param tableOptions resolved table options + * @param lookupShuffleEligible whether the table has the primary-key and bucket-key metadata + * required by the custom lookup shuffle + * @return the source unchanged for Flink 1.x + */ + public static FlinkTableSource maybeWithCustomShuffle( + FlinkTableSource source, ReadableConfig tableOptions, boolean lookupShuffleEligible) { + return source; + } +} diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java index ab64959abd..746548914a 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/catalog/FlinkTableFactory.java @@ -21,6 +21,7 @@ import org.apache.fluss.config.Configuration; import org.apache.fluss.config.TableConfig; import org.apache.fluss.flink.FlinkConnectorOptions; +import org.apache.fluss.flink.adapter.LookupShuffleSourceAdapter; import org.apache.fluss.flink.lake.LakeFlinkCatalog; import org.apache.fluss.flink.lake.LakeTableFactory; import org.apache.fluss.flink.sink.FlinkTableSink; @@ -150,26 +151,29 @@ public DynamicTableSource createDynamicTableSource(Context context) { tableOptions.get(FlinkConnectorOptions.SCAN_SPLIT_ASSIGNMENT_BATCH_SIZE); LeaseContext leaseContext = LeaseContext.fromConf(tableOptions); - return new FlinkTableSource( - toFlussTablePath(context.getObjectIdentifier()), - toFlussClientConfig( - context.getCatalogTable().getOptions(), context.getConfiguration()), - toFlussTableConfig(tableOptions), - tableOutputType, - primaryKeyIndexes, - bucketKeyIndexes, - partitionKeyIndexes, - isStreamingMode, - startupOptions, - tableOptions.get(FlinkConnectorOptions.LOOKUP_ASYNC), - tableOptions.get(FlinkConnectorOptions.LOOKUP_INSERT_IF_NOT_EXISTS), - cache, - partitionDiscoveryIntervalMs, - splitAssignmentBatchSize, - tableOptions.get(toFlinkOption(ConfigOptions.TABLE_DATALAKE_ENABLED)), - tableOptions.get(toFlinkOption(ConfigOptions.TABLE_MERGE_ENGINE)), - context.getCatalogTable().getOptions(), - leaseContext); + return LookupShuffleSourceAdapter.maybeWithCustomShuffle( + new FlinkTableSource( + toFlussTablePath(context.getObjectIdentifier()), + toFlussClientConfig( + context.getCatalogTable().getOptions(), context.getConfiguration()), + toFlussTableConfig(tableOptions), + tableOutputType, + primaryKeyIndexes, + bucketKeyIndexes, + partitionKeyIndexes, + isStreamingMode, + startupOptions, + tableOptions.get(FlinkConnectorOptions.LOOKUP_ASYNC), + tableOptions.get(FlinkConnectorOptions.LOOKUP_INSERT_IF_NOT_EXISTS), + cache, + partitionDiscoveryIntervalMs, + splitAssignmentBatchSize, + tableOptions.get(toFlinkOption(ConfigOptions.TABLE_DATALAKE_ENABLED)), + tableOptions.get(toFlinkOption(ConfigOptions.TABLE_MERGE_ENGINE)), + context.getCatalogTable().getOptions(), + leaseContext), + tableOptions, + primaryKeyIndexes.length > 0 && bucketKeyIndexes.length > 0); } @Override diff --git a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java index 24f3d94d67..16ea7f7b46 100644 --- a/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java +++ b/fluss-flink/fluss-flink-common/src/main/java/org/apache/fluss/flink/source/FlinkTableSource.java @@ -178,6 +178,9 @@ public class FlinkTableSource /** Watermark strategy that is pushed down by the Flink optimizer. */ @Nullable private WatermarkStrategy watermarkStrategy; + /** Stashed at getLookupRuntimeProvider time, used by Flink-2.x custom lookup shuffle. */ + @Nullable private LookupNormalizer lastLookupNormalizer; + public FlinkTableSource( TablePath tablePath, Configuration flussConfig, @@ -269,6 +272,44 @@ public FlinkTableSource( this.availableStatsColumns = computeAvailableStatsColumns(flussRowType); } + /** + * Copy constructor: the single source of truth for cloning this source's state. Used by {@link + * #copy()} and by version-specific subclasses (e.g. the Flink 2.x custom-lookup-shuffle + * variant) to wrap an already-built source while preserving all state. + */ + protected FlinkTableSource(FlinkTableSource other) { + this.tablePath = other.tablePath; + this.flussConfig = other.flussConfig; + this.tableOutputType = other.tableOutputType; + this.primaryKeyIndexes = other.primaryKeyIndexes; + this.bucketKeyIndexes = other.bucketKeyIndexes; + this.partitionKeyIndexes = other.partitionKeyIndexes; + this.streaming = other.streaming; + this.startupOptions = other.startupOptions; + this.lookupAsync = other.lookupAsync; + this.insertIfNotExists = other.insertIfNotExists; + this.cache = other.cache; + this.scanPartitionDiscoveryIntervalMs = other.scanPartitionDiscoveryIntervalMs; + this.splitPerAssignmentBatchSize = other.splitPerAssignmentBatchSize; + this.isDataLakeEnabled = other.isDataLakeEnabled; + this.leaseContext = other.leaseContext; + this.mergeEngineType = other.mergeEngineType; + this.tableConfig = other.tableConfig; + this.availableStatsColumns = other.availableStatsColumns; + this.tableOptions = other.tableOptions; + // mutable / push-down state + this.producedDataType = other.producedDataType; + this.projectedFields = other.projectedFields; + this.singleRowFilter = other.singleRowFilter; + this.modificationScanType = other.modificationScanType; + // Note: selectRowCount/limit are intentionally not carried over (reset on copy). + this.partitionFilters = other.partitionFilters; + this.lakeSource = other.lakeSource; + this.logRecordBatchFilter = other.logRecordBatchFilter; + this.watermarkStrategy = other.watermarkStrategy; + this.lastLookupNormalizer = other.lastLookupNormalizer; + } + @Override public ChangelogMode getChangelogMode() { if (!streaming) { @@ -474,6 +515,8 @@ public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext context) { partitionKeyIndexes, tableOutputType, projectedFields); + // Stash for SupportsLookupCustomShuffle (Flink 2.x). Harmless for Flink 1.x. + this.lastLookupNormalizer = lookupNormalizer; if (lookupAsync) { AsyncLookupFunction asyncLookupFunction = new FlinkAsyncLookupFunction( @@ -507,36 +550,36 @@ public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext context) { @Override public DynamicTableSource copy() { - FlinkTableSource source = - new FlinkTableSource( - tablePath, - flussConfig, - tableConfig, - tableOutputType, - primaryKeyIndexes, - bucketKeyIndexes, - partitionKeyIndexes, - streaming, - startupOptions, - lookupAsync, - insertIfNotExists, - cache, - scanPartitionDiscoveryIntervalMs, - splitPerAssignmentBatchSize, - isDataLakeEnabled, - mergeEngineType, - tableOptions, - leaseContext); - source.producedDataType = producedDataType; - source.projectedFields = projectedFields; - source.singleRowFilter = singleRowFilter; - source.modificationScanType = modificationScanType; - source.partitionFilters = partitionFilters; - source.lakeSource = lakeSource; - source.logRecordBatchFilter = logRecordBatchFilter; - source.watermarkStrategy = watermarkStrategy; - // Note: availableStatsColumns is already computed in the constructor - return source; + // Single source of truth for state copying: the copy constructor. It also carries over the + // stashed lookup normalizer, so a copy made between getLookupRuntimeProvider() and + // getPartitioner() (Flink 2.x custom lookup shuffle) does not silently disable the shuffle. + return new FlinkTableSource(this); + } + + // ---- accessors for version-specific (Flink 2.x) custom lookup shuffle ---- + + /** Returns the indexes of the bucket-key columns within the table output row. */ + protected int[] bucketKeyIndexes() { + return bucketKeyIndexes; + } + + /** Returns the Flink logical row type produced by this source. */ + protected org.apache.flink.table.types.logical.RowType tableOutputType() { + return tableOutputType; + } + + /** Returns the table config, e.g. to resolve the data lake format for bucketing. */ + protected TableConfig tableConfigInternal() { + return tableConfig; + } + + /** + * Returns the lookup normalizer stashed by the last {@link #getLookupRuntimeProvider}, or + * {@code null} if the lookup runtime provider has not been requested yet. + */ + @Nullable + protected LookupNormalizer lastLookupNormalizer() { + return lastLookupNormalizer; } @Override