feat: Supports topic partition increase.

src/main/java/com/google/cloud/pubsublite/spark/CachedPartitionCountReader.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed 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 com.google.cloud.pubsublite.spark;
+
+import com.google.cloud.pubsublite.AdminClient;
+import com.google.cloud.pubsublite.PartitionLookupUtils;
+import com.google.cloud.pubsublite.TopicPath;
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.concurrent.ThreadSafe;
+
+@ThreadSafe
+public class CachedPartitionCountReader implements PartitionCountReader {
+  private final AdminClient adminClient;
+  private final Supplier<Integer> supplier;
+
+  public CachedPartitionCountReader(AdminClient adminClient, TopicPath topicPath) {
+    this.adminClient = adminClient;
+    this.supplier =
+        Suppliers.memoizeWithExpiration(
+            () -> PartitionLookupUtils.numPartitions(topicPath, adminClient), 10, TimeUnit.SECONDS);
+  }
+
+  @Override
+  public void close() {
+    adminClient.close();
+  }
+
+  public int getPartitionCount() {
+    return supplier.get();
+  }
+}

src/main/java/com/google/cloud/pubsublite/spark/LimitingHeadOffsetReader.java

@@ -43,15 +43,18 @@ public class LimitingHeadOffsetReader implements PerTopicHeadOffsetReader {
 
   private final TopicStatsClient topicStatsClient;
   private final TopicPath topic;
-  private final long topicPartitionCount;
+  private final PartitionCountReader partitionCountReader;
   private final AsyncLoadingCache<Partition, Offset> cachedHeadOffsets;
 
   @VisibleForTesting
   public LimitingHeadOffsetReader(
-      TopicStatsClient topicStatsClient, TopicPath topic, long topicPartitionCount, Ticker ticker) {
+      TopicStatsClient topicStatsClient,
+      TopicPath topic,
+      PartitionCountReader partitionCountReader,
+      Ticker ticker) {
     this.topicStatsClient = topicStatsClient;
     this.topic = topic;
-    this.topicPartitionCount = topicPartitionCount;
+    this.partitionCountReader = partitionCountReader;
     this.cachedHeadOffsets =
         Caffeine.newBuilder()
             .ticker(ticker)
@@ -82,7 +85,7 @@ public void onSuccess(Cursor c) {
   @Override
   public PslSourceOffset getHeadOffset() {
     Set<Partition> keySet = new HashSet<>();
-    for (int i = 0; i < topicPartitionCount; i++) {
+    for (int i = 0; i < partitionCountReader.getPartitionCount(); i++) {
       keySet.add(Partition.of(i));
     }
     CompletableFuture<Map<Partition, Offset>> future = cachedHeadOffsets.getAll(keySet);
@@ -96,5 +99,6 @@ public PslSourceOffset getHeadOffset() {
   @Override
   public void close() {
     topicStatsClient.close();
+    partitionCountReader.close();
   }
 }

src/main/java/com/google/cloud/pubsublite/spark/MultiPartitionCommitterImpl.java

@@ -27,18 +27,22 @@
 import java.util.HashMap;
 import java.util.Map;
 
+/**
+ * A {@link MultiPartitionCommitter} that lazily adjusts for partition changes when {@link
+ * MultiPartitionCommitter#commit(PslSourceOffset)};
+ */
 public class MultiPartitionCommitterImpl implements MultiPartitionCommitter {
   private static final GoogleLogger log = GoogleLogger.forEnclosingClass();
 
+  private final CommitterFactory committerFactory;
   private final Map<Partition, Committer> committerMap = new HashMap<>();
 
   @VisibleForTesting
   MultiPartitionCommitterImpl(long topicPartitionCount, CommitterFactory committerFactory) {
+    this.committerFactory = committerFactory;
     for (int i = 0; i < topicPartitionCount; i++) {
       Partition p = Partition.of(i);
-      Committer committer = committerFactory.newCommitter(p);
-      committer.startAsync().awaitRunning();
-      committerMap.put(p, committer);
+      committerMap.put(p, createCommitter(p));
     }
   }
 
@@ -47,8 +51,24 @@ public synchronized void close() {
     committerMap.values().forEach(c -> c.stopAsync().awaitTerminated());
   }
 
+  /** Adjust committerMap based on the partitions that needs to be committed. */
+  public synchronized void update(PslSourceOffset offset) {
+    for (Partition p : offset.partitionOffsetMap().keySet()) {
+      if (!committerMap.containsKey(p)) {
+        committerMap.put(p, createCommitter(p));
+      }
+    }
+  }
+
+  private synchronized Committer createCommitter(Partition p) {
+    Committer committer = committerFactory.newCommitter(p);
+    committer.startAsync().awaitRunning();
+    return committer;
+  }
+
   @Override
   public synchronized void commit(PslSourceOffset offset) {
+    update(offset);
     offset
         .partitionOffsetMap()
         .forEach(

src/main/java/com/google/cloud/pubsublite/spark/PartitionCountReader.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed 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 com.google.cloud.pubsublite.spark;
+
+import java.io.Closeable;
+
+public interface PartitionCountReader extends Closeable {
+  int getPartitionCount();
+
+  @Override
+  void close();
+}

src/main/java/com/google/cloud/pubsublite/spark/PslContinuousReader.java

@@ -41,8 +41,9 @@ public class PslContinuousReader implements ContinuousReader {
   private final PartitionSubscriberFactory partitionSubscriberFactory;
   private final SubscriptionPath subscriptionPath;
   private final FlowControlSettings flowControlSettings;
-  private final long topicPartitionCount;
   private SparkSourceOffset startOffset;
+  private final PartitionCountReader partitionCountReader;
+  private final long topicPartitionCount;
 
   @VisibleForTesting
   public PslContinuousReader(
@@ -51,13 +52,14 @@ public PslContinuousReader(
       PartitionSubscriberFactory partitionSubscriberFactory,
       SubscriptionPath subscriptionPath,
       FlowControlSettings flowControlSettings,
-      long topicPartitionCount) {
+      PartitionCountReader partitionCountReader) {
     this.cursorClient = cursorClient;
     this.committer = committer;
     this.partitionSubscriberFactory = partitionSubscriberFactory;
     this.subscriptionPath = subscriptionPath;
     this.flowControlSettings = flowControlSettings;
-    this.topicPartitionCount = topicPartitionCount;
+    this.partitionCountReader = partitionCountReader;
+    this.topicPartitionCount = partitionCountReader.getPartitionCount();
   }
 
   @Override
@@ -126,4 +128,9 @@ public List<InputPartition<InternalRow>> planInputPartitions() {
     }
     return list;
   }
+
+  @Override
+  public boolean needsReconfiguration() {
+    return partitionCountReader.getPartitionCount() != topicPartitionCount;
+  }
 }

src/main/java/com/google/cloud/pubsublite/spark/PslDataSource.java

@@ -21,7 +21,6 @@
 import com.github.benmanes.caffeine.cache.Ticker;
 import com.google.auto.service.AutoService;
 import com.google.cloud.pubsublite.AdminClient;
-import com.google.cloud.pubsublite.PartitionLookupUtils;
 import com.google.cloud.pubsublite.SubscriptionPath;
 import com.google.cloud.pubsublite.TopicPath;
 import java.util.Objects;
@@ -55,17 +54,22 @@ public ContinuousReader createContinuousReader(
     PslDataSourceOptions pslDataSourceOptions =
         PslDataSourceOptions.fromSparkDataSourceOptions(options);
     SubscriptionPath subscriptionPath = pslDataSourceOptions.subscriptionPath();
-    long topicPartitionCount;
-    try (AdminClient adminClient = pslDataSourceOptions.newAdminClient()) {
-      topicPartitionCount = PartitionLookupUtils.numPartitions(subscriptionPath, adminClient);
+    TopicPath topicPath;
+    AdminClient adminClient = pslDataSourceOptions.newAdminClient();
+    try {
+      topicPath = TopicPath.parse(adminClient.getSubscription(subscriptionPath).get().getTopic());
+    } catch (Throwable t) {
+      throw toCanonical(t).underlying;
     }
+    PartitionCountReader partitionCountReader =
+        new CachedPartitionCountReader(adminClient, topicPath);
     return new PslContinuousReader(
         pslDataSourceOptions.newCursorClient(),
-        pslDataSourceOptions.newMultiPartitionCommitter(topicPartitionCount),
+        pslDataSourceOptions.newMultiPartitionCommitter(partitionCountReader.getPartitionCount()),
         pslDataSourceOptions.getSubscriberFactory(),
         subscriptionPath,
         Objects.requireNonNull(pslDataSourceOptions.flowControlSettings()),
-        topicPartitionCount);
+        partitionCountReader);
   }
 
   @Override
@@ -80,25 +84,25 @@ public MicroBatchReader createMicroBatchReader(
         PslDataSourceOptions.fromSparkDataSourceOptions(options);
     SubscriptionPath subscriptionPath = pslDataSourceOptions.subscriptionPath();
     TopicPath topicPath;
-    long topicPartitionCount;
-    try (AdminClient adminClient = pslDataSourceOptions.newAdminClient()) {
+    AdminClient adminClient = pslDataSourceOptions.newAdminClient();
+    try {
       topicPath = TopicPath.parse(adminClient.getSubscription(subscriptionPath).get().getTopic());
-      topicPartitionCount = PartitionLookupUtils.numPartitions(topicPath, adminClient);
     } catch (Throwable t) {
       throw toCanonical(t).underlying;
     }
+    PartitionCountReader partitionCountReader =
+        new CachedPartitionCountReader(adminClient, topicPath);
     return new PslMicroBatchReader(
         pslDataSourceOptions.newCursorClient(),
-        pslDataSourceOptions.newMultiPartitionCommitter(topicPartitionCount),
+        pslDataSourceOptions.newMultiPartitionCommitter(partitionCountReader.getPartitionCount()),
         pslDataSourceOptions.getSubscriberFactory(),
         new LimitingHeadOffsetReader(
             pslDataSourceOptions.newTopicStatsClient(),
             topicPath,
-            topicPartitionCount,
+            partitionCountReader,
             Ticker.systemTicker()),
         subscriptionPath,
         Objects.requireNonNull(pslDataSourceOptions.flowControlSettings()),
-        pslDataSourceOptions.maxMessagesPerBatch(),
-        topicPartitionCount);
+        pslDataSourceOptions.maxMessagesPerBatch());
   }
 }

src/main/java/com/google/cloud/pubsublite/spark/PslMicroBatchReader.java

@@ -19,6 +19,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 
+import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.SubscriptionPath;
 import com.google.cloud.pubsublite.cloudpubsub.FlowControlSettings;
 import com.google.cloud.pubsublite.internal.CursorClient;
@@ -34,14 +35,12 @@
 import org.apache.spark.sql.types.StructType;
 
 public class PslMicroBatchReader implements MicroBatchReader {
-
   private final CursorClient cursorClient;
   private final MultiPartitionCommitter committer;
   private final PartitionSubscriberFactory partitionSubscriberFactory;
   private final PerTopicHeadOffsetReader headOffsetReader;
   private final SubscriptionPath subscriptionPath;
   private final FlowControlSettings flowControlSettings;
-  private final long topicPartitionCount;
   private final long maxMessagesPerBatch;
   @Nullable private SparkSourceOffset startOffset = null;
   private SparkSourceOffset endOffset;
@@ -53,41 +52,45 @@ public PslMicroBatchReader(
       PerTopicHeadOffsetReader headOffsetReader,
       SubscriptionPath subscriptionPath,
       FlowControlSettings flowControlSettings,
-      long maxMessagesPerBatch,
-      long topicPartitionCount) {
+      long maxMessagesPerBatch) {
     this.cursorClient = cursorClient;
     this.committer = committer;
     this.partitionSubscriberFactory = partitionSubscriberFactory;
     this.headOffsetReader = headOffsetReader;
     this.subscriptionPath = subscriptionPath;
     this.flowControlSettings = flowControlSettings;
-    this.topicPartitionCount = topicPartitionCount;
     this.maxMessagesPerBatch = maxMessagesPerBatch;
   }
 
   @Override
   public void setOffsetRange(Optional<Offset> start, Optional<Offset> end) {
+    int currentTopicPartitionCount;
+    if (end.isPresent()) {
+      checkArgument(
+          end.get() instanceof SparkSourceOffset,
+          "end offset is not instance of SparkSourceOffset.");
+      endOffset = (SparkSourceOffset) end.get();
+      currentTopicPartitionCount = ((SparkSourceOffset) end.get()).getPartitionOffsetMap().size();
+    } else {
+      endOffset = PslSparkUtils.toSparkSourceOffset(headOffsetReader.getHeadOffset());
+      currentTopicPartitionCount = endOffset.getPartitionOffsetMap().size();
+    }
+
     if (start.isPresent()) {
       checkArgument(
           start.get() instanceof SparkSourceOffset,
           "start offset is not instance of SparkSourceOffset.");
       startOffset = (SparkSourceOffset) start.get();
     } else {
       startOffset =
-          PslSparkUtils.getSparkStartOffset(cursorClient, subscriptionPath, topicPartitionCount);
-    }
-    if (end.isPresent()) {
-      checkArgument(
-          end.get() instanceof SparkSourceOffset,
-          "end offset is not instance of SparkSourceOffset.");
-      endOffset = (SparkSourceOffset) end.get();
-    } else {
-      SparkSourceOffset headOffset =
-          PslSparkUtils.toSparkSourceOffset(headOffsetReader.getHeadOffset());
-      endOffset =
-          PslSparkUtils.getSparkEndOffset(
-              headOffset, startOffset, maxMessagesPerBatch, topicPartitionCount);
+          PslSparkUtils.getSparkStartOffset(
+              cursorClient, subscriptionPath, currentTopicPartitionCount);
     }
+
+    // Limit endOffset by maxMessagesPerBatch.
+    endOffset =
+        PslSparkUtils.getSparkEndOffset(
+            endOffset, startOffset, maxMessagesPerBatch, currentTopicPartitionCount);
   }
 
   @Override
@@ -126,23 +129,28 @@ public StructType readSchema() {
 
   @Override
   public List<InputPartition<InternalRow>> planInputPartitions() {
-    checkState(startOffset != null);
+    checkState(startOffset != null && endOffset != null);
+
     List<InputPartition<InternalRow>> list = new ArrayList<>();
-    for (SparkPartitionOffset offset : startOffset.getPartitionOffsetMap().values()) {
-      SparkPartitionOffset endPartitionOffset =
-          endOffset.getPartitionOffsetMap().get(offset.partition());
-      if (offset.equals(endPartitionOffset)) {
+    // Since this is called right after setOffsetRange, we could use partitions in endOffset as
+    // current partition count.
+    for (SparkPartitionOffset endPartitionOffset : endOffset.getPartitionOffsetMap().values()) {
+      Partition p = endPartitionOffset.partition();
+      SparkPartitionOffset startPartitionOffset =
+          startOffset.getPartitionOffsetMap().getOrDefault(p, SparkPartitionOffset.create(p, -1L));
+      if (startPartitionOffset.equals(endPartitionOffset)) {
         // There is no message to pull for this partition.
         continue;
       }
       PartitionSubscriberFactory partitionSubscriberFactory = this.partitionSubscriberFactory;
       SubscriberFactory subscriberFactory =
-          (consumer) -> partitionSubscriberFactory.newSubscriber(offset.partition(), consumer);
+          (consumer) ->
+              partitionSubscriberFactory.newSubscriber(endPartitionOffset.partition(), consumer);
       list.add(
           new PslMicroBatchInputPartition(
               subscriptionPath,
               flowControlSettings,
-              offset,
+              startPartitionOffset,
               endPartitionOffset,
               subscriberFactory));
     }