This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 4ff9dbe8389 Fix Kafka metadata for sparse partition recovery (#18687)
4ff9dbe8389 is described below
commit 4ff9dbe8389c1559a0c04c0db0049917bfeedfdd
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Jun 8 20:01:25 2026 -0700
Fix Kafka metadata for sparse partition recovery (#18687)
---
.../realtime/PinotLLCRealtimeSegmentManager.java | 10 +-
.../fakestream/FakeStreamMetadataProvider.java | 27 +++
.../kafka30/KafkaStreamMetadataProvider.java | 32 ++--
.../kafka30/KafkaPartitionLevelConsumerTest.java | 36 ++++
.../kafka30/KafkaStreamMetadataProviderTest.java | 188 +++++++++++++++++++++
.../kafka40/KafkaStreamMetadataProvider.java | 32 ++--
.../kafka40/KafkaPartitionLevelConsumerTest.java | 36 ++++
.../kafka40/KafkaStreamMetadataProviderTest.java | 188 +++++++++++++++++++++
.../pinot/spi/stream/StreamMetadataProvider.java | 64 ++-----
.../stream/PartitionGroupMetadataFetcherTest.java | 99 -----------
10 files changed, 531 insertions(+), 181 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java
index 46982aac0fb..0e27d156529 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java
@@ -2022,13 +2022,9 @@ public class PinotLLCRealtimeSegmentManager implements
PinotClusterConfigChangeL
streamConfig.setOffsetCriteria(OffsetCriteria.SMALLEST_OFFSET_CRITERIA);
// Kinesis shard-split flow requires us to pass
currentPartitionGroupConsumptionStatusList so that
- // we can check if its completely consumed
- // However the kafka implementation of computePartitionGroupMetadata()
breaks if we pass the current status
- // This leads to streamSmallestOffset set to null in selectStartOffset()
method
- // The overall dependency isn't clean and is causing the issue and
requires refactor
- // Temporarily, we are passing a boolean flag to indicate if we want to
use the current status
- // The kafka implementation of computePartitionGroupMetadata() will
ignore the current status
- // while the kinesis implementation will use it.
+ // we can check if it is completely consumed. Kafka needs stream offsets
for every partition in this path, so
+ // forceGetOffsetFromStream lets its metadata provider fetch offsets
from the stream instead of reusing
+ // current status offsets.
List<StreamMetadata> streamMetadataList = getNewStreamMetadataList(
streamConfigs, currentPartitionGroupConsumptionStatusList,
idealState, true);
streamConfig.setOffsetCriteria(originalOffsetCriteria);
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/realtime/impl/fakestream/FakeStreamMetadataProvider.java
b/pinot-core/src/test/java/org/apache/pinot/core/realtime/impl/fakestream/FakeStreamMetadataProvider.java
index 99e3d90749d..be662b53805 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/realtime/impl/fakestream/FakeStreamMetadataProvider.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/realtime/impl/fakestream/FakeStreamMetadataProvider.java
@@ -19,13 +19,17 @@
package org.apache.pinot.core.realtime.impl.fakestream;
import java.io.IOException;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.pinot.spi.stream.LongMsgOffset;
import org.apache.pinot.spi.stream.OffsetCriteria;
+import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
+import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.StreamConfig;
import org.apache.pinot.spi.stream.StreamMetadataProvider;
import org.apache.pinot.spi.stream.StreamPartitionMsgOffset;
@@ -60,6 +64,29 @@ public class FakeStreamMetadataProvider implements
StreamMetadataProvider {
}
}
+ @Override
+ public List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
+ List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses,
int timeoutMillis) {
+ Map<Integer, StreamPartitionMsgOffset> partitionIdToEndOffset =
+ new HashMap<>(partitionGroupConsumptionStatuses.size());
+ for (PartitionGroupConsumptionStatus partitionGroupConsumptionStatus :
partitionGroupConsumptionStatuses) {
+
partitionIdToEndOffset.put(partitionGroupConsumptionStatus.getStreamPartitionGroupId(),
+ partitionGroupConsumptionStatus.getEndOffset());
+ }
+
+ List<PartitionGroupMetadata> partitionGroupMetadataList = new
ArrayList<>(_numPartitions);
+ for (int partitionId = 0; partitionId < _numPartitions; partitionId++) {
+ if (partitionIdToEndOffset.containsKey(partitionId)) {
+ partitionGroupMetadataList.add(
+ new PartitionGroupMetadata(partitionId,
partitionIdToEndOffset.get(partitionId)));
+ } else {
+ partitionGroupMetadataList.add(new PartitionGroupMetadata(partitionId,
+ fetchStreamPartitionOffset(streamConfig.getOffsetCriteria(),
timeoutMillis)));
+ }
+ }
+ return partitionGroupMetadataList;
+ }
+
@Override
public boolean supportsOffsetLag() {
return true;
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
index ebc0a044ccd..77d76c1a434 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java
@@ -123,31 +123,37 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
public List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses,
int timeoutMillis)
throws IOException, java.util.concurrent.TimeoutException {
- if (_partitionIdSubset.isEmpty()) {
- return
StreamMetadataProvider.super.computePartitionGroupMetadata(clientId,
streamConfig,
- partitionGroupConsumptionStatuses, timeoutMillis);
- }
Map<Integer, StreamPartitionMsgOffset> partitionIdToEndOffset =
new HashMap<>(partitionGroupConsumptionStatuses.size());
for (PartitionGroupConsumptionStatus s :
partitionGroupConsumptionStatuses) {
partitionIdToEndOffset.put(s.getStreamPartitionGroupId(),
s.getEndOffset());
}
+
+ List<Integer> partitionIds;
+ if (_partitionIdSubset.isEmpty()) {
+ int partitionCount = fetchPartitionCount(timeoutMillis);
+ partitionIds = new ArrayList<>(partitionCount);
+ for (int partitionId = 0; partitionId < partitionCount; partitionId++) {
+ partitionIds.add(partitionId);
+ }
+ } else {
+ partitionIds = _partitionIdSubset;
+ }
+
StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
- List<PartitionGroupMetadata> result = new
ArrayList<>(_partitionIdSubset.size());
- for (Integer partitionId : _partitionIdSubset) {
- StreamPartitionMsgOffset endOffset =
partitionIdToEndOffset.get(partitionId);
- StreamPartitionMsgOffset startOffset;
- if (endOffset == null) {
+ List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ if (partitionIdToEndOffset.containsKey(partitionId)) {
+ result.add(new PartitionGroupMetadata(partitionId,
partitionIdToEndOffset.get(partitionId)));
+ } else {
try (StreamMetadataProvider partitionMetadataProvider =
streamConsumerFactory.createPartitionMetadataProvider(
StreamConsumerFactory.getUniqueClientId(clientId),
partitionId)) {
- startOffset = partitionMetadataProvider.fetchStreamPartitionOffset(
+ StreamPartitionMsgOffset startOffset =
partitionMetadataProvider.fetchStreamPartitionOffset(
streamConfig.getOffsetCriteria(), timeoutMillis);
+ result.add(new PartitionGroupMetadata(partitionId, startOffset));
}
- } else {
- startOffset = endOffset;
}
- result.add(new PartitionGroupMetadata(partitionId, startOffset));
}
return result;
}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java
index 96afe08e070..db5ee9bcfd5 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java
@@ -49,6 +49,7 @@ import org.apache.pinot.spi.stream.MessageBatch;
import org.apache.pinot.spi.stream.OffsetCriteria;
import org.apache.pinot.spi.stream.PartitionGroupConsumer;
import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
+import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.StreamConfig;
import org.apache.pinot.spi.stream.StreamConsumerFactory;
import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
@@ -282,6 +283,41 @@ public class KafkaPartitionLevelConsumerTest {
}
}
+ @Test
+ public void testComputePartitionGroupMetadataUsesKafkaPartitionIds()
+ throws Exception {
+ String streamType = "kafka";
+ String streamKafkaBrokerList = _kafkaBrokerAddress;
+ String clientId = "partition-id-client";
+ String tableNameWithType = "tableName_REALTIME";
+
+ Map<String, String> streamConfigMap = new HashMap<>();
+ streamConfigMap.put("streamType", streamType);
+ streamConfigMap.put("stream.kafka.topic.name",
TEST_TOPIC_SUBSET_PARTITION);
+ streamConfigMap.put("stream.kafka.broker.list", streamKafkaBrokerList);
+ streamConfigMap.put("stream.kafka.consumer.factory.class.name",
getKafkaConsumerFactoryName());
+ streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
+ StreamConfig streamConfig = new StreamConfig(tableNameWithType,
streamConfigMap);
+
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(4, 4, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(5, 5, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(6, 6, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(7, 7, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new KafkaStreamMetadataProvider(clientId, streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata(clientId,
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ }
+ }
+
@Test
public void testFetchMessages() {
String streamType = "kafka";
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
new file mode 100644
index 00000000000..973f338bbb7
--- /dev/null
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProviderTest.java
@@ -0,0 +1,188 @@
+/**
+ * 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.pinot.plugin.stream.kafka30;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Bytes;
+import org.apache.pinot.spi.stream.LongMsgOffset;
+import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
+import org.apache.pinot.spi.stream.PartitionGroupMetadata;
+import org.apache.pinot.spi.stream.StreamConfig;
+import org.apache.pinot.spi.stream.StreamConfigProperties;
+import org.apache.pinot.spi.stream.StreamMetadataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
+
+public class KafkaStreamMetadataProviderTest {
+ private static final ThreadLocal<Consumer<Bytes, Bytes>> MOCK_CONSUMER = new
ThreadLocal<>();
+
+ @Test
+ public void testComputePartitionGroupMetadataRecoversMissingLowPartitionId()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(11), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(13), "DONE"),
+ new PartitionGroupConsumptionStatus(4, 4, new LongMsgOffset(0), new
LongMsgOffset(14), "DONE"),
+ new PartitionGroupConsumptionStatus(5, 5, new LongMsgOffset(0), new
LongMsgOffset(15), "DONE"),
+ new PartitionGroupConsumptionStatus(6, 6, new LongMsgOffset(0), new
LongMsgOffset(16), "DONE"),
+ new PartitionGroupConsumptionStatus(7, 7, new LongMsgOffset(0), new
LongMsgOffset(17), "DONE"));
+
+ assertEquals(computePartitionIdsWithSizeBasedAlgorithm(currentStatuses,
8), List.of(0, 1, 3, 4, 5, 6, 7, 7));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata("client",
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+
assertEquals(partitionGroupMetadataList.get(2).getStartOffset().toString(),
"1002");
+
assertEquals(partitionGroupMetadataList.get(3).getStartOffset().toString(),
"13");
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataHandlesTopicExpansion()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(11), "DONE"),
+ new PartitionGroupConsumptionStatus(2, 2, new LongMsgOffset(0), new
LongMsgOffset(12), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(13), "DONE"));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata("client",
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ assertEquals(partitionGroupMetadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("10", "11", "12", "13",
"1004", "1005", "1006", "1007"));
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ private static StreamConfig getStreamConfig(String topicName) {
+ Map<String, String> streamConfigMap = new HashMap<>();
+ streamConfigMap.put("streamType", "kafka");
+ streamConfigMap.put("stream.kafka.topic.name", topicName);
+ streamConfigMap.put("stream.kafka.broker.list", "unused:9092");
+ streamConfigMap.put("stream.kafka.consumer.factory.class.name",
MockKafkaConsumerFactory.class.getName());
+ streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, "smallest");
+ streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
+ return new StreamConfig("tableName_REALTIME", streamConfigMap);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Consumer<Bytes, Bytes> mockConsumer(String topicName, int
partitionCount) {
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ List<PartitionInfo> partitionInfos = IntStream.range(0, partitionCount)
+ .mapToObj(partitionId -> new PartitionInfo(topicName, partitionId,
null, null, null))
+ .collect(Collectors.toList());
+ when(consumer.partitionsFor(eq(topicName),
any(Duration.class))).thenReturn(partitionInfos);
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ offsets.put(topicPartition, 1000L + topicPartition.partition());
+ }
+ return offsets;
+ });
+ when(consumer.endOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ offsets.put(topicPartition, 2000L + topicPartition.partition());
+ }
+ return offsets;
+ });
+ return consumer;
+ }
+
+ private static List<Integer> computePartitionIdsWithSizeBasedAlgorithm(
+ List<PartitionGroupConsumptionStatus> currentStatuses, int
partitionCount) {
+ List<Integer> partitionIds = currentStatuses.stream()
+ .map(PartitionGroupConsumptionStatus::getStreamPartitionGroupId)
+ .collect(Collectors.toList());
+ for (int partitionId = currentStatuses.size(); partitionId <
partitionCount; partitionId++) {
+ partitionIds.add(partitionId);
+ }
+ return partitionIds;
+ }
+
+ public static class MockKafkaConsumerFactory extends KafkaConsumerFactory {
+ @Override
+ public StreamMetadataProvider createPartitionMetadataProvider(String
clientId, int partition) {
+ return new MockKafkaStreamMetadataProvider(clientId, _streamConfig,
partition);
+ }
+
+ @Override
+ public StreamMetadataProvider createStreamMetadataProvider(String
clientId) {
+ return new MockKafkaStreamMetadataProvider(clientId, _streamConfig);
+ }
+ }
+
+ private static class MockKafkaStreamMetadataProvider extends
KafkaStreamMetadataProvider {
+ MockKafkaStreamMetadataProvider(String clientId, StreamConfig
streamConfig) {
+ super(clientId, streamConfig);
+ }
+
+ MockKafkaStreamMetadataProvider(String clientId, StreamConfig
streamConfig, int partition) {
+ super(clientId, streamConfig, partition);
+ }
+
+ @Override
+ protected Consumer<Bytes, Bytes> createConsumer(Properties consumerProp) {
+ return MOCK_CONSUMER.get();
+ }
+ }
+}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
index 4a9ff545b67..5605a58c363 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/main/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProvider.java
@@ -123,31 +123,37 @@ public class KafkaStreamMetadataProvider extends
KafkaPartitionLevelConnectionHa
public List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses,
int timeoutMillis)
throws IOException, java.util.concurrent.TimeoutException {
- if (_partitionIdSubset.isEmpty()) {
- return
StreamMetadataProvider.super.computePartitionGroupMetadata(clientId,
streamConfig,
- partitionGroupConsumptionStatuses, timeoutMillis);
- }
Map<Integer, StreamPartitionMsgOffset> partitionIdToEndOffset =
new HashMap<>(partitionGroupConsumptionStatuses.size());
for (PartitionGroupConsumptionStatus s :
partitionGroupConsumptionStatuses) {
partitionIdToEndOffset.put(s.getStreamPartitionGroupId(),
s.getEndOffset());
}
+
+ List<Integer> partitionIds;
+ if (_partitionIdSubset.isEmpty()) {
+ int partitionCount = fetchPartitionCount(timeoutMillis);
+ partitionIds = new ArrayList<>(partitionCount);
+ for (int partitionId = 0; partitionId < partitionCount; partitionId++) {
+ partitionIds.add(partitionId);
+ }
+ } else {
+ partitionIds = _partitionIdSubset;
+ }
+
StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
- List<PartitionGroupMetadata> result = new
ArrayList<>(_partitionIdSubset.size());
- for (Integer partitionId : _partitionIdSubset) {
- StreamPartitionMsgOffset endOffset =
partitionIdToEndOffset.get(partitionId);
- StreamPartitionMsgOffset startOffset;
- if (endOffset == null) {
+ List<PartitionGroupMetadata> result = new ArrayList<>(partitionIds.size());
+ for (Integer partitionId : partitionIds) {
+ if (partitionIdToEndOffset.containsKey(partitionId)) {
+ result.add(new PartitionGroupMetadata(partitionId,
partitionIdToEndOffset.get(partitionId)));
+ } else {
try (StreamMetadataProvider partitionMetadataProvider =
streamConsumerFactory.createPartitionMetadataProvider(
StreamConsumerFactory.getUniqueClientId(clientId),
partitionId)) {
- startOffset = partitionMetadataProvider.fetchStreamPartitionOffset(
+ StreamPartitionMsgOffset startOffset =
partitionMetadataProvider.fetchStreamPartitionOffset(
streamConfig.getOffsetCriteria(), timeoutMillis);
+ result.add(new PartitionGroupMetadata(partitionId, startOffset));
}
- } else {
- startOffset = endOffset;
}
- result.add(new PartitionGroupMetadata(partitionId, startOffset));
}
return result;
}
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaPartitionLevelConsumerTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaPartitionLevelConsumerTest.java
index 3584595f870..d43650cddb6 100644
---
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaPartitionLevelConsumerTest.java
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaPartitionLevelConsumerTest.java
@@ -50,6 +50,7 @@ import org.apache.pinot.spi.stream.MessageBatch;
import org.apache.pinot.spi.stream.OffsetCriteria;
import org.apache.pinot.spi.stream.PartitionGroupConsumer;
import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
+import org.apache.pinot.spi.stream.PartitionGroupMetadata;
import org.apache.pinot.spi.stream.StreamConfig;
import org.apache.pinot.spi.stream.StreamConsumerFactory;
import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider;
@@ -286,6 +287,41 @@ public class KafkaPartitionLevelConsumerTest {
}
}
+ @Test
+ public void testComputePartitionGroupMetadataUsesKafkaPartitionIds()
+ throws Exception {
+ String streamType = "kafka";
+ String streamKafkaBrokerList = _kafkaBrokerAddress;
+ String clientId = "partition-id-client";
+ String tableNameWithType = "tableName_REALTIME";
+
+ Map<String, String> streamConfigMap = new HashMap<>();
+ streamConfigMap.put("streamType", streamType);
+ streamConfigMap.put("stream.kafka.topic.name",
TEST_TOPIC_SUBSET_PARTITION);
+ streamConfigMap.put("stream.kafka.broker.list", streamKafkaBrokerList);
+ streamConfigMap.put("stream.kafka.consumer.factory.class.name",
getKafkaConsumerFactoryName());
+ streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
+ StreamConfig streamConfig = new StreamConfig(tableNameWithType,
streamConfigMap);
+
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(4, 4, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(5, 5, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(6, 6, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(7, 7, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new KafkaStreamMetadataProvider(clientId, streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata(clientId,
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ }
+ }
+
@Test
public void testFetchMessages() {
String streamType = "kafka";
diff --git
a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
new file mode 100644
index 00000000000..f3a34b57fe7
--- /dev/null
+++
b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-4.0/src/test/java/org/apache/pinot/plugin/stream/kafka40/KafkaStreamMetadataProviderTest.java
@@ -0,0 +1,188 @@
+/**
+ * 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.pinot.plugin.stream.kafka40;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Bytes;
+import org.apache.pinot.spi.stream.LongMsgOffset;
+import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus;
+import org.apache.pinot.spi.stream.PartitionGroupMetadata;
+import org.apache.pinot.spi.stream.StreamConfig;
+import org.apache.pinot.spi.stream.StreamConfigProperties;
+import org.apache.pinot.spi.stream.StreamMetadataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
+
+public class KafkaStreamMetadataProviderTest {
+ private static final ThreadLocal<Consumer<Bytes, Bytes>> MOCK_CONSUMER = new
ThreadLocal<>();
+
+ @Test
+ public void testComputePartitionGroupMetadataRecoversMissingLowPartitionId()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(11), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(13), "DONE"),
+ new PartitionGroupConsumptionStatus(4, 4, new LongMsgOffset(0), new
LongMsgOffset(14), "DONE"),
+ new PartitionGroupConsumptionStatus(5, 5, new LongMsgOffset(0), new
LongMsgOffset(15), "DONE"),
+ new PartitionGroupConsumptionStatus(6, 6, new LongMsgOffset(0), new
LongMsgOffset(16), "DONE"),
+ new PartitionGroupConsumptionStatus(7, 7, new LongMsgOffset(0), new
LongMsgOffset(17), "DONE"));
+
+ assertEquals(computePartitionIdsWithSizeBasedAlgorithm(currentStatuses,
8), List.of(0, 1, 3, 4, 5, 6, 7, 7));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata("client",
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+
assertEquals(partitionGroupMetadataList.get(2).getStartOffset().toString(),
"1002");
+
assertEquals(partitionGroupMetadataList.get(3).getStartOffset().toString(),
"13");
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ @Test
+ public void testComputePartitionGroupMetadataHandlesTopicExpansion()
+ throws Exception {
+ String topicName = "asset";
+ Consumer<Bytes, Bytes> consumer = mockConsumer(topicName, 8);
+ MOCK_CONSUMER.set(consumer);
+ try {
+ StreamConfig streamConfig = getStreamConfig(topicName);
+ List<PartitionGroupConsumptionStatus> currentStatuses = List.of(
+ new PartitionGroupConsumptionStatus(0, 0, new LongMsgOffset(0), new
LongMsgOffset(10), "DONE"),
+ new PartitionGroupConsumptionStatus(1, 1, new LongMsgOffset(0), new
LongMsgOffset(11), "DONE"),
+ new PartitionGroupConsumptionStatus(2, 2, new LongMsgOffset(0), new
LongMsgOffset(12), "DONE"),
+ new PartitionGroupConsumptionStatus(3, 3, new LongMsgOffset(0), new
LongMsgOffset(13), "DONE"));
+
+ try (KafkaStreamMetadataProvider streamMetadataProvider =
+ new MockKafkaStreamMetadataProvider("client", streamConfig)) {
+ List<PartitionGroupMetadata> partitionGroupMetadataList =
+ streamMetadataProvider.computePartitionGroupMetadata("client",
streamConfig, currentStatuses, 10000);
+
+
assertEquals(partitionGroupMetadataList.stream().map(PartitionGroupMetadata::getPartitionGroupId)
+ .collect(Collectors.toList()), List.of(0, 1, 2, 3, 4, 5, 6, 7));
+ assertEquals(partitionGroupMetadataList.stream().map(metadata ->
metadata.getStartOffset().toString())
+ .collect(Collectors.toList()), List.of("10", "11", "12", "13",
"1004", "1005", "1006", "1007"));
+ }
+ } finally {
+ MOCK_CONSUMER.remove();
+ }
+ }
+
+ private static StreamConfig getStreamConfig(String topicName) {
+ Map<String, String> streamConfigMap = new HashMap<>();
+ streamConfigMap.put("streamType", "kafka");
+ streamConfigMap.put("stream.kafka.topic.name", topicName);
+ streamConfigMap.put("stream.kafka.broker.list", "unused:9092");
+ streamConfigMap.put("stream.kafka.consumer.factory.class.name",
MockKafkaConsumerFactory.class.getName());
+ streamConfigMap.put("stream.kafka." +
StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA, "smallest");
+ streamConfigMap.put("stream.kafka.decoder.class.name", "decoderClass");
+ return new StreamConfig("tableName_REALTIME", streamConfigMap);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Consumer<Bytes, Bytes> mockConsumer(String topicName, int
partitionCount) {
+ Consumer<Bytes, Bytes> consumer = mock(Consumer.class);
+ List<PartitionInfo> partitionInfos = IntStream.range(0, partitionCount)
+ .mapToObj(partitionId -> new PartitionInfo(topicName, partitionId,
null, null, null))
+ .collect(Collectors.toList());
+ when(consumer.partitionsFor(eq(topicName),
any(Duration.class))).thenReturn(partitionInfos);
+ when(consumer.beginningOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ offsets.put(topicPartition, 1000L + topicPartition.partition());
+ }
+ return offsets;
+ });
+ when(consumer.endOffsets(any(Collection.class),
any(Duration.class))).thenAnswer(invocation -> {
+ Collection<TopicPartition> topicPartitions = invocation.getArgument(0);
+ Map<TopicPartition, Long> offsets = new HashMap<>();
+ for (TopicPartition topicPartition : topicPartitions) {
+ offsets.put(topicPartition, 2000L + topicPartition.partition());
+ }
+ return offsets;
+ });
+ return consumer;
+ }
+
+ private static List<Integer> computePartitionIdsWithSizeBasedAlgorithm(
+ List<PartitionGroupConsumptionStatus> currentStatuses, int
partitionCount) {
+ List<Integer> partitionIds = currentStatuses.stream()
+ .map(PartitionGroupConsumptionStatus::getStreamPartitionGroupId)
+ .collect(Collectors.toList());
+ for (int partitionId = currentStatuses.size(); partitionId <
partitionCount; partitionId++) {
+ partitionIds.add(partitionId);
+ }
+ return partitionIds;
+ }
+
+ public static class MockKafkaConsumerFactory extends KafkaConsumerFactory {
+ @Override
+ public StreamMetadataProvider createPartitionMetadataProvider(String
clientId, int partition) {
+ return new MockKafkaStreamMetadataProvider(clientId, _streamConfig,
partition);
+ }
+
+ @Override
+ public StreamMetadataProvider createStreamMetadataProvider(String
clientId) {
+ return new MockKafkaStreamMetadataProvider(clientId, _streamConfig);
+ }
+ }
+
+ private static class MockKafkaStreamMetadataProvider extends
KafkaStreamMetadataProvider {
+ MockKafkaStreamMetadataProvider(String clientId, StreamConfig
streamConfig) {
+ super(clientId, streamConfig);
+ }
+
+ MockKafkaStreamMetadataProvider(String clientId, StreamConfig
streamConfig, int partition) {
+ super(clientId, streamConfig, partition);
+ }
+
+ @Override
+ protected Consumer<Bytes, Bytes> createConsumer(Properties consumerProp) {
+ return MOCK_CONSUMER.get();
+ }
+ }
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMetadataProvider.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMetadataProvider.java
index 4b5666a7a35..6a742771a21 100644
---
a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMetadataProvider.java
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMetadataProvider.java
@@ -20,7 +20,6 @@ package org.apache.pinot.spi.stream;
import java.io.Closeable;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
@@ -78,60 +77,27 @@ public interface StreamMetadataProvider extends Closeable {
StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis)
throws TimeoutException;
- /**
- * Computes the list of {@link PartitionGroupMetadata} for the latest state
of the stream, using the current
- * {@link PartitionGroupConsumptionStatus}
- *
- * Default behavior is the one for the Kafka stream, where each partition
group contains only one partition
- * @param partitionGroupConsumptionStatuses list of {@link
PartitionGroupConsumptionStatus} for current partition
- * groups
- */
- default List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
+ /// Computes the list of [PartitionGroupMetadata] for the latest state of
the stream, using the current
+ /// [PartitionGroupConsumptionStatus].
+ ///
+ /// Implementations must define how stream partitions map to Pinot partition
groups.
+ ///
+ /// @param partitionGroupConsumptionStatuses list of
[PartitionGroupConsumptionStatus] for current partition groups
+ List<PartitionGroupMetadata> computePartitionGroupMetadata(String clientId,
StreamConfig streamConfig,
List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses,
int timeoutMillis)
- throws IOException, TimeoutException {
- int partitionCount = fetchPartitionCount(timeoutMillis);
- List<PartitionGroupMetadata> newPartitionGroupMetadataList = new
ArrayList<>(partitionCount);
-
- // Add a PartitionGroupMetadata into the list, foreach partition already
present in current.
- // Setting endOffset (exclusive) as the startOffset for new partition
group.
- // If partition group is still in progress, this value will be null
- for (PartitionGroupConsumptionStatus
currentPartitionGroupConsumptionStatus : partitionGroupConsumptionStatuses) {
- newPartitionGroupMetadataList.add(
- new
PartitionGroupMetadata(currentPartitionGroupConsumptionStatus.getStreamPartitionGroupId(),
- currentPartitionGroupConsumptionStatus.getEndOffset()));
- }
- // Add PartitionGroupMetadata for new partitions
- // Use offset criteria from stream config
- StreamConsumerFactory streamConsumerFactory =
StreamConsumerFactoryProvider.create(streamConfig);
- for (int i = partitionGroupConsumptionStatuses.size(); i < partitionCount;
i++) {
- try (StreamMetadataProvider partitionMetadataProvider =
streamConsumerFactory.createPartitionMetadataProvider(
- StreamConsumerFactory.getUniqueClientId(clientId), i)) {
- StreamPartitionMsgOffset streamPartitionMsgOffset =
-
partitionMetadataProvider.fetchStreamPartitionOffset(streamConfig.getOffsetCriteria(),
timeoutMillis);
- newPartitionGroupMetadataList.add(new PartitionGroupMetadata(i,
streamPartitionMsgOffset));
- }
- }
- return newPartitionGroupMetadataList;
- }
+ throws IOException, TimeoutException;
- /**
- * @param forceGetOffsetFromStream - the flag is a workaround to not use
partitionGroupConsumptionStatuses.
- * This is required because
PinotLLCRealtimeSegmentManager.selectStartOffset()
- * actually requires the offsets from the
stream, but was originally relying on
- * passing an empty
partitionGroupConsumptionStatuses to the method.
- * The change for <a
href="https://github.com/apache/pinot/issues/15608";>...</a>
- * required to pass the actual
partitionGroupConsumptionStatuses
- * TODO - Remove the flag and fix the
clients calling computePartitionGroupMetadata()
- */
+ /// Computes partition group metadata with the option to force offsets to be
fetched from the stream instead of
+ /// derived from [PartitionGroupConsumptionStatus].
+ ///
+ /// @param forceGetOffsetFromStream whether the implementation should ignore
current consumption statuses when
+ /// selecting start offsets
default List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
List<PartitionGroupConsumptionStatus> partitionGroupConsumptionStatuses,
int timeoutMillis,
boolean forceGetOffsetFromStream)
throws IOException, TimeoutException {
- if (forceGetOffsetFromStream) {
- return computePartitionGroupMetadata(clientId, streamConfig,
Collections.emptyList(), timeoutMillis);
- } else {
- return computePartitionGroupMetadata(clientId, streamConfig,
partitionGroupConsumptionStatuses, timeoutMillis);
- }
+ return computePartitionGroupMetadata(clientId, streamConfig,
+ forceGetOffsetFromStream ? Collections.emptyList() :
partitionGroupConsumptionStatuses, timeoutMillis);
}
default Map<String, PartitionLagState> getCurrentPartitionLagState(
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java
index 3a175ef7fcd..4cab96d169c 100644
---
a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java
+++
b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java
@@ -500,103 +500,4 @@ public class PartitionGroupMetadataFetcherTest {
when(streamConfig.getTableNameWithType()).thenReturn(tableName);
return streamConfig;
}
-
- private static final class DefaultComputeOnlyMetadataProvider implements
StreamMetadataProvider {
- private int _fetchPartitionCountCalls;
-
- @Override
- public int fetchPartitionCount(long timeoutMillis) {
- _fetchPartitionCountCalls++;
- return 1;
- }
-
- @Override
- public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
- throw new UnsupportedOperationException("Should not be called");
- }
-
- @Override
- public boolean supportsOffsetLag() {
- return false;
- }
-
- @Override
- public void close() {
- }
-
- int getFetchPartitionCountCalls() {
- return _fetchPartitionCountCalls;
- }
- }
-
- private static final class OverriddenComputeMetadataProvider implements
StreamMetadataProvider {
- private int _fetchPartitionCountCalls;
- private final StreamPartitionMsgOffset _offset =
mock(StreamPartitionMsgOffset.class);
-
- @Override
- public int fetchPartitionCount(long timeoutMillis) {
- _fetchPartitionCountCalls++;
- return 3;
- }
-
- @Override
- public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
- throw new UnsupportedOperationException("Should not be called");
- }
-
- @Override
- public List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
- List<PartitionGroupConsumptionStatus>
partitionGroupConsumptionStatuses, int timeoutMillis,
- boolean forceGetOffsetFromStream) {
- return Collections.singletonList(new PartitionGroupMetadata(0, _offset));
- }
-
- @Override
- public boolean supportsOffsetLag() {
- return false;
- }
-
- @Override
- public void close() {
- }
-
- int getFetchPartitionCountCalls() {
- return _fetchPartitionCountCalls;
- }
- }
-
- private static final class OverriddenFourArgComputeMetadataProvider
implements StreamMetadataProvider {
- private int _fetchPartitionCountCalls;
- private final StreamPartitionMsgOffset _offset =
mock(StreamPartitionMsgOffset.class);
-
- @Override
- public int fetchPartitionCount(long timeoutMillis) {
- _fetchPartitionCountCalls++;
- return 3;
- }
-
- @Override
- public StreamPartitionMsgOffset fetchStreamPartitionOffset(OffsetCriteria
offsetCriteria, long timeoutMillis) {
- throw new UnsupportedOperationException("Should not be called");
- }
-
- @Override
- public List<PartitionGroupMetadata> computePartitionGroupMetadata(String
clientId, StreamConfig streamConfig,
- List<PartitionGroupConsumptionStatus>
partitionGroupConsumptionStatuses, int timeoutMillis) {
- return Collections.singletonList(new PartitionGroupMetadata(0, _offset));
- }
-
- @Override
- public boolean supportsOffsetLag() {
- return false;
- }
-
- @Override
- public void close() {
- }
-
- int getFetchPartitionCountCalls() {
- return _fetchPartitionCountCalls;
- }
- }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
