This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-12584-ebb2cc3033dd1689c535615242820150174c1c1e in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit e649b2a319c9234ecf8a47bae219f7e65e2c26c2 Author: hutiefang76 <[email protected]> AuthorDate: Wed Oct 28 04:06:28 2026 +0000 [Fix][Connector-V2] Preserve numeric JDBC partition bounds (#12584) Co-authored-by: zhiweiniu <[email protected]> --- docs/en/connectors/source/Jdbc.md | 2 + docs/zh/connectors/source/Jdbc.md | 2 + .../JdbcNumericBetweenParametersProvider.java | 54 +++--- .../JdbcNumericBetweenParametersProviderTest.java | 184 +++++++++++++++++++++ .../jdbc/source/FixedChunkSplitterTest.java | 168 +++++++++++++++++++ .../seatunnel/jdbc/JdbcMysqlSplitIT.java | 82 +++++++++ 6 files changed, 463 insertions(+), 29 deletions(-) diff --git a/docs/en/connectors/source/Jdbc.md b/docs/en/connectors/source/Jdbc.md index 4ee50242d8..f4c20f3150 100644 --- a/docs/en/connectors/source/Jdbc.md +++ b/docs/en/connectors/source/Jdbc.md @@ -321,6 +321,8 @@ For a full-table snapshot, configure `table_path` and normally leave split-key d Only the top-level combination of `query` and `partition_column` selects the legacy fixed splitter; `partition_num` then controls the number of splits. The partition column must be present in the query result. Optional lower and upper bounds avoid extra `MIN`/`MAX` queries, but incorrect bounds can omit source rows, so use them only when the full data range is known. +For fixed splitting of integer keys, including BIGINT and zero-scale DECIMAL, the partitions cover the numeric range without gaps or overlaps. An uneven division by `partition_num` preserves the remaining values, and each interval stays within the lower and upper bounds. + Entries inside `table_list` continue to use dynamic splitting even when they include `query` or partition settings. Do not expect `split.size` to affect the top-level fixed partition mode. ### Query and primary-key caution diff --git a/docs/zh/connectors/source/Jdbc.md b/docs/zh/connectors/source/Jdbc.md index 8f6fa0c37b..7d06080bb5 100644 --- a/docs/zh/connectors/source/Jdbc.md +++ b/docs/zh/connectors/source/Jdbc.md @@ -317,6 +317,8 @@ numeric_special_value_as_null = true 只有顶层同时配置 `query` 和 `partition_column` 时才会使用旧版 fixed splitter,再通过 `partition_num` 控制分片数。分区列必须包含在 query 结果中。配置上下界可以避免额外执行 `MIN`/`MAX` 查询,但错误的边界会漏读源数据,因此只有确认完整数据范围时才应设置。 +对 BIGINT、零标度 DECIMAL 等整数键使用固定分片时,各分片完整覆盖数值范围,不产生间隙或重叠。范围不能被 `partition_num` 整除时,剩余值仍会保留,各区间均位于上下界之内。 + `table_list` 中的条目即使配置了 `query` 或分区参数,仍然使用动态分片。`split.size` 不影响顶层 fixed partition 模式。 ### Query 与主键注意事项 diff --git a/seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProvider.java b/seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProvider.java index 4bf8834f6d..c514450032 100644 --- a/seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProvider.java +++ b/seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProvider.java @@ -44,8 +44,9 @@ public class JdbcNumericBetweenParametersProvider implements JdbcParameterValues private final BigDecimal minVal; private final BigDecimal maxVal; - private long batchSize; + private BigDecimal batchSize; private int batchNum; + private BigDecimal largeBatchCount; /** * NumericBetweenParametersProviderJdbc constructor. @@ -78,16 +79,15 @@ public class JdbcNumericBetweenParametersProvider implements JdbcParameterValues checkArgument(batchSize > 0, "Batch size must be positive"); BigDecimal maxElemCount = (maxVal.subtract(minVal)).add(BigDecimal.valueOf(1)); - if (BigDecimal.valueOf(batchSize).compareTo(maxElemCount) > 0) { - batchSize = maxElemCount.longValue(); + BigDecimal requestedBatchSize = BigDecimal.valueOf(batchSize); + if (requestedBatchSize.compareTo(maxElemCount) > 0) { + requestedBatchSize = maxElemCount; } - this.batchSize = batchSize; - this.batchNum = - new Double( - Math.ceil( - (maxElemCount.divide(BigDecimal.valueOf(batchSize))) - .doubleValue())) - .intValue(); + int numberOfBatches = + maxElemCount.divide(requestedBatchSize, 0, RoundingMode.CEILING).intValueExact(); + this.batchSize = requestedBatchSize; + this.batchNum = numberOfBatches; + this.largeBatchCount = BigDecimal.valueOf(this.batchNum); return this; } @@ -99,36 +99,32 @@ public class JdbcNumericBetweenParametersProvider implements JdbcParameterValues batchNum = maxElemCount.intValue(); } this.batchNum = batchNum; - // For the presence of a decimal we take the integer up - this.batchSize = - (maxElemCount.divide(BigDecimal.valueOf(batchNum), 2, RoundingMode.HALF_UP)) - .setScale(0, RoundingMode.CEILING) - .longValue(); + this.batchSize = maxElemCount.divide(BigDecimal.valueOf(batchNum), 0, RoundingMode.CEILING); + // Distribute the remainder without rounding away keys or exceeding the upper bound. + this.largeBatchCount = + maxElemCount.subtract( + this.batchSize + .subtract(BigDecimal.valueOf(1)) + .multiply(BigDecimal.valueOf(batchNum))); return this; } @Override public Serializable[][] getParameterValues() { checkState( - batchSize > 0, + batchSize != null && batchSize.compareTo(BigDecimal.ZERO) > 0 && batchNum > 0, "Batch size and batch number must be positive. Have you called `ofBatchSize` or `ofBatchNum`?"); - BigDecimal maxElemCount = (maxVal.subtract(minVal)).add(BigDecimal.valueOf(1)); - BigDecimal bigBatchNum = - maxElemCount - .subtract(BigDecimal.valueOf(batchSize - 1)) - .multiply(BigDecimal.valueOf(batchNum)); - Serializable[][] parameters = new Serializable[batchNum][2]; BigDecimal start = minVal; for (int i = 0; i < batchNum; i++) { - BigDecimal end = - start.add(BigDecimal.valueOf(batchSize)) - .subtract(BigDecimal.valueOf(1)) - .subtract( - BigDecimal.valueOf(i).compareTo(bigBatchNum) >= 0 - ? BigDecimal.ONE - : BigDecimal.ZERO); + BigDecimal end = start.add(batchSize).subtract(BigDecimal.valueOf(1)); + if (BigDecimal.valueOf(i).compareTo(largeBatchCount) >= 0) { + end = end.subtract(BigDecimal.valueOf(1)); + } + if (end.compareTo(maxVal) > 0) { + end = maxVal; + } parameters[i] = new BigDecimal[] {start, end}; start = end.add(BigDecimal.valueOf(1)); } diff --git a/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProviderTest.java b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProviderTest.java new file mode 100644 index 0000000000..40161f94ee --- /dev/null +++ b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/split/JdbcNumericBetweenParametersProviderTest.java @@ -0,0 +1,184 @@ +/* + * 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.seatunnel.connectors.seatunnel.jdbc.internal.split; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; + +import java.io.Serializable; +import java.math.BigDecimal; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class JdbcNumericBetweenParametersProviderTest { + + private static final BigDecimal ONE = BigDecimal.ONE; + private static final BigDecimal DECIMAL_38_MAX = + new BigDecimal("99999999999999999999999999999999999999"); + + @Test + void splitsRangeByFetchSize() { + Serializable[][] actual = + new JdbcNumericBetweenParametersProvider(300, ONE, BigDecimal.valueOf(1000)) + .getParameterValues(); + BigDecimal[][] expected = { + {ONE, BigDecimal.valueOf(300)}, + {BigDecimal.valueOf(301), BigDecimal.valueOf(600)}, + {BigDecimal.valueOf(601), BigDecimal.valueOf(900)}, + {BigDecimal.valueOf(901), BigDecimal.valueOf(1000)} + }; + assertPairs(expected, actual); + } + + @Test + void clampsBatchSizeToSingletonRange() { + Serializable[][] actual = + new JdbcNumericBetweenParametersProvider( + BigDecimal.valueOf(5), BigDecimal.valueOf(5)) + .ofBatchSize(3) + .getParameterValues(); + BigDecimal[][] expected = {{BigDecimal.valueOf(5), BigDecimal.valueOf(5)}}; + assertPairs(expected, actual); + } + + @Test + void clampsBatchNumToElementCount() { + Serializable[][] actual = + new JdbcNumericBetweenParametersProvider( + BigDecimal.valueOf(-2), BigDecimal.valueOf(2)) + .ofBatchNum(10) + .getParameterValues(); + BigDecimal[][] expected = { + {BigDecimal.valueOf(-2), BigDecimal.valueOf(-2)}, + {BigDecimal.valueOf(-1), BigDecimal.valueOf(-1)}, + {BigDecimal.ZERO, BigDecimal.ZERO}, + {ONE, ONE}, + {BigDecimal.valueOf(2), BigDecimal.valueOf(2)} + }; + assertPairs(expected, actual); + } + + @Test + void keepsExactBoundsForDecimal38Range() { + Serializable[][] actual = + new JdbcNumericBetweenParametersProvider(ONE, DECIMAL_38_MAX) + .ofBatchNum(3) + .getParameterValues(); + assertEquals(3, actual.length); + assertEquals(ONE, actual[0][0]); + assertEquals(DECIMAL_38_MAX, actual[2][1]); + for (int i = 1; i < actual.length; i++) { + assertEquals(((BigDecimal) actual[i - 1][1]).add(ONE), actual[i][0]); + assertTrue(((BigDecimal) actual[i][0]).compareTo((BigDecimal) actual[i][1]) <= 0); + } + } + + @Test + void rejectsNonPositiveBatchSize() { + assertMessage( + IllegalArgumentException.class, + "Batch size must be positive", + () -> new JdbcNumericBetweenParametersProvider(ONE, BigDecimal.TEN).ofBatchSize(0)); + } + + @Test + void rejectsNonPositiveBatchCount() { + assertMessage( + IllegalArgumentException.class, + "Batch number must be positive", + () -> new JdbcNumericBetweenParametersProvider(ONE, BigDecimal.TEN).ofBatchNum(0)); + } + + @Test + void rejectsReversedBounds() { + assertMessage( + IllegalArgumentException.class, + "minVal must not be larger than maxVal", + () -> new JdbcNumericBetweenParametersProvider(BigDecimal.TEN, ONE)); + } + + @Test + void rejectsUnconfiguredState() { + assertMessage( + IllegalStateException.class, + "Batch size and batch number must be positive. Have you called `ofBatchSize` or `ofBatchNum`?", + () -> + new JdbcNumericBetweenParametersProvider(ONE, BigDecimal.TEN) + .getParameterValues()); + } + + @Test + void preservesFractionalBoundsWhenClampingToOneBatch() { + BigDecimal min = new BigDecimal("0.1"); + BigDecimal max = new BigDecimal("0.2"); + assertPairs( + new BigDecimal[][] {{min, max}}, + new JdbcNumericBetweenParametersProvider(min, max) + .ofBatchNum(2) + .getParameterValues()); + } + + @Test + void replacesPreviousPartitionMode() { + JdbcNumericBetweenParametersProvider provider = + new JdbcNumericBetweenParametersProvider(ONE, BigDecimal.TEN); + provider.ofBatchNum(3).ofBatchSize(4); + assertPairs( + new BigDecimal[][] { + {ONE, BigDecimal.valueOf(4)}, + {BigDecimal.valueOf(5), BigDecimal.valueOf(8)}, + {BigDecimal.valueOf(9), BigDecimal.TEN} + }, + provider.getParameterValues()); + provider.ofBatchNum(3); + assertPairs( + new BigDecimal[][] { + {ONE, BigDecimal.valueOf(4)}, + {BigDecimal.valueOf(5), BigDecimal.valueOf(7)}, + {BigDecimal.valueOf(8), BigDecimal.TEN} + }, + provider.getParameterValues()); + } + + @Test + void rejectsTooManyBatchesWithoutChangingConfiguration() { + BigDecimal min = BigDecimal.valueOf(Long.MIN_VALUE); + BigDecimal max = BigDecimal.valueOf(Long.MAX_VALUE); + JdbcNumericBetweenParametersProvider provider = + new JdbcNumericBetweenParametersProvider(min, max).ofBatchNum(2); + assertMessage(ArithmeticException.class, "Overflow", () -> provider.ofBatchSize(1)); + assertPairs( + new BigDecimal[][] {{min, BigDecimal.valueOf(-1)}, {BigDecimal.ZERO, max}}, + provider.getParameterValues()); + } + + private static void assertPairs(BigDecimal[][] expected, Serializable[][] actual) { + assertEquals(expected.length, actual.length); + for (int i = 0; i < expected.length; i++) { + assertArrayEquals(expected[i], actual[i], "pair " + i); + } + } + + private static void assertMessage( + Class<? extends Throwable> type, String message, Executable executable) { + assertEquals(message, assertThrows(type, executable).getMessage()); + } +} diff --git a/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/FixedChunkSplitterTest.java b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/FixedChunkSplitterTest.java index cdd4bb8e99..c1b230c2c9 100644 --- a/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/FixedChunkSplitterTest.java +++ b/seatunnel-connectors-v2/connector-jdbc/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/FixedChunkSplitterTest.java @@ -26,6 +26,7 @@ import org.apache.seatunnel.connectors.seatunnel.jdbc.config.JdbcConnectionConfi import org.apache.seatunnel.connectors.seatunnel.jdbc.config.JdbcSourceConfig; import org.apache.seatunnel.connectors.seatunnel.jdbc.exception.JdbcConnectorException; +import org.duckdb.DuckDBDriver; import org.junit.jupiter.api.Test; import lombok.extern.slf4j.Slf4j; @@ -34,12 +35,27 @@ import java.lang.reflect.InvocationHandler; import java.lang.reflect.Method; import java.lang.reflect.Proxy; import java.math.BigDecimal; +import java.sql.Connection; import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.LongStream; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -201,4 +217,156 @@ public class FixedChunkSplitterTest { handler); } } + + private static final String DUCKDB_URL = "jdbc:duckdb:"; + + @Test + public void testNumericSplitsReadEveryRowExactlyOnce() throws Exception { + try (Connection connection = new DuckDBDriver().connect(DUCKDB_URL, new Properties()); + DuckDbFixedChunkSplitter splitter = + new DuckDbFixedChunkSplitter(duckDbConfig(), connection)) { + try (Statement statement = connection.createStatement()) { + statement.execute("CREATE TABLE t (id BIGINT)"); + statement.execute("INSERT INTO t SELECT * FROM range(1, 2002)"); + } + + JdbcSourceTable table = + JdbcSourceTable.builder() + .tablePath(TablePath.of("main", "t")) + .query("SELECT id FROM t") + .partitionColumn("id") + .partitionStart("1") + .partitionEnd("2001") + .partitionNumber(1000) + .build(); + + Collection<JdbcSourceSplit> splits = splitter.createSplits(table, splitKeyType()); + List<Long> ids = readAllIds(splitter, splits); + + assertEquals(2001, ids.size(), "every row must be read exactly once"); + List<Long> distinct = new ArrayList<>(new LinkedHashSet<>(ids)); + assertEquals(2001, distinct.size(), "no row may be read twice"); + Collections.sort(distinct); + assertEquals( + LongStream.rangeClosed(1, 2001).boxed().collect(Collectors.toList()), distinct); + } + } + + @Test + public void testSignedBigIntBoundariesAreSplitWithoutNarrowing() throws Exception { + try (Connection connection = new DuckDBDriver().connect(DUCKDB_URL, new Properties()); + DuckDbFixedChunkSplitter splitter = + new DuckDbFixedChunkSplitter(duckDbConfig(), connection)) { + long[] boundaries = {Long.MIN_VALUE, -1L, 0L, Long.MAX_VALUE}; + try (Statement statement = connection.createStatement()) { + statement.execute("CREATE TABLE t (id BIGINT)"); + } + try (PreparedStatement insert = + connection.prepareStatement("INSERT INTO t (id) VALUES (?)")) { + for (long value : boundaries) { + insert.setLong(1, value); + insert.addBatch(); + } + insert.executeBatch(); + } + + JdbcSourceTable table = + JdbcSourceTable.builder() + .tablePath(TablePath.of("main", "t")) + .query("SELECT id FROM t") + .partitionColumn("id") + .partitionStart(String.valueOf(Long.MIN_VALUE)) + .partitionEnd(String.valueOf(Long.MAX_VALUE)) + .partitionNumber(2) + .build(); + + Collection<JdbcSourceSplit> splits = splitter.createSplits(table, splitKeyType()); + List<Long> ids = readAllIds(splitter, splits); + + assertEquals(4, ids.size(), "each signed BIGINT boundary must be read exactly once"); + List<Long> sorted = new ArrayList<>(ids); + Collections.sort(sorted); + assertEquals(Arrays.asList(Long.MIN_VALUE, -1L, 0L, Long.MAX_VALUE), sorted); + } + } + + @Test + public void testConfiguredBoundsCoverRangeAndExcludeOutsiders() throws Exception { + try (Connection connection = new DuckDBDriver().connect(DUCKDB_URL, new Properties()); + DuckDbFixedChunkSplitter splitter = + new DuckDbFixedChunkSplitter(duckDbConfig(), connection)) { + try (Statement statement = connection.createStatement()) { + statement.execute("CREATE TABLE t (id BIGINT)"); + statement.execute("INSERT INTO t SELECT * FROM range(0, 12)"); + } + + JdbcSourceTable table = + JdbcSourceTable.builder() + .tablePath(TablePath.of("main", "t")) + .query("SELECT id FROM t") + .partitionColumn("id") + .partitionStart("1") + .partitionEnd("10") + .partitionNumber(3) + .build(); + + Collection<JdbcSourceSplit> splits = splitter.createSplits(table, splitKeyType()); + List<Long> ids = readAllIds(splitter, splits); + assertEquals(10, ids.size()); + Set<Long> distinct = new HashSet<>(ids); + + assertEquals( + LongStream.rangeClosed(1, 10).boxed().collect(Collectors.toSet()), + distinct, + "the configured bounds must cover 1..10"); + assertFalse( + distinct.contains(0L), "values below the configured start must be excluded"); + assertFalse(distinct.contains(11L), "values above the configured end must be excluded"); + } + } + + private static List<Long> readAllIds(ChunkSplitter splitter, Collection<JdbcSourceSplit> splits) + throws SQLException { + List<Long> ids = new ArrayList<>(); + for (JdbcSourceSplit split : splits) { + try (PreparedStatement statement = + splitter.generateSplitStatement(split, TableSchema.builder().build()); + ResultSet resultSet = statement.executeQuery()) { + while (resultSet.next()) { + ids.add(resultSet.getLong(1)); + } + } + } + return ids; + } + + private static JdbcSourceConfig duckDbConfig() { + return JdbcSourceConfig.builder() + .jdbcConnectionConfig( + JdbcConnectionConfig.builder() + .url(DUCKDB_URL) + .driverName("org.duckdb.DuckDBDriver") + .build()) + .build(); + } + + private static SeaTunnelRowType splitKeyType() { + return new SeaTunnelRowType( + new String[] {"id"}, new SeaTunnelDataType<?>[] {BasicType.LONG_TYPE}); + } + + private static final class DuckDbFixedChunkSplitter extends FixedChunkSplitter { + + private final Connection connection; + + private DuckDbFixedChunkSplitter(JdbcSourceConfig config, Connection connection) { + super(config); + this.connection = connection; + } + + @Override + protected PreparedStatement createPreparedStatement(String sql) throws SQLException { + return connection.prepareStatement(sql); + } + } } diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-jdbc-e2e/connector-jdbc-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-jdbc-e2e/connector-jdbc-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java index e32dc8b593..fe850d5f9c 100644 --- a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-jdbc-e2e/connector-jdbc-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java +++ b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-jdbc-e2e/connector-jdbc-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java @@ -56,7 +56,9 @@ import java.sql.Connection; import java.sql.Date; import java.sql.DriverManager; import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.Statement; import java.sql.Timestamp; import java.time.Duration; import java.time.LocalDate; @@ -65,8 +67,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -791,4 +795,82 @@ public class JdbcMysqlSplitIT extends TestSuiteBase implements TestResource { } } } + + @Test + public void testFixedNumericSplitsReadAllRows() throws Exception { + try (Connection connection = getJdbcConnection(); + Statement statement = connection.createStatement()) { + statement.execute( + "CREATE TABLE " + + MYSQL_DATABASE + + ".fixed_numeric_uneven (id BIGINT PRIMARY KEY)"); + try { + try (PreparedStatement insert = + connection.prepareStatement( + "INSERT INTO " + + MYSQL_DATABASE + + ".fixed_numeric_uneven (id) VALUES (?)")) { + for (int id = 0; id <= 2002; id++) { + insert.setLong(1, id); + insert.addBatch(); + } + insert.executeBatch(); + } + + Map<String, Object> configMap = new HashMap<>(); + configMap.put("url", mysqlUrlInfo.getUrlWithDatabase().get()); + configMap.put("driver", "com.mysql.cj.jdbc.Driver"); + configMap.put("username", MYSQL_USERNAME); + configMap.put("password", MYSQL_PASSWORD); + configMap.put( + "query", "SELECT id FROM " + MYSQL_DATABASE + ".fixed_numeric_uneven"); + configMap.put("partition_column", "id"); + + TablePath tablePath = TablePath.of(MYSQL_DATABASE, "fixed_numeric_uneven"); + try (MySqlCatalog mySqlCatalog = + new MySqlCatalog( + "mysql", + MYSQL_USERNAME, + MYSQL_PASSWORD, + mysqlUrlInfo, + null); + FixedChunkSplitter splitter = getFixedChunkSplitter(configMap)) { + mySqlCatalog.open(); + CatalogTable table = mySqlCatalog.getTable(tablePath); + JdbcSourceTable jdbcSourceTable = + JdbcSourceTable.builder() + .tablePath(tablePath) + .catalogTable(table) + .query( + "SELECT id FROM " + + MYSQL_DATABASE + + ".fixed_numeric_uneven") + .partitionColumn("id") + .partitionNumber(1000) + .partitionStart("1") + .partitionEnd("2001") + .build(); + + Set<Long> uniqueIds = new HashSet<>(); + for (JdbcSourceSplit split : splitter.generateSplits(jdbcSourceTable)) { + try (PreparedStatement splitStatement = + splitter.generateSplitStatement( + split, table.getTableSchema()); + ResultSet resultSet = splitStatement.executeQuery()) { + while (resultSet.next()) { + long id = resultSet.getLong(1); + Assertions.assertTrue( + id >= 1 && id <= 2001, "id out of range: " + id); + Assertions.assertTrue(uniqueIds.add(id), "duplicated id: " + id); + } + } + } + Assertions.assertEquals(2001, uniqueIds.size()); + } + } finally { + statement.execute( + "DROP TABLE IF EXISTS " + MYSQL_DATABASE + ".fixed_numeric_uneven"); + } + } + } }
