offsets = new HashMap<>();
diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactory.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactory.java
index c5a25cf..56fa1d7 100644
--- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactory.java
+++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactory.java
@@ -99,7 +99,7 @@
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateDeliveryGuarantee;
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateTableSinkOptions;
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateTableSourceOptions;
-import static com.pinterest.flink.streaming.connectors.psc.table.PscTableCommonUtils.shouldApplyRescale;
+import static com.pinterest.flink.streaming.connectors.psc.table.PscTableCommonUtils.getEffectiveSourceParallelism;
/**
* Factory for creating configured instances of {@link PscDynamicSource} and {@link
@@ -220,22 +220,24 @@ public DynamicTableSource createDynamicTableSource(Context context) {
final String keyPrefix = tableOptions.getOptional(KEY_FIELDS_PREFIX).orElse(null);
+ final boolean shouldRescale =tableOptions.get(SCAN_ENABLE_RESCALE);
// Read scan.parallelism configuration
final Integer scanParallelism = tableOptions.getOptional(SCAN_PARALLELISM).orElse(null);
- // Log scan parallelism configuration if set
- if (scanParallelism != null) {
- LOG.info("scan.parallelism configured: {} (will override partition-based parallelism)",
- scanParallelism);
- }
-
- // Determine if rescale should be applied based on scan.parallelism (if set) or global default
- final boolean shouldRescale = shouldApplyRescale(
- tableOptions,
+ final Integer effectiveParallelism = shouldRescale ?
+ getEffectiveSourceParallelism(
context.getConfiguration(),
getSourceTopicUris(tableOptions),
properties,
- scanParallelism);
+ scanParallelism)
+ : -1;
+
+
+ // Log scan parallelism configuration if set
+ if (shouldRescale && effectiveParallelism > 0 ) {
+ LOG.info("Psc will use parallelism={} for source operator",
+ effectiveParallelism);
+ }
// Get rate limit configuration
final Double rateLimitRecordsPerSecond = tableOptions.getOptional(SCAN_RATE_LIMIT).orElse(null);
@@ -266,7 +268,7 @@ public DynamicTableSource createDynamicTableSource(Context context) {
tableOptions.getOptional(SOURCE_UID_PREFIX).orElse(null),
shouldRescale,
rateLimitRecordsPerSecond,
- scanParallelism);
+ effectiveParallelism);
}
@Override
diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtils.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtils.java
index 558d379..6f457c2 100644
--- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtils.java
+++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtils.java
@@ -24,6 +24,7 @@
import com.pinterest.psc.metadata.TopicUriMetadata;
import com.pinterest.psc.metadata.client.PscMetadataClient;
import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.configuration.CoreOptions;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.table.api.config.ExecutionConfigOptions;
import org.slf4j.Logger;
@@ -91,77 +92,60 @@ static synchronized void resetProvider() {
}
/**
- * Determines whether rescale() should be applied based on:
- * 1. scan.enable-rescale flag must be true
- * 2. Effective parallelism (scan.parallelism or global default) > partition count
- *
- * When scan.parallelism is set, it takes precedence over global default for
- * rescale decision logic. This ensures rescale is only applied when the user's
- * intended source parallelism exceeds partition count.
+ * Determines the effective parallelism for the source by walking the following
+ * fallback chain and returning the first source that yields a positive value:
+ *
+ * - {@code scan.parallelism} (table-level override) — used if non-null and not -1
+ * - {@code table.exec.resource.default-parallelism} — used if set and not -1
+ * - Kafka partition count for the source topics (queried via PSC metadata client)
+ *
*
- * @param tableOptions User's table configuration options
- * @param globalConfig Global Flink configuration
- * @param topicUris List of topic URIs to query for partition counts
- * @param pscProperties PSC properties for metadata client connection
+ * If none of the above produces a positive value (e.g. the partition-count query
+ * fails), this method returns {@code -1} to signal "unknown".
+ *
+ * @param globalConfig Global Flink configuration (read for table.exec.resource.default-parallelism)
+ * @param topicUris List of topic URIs used to query partition count
+ * @param pscProperties PSC properties for metadata client connection
* @param scanParallelism Optional explicit scan.parallelism configuration
- * @return true if rescale should be applied, false otherwise
+ * @return Effective parallelism, or -1 if it cannot be determined
*/
- public static boolean shouldApplyRescale(
- ReadableConfig tableOptions,
+ public static int getEffectiveSourceParallelism(
ReadableConfig globalConfig,
List topicUris,
Properties pscProperties,
@Nullable Integer scanParallelism) {
-
- // First check if rescale is enabled by user
- if (!tableOptions.get(SCAN_ENABLE_RESCALE)) {
- return false;
- }
- // Determine effective parallelism: scan.parallelism takes precedence over global default
- Integer effectiveParallelism;
- String parallelismSource;
-
+ // 1) scan.parallelism
if (scanParallelism != null && scanParallelism > 0) {
- effectiveParallelism = scanParallelism;
- parallelismSource = PscConnectorOptions.SCAN_PARALLELISM.key();
- } else {
- effectiveParallelism = globalConfig.get(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM);
- parallelismSource = ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM.key();
+ LOG.info("Effective source parallelism = {} (source: {})",
+ scanParallelism, PscConnectorOptions.SCAN_PARALLELISM.key());
+ return scanParallelism;
}
-
- // If parallelism is not set or invalid, cannot determine
- if (effectiveParallelism == null || effectiveParallelism <= 0) {
- LOG.info("scan.enable-rescale is true, but effective parallelism from {} is {} (not set or invalid). " +
- "Rescale will not be applied.",
- parallelismSource, effectiveParallelism);
- return false;
+
+ // 2) table.exec.resource.default-parallelism
+ Integer tableExecParallelism =
+ globalConfig.get(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM);
+ if (tableExecParallelism != null && tableExecParallelism > 0) {
+ LOG.info("Effective source parallelism = {} (source: {})",
+ tableExecParallelism,
+ ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM.key());
+ return tableExecParallelism;
}
- // Query partition count using the configured provider (mockable for testing)
+ // 3) Kafka partition count
int partitionCount = partitionCountProvider.getPartitionCount(topicUris, pscProperties);
-
- // If partition count couldn't be determined, don't apply rescale (fail-safe)
- if (partitionCount <= 0) {
- LOG.warn("scan.enable-rescale is true, but partition count could not be determined. " +
- "Rescale will not be applied.");
- return false;
+ if (partitionCount > 0) {
+ LOG.info("Effective source parallelism = {} (source: kafka partition count)",
+ partitionCount);
+ return partitionCount;
}
- // Apply rescale only if effective parallelism exceeds partition count
- boolean shouldRescale = effectiveParallelism > partitionCount;
-
- if (shouldRescale) {
- LOG.info("Applying rescale(): {} ({}) > partition count ({}). " +
- "Data will be redistributed to fully utilize downstream operators.",
- parallelismSource, effectiveParallelism, partitionCount);
- } else {
- LOG.info("Skipping rescale(): {} ({}) <= partition count ({}). " +
- "No shuffle needed as source parallelism matches or is less than partition count.",
- parallelismSource, effectiveParallelism, partitionCount);
- }
-
- return shouldRescale;
+ // 4) Unknown
+ LOG.warn("Could not determine effective source parallelism: {} is unset/-1, {} is unset/-1, " +
+ "and partition count could not be retrieved. Returning -1.",
+ PscConnectorOptions.SCAN_PARALLELISM.key(),
+ ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM.key());
+ return -1;
}
/**
diff --git a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactory.java b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactory.java
index 0805d0a..ccae228 100644
--- a/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactory.java
+++ b/psc-flink/src/main/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactory.java
@@ -85,7 +85,7 @@
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateScanBoundedMode;
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateConsumerClientOptions;
import static com.pinterest.flink.streaming.connectors.psc.table.PscConnectorOptionsUtil.validateProducerClientOptions;
-import static com.pinterest.flink.streaming.connectors.psc.table.PscTableCommonUtils.shouldApplyRescale;
+import static com.pinterest.flink.streaming.connectors.psc.table.PscTableCommonUtils.getEffectiveSourceParallelism;
/** Upsert-Psc factory. */
public class UpsertPscDynamicTableFactory
@@ -161,22 +161,24 @@ public DynamicTableSource createDynamicTableSource(Context context) {
final PscConnectorOptionsUtil.BoundedOptions boundedOptions = getBoundedOptions(tableOptions);
+ final boolean shouldRescale =tableOptions.get(SCAN_ENABLE_RESCALE);
// Read scan.parallelism configuration
final Integer scanParallelism = tableOptions.getOptional(SCAN_PARALLELISM).orElse(null);
- // Log scan parallelism configuration if set
- if (scanParallelism != null) {
- LOG.info("scan.parallelism configured: {} (will override partition-based parallelism)",
- scanParallelism);
- }
-
- // Determine if rescale should be applied based on scan.parallelism (if set) or global default
- final boolean shouldRescale = shouldApplyRescale(
- tableOptions,
+ final Integer effectiveParallelism = shouldRescale ?
+ getEffectiveSourceParallelism(
context.getConfiguration(),
getSourceTopicUris(tableOptions),
properties,
- scanParallelism);
+ scanParallelism)
+ : -1;
+
+
+ // Log scan parallelism configuration if set
+ if (shouldRescale && effectiveParallelism > 0 ) {
+ LOG.info("Psc will use parallelism={} for source operator",
+ effectiveParallelism);
+ }
// Get rate limit configuration
final Double rateLimitRecordsPerSecond = tableOptions.getOptional(SCAN_RATE_LIMIT).orElse(null);
@@ -208,7 +210,7 @@ public DynamicTableSource createDynamicTableSource(Context context) {
tableOptions.getOptional(SOURCE_UID_PREFIX).orElse(null),
shouldRescale,
rateLimitRecordsPerSecond,
- scanParallelism);
+ effectiveParallelism);
}
@Override
diff --git a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java
index 917cb38..e747e49 100644
--- a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java
+++ b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscDynamicTableFactoryTest.java
@@ -52,8 +52,12 @@
import org.apache.flink.streaming.api.transformations.PartitionTransformation;
import org.apache.flink.streaming.api.transformations.SourceTransformation;
import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.ValidationException;
+import org.apache.flink.table.api.config.ExecutionConfigOptions;
+import org.apache.flink.table.catalog.CatalogTable;
import org.apache.flink.table.catalog.Column;
+import org.apache.flink.table.catalog.ResolvedCatalogTable;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.catalog.UniqueConstraint;
import org.apache.flink.table.catalog.WatermarkSpec;
@@ -1438,6 +1442,34 @@ private void addScanParallelismConfig(Map options, int paralleli
options.put(PscConnectorOptions.SCAN_PARALLELISM.key(), String.valueOf(parallelism));
}
+ /**
+ * Builds a PSC dynamic table source using {@link FactoryUtil#createDynamicTableSource} with
+ * a caller-supplied global {@link Configuration}. {@link FactoryMocks#createTableSource}
+ * hard-codes {@code new Configuration()}, so this overload is needed when a test must
+ * exercise factory behavior that reads from the global config (e.g. when the factory calls
+ * {@link PscTableCommonUtils#getEffectiveSourceParallelism} and we need to set
+ * {@code table.exec.resource.default-parallelism}).
+ */
+ private DynamicTableSource createTableSourceWithGlobalConfig(
+ ResolvedSchema schema,
+ Map tableOptions,
+ Configuration globalConfig) {
+ return FactoryUtil.createDynamicTableSource(
+ null,
+ FactoryMocks.IDENTIFIER,
+ new ResolvedCatalogTable(
+ CatalogTable.of(
+ Schema.newBuilder().fromResolvedSchema(schema).build(),
+ "mock source",
+ Collections.emptyList(),
+ tableOptions),
+ schema),
+ Collections.emptyMap(),
+ globalConfig,
+ Thread.currentThread().getContextClassLoader(),
+ false);
+ }
+
@Test
public void testOperatorChainingWithRateLimitOnly() {
// When only rate limiting is enabled (rescale disabled), verify rate limiter is applied
@@ -1469,36 +1501,50 @@ public void testOperatorChainingWithRateLimitOnly() {
@Test
public void testOperatorChainingWithRescaleAndRateLimit() {
- // When both rescale and rate limiting are enabled, verify configuration
+ // Verifies the new operator chain when both rescale and rate limiting are enabled:
+ // Source -> PscRateLimit -> rescale (terminal PartitionTransformation)
+ // scan.parallelism (tier 1) determines the source/rate-limit parallelism, capped
+ // by the env parallelism in produceTransformationFromSource.
final Map modifiedOptions =
getModifiedOptions(
getBasicSourceOptions(),
options -> {
addRescaleConfig(options, true);
addRateLimitConfig(options, 5000.0);
- addScanParallelismConfig(options, 100); // Explicitly set higher than partition count
+ addScanParallelismConfig(options, 100);
});
-
+
final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
-
+
final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
-
- // Verify configuration is correctly set
- assertThat(pscSource.rateLimitRecordsPerSecond).isNotNull();
+ assertThat(pscSource.enableRescale).isTrue();
assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(5000.0);
assertThat(pscSource.scanParallelism).isEqualTo(100);
-
- // Get transformation using helper method (reuses same source)
- final Transformation transformation = produceTransformationFromSource(pscSource, 10);
-
- // The final transformation should be the rate limiter
- assertThat(transformation).isNotNull();
- assertThat(transformation.getName()).contains("PscRateLimit");
-
- // In test environment without real partitions, runtime rescale decision may differ
- // The key verification is that configuration is properly passed through
- // In production, with actual partition count > scan.parallelism, rescale would be applied
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // Terminal is the rescale PartitionTransformation.
+ assertThat(terminal).isNotNull();
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ // Terminal's input is the rate-limit operator.
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ // The rate-limit operator wraps the Kafka source.
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Source and rate-limit are both pinned to min(scanParallelism, env).
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism);
}
@Test
@@ -1534,7 +1580,10 @@ public void testOperatorChainingWithRescaleWithoutRateLimit() {
@Test
public void testRateLimiterParallelismConfiguration() {
- // Verify that scan.parallelism and rate limiting configurations are properly stored
+ // Verifies that scan.parallelism and rate-limit options are wired into the
+ // PscDynamicSource and that the rate-limit operator is inserted between the
+ // source and the rescale, with its parallelism matching the (capped) source
+ // parallelism.
final Map modifiedOptions =
getModifiedOptions(
getBasicSourceOptions(),
@@ -1543,23 +1592,31 @@ public void testRateLimiterParallelismConfiguration() {
addScanParallelismConfig(options, 200);
addRateLimitConfig(options, 10000.0);
});
-
+
final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
-
- // Verify all configuration values are correctly stored in the source
+
assertThat(pscSource.scanParallelism).isEqualTo(200);
assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(10000.0);
-
- // Get transformation using helper method (reuses same source)
- final Transformation transformation = produceTransformationFromSource(pscSource, 10);
-
- // Rate limiter operator should be present
- assertThat(transformation).isNotNull();
- assertThat(transformation.getName()).contains("PscRateLimit");
-
- // Note: See testRescaleAndRateLimitChain() and testSkipsRescaleWhenNotNeeded()
- // for comprehensive partition count testing with mocked providers.
+ assertThat(pscSource.enableRescale).isTrue();
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // Terminal is the rescale PartitionTransformation; its input is the rate-limit map.
+ assertThat(terminal).isNotNull();
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ // Rate-limit parallelism follows the upstream source parallelism, which is
+ // capped at min(scanParallelism, env).
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism);
}
@Test
@@ -1603,13 +1660,21 @@ public void testRescaleCreatesPartitionTransformation() {
@Test
public void testRescaleAndRateLimitChain() {
- // Verifies complete operator chain: Source → Rescale → RateLimit
-
+ // Verifies the operator chain: Source -> RateLimit -> Rescale.
+ // Rescale is now applied AFTER the rate limiter, so the terminal transformation
+ // is a PartitionTransformation (the rescale), whose input is the rate-limit
+ // OneInputTransformation, whose input is the SourceTransformation.
+ // Also verifies that source/rate-limit parallelism is capped at env parallelism:
+ // sourceParallelism = min(scanParallelism, env.getParallelism()) = min(100, 10) = 10
+ // rateLimiterParallelism = sourceStream.getParallelism() = 10
+
try {
- // Mock partition count = 20
+ // Mock partition count = 20 (irrelevant to the gate now; rescale is gated only
+ // by scan.enable-rescale, but kept here so the resolver in PscTableCommonUtils
+ // is deterministic if it ever falls through to tier 3).
PscTableCommonUtils.setProviderForTest((topicUris, props) -> 20);
-
- // Create source with scan.parallelism = 100 (> 20), rescale and rate limiting enabled
+
+ // Create source with scan.parallelism = 100, rescale and rate limiting enabled
final Map modifiedOptions =
getModifiedOptions(
getBasicSourceOptions(),
@@ -1618,42 +1683,66 @@ public void testRescaleAndRateLimitChain() {
addScanParallelismConfig(options, 100);
addRateLimitConfig(options, 5000.0);
});
-
+
final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
-
+
final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
-
- // Verify configuration
+
+ // Verify configuration carried into the source.
+ // (PscDynamicSource.scanParallelism now holds the effective parallelism
+ // resolved by the factory; with scan.parallelism set, tier 1 returns 100.)
assertThat(pscSource.scanParallelism).isEqualTo(100);
assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(5000.0);
-
- // Get transformation - should apply rescale since 100 > 20
- final Transformation transformation = produceTransformationFromSource(pscSource, 10);
-
- // The final transformation should be the rate limiter
- assertThat(transformation).isNotNull();
- assertThat(transformation.getName()).contains("PscRateLimit");
- assertThat(transformation.getParallelism()).isEqualTo(100);
-
- // The input to rate limiter should be PartitionTransformation (rescale)
- assertThat(transformation.getInputs()).isNotEmpty();
- Transformation inputTransformation = transformation.getInputs().get(0);
- assertThat(inputTransformation).isInstanceOf(PartitionTransformation.class);
+ assertThat(pscSource.enableRescale).isTrue();
+
+ // Build the transformation with env parallelism = 10.
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // 1) Terminal is the rescale operator (PartitionTransformation).
+ assertThat(terminal).isNotNull();
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ // 2) Terminal's input is the rate-limit map.
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ // Rate-limit parallelism should match the (capped) source parallelism:
+ // min(scanParallelism=100, env=10) = 10.
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism);
+
+ // 3) The rate-limit map's input is the Kafka source.
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Source operator parallelism should be capped at env parallelism.
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism);
} finally {
PscTableCommonUtils.resetProvider();
}
}
@Test
- public void testSkipsRescaleWhenNotNeeded() {
- // Verifies rescale is NOT applied when scan.parallelism <= partition count
-
+ public void testRescaleAlwaysAppliedWhenEnabled() {
+ // Verifies that rescale is applied whenever scan.enable-rescale=true, regardless
+ // of how scan.parallelism compares to partition count. The factory now passes
+ // shouldRescale through directly and computes the source parallelism via
+ // PscTableCommonUtils.getEffectiveSourceParallelism(), and PscDynamicSource
+ // unconditionally wraps the stream in rescale() when enableRescale is true.
+
try {
- // Mock partition count = 50
+ // Mock partition count = 50. In the OLD logic this would have caused
+ // shouldApplyRescale() to return false (since scan.parallelism 20 <= 50);
+ // in the NEW logic it must NOT prevent rescale from being applied.
PscTableCommonUtils.setProviderForTest((topicUris, props) -> 50);
-
- // Create source with scan.parallelism = 20 (< 50), rescale enabled
+
+ // scan.parallelism = 20, rescale enabled, no rate limiting.
final Map modifiedOptions =
getModifiedOptions(
getBasicSourceOptions(),
@@ -1661,23 +1750,32 @@ public void testSkipsRescaleWhenNotNeeded() {
addRescaleConfig(options, true);
addScanParallelismConfig(options, 20);
});
-
+
final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
-
+
final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
-
- // Verify configuration
+ assertThat(pscSource.enableRescale).isTrue();
assertThat(pscSource.scanParallelism).isEqualTo(20);
-
- // Get transformation - should NOT apply rescale since 20 < 50
- final Transformation transformation = produceTransformationFromSource(pscSource, 10);
-
- // With scan.parallelism (20) < partition count (50), rescale is NOT applied
- // The transformation should be SourceTransformation, not PartitionTransformation
- assertThat(transformation).isNotNull();
- assertThat(transformation).isInstanceOf(SourceTransformation.class);
- assertThat(transformation).isNotInstanceOf(PartitionTransformation.class);
+
+ // Build the transformation with env parallelism = 10 (smaller than scan.parallelism).
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // 1) Terminal is the rescale operator (PartitionTransformation), NOT the bare source.
+ assertThat(terminal).isNotNull();
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ // 2) The rescale's input is the Kafka SourceTransformation (no rate limiter in this scenario).
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // 3) Source operator parallelism is capped at min(scanParallelism, env).
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism);
} finally {
PscTableCommonUtils.resetProvider();
}
@@ -1685,14 +1783,18 @@ public void testSkipsRescaleWhenNotNeeded() {
@Test
public void testRescaleAndRateLimitWithDifferentParallelism() {
- // Verifies PartitionTransformation and correct rate limiter parallelism
- // when scan.parallelism > partition count with both rescale and rate limiting
-
+ // Verifies that when scan.parallelism is larger than env parallelism, both the
+ // source and the rate-limit operator are capped at env parallelism. The chain is:
+ // Source(parallelism=min(scanParallelism, env))
+ // -> PscRateLimit(parallelism=sourceStream.getParallelism())
+ // -> rescale (terminal PartitionTransformation)
+
try {
- // Mock partition count = 15
+ // Mock partition count = 15. In the new logic this is not consulted because
+ // scan.parallelism (tier 1) wins in getEffectiveSourceParallelism.
PscTableCommonUtils.setProviderForTest((topicUris, props) -> 15);
-
- // Create source with scan.parallelism = 80 (> 15), rescale and rate limiting
+
+ // scan.parallelism = 80 (> env parallelism = 10), rescale + rate limit enabled.
final Map modifiedOptions =
getModifiedOptions(
getBasicSourceOptions(),
@@ -1701,35 +1803,343 @@ public void testRescaleAndRateLimitWithDifferentParallelism() {
addScanParallelismConfig(options, 80);
addRateLimitConfig(options, 10000.0);
});
-
+
final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
-
+
final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
-
- // Verify configuration
assertThat(pscSource.scanParallelism).isEqualTo(80);
assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(10000.0);
assertThat(pscSource.enableRescale).isTrue();
-
- // Get transformation with global parallelism = 10 (different from scan.parallelism)
- final Transformation transformation = produceTransformationFromSource(pscSource, 10);
-
- // The final transformation should be the rate limiter
- assertThat(transformation).isNotNull();
- assertThat(transformation.getName()).contains("PscRateLimit");
- // Rate limiter parallelism should be scan.parallelism (80), not global (10)
- assertThat(transformation.getParallelism()).isEqualTo(80);
-
- // The input to rate limiter should be PartitionTransformation (rescale)
- assertThat(transformation.getInputs()).isNotEmpty();
- Transformation inputTransformation = transformation.getInputs().get(0);
- assertThat(inputTransformation).isInstanceOf(PartitionTransformation.class);
+
+ // Build the transformation with env parallelism = 10 (smaller than scan.parallelism).
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // 1) Terminal is the rescale PartitionTransformation.
+ assertThat(terminal).isNotNull();
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ // 2) Terminal's input is the rate-limit map.
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ // Rate-limit and source are both capped at env parallelism (NOT 80).
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism);
+
+ // 3) The rate-limit map's input is the Kafka source, also at the capped parallelism.
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism);
} finally {
PscTableCommonUtils.resetProvider();
}
}
+ @Test
+ public void testSourceParallelismCappedByEnvParallelism() {
+ // When scan.parallelism > env parallelism, PscDynamicSource.produceDataStream() caps
+ // the source operator at env.getParallelism() so we never exceed the available slots.
+ // Chain (no rate limiter): Source(parallelism=env) -> rescale (PartitionTransformation).
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> {
+ addRescaleConfig(options, true);
+ addScanParallelismConfig(options, 80);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ // scan.parallelism (tier 1) wins in getEffectiveSourceParallelism and is stored on the source.
+ assertThat(pscSource.scanParallelism).isEqualTo(80);
+ assertThat(pscSource.enableRescale).isTrue();
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // Terminal is the rescale PartitionTransformation; its input is the Kafka source.
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Cap kicks in: source parallelism = env (10), NOT scan.parallelism (80).
+ assertThat(sourceOp.getParallelism()).isEqualTo(envParallelism);
+ assertThat(sourceOp.getParallelism()).isLessThan(pscSource.scanParallelism);
+ }
+
+ @Test
+ public void testSourceParallelismUsesScanWhenSmallerThanEnv() {
+ // When scan.parallelism < env parallelism, the cap does NOT change the value:
+ // source parallelism = scan.parallelism. This is the symmetric counterpart to
+ // testSourceParallelismCappedByEnvParallelism.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> {
+ addRescaleConfig(options, true);
+ addScanParallelismConfig(options, 4);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.scanParallelism).isEqualTo(4);
+ assertThat(pscSource.enableRescale).isTrue();
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // Terminal is the rescale PartitionTransformation; its input is the Kafka source.
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Source parallelism = scan.parallelism (4), since it is already <= env (10).
+ assertThat(sourceOp.getParallelism()).isEqualTo(pscSource.scanParallelism);
+ assertThat(sourceOp.getParallelism()).isLessThan(envParallelism);
+ }
+
+ @Test
+ public void testEffectiveParallelismFallsBackToPartitionCount() {
+ // With neither scan.parallelism nor table.exec.resource.default-parallelism set,
+ // PscTableCommonUtils.getEffectiveSourceParallelism() must fall through to the
+ // Kafka partition count (tier 3). The factory stores that resolved value on
+ // PscDynamicSource.scanParallelism, so we can assert it directly.
+ try {
+ final int partitionCount = 7;
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> partitionCount);
+
+ // Only rescale enabled - no scan.parallelism, no rate limit.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> addRescaleConfig(options, true));
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ // tier 1 (scan.parallelism) and tier 2 (table.exec.resource.default-parallelism)
+ // are unset, so the partition-count fallback wins.
+ assertThat(pscSource.scanParallelism).isEqualTo(partitionCount);
+
+ // Verify the resolved value is applied to the source operator at stream-build time.
+ final int envParallelism = 20;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+ assertThat(sourceOp.getParallelism())
+ .isEqualTo(Math.min(partitionCount, envParallelism));
+ } finally {
+ PscTableCommonUtils.resetProvider();
+ }
+ }
+
+ @Test
+ public void testEffectiveParallelismFallsBackToTableExec() {
+ // Tier 2 of getEffectiveSourceParallelism(): when scan.parallelism is not set but
+ // table.exec.resource.default-parallelism is, the factory must resolve the source
+ // parallelism to that value AND must NOT fall through to the partition-count
+ // provider (tier 3). We inject a non-default value via a custom global Configuration
+ // and install a partition-count mock with a distinguishable value to assert
+ // tier 3 is never consulted.
+ try {
+ final int sentinelPartitionCount = 999; // distinct from the tier-2 value
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> sentinelPartitionCount);
+
+ final int tableExecParallelism = 12;
+ final Configuration globalConfig = new Configuration();
+ globalConfig.set(
+ ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM,
+ tableExecParallelism);
+
+ // Only rescale enabled. No scan.parallelism, no rate limit.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> addRescaleConfig(options, true));
+
+ final DynamicTableSource actualSource =
+ createTableSourceWithGlobalConfig(SCHEMA, modifiedOptions, globalConfig);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ // tier 2 wins: resolved effective parallelism equals table.exec.resource.default-parallelism,
+ // NOT the partition-count sentinel.
+ assertThat(pscSource.scanParallelism).isEqualTo(tableExecParallelism);
+ assertThat(pscSource.scanParallelism).isNotEqualTo(sentinelPartitionCount);
+
+ // The resolved value flows through to the source operator at stream-build time,
+ // capped by env parallelism.
+ final int envParallelism = 20;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+ assertThat(sourceOp.getParallelism())
+ .isEqualTo(Math.min(tableExecParallelism, envParallelism));
+ } finally {
+ PscTableCommonUtils.resetProvider();
+ }
+ }
+
+ @Test
+ public void testRescaleWithUnknownEffectiveParallelism() {
+ // When all three tiers of getEffectiveSourceParallelism() fail to produce a positive
+ // value, the factory stores -1 on PscDynamicSource.scanParallelism. PscDynamicSource
+ // must still honor scan.enable-rescale=true (i.e. apply rescale()), but it must NOT
+ // call setParallelism() on the source operator -- the source inherits env parallelism
+ // unchanged.
+ try {
+ // Tier 3 failure: partition-count provider returns -1 (signals "could not determine").
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> -1);
+
+ // No scan.parallelism (tier 1 unset). FactoryMocks.createTableSource passes an empty
+ // Configuration, so table.exec.resource.default-parallelism defaults to -1 (tier 2
+ // unset). Combined with the tier-3 failure above, getEffectiveSourceParallelism()
+ // must return -1.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> addRescaleConfig(options, true));
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ // All three tiers failed -> -1 ("unknown") is propagated to the source.
+ assertThat(pscSource.scanParallelism).isEqualTo(-1);
+
+ // Build the transformation with env parallelism = 5.
+ final int envParallelism = 5;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // 1) Rescale is STILL applied (terminal is PartitionTransformation) because
+ // scan.enable-rescale=true is treated as authoritative -- the unknown
+ // effective parallelism only suppresses the explicit setParallelism call.
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ // 2) The rescale's input is the bare SourceTransformation (no rate limiter).
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // 3) Source operator inherits env parallelism (no setParallelism call was made
+ // because scanParallelism <= 0). This is the key behavioral assertion for
+ // the unknown case: we do NOT pin the source to -1, we leave it as-is.
+ assertThat(sourceOp.getParallelism()).isEqualTo(envParallelism);
+ } finally {
+ PscTableCommonUtils.resetProvider();
+ }
+ }
+
+ @Test
+ public void testRescaleAfterRateLimitOrder() {
+ // Pins the operator ordering established in PscDynamicSource.produceDataStream():
+ // Source -> PscRateLimit (rate limiter) -> rescale (terminal)
+ // i.e. rescale is the OUTERMOST operator, applied AFTER the rate limiter,
+ // not before it. This guards against accidental reordering.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> {
+ addRescaleConfig(options, true);
+ addRateLimitConfig(options, 1000.0);
+ addScanParallelismConfig(options, 20);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(1000.0);
+ assertThat(pscSource.scanParallelism).isEqualTo(20);
+
+ final Transformation terminal = produceTransformationFromSource(pscSource, 10);
+
+ // 1) Outermost / terminal op is the rescale.
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getName()).doesNotContain("PscRateLimit");
+
+ // 2) Immediately below rescale is the rate limiter (NOT the bare source).
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+ assertThat(rateLimitOp).isNotInstanceOf(PartitionTransformation.class);
+ assertThat(rateLimitOp).isNotInstanceOf(SourceTransformation.class);
+
+ // 3) Below the rate limiter is the Kafka source. This rules out the swapped
+ // ordering (Source -> rescale -> rate limiter), in which case the rate
+ // limiter's input would be a PartitionTransformation.
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+ }
+
+ @Test
+ public void testRateLimiterParallelismMatchesSourceParallelism() {
+ // The rate limiter is chained via DataStream.map(...) on the source stream, so its
+ // parallelism must equal the source stream's parallelism (the capped value),
+ // never scan.parallelism directly when scan.parallelism > env parallelism.
+ final Map modifiedOptions =
+ getModifiedOptions(
+ getBasicSourceOptions(),
+ options -> {
+ addRescaleConfig(options, true);
+ addRateLimitConfig(options, 2000.0);
+ addScanParallelismConfig(options, 50);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SCHEMA, modifiedOptions);
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.scanParallelism).isEqualTo(50);
+ assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(2000.0);
+
+ final int envParallelism = 8;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Both must equal min(scanParallelism, env). The rate limiter parallelism is
+ // explicitly read from sourceStream.getParallelism() in PscDynamicSource, so it
+ // is locked to the source operator's parallelism.
+ final int expectedParallelism = Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(sourceOp.getParallelism());
+ }
+
@Test
public void testTableSinkAutoCompleteSchemaRegistrySubject() {
// only format
@@ -2217,7 +2627,11 @@ private static PscDynamicSource createExpectedScanSource(
null,
false,
null, // rateLimitRecordsPerSecond
- null); // scanParallelism
+ -1); // scanParallelism: PscDynamicTableFactory passes
+ // `shouldRescale ? getEffectiveSourceParallelism(...) : -1`, so when
+ // rescale is disabled (the default for these expected-source builders)
+ // the actual source carries Integer(-1), not null. We mirror that here
+ // so PscDynamicSource.equals()'s scanParallelism check passes.
}
private static PscDynamicSink createExpectedSink(
diff --git a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtilsTest.java b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtilsTest.java
index c709e64..118af89 100644
--- a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtilsTest.java
+++ b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/PscTableCommonUtilsTest.java
@@ -66,299 +66,105 @@ public void tearDown() {
}
// ============================================
- // Tests for rescale enabled flag checks
+ // Tests for getEffectiveSourceParallelism()
+ // Precedence: scan.parallelism > table.exec.resource.default-parallelism > kafka partition count
// ============================================
@Test
- public void testShouldNotRescaleWhenDisabled() {
- // Given: rescale is disabled
- tableOptions.set(SCAN_ENABLE_RESCALE, false);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 10);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is not applied
- assertThat(result).isFalse();
- }
+ public void testEffectiveParallelismFromScanParallelism() {
+ // Given: scan.parallelism = 12, others would also resolve but should be ignored
+ globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4);
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 7);
- // ============================================
- // Tests for scan.parallelism with mocked partition counts
- // ============================================
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, 12);
- @Test
- public void testShouldRescaleWhenScanParallelismExceedsPartitionCount() {
- // Given: scan.parallelism = 10, partition count = 5, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 2); // Should be ignored
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 5);
-
- // When: shouldApplyRescale is called with scan.parallelism = 10
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 10);
-
- // Then: rescale is applied (10 > 5)
- assertThat(result).isTrue();
+ // Then: scan.parallelism wins
+ assertThat(parallelism).isEqualTo(12);
}
@Test
- public void testShouldNotRescaleWhenScanParallelismLessThanPartitionCount() {
- // Given: scan.parallelism = 5, partition count = 20, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 100); // Should be ignored
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 20);
-
- // When: shouldApplyRescale is called with scan.parallelism = 5
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 5);
-
- // Then: rescale is not applied (5 < 20)
- assertThat(result).isFalse();
- }
+ public void testEffectiveParallelismFallsThroughWhenScanParallelismIsNull() {
+ // Given: scan.parallelism not set; table.exec set to 4; partition count = 7
+ globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4);
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 7);
- @Test
- public void testShouldNotRescaleWhenScanParallelismEqualsPartitionCount() {
- // Given: scan.parallelism = 10, partition count = 10, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 10);
-
- // When: shouldApplyRescale is called with scan.parallelism = 10
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 10);
-
- // Then: rescale is not applied (10 == 10)
- assertThat(result).isFalse();
- }
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, null);
- // ============================================
- // Tests for global default parallelism with mocked partition counts
- // ============================================
-
- @Test
- public void testShouldRescaleWhenGlobalParallelismExceedsPartitionCount() {
- // Given: global parallelism = 50, partition count = 10, no scan.parallelism, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 50);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 10);
-
- // When: shouldApplyRescale is called with no scan.parallelism
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is applied (50 > 10)
- assertThat(result).isTrue();
+ // Then: table.exec wins (4)
+ assertThat(parallelism).isEqualTo(4);
}
@Test
- public void testShouldNotRescaleWhenGlobalParallelismLessThanPartitionCount() {
- // Given: global parallelism = 5, partition count = 20, no scan.parallelism, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 5);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 20);
-
- // When: shouldApplyRescale is called with no scan.parallelism
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is not applied (5 < 20)
- assertThat(result).isFalse();
- }
+ public void testEffectiveParallelismFallsThroughWhenScanParallelismIsMinusOne() {
+ // Given: scan.parallelism = -1 (unset sentinel)
+ globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 4);
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 7);
- // ============================================
- // Tests for invalid parallelism configurations
- // ============================================
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, -1);
- @Test
- public void testShouldNotRescaleWhenScanParallelismIsZero() {
- // Given: scan.parallelism = 0 (invalid), rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 10);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 5);
-
- // When: shouldApplyRescale is called with scan.parallelism = 0
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 0);
-
- // Then: falls back to global parallelism, rescale is applied (10 > 5)
- assertThat(result).isTrue();
+ // Then: -1 is treated as unset; table.exec wins (4)
+ assertThat(parallelism).isEqualTo(4);
}
@Test
- public void testShouldNotRescaleWhenScanParallelismIsNegative() {
- // Given: scan.parallelism = -1 (invalid), rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 10);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 5);
-
- // When: shouldApplyRescale is called with scan.parallelism = -1
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, -1);
-
- // Then: falls back to global parallelism, rescale is applied (10 > 5)
- assertThat(result).isTrue();
- }
+ public void testEffectiveParallelismFallsThroughToPartitionCount() {
+ // Given: scan.parallelism unset; table.exec = -1; partition count = 7
+ globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, -1);
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 7);
- @Test
- public void testShouldNotRescaleWhenNoParallelismConfigured() {
- // Given: no scan.parallelism, no global parallelism, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- // globalConfig has no default parallelism set (returns null)
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 10);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is not applied (no valid parallelism to compare)
- assertThat(result).isFalse();
- }
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, null);
- // ============================================
- // Tests for partition count edge cases
- // ============================================
-
- @Test
- public void testShouldNotRescaleWhenPartitionCountCannotBeDetermined() {
- // Given: partition count = -1 (cannot be determined), rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 10);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> -1);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is not applied (fail-safe behavior)
- assertThat(result).isFalse();
+ // Then: kafka partition count is used
+ assertThat(parallelism).isEqualTo(7);
}
@Test
- public void testShouldNotRescaleWhenPartitionCountIsZero() {
- // Given: partition count = 0 (invalid), rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 10);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 0);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
-
- // Then: rescale is not applied (fail-safe behavior)
- assertThat(result).isFalse();
- }
+ public void testEffectiveParallelismFallsThroughToPartitionCountWhenTableExecIsUnset() {
+ // Given: scan.parallelism unset; table.exec not configured at all (returns null/default)
+ // Note: globalConfig has no value for TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 9);
- // ============================================
- // Tests for scan.parallelism precedence
- // ============================================
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, null);
- @Test
- public void testScanParallelismTakesPrecedenceOverGlobalParallelism() {
- // Given: scan.parallelism = 100, global parallelism = 5, partition count = 10, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 5);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 10);
-
- // When: shouldApplyRescale is called with scan.parallelism = 100
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 100);
-
- // Then: rescale is applied based on scan.parallelism (100 > 10), not global (5 < 10)
- assertThat(result).isTrue();
+ // Then: kafka partition count is used
+ assertThat(parallelism).isEqualTo(9);
}
- // ============================================
- // Tests for multi-topic scenarios
- // ============================================
-
@Test
- public void testShouldRescaleWithMultipleTopics() {
- // Given: multiple topics, min partition count = 8, scan.parallelism = 20, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- List multipleTopics = Arrays.asList(
- "plaintext:kafka:local:test-cluster:/topic1",
- "plaintext:kafka:local:test-cluster:/topic2"
- );
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 8);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, multipleTopics, pscProperties, 20);
-
- // Then: rescale is applied (20 > 8)
- assertThat(result).isTrue();
- }
+ public void testEffectiveParallelismReturnsMinusOneWhenAllSourcesFail() {
+ // Given: scan.parallelism unset; table.exec = -1; partition count provider returns -1
+ globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, -1);
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> -1);
- // ============================================
- // Tests for high parallelism scenarios
- // ============================================
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, null);
- @Test
- public void testShouldRescaleWithHighParallelismAndLowPartitionCount() {
- // Given: scan.parallelism = 2000, partition count = 10, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 10);
-
- // When: shouldApplyRescale is called with high parallelism
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 2000);
-
- // Then: rescale is applied (2000 > 10)
- assertThat(result).isTrue();
+ // Then: -1 (unknown)
+ assertThat(parallelism).isEqualTo(-1);
}
@Test
- public void testShouldNotRescaleWithHighPartitionCountAndLowParallelism() {
- // Given: scan.parallelism = 2, partition count = 1000, rescale enabled
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
-
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 1000);
-
- // When: shouldApplyRescale is called
- boolean result = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, 2);
-
- // Then: rescale is not applied (2 < 1000)
- assertThat(result).isFalse();
- }
+ public void testEffectiveParallelismReturnsMinusOneWhenPartitionCountIsZero() {
+ // Given: scan.parallelism null; table.exec unset; partition count provider returns 0
+ PscTableCommonUtils.setProviderForTest((topicUris, props) -> 0);
- // ============================================
- // Tests for provider reset mechanism
- // ============================================
+ // When
+ int parallelism = PscTableCommonUtils.getEffectiveSourceParallelism(
+ globalConfig, topicUris, pscProperties, null);
- @Test
- public void testProviderResetRestoresDefaultBehavior() {
- // Given: custom provider is set
- PscTableCommonUtils.setProviderForTest((topicUris, props) -> 42);
-
- tableOptions.set(SCAN_ENABLE_RESCALE, true);
- globalConfig.set(ExecutionConfigOptions.TABLE_EXEC_RESOURCE_DEFAULT_PARALLELISM, 100);
-
- // Verify custom provider works
- boolean resultWithMock = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
- assertThat(resultWithMock).isTrue(); // 100 > 42
-
- // When: provider is reset
- PscTableCommonUtils.resetProvider();
-
- // Then: default behavior is restored (returns -1 in unit test environment)
- boolean resultAfterReset = PscTableCommonUtils.shouldApplyRescale(
- tableOptions, globalConfig, topicUris, pscProperties, null);
- assertThat(resultAfterReset).isFalse(); // partition count = -1, fail-safe
+ // Then: 0 is invalid → -1 (unknown)
+ assertThat(parallelism).isEqualTo(-1);
}
// ============================================
diff --git a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java
index 2afa561..41ab741 100644
--- a/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java
+++ b/psc-flink/src/test/java/com/pinterest/flink/streaming/connectors/psc/table/UpsertPscDynamicTableFactoryTest.java
@@ -43,6 +43,7 @@
import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
+import org.apache.flink.streaming.api.transformations.PartitionTransformation;
import org.apache.flink.streaming.api.transformations.SourceTransformation;
import org.apache.flink.streaming.runtime.operators.sink.SinkWriterOperatorFactory;
import org.apache.flink.table.api.DataTypes;
@@ -71,6 +72,7 @@
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.VarCharType;
import org.apache.flink.util.TestLogger;
+import org.junit.After;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -158,6 +160,16 @@ public class UpsertPscDynamicTableFactoryTest extends TestLogger {
@Rule public ExpectedException thrown = ExpectedException.none();
+ /**
+ * Reset partition count provider after each test to prevent test pollution.
+ * Tests that mock the provider via {@link PscTableCommonUtils#setProviderForTest}
+ * must not affect other tests.
+ */
+ @After
+ public void tearDown() {
+ PscTableCommonUtils.resetProvider();
+ }
+
@Test
public void testTableSource() {
final DataType producedDataType = SOURCE_SCHEMA.toPhysicalRowDataType();
@@ -700,6 +712,194 @@ public void testExactlyOnceGuaranteeWithoutTransactionalIdPrefix() {
createTableSink(SINK_SCHEMA, modifiedOptions);
}
+ // --------------------------------------------------------------------------------------------
+ // Rescale / parallelism tests for the upsert source
+ // --------------------------------------------------------------------------------------------
+
+ @Test
+ public void testUpsertSourceWithRescaleEnabledProducesPartitionTransformation() {
+ // With scan.enable-rescale=true on the upsert source, the produced stream must end
+ // in a PartitionTransformation (rescale()), and the upstream operator must be the
+ // bare Kafka source (no rate limiter set).
+ final Map options = getModifiedOptions(
+ getFullSourceOptions(),
+ opts -> {
+ addRescaleConfig(opts, true);
+ addScanParallelismConfig(opts, 50);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SOURCE_SCHEMA, options);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ // scan.parallelism (tier 1) wins in getEffectiveSourceParallelism and is stored on
+ // the source by UpsertPscDynamicTableFactory.
+ assertThat(pscSource.scanParallelism).isEqualTo(50);
+
+ final int envParallelism = 20;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+ assertThat(sourceOp.getParallelism())
+ .isEqualTo(Math.min(pscSource.scanParallelism, envParallelism));
+ }
+
+ @Test
+ public void testUpsertSourceWithRescaleAndRateLimitOrder() {
+ // Pins the operator ordering for the upsert source path:
+ // Source -> PscRateLimit -> rescale (terminal PartitionTransformation)
+ // This must match the non-upsert factory's behavior.
+ final Map options = getModifiedOptions(
+ getFullSourceOptions(),
+ opts -> {
+ addRescaleConfig(opts, true);
+ addScanParallelismConfig(opts, 20);
+ addRateLimitConfig(opts, 1000.0);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SOURCE_SCHEMA, options);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isTrue();
+ assertThat(pscSource.rateLimitRecordsPerSecond).isEqualTo(1000.0);
+ assertThat(pscSource.scanParallelism).isEqualTo(20);
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // 1) Outermost / terminal op is the rescale.
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getName()).doesNotContain("PscRateLimit");
+
+ // 2) Immediately below rescale is the rate limiter (NOT the bare source).
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation rateLimitOp = terminal.getInputs().get(0);
+ assertThat(rateLimitOp.getName()).contains("PscRateLimit");
+
+ // 3) Below the rate limiter is the Kafka source.
+ assertThat(rateLimitOp.getInputs()).isNotEmpty();
+ final Transformation sourceOp = rateLimitOp.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Source and rate-limit are both pinned to min(scanParallelism, env).
+ final int expectedSourceParallelism =
+ Math.min(pscSource.scanParallelism, envParallelism);
+ assertThat(sourceOp.getParallelism()).isEqualTo(expectedSourceParallelism);
+ assertThat(rateLimitOp.getParallelism()).isEqualTo(expectedSourceParallelism);
+ }
+
+ @Test
+ public void testUpsertSourceParallelismCappedByEnv() {
+ // When scan.parallelism (80) exceeds env.getParallelism() (10), the source
+ // operator parallelism must be capped at env parallelism (10) by
+ // PscDynamicSource.produceDataStream(), even though pscSource.scanParallelism
+ // continues to report the configured value (80).
+ final Map options = getModifiedOptions(
+ getFullSourceOptions(),
+ opts -> {
+ addRescaleConfig(opts, true);
+ addScanParallelismConfig(opts, 80);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SOURCE_SCHEMA, options);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ // Configured scan.parallelism is preserved on the source instance.
+ assertThat(pscSource.scanParallelism).isEqualTo(80);
+ assertThat(pscSource.enableRescale).isTrue();
+
+ final int envParallelism = 10;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ assertThat(terminal).isInstanceOf(PartitionTransformation.class);
+ assertThat(terminal.getInputs()).isNotEmpty();
+ final Transformation sourceOp = terminal.getInputs().get(0);
+ assertThat(sourceOp).isInstanceOf(SourceTransformation.class);
+
+ // Cap kicks in: source parallelism = env (10), NOT scan.parallelism (80).
+ assertThat(sourceOp.getParallelism()).isEqualTo(envParallelism);
+ assertThat(sourceOp.getParallelism()).isLessThan(pscSource.scanParallelism);
+ }
+
+ @Test
+ public void testUpsertSourceWithRescaleDisabled() {
+ // With scan.enable-rescale=false the upsert factory must skip
+ // getEffectiveSourceParallelism() entirely and store -1 on the source.
+ // No rescale() should be applied, and (in the absence of a rate limiter) the
+ // terminal transformation should be the bare SourceTransformation.
+ final Map options = getModifiedOptions(
+ getFullSourceOptions(),
+ opts -> {
+ addRescaleConfig(opts, false);
+ // scan.parallelism is set but should NOT be consulted because rescale=false.
+ addScanParallelismConfig(opts, 20);
+ });
+
+ final DynamicTableSource actualSource = createTableSource(SOURCE_SCHEMA, options);
+ assertThat(actualSource).isInstanceOf(PscDynamicSource.class);
+
+ final PscDynamicSource pscSource = (PscDynamicSource) actualSource;
+ assertThat(pscSource.enableRescale).isFalse();
+ // When rescale is disabled the factory short-circuits to -1 regardless of scan.parallelism.
+ assertThat(pscSource.scanParallelism).isEqualTo(-1);
+
+ final int envParallelism = 5;
+ final Transformation terminal =
+ produceTransformationFromSource(pscSource, envParallelism);
+
+ // No rescale, no rate limiter => terminal is the bare Kafka source.
+ assertThat(terminal).isInstanceOf(SourceTransformation.class);
+ assertThat(terminal).isNotInstanceOf(PartitionTransformation.class);
+ // Source inherits env parallelism (no explicit setParallelism call).
+ assertThat(terminal.getParallelism()).isEqualTo(envParallelism);
+ }
+
+ /**
+ * Helper method to create a transformation from a {@link PscDynamicSource}.
+ * Reduces repetitive code in operator chaining tests.
+ *
+ * @param pscSource the PSC dynamic source
+ * @param globalParallelism global parallelism for the execution environment
+ * @return the final transformation in the operator chain
+ */
+ private Transformation produceTransformationFromSource(
+ PscDynamicSource pscSource, int globalParallelism) {
+ final ScanTableSource.ScanRuntimeProvider runtimeProvider =
+ pscSource.getScanRuntimeProvider(ScanRuntimeProviderContext.INSTANCE);
+ assertThat(runtimeProvider).isInstanceOf(DataStreamScanProvider.class);
+
+ final DataStreamScanProvider dataStreamProvider = (DataStreamScanProvider) runtimeProvider;
+ final StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment();
+ env.setParallelism(globalParallelism);
+
+ return dataStreamProvider.produceDataStream(n -> Optional.empty(), env).getTransformation();
+ }
+
+ /** Helper method to add rescale configuration to options map. */
+ private static void addRescaleConfig(Map options, boolean enableRescale) {
+ options.put(PscConnectorOptions.SCAN_ENABLE_RESCALE.key(), String.valueOf(enableRescale));
+ }
+
+ /** Helper method to add rate limit configuration to options map. */
+ private static void addRateLimitConfig(Map options, double rateLimit) {
+ options.put(PscConnectorOptions.SCAN_RATE_LIMIT.key(), String.valueOf(rateLimit));
+ }
+
+ /** Helper method to add scan.parallelism configuration to options map. */
+ private static void addScanParallelismConfig(Map options, int parallelism) {
+ options.put(PscConnectorOptions.SCAN_PARALLELISM.key(), String.valueOf(parallelism));
+ }
+
// --------------------------------------------------------------------------------------------
// Utilities
// --------------------------------------------------------------------------------------------
@@ -837,6 +1037,12 @@ private PscDynamicSource createExpectedScanSource(
String keyPrefix,
String topicUri,
Properties properties) {
+ // NOTE: UpsertPscDynamicTableFactory now ALWAYS passes a non-null effectiveParallelism
+ // to PscDynamicSource (it computes `shouldRescale ? getEffectiveSourceParallelism(...) : -1`).
+ // The backward-compatible 16-arg PscDynamicSource ctor delegates with scanParallelism=null,
+ // which would no longer match the factory output. We therefore use the full ctor here and
+ // mirror the factory's defaults explicitly: sourceUidPrefix=null, enableRescale=false,
+ // rateLimitRecordsPerSecond=null, scanParallelism=-1.
return new PscDynamicSource(
producedDataType,
keyDecodingFormat,
@@ -854,7 +1060,11 @@ private PscDynamicSource createExpectedScanSource(
Collections.emptyMap(),
0,
true,
- FactoryMocks.IDENTIFIER.asSummaryString());
+ FactoryMocks.IDENTIFIER.asSummaryString(),
+ null,
+ false,
+ null,
+ -1);
}
private static PscDynamicSink createExpectedSink(
diff --git a/psc-integration-test/pom.xml b/psc-integration-test/pom.xml
index 2e8245c..9aa66ad 100644
--- a/psc-integration-test/pom.xml
+++ b/psc-integration-test/pom.xml
@@ -5,7 +5,7 @@
psc-java-oss
com.pinterest.psc
- 4.2.1
+ 4.2.2-RC1
../pom.xml
4.0.0
diff --git a/psc-logging/pom.xml b/psc-logging/pom.xml
index 5da21c2..3b1f601 100644
--- a/psc-logging/pom.xml
+++ b/psc-logging/pom.xml
@@ -5,7 +5,7 @@
psc-java-oss
com.pinterest.psc
- 4.2.1
+ 4.2.2-RC1
../pom.xml
4.0.0
diff --git a/psc/pom.xml b/psc/pom.xml
index 0865f3e..de78096 100644
--- a/psc/pom.xml
+++ b/psc/pom.xml
@@ -5,7 +5,7 @@
psc-java-oss
com.pinterest.psc
- 4.2.1
+ 4.2.2-RC1
../pom.xml
4.0.0