pabloem commented on a change in pull request #15848:
URL: https://github.com/apache/beam/pull/15848#discussion_r798090176
##########
File path:
sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcUtil.java
##########
@@ -334,30 +342,188 @@ private static Calendar
withTimestampAndTimezone(DateTime dateTime) {
}
/** Create partitions on a table. */
- static class PartitioningFn extends DoFn<KV<Integer, KV<Long, Long>>,
KV<String, Long>> {
+ static class PartitioningFn<T> extends DoFn<KV<Integer, KV<T, T>>, KV<T, T>>
{
+ private static final Logger LOG =
LoggerFactory.getLogger(PartitioningFn.class);
+ final TypeDescriptor<T> partitioningColumnType;
+
+ PartitioningFn(TypeDescriptor<T> partitioningColumnType) {
+ this.partitioningColumnType = partitioningColumnType;
+ }
+
@ProcessElement
public void processElement(ProcessContext c) {
- Integer numPartitions = c.element().getKey();
- Long lowerBound = c.element().getValue().getKey();
- Long upperBound = c.element().getValue().getValue();
- if (lowerBound > upperBound) {
- throw new RuntimeException(
- String.format(
- "Lower bound [%s] is higher than upper bound [%s]",
lowerBound, upperBound));
- }
- long stride = (upperBound - lowerBound) / numPartitions + 1;
- for (long i = lowerBound; i < upperBound - stride; i += stride) {
- String range = String.format("%s,%s", i, i + stride);
- KV<String, Long> kvRange = KV.of(range, 1L);
- c.output(kvRange);
- }
- if (upperBound - lowerBound > stride * (numPartitions - 1)) {
- long indexFrom = (numPartitions - 1) * stride;
- long indexTo = upperBound + 1;
- String range = String.format("%s,%s", indexFrom, indexTo);
- KV<String, Long> kvRange = KV.of(range, 1L);
- c.output(kvRange);
+ T lowerBound = c.element().getValue().getKey();
+ T upperBound = c.element().getValue().getValue();
+ JdbcReadWithPartitionsHelper<T> helper =
+ (JdbcReadWithPartitionsHelper<T>)
PRESET_HELPERS.get(partitioningColumnType.getRawType());
+ List<KV<T, T>> ranges =
+ Lists.newArrayList(helper.calculateRanges(lowerBound, upperBound,
c.element().getKey()));
+ LOG.warn("Total of {} ranges: {}", ranges.size(), ranges);
+ for (KV<T, T> e : ranges) {
+ c.output(e);
}
}
}
+
+ public static final Map<Class<?>, JdbcReadWithPartitionsHelper<?>>
PRESET_HELPERS =
+ ImmutableMap.of(
+ String.class,
Review comment:
I agree that this will not partition ranges equally, but it's hard to
partition intelligently without querying the database ahead of time.
I think the best strategy to improve the automatic partitioning will be to
make ReadFn a splittable DoFn where a range can be partitioned at runtime. What
do you think about that?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
