This is an automated email from the ASF dual-hosted git repository.
FrankChen021 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 3469a0fdaa4 feat: support ingesting Kafka messages from selected
partitions (#20474)
3469a0fdaa4 is described below
commit 3469a0fdaa4cb3857e00f51fbc7d490c3305d99f
Author: Frank Chen <[email protected]>
AuthorDate: Sun Oct 11 15:03:16 2026 +0800
feat: support ingesting Kafka messages from selected partitions (#20474)
* feat: support ingesting Kafka messages from selected partitions
* fix(web-console): validate Kafka partition IDs against the server's
accepted range and string forms
* address review: keep partition ID field visible, stop coercing IDs, drop
unused constructors
* address review: filter stored offsets in getOffsetsFromMetadataStorage
and drop the extra hook
* address review: report lag only for partitions a bounded supervisor tracks
---
docs/ingestion/kafka-ingestion.md | 41 +++
.../druid/indexing/kafka/KafkaIndexTask.java | 3 +-
.../druid/indexing/kafka/KafkaRecordSupplier.java | 27 +-
.../druid/indexing/kafka/KafkaSamplerSpec.java | 3 +-
.../kafka/supervisor/KafkaIOConfigBuilder.java | 12 +-
.../indexing/kafka/supervisor/KafkaSupervisor.java | 127 +++++++--
.../kafka/supervisor/KafkaSupervisorIOConfig.java | 34 ++-
.../kafka/supervisor/KafkaSupervisorSpec.java | 3 +-
.../kafka/KafkaRecordSupplierHeaderFilterTest.java | 18 +-
.../indexing/kafka/KafkaRecordSupplierTest.java | 47 +++
.../druid/indexing/kafka/KafkaSamplerSpecTest.java | 25 ++
.../simulate/EmbeddedKafkaSupervisorTest.java | 52 ++++
.../supervisor/KafkaSupervisorIOConfigTest.java | 168 ++++++++++-
.../kafka/supervisor/KafkaSupervisorSpecTest.java | 30 ++
.../kafka/supervisor/KafkaSupervisorTest.java | 314 ++++++++++++++++++++-
.../supervisor/SeekableStreamSupervisor.java | 28 ++
.../druid/jackson/StrictIntegerDeserializer.java | 59 ++++
.../jackson/StrictIntegerDeserializerTest.java | 89 ++++++
.../ingestion-spec/ingestion-spec.spec.ts | 62 ++++
.../druid-models/ingestion-spec/ingestion-spec.tsx | 38 +++
20 files changed, 1138 insertions(+), 42 deletions(-)
diff --git a/docs/ingestion/kafka-ingestion.md
b/docs/ingestion/kafka-ingestion.md
index 91f1706f134..3b07df3495b 100644
--- a/docs/ingestion/kafka-ingestion.md
+++ b/docs/ingestion/kafka-ingestion.md
@@ -121,6 +121,7 @@ For configuration properties shared across all streaming
ingestion methods, refe
|Property|Type|Description|Required|Default|
|--------|----|-----------|--------|-------|
|`topic`|String|The Kafka topic to read from. Note that once this value is
established for a supervisor, updating it is not supported. To ingest data from
multiple topic, use `topicPattern`. |Yes if `topicPattern` isn't set.||
+|`partitionIds`|Array of integers|Optional nonempty set of nonnegative
partition IDs for a single `topic`. Omit or set to null to read all partitions.
Cannot be combined with `topicPattern` or `boundedStreamConfig`. See [Ingest
selected partitions](#ingest-selected-partitions).|No|null|
|`topicPattern`|String|Multiple Kafka topics to read from, passed as a regex
pattern. See [Ingest from multiple topics](#ingest-from-multiple-topics) for
more information.|Yes if `topic` isn't set.||
|`consumerProperties`|String, Object|A map of properties to pass to the Kafka
consumer. See [Consumer properties](#consumer-properties) for details.|Yes. At
the minimum, you must set the `bootstrap.servers` property to establish the
initial connection to the Kafka cluster.||
|`pollTimeout`|Long|The length of time to wait for the Kafka consumer to poll
records, in milliseconds.|No|100|
@@ -600,6 +601,46 @@ For configuration properties shared across all streaming
ingestion methods, refe
|`numPersistThreads`|Integer|The number of threads to use to create and
persist incremental segments on the disk. Higher ingestion data throughput
results in a larger number of incremental segments, causing significant CPU
time to be spent on the creation of the incremental segments on the disk. For
datasources with number of columns running into hundreds or thousands, creation
of the incremental segments may take up significant time, in the order of
multiple seconds. In both of these sc [...]
|`streamingPartitionsSpec`|Object|Configures query-time segment pruning for
streaming-ingested segments. A typed object with an optional `type` (defaults
to `dim_value_set`), `partitionDimensions` (List of String), the dimensions
whose observed values each segment records so the broker can skip segments that
can't match a query filter, and an optional `maxValuesPerDimension` (Integer)
cap on the distinct values recorded per dimension per segment. See [Streaming
partitions spec](#streamin [...]
+## Ingest selected partitions
+
+Set `partitionIds` in the supervisor's `ioConfig` to consume only specific
partitions:
+
+```json
+{
+ "type": "kafka",
+ "topic": "events",
+ "partitionIds": [0, 2, 5],
+ "consumerProperties": {"bootstrap.servers": "localhost:9092"},
+ "useEarliestOffset": true
+}
+```
+
+This fragment can be used to investigate production data in a separate
diagnostic supervisor and datasource without ingesting the entire topic again.
+The diagnostic supervisor maintains its own offsets.
+Selected partitions are not necessarily a representative sample, and selecting
partitions does not limit their throughput.
+Additional ingestion still consumes Kafka and Druid resources.
+
+Druid validates IDs against Kafka `partitionsFor(topic)` metadata during
discovery, including discovery for sampling and monitoring.
+Spec parsing does not contact Kafka.
+Missing IDs produce a supervisor error and are retried; a new supervisor does
not start with only part of its selection.
+Empty lists, negative IDs, and null elements are invalid.
+Duplicate IDs are normalized and list order does not matter.
+
+New topic partitions outside the selection are ignored.
+Omit `partitionIds` or set it to null to restore discovery of all partitions.
+Selected partitions are assigned to task groups by their position in the
sorted selection: `[0, 3, 6]` with three tasks uses three groups.
+
+Submitting a supervisor spec that changes or removes `partitionIds` requests
managed readers to finish and publish, then restarts the supervisor with the
updated selection.
+Supervisors with an explicit selection also replace incompatible adopted
readers on restart.
+Narrowing `partitionIds` does not delete previously ingested rows from
partitions removed from the selection.
+Already publishing tasks may finish; the update does not provide an
instantaneous cutoff.
+Retained and re-added partitions resume from saved offsets when available.
+Otherwise, the existing `useEarliestOffset` behavior applies.
+
+Scoped reset requests reject IDs outside the selection.
+If an excluded partition's saved offset has expired and automatic offset reset
is disabled, first re-add the partition, then issue a scoped reset after the
unavailable-offset error, or set an explicitly chosen available offset using
`resetOffsets`.
+A full supervisor reset still clears all saved offsets, including those of
excluded partitions.
+
## Deployment notes on Kafka partitions and Druid segments
Druid assigns Kafka partitions to each Kafka indexing task. A task writes the
events it consumes from Kafka into a single segment for the segment granularity
interval until it reaches one of the following limits: `maxRowsPerSegment`,
`maxTotalRows`, or `intermediateHandoffPeriod`. At this point, the task creates
a new partition for this segment granularity to contain subsequent events.
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java
index be50b9554a1..e15753216d0 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaIndexTask.java
@@ -122,7 +122,8 @@ public class KafkaIndexTask extends
SeekableStreamIndexTask<KafkaTopicPartition,
kafkaIndexTaskIOConfig.getConfigOverrides(),
kafkaIndexTaskIOConfig.isMultiTopic(),
this::getMetricBuilder,
- kafkaIndexTaskIOConfig.getHeaderBasedFilterConfig()
+ kafkaIndexTaskIOConfig.getHeaderBasedFilterConfig(),
+ null
);
if (toolbox.getMonitorScheduler() != null) {
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java
index 2c1ea896af6..c20aa81f2c7 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaRecordSupplier.java
@@ -60,6 +60,7 @@ import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
+import java.util.TreeSet;
import java.util.concurrent.Callable;
import java.util.function.Supplier;
import java.util.regex.Pattern;
@@ -83,6 +84,8 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
private boolean closed;
private final boolean multiTopic;
+ @Nullable
+ private final Set<Integer> partitionIds;
@Nullable
private final KafkaHeaderBasedFilterEvaluator headerFilterEvaluator;
@@ -101,7 +104,7 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
@Nullable Supplier<ServiceMetricEvent.Builder> metricBuilderSupplier
)
{
- this(consumerProperties, sortingMapper, configOverrides, multiTopic,
metricBuilderSupplier, null);
+ this(consumerProperties, sortingMapper, configOverrides, multiTopic,
metricBuilderSupplier, null, null);
}
public KafkaRecordSupplier(
@@ -110,14 +113,16 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
KafkaConfigOverrides configOverrides,
boolean multiTopic,
@Nullable Supplier<ServiceMetricEvent.Builder> metricBuilderSupplier,
- @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig
+ @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig,
+ @Nullable Set<Integer> partitionIds
)
{
this(
getKafkaConsumer(sortingMapper, consumerProperties, configOverrides),
multiTopic,
metricBuilderSupplier,
- headerBasedFilterConfig
+ headerBasedFilterConfig,
+ partitionIds
);
}
@@ -128,7 +133,7 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
@Nullable Supplier<ServiceMetricEvent.Builder> metricBuilderSupplier
)
{
- this(consumer, multiTopic, metricBuilderSupplier, null);
+ this(consumer, multiTopic, metricBuilderSupplier, null, null);
}
@VisibleForTesting
@@ -136,11 +141,13 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
KafkaConsumer<byte[], byte[]> consumer,
boolean multiTopic,
@Nullable Supplier<ServiceMetricEvent.Builder> metricBuilderSupplier,
- @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig
+ @Nullable KafkaHeaderBasedFilterConfig headerBasedFilterConfig,
+ @Nullable Set<Integer> partitionIds
)
{
this.consumer = consumer;
this.multiTopic = multiTopic;
+ this.partitionIds = partitionIds == null ? null : Set.copyOf(partitionIds);
this.monitor = new KafkaConsumerMonitor(consumer, metricBuilderSupplier);
this.headerFilterEvaluator = headerBasedFilterConfig != null
? new KafkaHeaderBasedFilterEvaluator(headerBasedFilterConfig) : null;
@@ -327,7 +334,17 @@ public class KafkaRecordSupplier implements
RecordSupplier<KafkaTopicPartition,
+ " Check that the topic exists in Kafka
cluster", stream);
}
}
+
+ if (partitionIds != null) {
+ final Set<Integer> missing = new TreeSet<>(partitionIds);
+ allPartitions.forEach(partition ->
missing.remove(partition.partition()));
+ if (!missing.isEmpty()) {
+ throw InvalidInput.exception("Requested partition IDs [%s] do not
exist in topic [%s]", missing, stream);
+ }
+ }
return allPartitions.stream()
+ // Keep all discovered partitions unless a fixed
selection is configured.
+ .filter(p -> partitionIds == null ||
partitionIds.contains(p.partition()))
.map(p -> new KafkaTopicPartition(multiTopic,
p.topic(), p.partition()))
.collect(Collectors.toSet());
});
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java
index 26164a1df7c..751a2144a07 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/KafkaSamplerSpec.java
@@ -76,7 +76,8 @@ public class KafkaSamplerSpec extends
SeekableStreamSamplerSpec
kafkaSupervisorIOConfig.isMultiTopic(),
null,
// Apply the same header-based filter during sampling so sampled
records match what ingestion tasks keep.
- kafkaSupervisorIOConfig.getHeaderBasedFilterConfig()
+ kafkaSupervisorIOConfig.getHeaderBasedFilterConfig(),
+ kafkaSupervisorIOConfig.getPartitionIds()
);
}
finally {
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java
index 60990e3a9fd..852fa358a31 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaIOConfigBuilder.java
@@ -25,6 +25,7 @@ import
org.apache.druid.indexing.seekablestream.extension.KafkaConfigOverrides;
import
org.apache.druid.indexing.seekablestream.supervisor.SupervisorIOConfigBuilder;
import java.util.Map;
+import java.util.Set;
/**
* Builder for {@link KafkaSupervisorIOConfig}.
@@ -33,6 +34,7 @@ public class KafkaIOConfigBuilder extends
SupervisorIOConfigBuilder<KafkaIOConfi
{
private String topic;
private String topicPattern;
+ private Set<Integer> partitionIds;
private Map<String, Object> consumerProperties;
private Long pollTimeout;
private KafkaConfigOverrides configOverrides;
@@ -45,6 +47,12 @@ public class KafkaIOConfigBuilder extends
SupervisorIOConfigBuilder<KafkaIOConfi
return this;
}
+ public KafkaIOConfigBuilder withPartitionIds(Set<Integer> partitionIds)
+ {
+ this.partitionIds = partitionIds;
+ return this;
+ }
+
public KafkaIOConfigBuilder withTopicPattern(String topicPattern)
{
this.topicPattern = topicPattern;
@@ -105,6 +113,7 @@ public class KafkaIOConfigBuilder extends
SupervisorIOConfigBuilder<KafkaIOConfi
copyFromBase(io);
this.topic = io.getTopic();
this.topicPattern = io.getTopicPattern();
+ this.partitionIds = io.getPartitionIds();
this.consumerProperties = io.getConsumerProperties();
this.pollTimeout = io.getPollTimeout();
this.configOverrides = io.getConfigOverrides();
@@ -140,7 +149,8 @@ public class KafkaIOConfigBuilder extends
SupervisorIOConfigBuilder<KafkaIOConfi
stopTaskCount,
emitTimeLagMetrics,
serverPriorityToReplicas,
- boundedStreamConfig
+ boundedStreamConfig,
+ partitionIds
);
}
}
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java
index 7e019266e34..ef9b87573a5 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisor.java
@@ -29,6 +29,7 @@ import org.apache.druid.common.utils.IdUtils;
import org.apache.druid.data.input.kafka.KafkaRecordEntity;
import org.apache.druid.data.input.kafka.KafkaTopicPartition;
import org.apache.druid.error.DruidException;
+import org.apache.druid.error.InvalidInput;
import org.apache.druid.indexing.common.task.Task;
import org.apache.druid.indexing.common.task.TaskResource;
import org.apache.druid.indexing.kafka.KafkaDataSourceMetadata;
@@ -97,6 +98,12 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
private static final Long END_OF_PARTITION = Long.MAX_VALUE;
private final Pattern pattern;
+
+ /// Position of each selected partition ID in the sorted `partitionIds`,
e.g. `[0, 3, 6]` gives `{0=0, 3=1, 6=2}`.
+ /// Used to spread selected partitions evenly across task groups (`rank %
taskCount`), where
+ /// `partition % taskCount` could put them all in one group. Empty when no
selection is configured.
+ private final Map<Integer, Integer> selectedPartitionRanks;
+
private volatile Map<KafkaTopicPartition, Long> partitionToTimeLag;
private final KafkaSupervisorSpec spec;
@@ -128,6 +135,18 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
this.spec = spec;
this.pattern = getIoConfig().isMultiTopic() ?
Pattern.compile(getIoConfig().getStream()) : null;
+ final Set<Integer> selected = getIoConfig().getPartitionIds();
+ if (selected == null) {
+ this.selectedPartitionRanks = Map.of();
+ } else {
+ final List<Integer> sorted = new ArrayList<>(selected);
+ Collections.sort(sorted);
+ final Map<Integer, Integer> ranks = new HashMap<>();
+ for (int i = 0; i < sorted.size(); i++) {
+ ranks.put(sorted.get(i), i);
+ }
+ this.selectedPartitionRanks = Map.copyOf(ranks);
+ }
}
@Override
@@ -140,7 +159,8 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
ioConfig.getConfigOverrides(),
ioConfig.isMultiTopic(),
null,
- ioConfig.getHeaderBasedFilterConfig()
+ ioConfig.getHeaderBasedFilterConfig(),
+ ioConfig.getPartitionIds()
);
}
@@ -151,10 +171,55 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
if (partitionId.isMultiTopicPartition()) {
return Math.abs(31 * partitionId.topic().hashCode() +
partitionId.partition()) % taskCount;
} else {
+ final Integer rank = selectedPartitionRanks.get(partitionId.partition());
+ if (rank != null) {
+ return rank % taskCount;
+ }
+ // An adopted task can still refer to a partition excluded by this
supervisor.
return partitionId.partition() % taskCount;
}
}
+ @Override
+ protected boolean isTaskPartitionSetCurrent(int taskGroupId,
Set<KafkaTopicPartition> taskPartitions)
+ {
+ // Removing a selection is handled by the spec-update restart, which
requests existing readers to publish.
+ // All-partition supervisors must still adopt readers that predate
ordinary topic partition growth.
+ return selectedPartitionRanks.isEmpty() ||
taskPartitions.equals(partitionGroups.get(taskGroupId));
+ }
+
+ /**
+ * With {@code partitionIds} set, rejects resets naming a different topic or
any partition outside the selection.
+ * A full reset (null metadata) is always allowed.
+ */
+ @Override
+ protected void validatePartitionReset(@Nullable DataSourceMetadata metadata)
+ {
+ final Set<Integer> selected = getIoConfig().getPartitionIds();
+ if (selected == null || metadata == null) {
+ return;
+ }
+
+ if (!(metadata instanceof KafkaDataSourceMetadata kafkaMetadata)) {
+ throw InvalidInput.exception("Partition reset requires Kafka metadata
for supervisor [%s]", spec.getId());
+ }
+
+ final SeekableStreamSequenceNumbers<KafkaTopicPartition, Long> sequences =
kafkaMetadata.getSeekableStreamSequenceNumbers();
+ if (sequences == null ||
!getIoConfig().getStream().equals(sequences.getStream())) {
+ throw InvalidInput.exception("Partition reset must specify offsets for
topic [%s]", getIoConfig().getStream());
+ }
+
+ final Set<KafkaTopicPartition> excluded =
sequences.getPartitionSequenceNumberMap().keySet().stream()
+ .filter(p ->
p.isMultiTopicPartition() || !selected.contains(p.partition()))
+
.collect(Collectors.toSet());
+ if (!excluded.isEmpty()) {
+ throw InvalidInput.exception(
+ "Cannot reset excluded or topic-qualified partitions [%s] for
supervisor [%s] with partition IDs [%s]",
+ excluded, spec.getId(), selected
+ );
+ }
+ }
+
@Override
protected boolean checkSourceMetadataMatch(DataSourceMetadata metadata)
{
@@ -270,6 +335,19 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
return taskList;
}
+ /** Keeps offsets for the configured partitions only. */
+ @Override
+ protected Map<KafkaTopicPartition, Long> getHighestCurrentOffsets()
+ {
+ final Set<Integer> selected = getIoConfig().getPartitionIds();
+ if (selected == null) {
+ return super.getHighestCurrentOffsets();
+ }
+ final Map<KafkaTopicPartition, Long> offsets = new
HashMap<>(super.getHighestCurrentOffsets());
+ offsets.keySet().removeIf(partition ->
!selected.contains(partition.partition()));
+ return offsets;
+ }
+
@Override
protected Map<KafkaTopicPartition, Long> getPartitionRecordLag()
{
@@ -443,6 +521,25 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
return aggregatePartitionLags(partitionRecordLag);
}
+ /**
+ * Partitions to report lag for. A bounded supervisor, such as a backfill,
only assigns the partitions in its bounded
+ * config, which can be a subset of the topic. Reading positions of any
other partition would fail.
+ * Must be called while holding the record supplier lock.
+ */
+ private Set<KafkaTopicPartition> getPartitionsForLag()
+ {
+ if (getIoConfig().isBounded()) {
+ return new HashSet<>(partitionIds);
+ }
+ try {
+ return recordSupplier.getPartitionIds(getIoConfig().getStream());
+ }
+ catch (Exception e) {
+ log.warn("Could not fetch partitions for topic/stream [%s]",
getIoConfig().getStream());
+ throw new StreamException(e);
+ }
+ }
+
/**
* This method is similar to updatePartitionLagFromStream
* but also determines time lag. Once this method has been
@@ -454,14 +551,7 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
getRecordSupplierLock().lock();
try {
- Set<KafkaTopicPartition> partitionIds;
- try {
- partitionIds =
recordSupplier.getPartitionIds(getIoConfig().getStream());
- }
- catch (Exception e) {
- log.warn("Could not fetch partitions for topic/stream [%s]",
getIoConfig().getStream());
- throw new StreamException(e);
- }
+ final Set<KafkaTopicPartition> partitionIds = getPartitionsForLag();
final Set<StreamPartition<KafkaTopicPartition>> partitions = partitionIds
.stream()
@@ -563,14 +653,7 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
getRecordSupplierLock().lock();
try {
- Set<KafkaTopicPartition> partitionIds;
- try {
- partitionIds =
recordSupplier.getPartitionIds(getIoConfig().getStream());
- }
- catch (Exception e) {
- log.warn("Could not fetch partitions for topic/stream [%s]",
getIoConfig().getStream());
- throw new StreamException(e);
- }
+ final Set<KafkaTopicPartition> partitionIds = getPartitionsForLag();
Set<StreamPartition<KafkaTopicPartition>> partitions = partitionIds
.stream()
@@ -632,7 +715,8 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
/**
* Gets the offsets as stored in the metadata store. The map returned will
only contain
- * offsets from topic partitions that match the current supervisor config
stream. This
+ * offsets from topic partitions that match the current supervisor config
stream, and from the partitions
+ * selected by {@code partitionIds} if it is set. This
* override is needed because in the case of multi-topic, a user could have
updated the supervisor
* config from single topic to multi-topic, where the new multi-topic
pattern regex matches the
* old config single topic. Without this override, the previously stored
metadata for the single
@@ -678,6 +762,13 @@ public class KafkaSupervisor extends
SeekableStreamSupervisor<KafkaTopicPartitio
partitionOffsets.put(matchingTopicPartition, value);
}
});
+
+ // Keep offsets for the configured partitions only. Offsets of
excluded partitions stay in metadata storage.
+ final Set<Integer> selected = getIoConfig().getPartitionIds();
+ if (selected != null) {
+ partitionOffsets.keySet().removeIf(partition ->
!selected.contains(partition.partition()));
+ }
+
return partitionOffsets;
}
}
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java
index bb818fb5d4b..3ad41d0a575 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfig.java
@@ -20,8 +20,11 @@
package org.apache.druid.indexing.kafka.supervisor;
import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSortedSet;
import org.apache.druid.common.config.Configs;
import org.apache.druid.data.input.InputFormat;
import org.apache.druid.error.InvalidInput;
@@ -31,6 +34,7 @@ import
org.apache.druid.indexing.seekablestream.supervisor.IdleConfig;
import org.apache.druid.indexing.seekablestream.supervisor.LagAggregator;
import
org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisorIOConfig;
import
org.apache.druid.indexing.seekablestream.supervisor.autoscaler.AutoScalerConfig;
+import org.apache.druid.jackson.StrictIntegerDeserializer;
import org.apache.druid.java.util.common.StringUtils;
import org.joda.time.DateTime;
import org.joda.time.Period;
@@ -38,6 +42,7 @@ import org.joda.time.Period;
import javax.annotation.Nullable;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
public class KafkaSupervisorIOConfig extends SeekableStreamSupervisorIOConfig
{
@@ -55,6 +60,8 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
private final KafkaConfigOverrides configOverrides;
private final String topic;
private final String topicPattern;
+ @Nullable
+ private final ImmutableSortedSet<Integer> partitionIds;
private final boolean emitTimeLagMetrics;
private final KafkaHeaderBasedFilterConfig headerBasedFilterConfig;
@@ -83,7 +90,9 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
@JsonProperty("stopTaskCount") Integer stopTaskCount,
@Nullable @JsonProperty("emitTimeLagMetrics") Boolean emitTimeLagMetrics,
@Nullable @JsonProperty("serverPriorityToReplicas") Map<Integer,
Integer> serverPriorityToReplicas,
- @Nullable @JsonProperty("boundedStreamConfig") BoundedStreamConfig
boundedStreamConfig
+ @Nullable @JsonProperty("boundedStreamConfig") BoundedStreamConfig
boundedStreamConfig,
+ @Nullable @JsonProperty("partitionIds")
+ @JsonDeserialize(contentUsing = StrictIntegerDeserializer.class)
Set<Integer> partitionIds
)
{
super(
@@ -117,6 +126,15 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
this.configOverrides = configOverrides;
this.topic = topic;
this.topicPattern = topicPattern;
+ if (partitionIds != null) {
+ if (partitionIds.isEmpty() || partitionIds.stream().anyMatch(id -> id ==
null || id < 0)) {
+ throw InvalidInput.exception("partitionIds must contain nonnegative
partition IDs and must not be empty");
+ }
+ if (topicPattern != null || boundedStreamConfig != null) {
+ throw InvalidInput.exception("partitionIds requires a single topic and
cannot be combined with boundedStreamConfig");
+ }
+ }
+ this.partitionIds = partitionIds == null ? null :
ImmutableSortedSet.copyOf(partitionIds);
this.emitTimeLagMetrics = Configs.valueOrDefault(emitTimeLagMetrics,
false);
}
@@ -174,7 +192,8 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
stopTaskCount,
emitTimeLagMetrics,
serverPriorityToReplicas,
- boundedStreamConfig
+ boundedStreamConfig,
+ null
);
}
@@ -196,6 +215,14 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
return topicPattern;
}
+ @Nullable
+ @JsonProperty
+ @JsonInclude(JsonInclude.Include.NON_NULL)
+ public Set<Integer> getPartitionIds()
+ {
+ return partitionIds;
+ }
+
@JsonProperty
public Map<String, Object> getConsumerProperties()
{
@@ -248,6 +275,7 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
return "KafkaSupervisorIOConfig{" +
"topic='" + getTopic() + '\'' +
"topicPattern='" + getTopicPattern() + '\'' +
+ ", partitionIds=" + partitionIds +
", replicas=" + getReplicas() +
", taskCount=" + getTaskCount() +
", taskDuration=" + getTaskDuration() +
@@ -299,6 +327,7 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
&& Objects.equals(configOverrides, that.configOverrides)
&& Objects.equals(topic, that.topic)
&& Objects.equals(topicPattern, that.topicPattern)
+ && Objects.equals(partitionIds, that.partitionIds)
&& Objects.equals(headerBasedFilterConfig,
that.headerBasedFilterConfig);
}
@@ -312,6 +341,7 @@ public class KafkaSupervisorIOConfig extends
SeekableStreamSupervisorIOConfig
configOverrides,
topic,
topicPattern,
+ partitionIds,
emitTimeLagMetrics,
headerBasedFilterConfig
);
diff --git
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java
index b1d3721555b..e755a8f57c2 100644
---
a/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java
+++
b/extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpec.java
@@ -206,7 +206,8 @@ public class KafkaSupervisorSpec extends
SeekableStreamSupervisorSpec
ioConfig.getStopTaskCount(),
ioConfig.isEmitTimeLagMetrics(),
ioConfig.getServerPriorityToReplicas(),
- boundedStreamConfig
+ boundedStreamConfig,
+ null
);
return new KafkaSupervisorSpec(
backfillId,
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java
index 9bd8fdbb7b5..674e64a3400 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierHeaderFilterTest.java
@@ -99,7 +99,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
ConsumerRecord<byte[], byte[]> prodRecord = createRecord("topic", 0, 100L,
headers("environment", "production"));
@@ -141,7 +141,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
ConsumerRecord<byte[], byte[]> prodRecord = createRecord("topic", 0, 100L,
headers("environment", "production"));
@@ -176,7 +176,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("service",
Arrays.asList("user-service", "payment-service"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
ConsumerRecord<byte[], byte[]> userServiceRecord = createRecord("topic",
0, 100L,
headers("service", "user-service"));
@@ -218,7 +218,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter serviceFilter = new InDimFilter("service",
Arrays.asList("user-service", "payment-service"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(serviceFilter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
ConsumerRecord<byte[], byte[]> userServiceRecord = createRecord("topic",
0, 100L,
headers("service", "user-service"));
@@ -260,7 +260,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
// First poll
ConsumerRecord<byte[], byte[]> prodRecord1 = createRecord("topic", 0, 100L,
@@ -305,7 +305,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
EasyMock.expect(mockConsumer.poll(EasyMock.anyObject(Duration.class)))
.andReturn(createConsumerRecords(Collections.emptyList()));
@@ -326,7 +326,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
// All records have "staging" environment - none should pass the
"production" filter
ConsumerRecord<byte[], byte[]> stagingRecord1 = createRecord("topic", 0,
100L,
@@ -367,7 +367,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter);
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, false, null,
headerFilter, null);
ConsumerRecord<byte[], byte[]> prodRecord = createRecord("topic", 0, 100L,
headers("environment", "production"));
@@ -407,7 +407,7 @@ public class KafkaRecordSupplierHeaderFilterTest
InDimFilter filter = new InDimFilter("environment",
Collections.singletonList("production"), null);
KafkaHeaderBasedFilterConfig headerFilter = new
KafkaHeaderBasedFilterConfig(filter, null, null);
- recordSupplier = new KafkaRecordSupplier(mockConsumer, true, null,
headerFilter); // multiTopic = true
+ recordSupplier = new KafkaRecordSupplier(mockConsumer, true, null,
headerFilter, null); // multiTopic = true
ConsumerRecord<byte[], byte[]> topic1Record = createRecord("topic1", 0,
100L,
headers("environment", "production"));
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java
index c09ee5f6c71..bd5df4a1780 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaRecordSupplierTest.java
@@ -39,6 +39,8 @@ import org.apache.druid.metadata.DynamicConfigProvider;
import org.apache.druid.metadata.MapStringDynamicConfigProvider;
import org.apache.druid.query.DruidMetrics;
import org.apache.druid.segment.TestHelper;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.NewPartitions;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
@@ -311,6 +313,44 @@ public class KafkaRecordSupplierTest
Assertions.assertEquals(0, diff.size(), diff.toString());
}
+ @Test
+ public void testGetPartitionIdsWithSelectionIgnoresUnselectedPartitions()
throws Exception
+ {
+ // The topic has partitions 0 and 1.
+ insertData();
+ final KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier(
+ KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null,
null, Set.of(PARTITION_1.partition()));
+
+ Assertions.assertEquals(Set.of(PARTITION_1),
recordSupplier.getPartitionIds(TOPIC));
+
+ // A partition added later that is outside the selection is still ignored.
+ increasePartitions(TOPIC, 4);
+
+ Assertions.assertEquals(Set.of(PARTITION_1),
recordSupplier.getPartitionIds(TOPIC));
+ recordSupplier.close();
+ }
+
+ @Test
+ public void testGetPartitionIdsFailsWhenSelectedPartitionIsMissing() throws
Exception
+ {
+ // The topic has partitions 0 and 1 only.
+ insertData();
+ final KafkaRecordSupplier recordSupplier = new KafkaRecordSupplier(
+ KAFKA_SERVER.consumerProperties(), OBJECT_MAPPER, null, false, null,
null, Set.of(0, 3));
+
+ final StreamException error =
Assertions.assertThrows(StreamException.class, () ->
recordSupplier.getPartitionIds(TOPIC));
+ Assertions.assertTrue(error.getMessage().contains("[3]"),
error.getMessage());
+
+ // Discovery recovers once the missing partition exists.
+ increasePartitions(TOPIC, 4);
+
+ Assertions.assertEquals(
+ Set.of(PARTITION_0, new KafkaTopicPartition(false, null, 3)),
+ recordSupplier.getPartitionIds(TOPIC)
+ );
+ recordSupplier.close();
+ }
+
@Test
public void testSupplierSetupCustomDeserializer() throws ExecutionException,
InterruptedException
{
@@ -817,6 +857,13 @@ public class KafkaRecordSupplierTest
}
}
+ private static void increasePartitions(String topic, int totalPartitions)
throws Exception
+ {
+ try (Admin admin = KAFKA_SERVER.newAdminClient()) {
+ admin.createPartitions(Map.of(topic,
NewPartitions.increaseTo(totalPartitions))).all().get();
+ }
+ }
+
private void insertData() throws ExecutionException, InterruptedException
{
try (final KafkaProducer<byte[], byte[]> kafkaProducer =
KAFKA_SERVER.newProducer()) {
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java
index d91fc68bda8..478f742c931 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaSamplerSpecTest.java
@@ -59,6 +59,7 @@ import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.function.Consumer;
import java.util.regex.Pattern;
@@ -123,6 +124,30 @@ public class KafkaSamplerSpecTest extends
InitializedNullHandlingTest
kafkaServer.close();
}
+ @Test
+ public void testSampleOnlySelectedPartition()
+ {
+ final String topic = TOPIC + "_selected";
+ insertData(List.of(
+ new ProducerRecord<>(topic, 0, null, jb("2008", "excluded", "y", "10",
"20.0", "1.0")),
+ new ProducerRecord<>(topic, 1, null, jb("2008", "selected", "y", "10",
"20.0", "1.0"))
+ ));
+ final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder()
+ .withDataSchema(DATA_SCHEMA)
+ .withIoConfig(io ->
io.withJsonInputFormat().withConsumerProperties(kafkaServer.consumerProperties())
+
.withUseEarliestSequenceNumber(true).withPartitionIds(Set.of(1)))
+ .build(DATASOURCE, topic);
+ final SamplerResponse response = new KafkaSamplerSpec(
+ spec,
+ new SamplerConfig(1, 5_000, null, null),
+ new InputSourceSampler(OBJECT_MAPPER),
+ OBJECT_MAPPER
+ ).sample();
+
+ Assertions.assertEquals(1, response.getNumRowsIndexed());
+ Assertions.assertEquals("selected",
response.getData().get(0).getParsed().get("dim1"));
+ }
+
@Test
public void testSample()
{
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java
index ea0a3ca49ff..f407b13fa98 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedKafkaSupervisorTest.java
@@ -58,6 +58,7 @@ import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
public class EmbeddedKafkaSupervisorTest extends EmbeddedClusterTestBase
@@ -112,6 +113,57 @@ public class EmbeddedKafkaSupervisorTest extends
EmbeddedClusterTestBase
}
}
+ @Test
+ public void testSupervisorIngestsOnlySelectedPartitions()
+ {
+ final String topic = IdUtils.getRandomId();
+ kafkaServer.createTopicWithPartitions(topic, 4);
+ final KafkaSupervisorSpec spec = newKafkaSupervisor()
+ .withIoConfig(io ->
io.withPartitionIds(Set.of(1)).withTaskCount(2).withStopTaskCount(1))
+ .build(dataSource, topic);
+ cluster.callApi().postSupervisor(spec);
+
+ // One record per partition, but only partition 1 is selected.
+ kafkaServer.produceRecordsToTopic(partitionRecords(topic, 4));
+ waitForSelectedRows(dataSource, 1);
+
+ Assertions.assertEquals("1", cluster.runSql("SELECT COUNT(*) FROM %s",
dataSource));
+ Assertions.assertEquals("0", cluster.runSql("SELECT COUNT(*) FROM %s WHERE
item != 'p1'", dataSource));
+
+ // Replacing the supervisor with a wider selection picks up partition 3
but still skips partition 2.
+ final KafkaSupervisorSpec expanded = newKafkaSupervisor()
+ .withIoConfig(io -> io.withPartitionIds(Set.of(1,
3)).withTaskCount(2).withStopTaskCount(1))
+ .build(dataSource, topic);
+ cluster.callApi().postSupervisor(expanded);
+ waitForSelectedRows(dataSource, 2);
+
+ Assertions.assertEquals("2", cluster.runSql("SELECT COUNT(*) FROM %s",
dataSource));
+ Assertions.assertEquals("0", cluster.runSql("SELECT COUNT(*) FROM %s WHERE
item = 'p2'", dataSource));
+ cluster.callApi().postSupervisor(expanded.createSuspendedSpec());
+ }
+
+ private static List<ProducerRecord<byte[], byte[]>> partitionRecords(String
topic, int partitions)
+ {
+ final List<ProducerRecord<byte[], byte[]>> records = new ArrayList<>();
+ for (int partition = 0; partition < partitions; partition++) {
+ records.add(new ProducerRecord<>(
+ topic,
+ partition,
+ null,
+ StringUtils.toUtf8("2025-06-01T00:00:00Z,p" + partition)
+ ));
+ }
+ return records;
+ }
+
+ private void waitForSelectedRows(String source, int count)
+ {
+ indexer.latchableEmitter().waitForEventAggregate(
+ event ->
event.hasMetricName("ingest/events/processed").hasDimension(DruidMetrics.DATASOURCE,
source),
+ agg -> agg.hasSumAtLeast(count)
+ );
+ }
+
@Test
public void test_runKafkaSupervisor()
{
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java
index 50486ed0743..6c758928400 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorIOConfigTest.java
@@ -24,13 +24,16 @@ import com.fasterxml.jackson.databind.JsonMappingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedSet;
import nl.jqno.equalsverifier.EqualsVerifier;
import nl.jqno.equalsverifier.Warning;
import org.apache.druid.data.input.InputFormat;
+import org.apache.druid.data.input.impl.TimestampSpec;
import org.apache.druid.error.DruidException;
import org.apache.druid.indexing.kafka.KafkaConsumerConfigs;
import org.apache.druid.indexing.kafka.KafkaIndexTaskModule;
import org.apache.druid.indexing.kafka.KafkaRecordSupplier;
+import
org.apache.druid.indexing.overlord.supervisor.SupervisorSpecUpdateAction;
import org.apache.druid.indexing.seekablestream.extension.KafkaConfigOverrides;
import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig;
import org.apache.druid.indexing.seekablestream.supervisor.IdleConfig;
@@ -39,6 +42,7 @@ import
org.apache.druid.indexing.seekablestream.supervisor.autoscaler.AutoScaler
import
org.apache.druid.indexing.seekablestream.supervisor.autoscaler.LagBasedAutoScalerConfig;
import org.apache.druid.jackson.DefaultObjectMapper;
import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.StringUtils;
import org.joda.time.Duration;
import org.joda.time.Period;
import org.junit.jupiter.api.Assertions;
@@ -48,6 +52,7 @@ import org.junit.jupiter.api.function.Executable;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
+import java.util.Set;
import static org.easymock.EasyMock.createMock;
@@ -61,6 +66,160 @@ public class KafkaSupervisorIOConfigTest
mapper.registerModules(new KafkaIndexTaskModule().getJacksonModules());
}
+ @Test
+ public void testPartitionSelectionSerdeAndUpdate() throws Exception
+ {
+ // Duplicates and ordering are normalized into a sorted set.
+ final KafkaSupervisorIOConfig config = mapper.readValue(
+
"{\"topic\":\"events\",\"consumerProperties\":{\"bootstrap.servers\":\"localhost:9092\"},"
+ + "\"partitionIds\":[5,0,2,0]}",
+ KafkaSupervisorIOConfig.class
+ );
+ Assertions.assertEquals(Set.of(0, 2, 5), config.getPartitionIds());
+
+ // Serialized in sorted order; survives builder copy and JSON round trip;
the selection is immutable.
+ Assertions.assertEquals("[0,2,5]",
mapper.writeValueAsString(config.getPartitionIds()));
+ Assertions.assertEquals(config, config.toBuilder().build());
+ Assertions.assertEquals(config,
mapper.readValue(mapper.writeValueAsString(config),
KafkaSupervisorIOConfig.class));
+ Assertions.assertThrows(UnsupportedOperationException.class, () ->
config.getPartitionIds().add(7));
+
+ // A different selection makes the config unequal.
+ final KafkaSupervisorIOConfig changed =
config.toBuilder().withPartitionIds(Set.of(0, 2)).build();
+ Assertions.assertNotEquals(config, changed);
+
+ // Changing the selection restarts the supervisor and its tasks.
+ final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder()
+ .withDataSchema(schema -> schema.withTimestamp(new
TimestampSpec("timestamp", "auto", null)))
+ .withIoConfig(io -> io.copyFrom(config))
+ .build("diagnostic", "events");
+ Assertions.assertEquals(
+ SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS,
+ spec.getActionOnUpdateTo(spec.toBuilder().ioConfig(changed).build())
+ );
+
+ // Removing the selection (back to all partitions) also restarts the
supervisor and its tasks.
+ final KafkaSupervisorSpec allPartitions = spec.toBuilder()
+ .ioConfig(config.toBuilder().withPartitionIds(null).build()).build();
+
Assertions.assertNull(allPartitions.getSpec().getIOConfig().getPartitionIds());
+ Assertions.assertEquals(
+ SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS,
+ spec.getActionOnUpdateTo(allPartitions)
+ );
+ }
+
+ @Test
+ public void testLegacyConstructorDefaultsToAllPartitions()
+ {
+ final KafkaSupervisorIOConfig config = new KafkaSupervisorIOConfig(
+ "events",
+ null,
+ null,
+ null,
+ null,
+ null,
+ Map.of("bootstrap.servers", "localhost:9092"),
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null,
+ null
+ );
+
+ Assertions.assertNull(config.getPartitionIds());
+ Assertions.assertFalse(mapper.valueToTree(config).has("partitionIds"));
+ Assertions.assertEquals(config, config.toBuilder().build());
+ }
+
+ @Test
+ public void testPartitionSelectionSurvivesAutoscalerUpdates() throws
Exception
+ {
+ final LagBasedAutoScalerConfig autoscaler = mapper.readValue(
+
"{\"enableTaskAutoScaler\":true,\"taskCountMin\":1,\"taskCountMax\":4}",
+ LagBasedAutoScalerConfig.class
+ );
+ final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder()
+ .withDataSchema(schema -> schema.withTimestamp(TimestampSpec.DEFAULT))
+ .withIoConfig(io ->
io.withConsumerProperties(Map.of("bootstrap.servers", "localhost:9092"))
+ .withPartitionIds(Set.of(0,
2)).withAutoScalerConfig(autoscaler))
+ .build("diagnostic", "events");
+ final KafkaSupervisorSpec scaled = spec.toBuilder().taskCount(2).build();
+
+ Assertions.assertEquals(Set.of(0, 2),
scaled.getSpec().getIOConfig().getPartitionIds());
+ Assertions.assertEquals(SupervisorSpecUpdateAction.NONE,
spec.getActionOnUpdateTo(scaled));
+ final KafkaSupervisorSpec narrowed = scaled.toBuilder()
+
.ioConfig(scaled.getSpec().getIOConfig().toBuilder().withPartitionIds(Set.of(0)).build()).build();
+
+ Assertions.assertEquals(
+ SupervisorSpecUpdateAction.RESTART_SUPERVISOR_AND_TASKS,
+ spec.getActionOnUpdateTo(narrowed)
+ );
+ Assertions.assertEquals(Set.of(0),
narrowed.getSpec().getIOConfig().getPartitionIds());
+ }
+
+ @Test
+ public void testInvalidPartitionIds()
+ {
+ for (final String ids : new String[]{"[]", "[-1]", "[null]", "[0.5]",
"[1.0]", "[\"abc\"]", "[true]", "[2147483648]"}) {
+ Assertions.assertThrows(
+ JsonMappingException.class,
+ () -> mapper.readValue(
+ StringUtils.format(
+ """
+ {
+ "topic": "events",
+ "consumerProperties": {"bootstrap.servers":
"localhost:9092"},
+ "partitionIds": %s
+ }
+ """,
+ ids
+ ),
+ KafkaSupervisorIOConfig.class
+ )
+ );
+ }
+ }
+
+ @Test
+ public void testPartitionIdsRequireSingleTopicWithoutBoundedStream()
+ {
+ final String expectedMessage =
+ "partitionIds requires a single topic and cannot be combined with
boundedStreamConfig";
+
+ // Rejected with topicPattern, which reads multiple topics.
+ final DruidException withTopicPattern = Assertions.assertThrows(
+ DruidException.class,
+ () -> new KafkaIOConfigBuilder()
+ .withTopicPattern("events.*")
+ .withConsumerProperties(Map.of("bootstrap.servers",
"localhost:9092"))
+ .withPartitionIds(Set.of(0))
+ .build()
+ );
+ Assertions.assertEquals(expectedMessage, withTopicPattern.getMessage());
+
+ // Rejected with boundedStreamConfig.
+ final DruidException withBoundedStream = Assertions.assertThrows(
+ DruidException.class,
+ () -> new KafkaIOConfigBuilder()
+ .withTopic("events")
+ .withConsumerProperties(Map.of("bootstrap.servers",
"localhost:9092"))
+ .withPartitionIds(Set.of(0))
+ .withBoundedStreamConfig(new BoundedStreamConfig(Map.of(0, 0L),
Map.of(0, 10L)))
+ .build()
+ );
+ Assertions.assertEquals(expectedMessage, withBoundedStream.getMessage());
+ }
+
@Test
public void testSerdeWithDefaults() throws Exception
{
@@ -81,6 +240,7 @@ public class KafkaSupervisorIOConfigTest
Assertions.assertEquals("my-topic", config.getTopic());
Assertions.assertNull(config.getTopicPattern());
+ Assertions.assertNull(config.getPartitionIds());
Assertions.assertEquals(1, (int) config.getReplicas());
Assertions.assertEquals(1, (int) config.getTaskCount());
Assertions.assertNull(config.getStopTaskCount());
@@ -353,6 +513,7 @@ public class KafkaSupervisorIOConfigTest
null,
false,
null,
+ null,
null
);
String ioConfig = mapper.writeValueAsString(kafkaSupervisorIOConfig);
@@ -391,6 +552,7 @@ public class KafkaSupervisorIOConfigTest
null,
false,
null,
+ null,
null
);
Assertions.assertEquals(1, kafkaSupervisorIOConfig.getTaskCount());
@@ -457,6 +619,7 @@ public class KafkaSupervisorIOConfigTest
null,
false,
null,
+ null,
null
);
}
@@ -495,6 +658,7 @@ public class KafkaSupervisorIOConfigTest
null,
false,
null,
+ null,
null
);
String ioConfig = mapper.writeValueAsString(kafkaSupervisorIOConfig);
@@ -621,7 +785,8 @@ public class KafkaSupervisorIOConfigTest
null,
false,
null,
- boundedConfig
+ boundedConfig,
+ null
);
String json = mapper.writeValueAsString(original);
@@ -687,6 +852,7 @@ public class KafkaSupervisorIOConfigTest
.withIgnoredFields("taskCountExplicit", "autoScalerEnabled")
.suppress(Warning.NONFINAL_FIELDS)
.withPrefabValues(Optional.class, Optional.of("a"),
Optional.of("b"))
+ .withPrefabValues(ImmutableSortedSet.class,
ImmutableSortedSet.of(0), ImmutableSortedSet.of(1))
.withPrefabValues(InputFormat.class,
createMock(InputFormat.class), createMock(InputFormat.class))
.withPrefabValues(AutoScalerConfig.class,
createMock(AutoScalerConfig.class), createMock(AutoScalerConfig.class))
.withPrefabValues(LagAggregator.class,
createMock(LagAggregator.class), createMock(LagAggregator.class))
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpecTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpecTest.java
index e0012c8c53f..37fb268b54f 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpecTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorSpecTest.java
@@ -50,6 +50,7 @@ import org.junit.jupiter.api.Test;
import java.io.IOException;
import java.util.Map;
+import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -577,6 +578,35 @@ public class KafkaSupervisorSpecTest
Assertions.assertEquals(boundedStreamConfig,
backfill.getSpec().getIOConfig().getBoundedStreamConfig());
}
+ @Test
+ public void testCreateBackfillSpecDropsPartitionSelection()
+ {
+ final KafkaSupervisorSpec spec = new KafkaSupervisorSpecBuilder()
+ .withDataSchema(
+ schema -> schema
+ .withTimestamp(TimestampSpec.DEFAULT)
+ .withAggregators(new CountAggregatorFactory("rows"))
+ .withGranularity(new
UniformGranularitySpec(Granularities.HOUR, Granularities.NONE, null))
+ )
+ .withIoConfig(
+ ioConfig -> ioConfig
+ .withJsonInputFormat()
+ .withConsumerProperties(Map.of("bootstrap.servers",
"localhost:9092"))
+ .withPartitionIds(Set.of(0))
+ )
+ .build("testDs", "metrics");
+ final BoundedStreamConfig boundedStreamConfig = new
BoundedStreamConfig(Map.of("0", 100L), Map.of("0", 500L));
+
+ final KafkaSupervisorSpec backfill =
spec.createBackfillSpec("backfill-id", boundedStreamConfig, 1);
+
+ // The bounded backfill reads its own range, so it must not inherit the
selection.
+ Assertions.assertNull(backfill.getSpec().getIOConfig().getPartitionIds());
+ Assertions.assertEquals(boundedStreamConfig,
backfill.getSpec().getIOConfig().getBoundedStreamConfig());
+
+ // The source spec keeps its selection.
+ Assertions.assertEquals(Set.of(0),
spec.getSpec().getIOConfig().getPartitionIds());
+ }
+
private KafkaSupervisorSpec getSpec(String topic, String topicPattern)
{
KafkaSupervisorSpecBuilder builder = new KafkaSupervisorSpecBuilder()
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java
index b64bd19e4f4..fc4a181c5cd 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/supervisor/KafkaSupervisorTest.java
@@ -21,6 +21,7 @@ package org.apache.druid.indexing.kafka.supervisor;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.jsontype.NamedType;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
@@ -34,6 +35,7 @@ import org.apache.druid.data.input.impl.StringDimensionSchema;
import org.apache.druid.data.input.impl.TimestampSpec;
import org.apache.druid.data.input.kafka.KafkaRecordEntity;
import org.apache.druid.data.input.kafka.KafkaTopicPartition;
+import org.apache.druid.error.DruidException;
import org.apache.druid.indexer.TaskLocation;
import org.apache.druid.indexer.TaskStatus;
import org.apache.druid.indexer.granularity.UniformGranularitySpec;
@@ -69,6 +71,7 @@ import
org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskRunner.St
import
org.apache.druid.indexing.seekablestream.SeekableStreamIndexTaskTuningConfig;
import
org.apache.druid.indexing.seekablestream.SeekableStreamStartSequenceNumbers;
import org.apache.druid.indexing.seekablestream.common.RecordSupplier;
+import org.apache.druid.indexing.seekablestream.common.StreamException;
import org.apache.druid.indexing.seekablestream.supervisor.BoundedStreamConfig;
import org.apache.druid.indexing.seekablestream.supervisor.IdleConfig;
import
org.apache.druid.indexing.seekablestream.supervisor.SeekableStreamSupervisor;
@@ -181,6 +184,7 @@ public class KafkaSupervisorTest extends EasyMockSupport
private String topic;
private String topicPattern;
private boolean multiTopic;
+ private Set<Integer> selectedPartitionIds;
private RowIngestionMetersFactory rowIngestionMetersFactory;
private StubServiceEmitter serviceEmitter;
private SupervisorStateManagerConfig supervisorConfig;
@@ -242,6 +246,7 @@ public class KafkaSupervisorTest extends EasyMockSupport
topic = getTopic();
topicPattern = getTopicPattern();
topicPostfix++;
+ selectedPartitionIds = null;
multiTopic = false; // assign to true in test if you wish to test
multi-topic
rowIngestionMetersFactory = new TestUtils().getRowIngestionMetersFactory();
serviceEmitter = new StubServiceEmitter("KafkaSupervisorTest",
"localhost");
@@ -523,6 +528,99 @@ public class KafkaSupervisorTest extends EasyMockSupport
Assertions.assertEquals(Integer.valueOf(20),
taskList.get(2).getServerPriority());
}
+ @Test
+ public void testSelectedPartitionAssignmentAndLag() throws Exception
+ {
+ selectedPartitionIds = Set.of(0, 2);
+ supervisor = getTestableSupervisor(1, 2, true, "PT1H", null, null);
+ addSomeEvents(1);
+ final Capture<KafkaIndexTask> captured =
Capture.newInstance(CaptureType.ALL);
+
EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
+
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes();
+
EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes();
+
EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(null).anyTimes();
+
EasyMock.expect(taskQueue.add(EasyMock.capture(captured))).andReturn(true).times(2);
+ replayAll();
+ supervisor.start();
+ // A previous generation had only partition 0. Rebuild the group, not just
its tasks.
+ supervisor.addTaskGroupToActivelyReadingTaskGroup(
+ 0,
+ singlePartitionMap(topic, 0, 0L),
+ null,
+ null,
+ Set.of(),
+ Set.of(),
+ null
+ );
+ supervisor.runInternal();
+ final Set<KafkaTopicPartition> selected = singlePartitionMap(topic, 0, 0L,
2, 0L).keySet();
+
+ Assertions.assertEquals(
+ Set.of(Set.of(new KafkaTopicPartition(false, topic, 0)), Set.of(new
KafkaTopicPartition(false, topic, 2))),
+ captured.getValues().stream()
+ .map(task ->
task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet())
+ .collect(Collectors.toSet())
+ );
+ captured.getValues().forEach(task -> Assertions.assertEquals(
+
task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet(),
+
task.getIOConfig().getEndSequenceNumbers().getPartitionSequenceNumberMap().keySet()
+ ));
+
+ Assertions.assertEquals(2, supervisor.getPartitionCount());
+ supervisor.updatePartitionLagFromStream();
+
+ Assertions.assertEquals(selected,
supervisor.getPartitionRecordLag().keySet());
+ verifyAll();
+ }
+
+ @Test
+ public void testSparseSelectedPartitionGrouping()
+ {
+ selectedPartitionIds = Set.of(0, 3, 6);
+ supervisor = getTestableSupervisor(1, 3, true, "PT1H", null, null);
+
+ Assertions.assertEquals(0, supervisor.getTaskGroupIdForPartition(new
KafkaTopicPartition(false, topic, 0)));
+ Assertions.assertEquals(1, supervisor.getTaskGroupIdForPartition(new
KafkaTopicPartition(false, topic, 3)));
+ Assertions.assertEquals(2, supervisor.getTaskGroupIdForPartition(new
KafkaTopicPartition(false, topic, 6)));
+ // Older tasks may still contain excluded partitions while being replaced.
+ Assertions.assertEquals(1, supervisor.getTaskGroupIdForPartition(new
KafkaTopicPartition(false, topic, 7)));
+ }
+
+ @Test
+ public void testOffsetsFromMetadataStorageIgnoreUnselectedPartitions()
+ {
+ selectedPartitionIds = Set.of(0);
+ supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null);
+ // Partition 1 is outside the selection. Its saved offset stays in
metadata storage but is not reported.
+ final Map<KafkaTopicPartition, Long> storedOffsets =
singlePartitionMap(topic, 0, 10L, 1, 20L);
+
EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(
+ new KafkaDataSourceMetadata(new
SeekableStreamEndSequenceNumbers<>(topic, storedOffsets))
+ );
+ replayAll();
+
+ Assertions.assertEquals(singlePartitionMap(topic, 0, 10L),
supervisor.getOffsetsFromMetadataStorage());
+ verifyAll();
+ }
+
+ @Test
+ public void testMissingSelectedPartitionBlocksStartup() throws Exception
+ {
+ selectedPartitionIds = Set.of(0, 99);
+ supervisor = getTestableSupervisor(1, 2, true, "PT1H", null, null);
+ addSomeEvents(1);
+
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes();
+ replayAll();
+ supervisor.start();
+ supervisor.runInternal();
+ supervisor.runInternal();
+
+ Assertions.assertTrue(supervisor.getPartitionGroups().isEmpty());
+
Assertions.assertFalse(supervisor.getStateManager().isAtLeastOneSuccessfulRun());
+
Assertions.assertTrue(supervisor.getStateManager().getExceptionEvents().toString().contains("99"));
+ Assertions.assertThrows(StreamException.class, () ->
supervisor.getPartitionCount());
+ verifyAll();
+ }
+
@Test
public void testNoInitialState() throws Exception
{
@@ -2879,8 +2977,6 @@ public class KafkaSupervisorTest extends EasyMockSupport
supervisor.start();
supervisor.updateCurrentAndLatestOffsets();
supervisor.runInternal();
- verifyAll();
-
Thread.sleep(100);
supervisor.updateCurrentAndLatestOffsets();
supervisor.runInternal();
@@ -2896,6 +2992,42 @@ public class KafkaSupervisorTest extends EasyMockSupport
Assertions.assertEquals(SupervisorStateManager.BasicState.IDLE,
supervisor.getState());
}
+ @Test
+ public void testSelectedSupervisorIsIdleWithExcludedStoredOffsets() throws
Exception
+ {
+ selectedPartitionIds = Set.of(0);
+ supervisor = getTestableSupervisorForIdleBehaviour(
+ 1,
+ 2,
+ true,
+ "PT10S",
+ null,
+ null,
+ false,
+ new IdleConfig(true, 200L)
+ );
+ addSomeEvents(1);
+
EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
+
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes();
+
EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes();
+ taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class),
EasyMock.anyObject(Executor.class));
+ // Saved offsets include partitions 1 and 2, which are outside the
selection.
+
EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(
+ new KafkaDataSourceMetadata(
+ new SeekableStreamEndSequenceNumbers<>(topic,
singlePartitionMap(topic, 0, 2L, 1, 2L, 2, 2L))
+ )
+ ).anyTimes();
+
EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes();
+ replayAll();
+ supervisor.start();
+ supervisor.updateCurrentAndLatestOffsets();
+
+ supervisor.runInternal();
+
+ Assertions.assertEquals(SupervisorStateManager.BasicState.IDLE,
supervisor.getState());
+ verifyAll();
+ }
+
@Test
public void testSupervisorNotIdleIfStreamInactiveWhenSuspended() throws
Exception
{
@@ -4809,6 +4941,153 @@ public class KafkaSupervisorTest extends EasyMockSupport
verifyAll();
}
+ @Test
+ public void testRestartAdoptsReaderAfterTopicGainsPartitions() throws
Exception
+ {
+ supervisor = getTestableSupervisor(1, 1, true, "PT1H", null, null);
+ addSomeEvents(1);
+ // The task predates partition 1; restart must adopt it without discarding
its uncommitted work.
+ final Map<KafkaTopicPartition, Long> originalOffsets =
singlePartitionMap(topic, 0, 0L, 2, 0L);
+ final KafkaIndexTask existingTask = createKafkaIndexTask(
+ "existing",
+ DATASOURCE,
+ 0,
+ new SeekableStreamStartSequenceNumbers<>(topic, originalOffsets,
Set.of()),
+ new SeekableStreamEndSequenceNumbers<>(topic,
singlePartitionMap(topic, 0, Long.MAX_VALUE, 2, Long.MAX_VALUE)),
+ null,
+ null,
+ supervisor.getTuningConfig()
+ );
+ final TreeMap<Integer, Map<KafkaTopicPartition, Long>> checkpoints = new
TreeMap<>();
+ checkpoints.put(0, originalOffsets);
+
EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
+
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.absent()).anyTimes();
+
EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of("existing",
existingTask)).anyTimes();
+
EasyMock.expect(taskStorage.getStatus("existing")).andReturn(Optional.of(TaskStatus.running("existing"))).anyTimes();
+
EasyMock.expect(taskStorage.getTask("existing")).andReturn(Optional.of(existingTask)).anyTimes();
+
EasyMock.expect(taskClient.getStatusAsync("existing")).andReturn(Futures.immediateFuture(Status.READING));
+
EasyMock.expect(taskClient.getStartTimeAsync("existing")).andReturn(Futures.immediateFuture(DateTimes.nowUtc()));
+ EasyMock.expect(taskClient.getCheckpointsAsync(EasyMock.eq("existing"),
EasyMock.anyBoolean()))
+ .andReturn(Futures.immediateFuture(checkpoints));
+
EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE))
+ .andReturn(new KafkaDataSourceMetadata(null)).anyTimes();
+ replayAll();
+ supervisor.start();
+ supervisor.runInternal();
+ verifyAll();
+
+
Assertions.assertTrue(supervisor.getStateManager().isAtLeastOneSuccessfulRun());
+ Assertions.assertEquals(3, supervisor.getPartitionGroups().get(0).size());
+ Assertions.assertEquals(1, supervisor.getActiveTaskGroupsCount());
+ // Strict mocks reject stop/pause/publish calls and replacement
submissions in the initial restart cycle.
+ }
+
+ @Test
+ public void testSelectedPartitionTaskCompatibility()
+ {
+ selectedPartitionIds = Set.of(0, 2);
+ final KafkaSupervisor selected = createSupervisor(
+ 1,
+ 2,
+ true,
+ "PT1H",
+ null,
+ null,
+ false,
+ kafkaHost,
+ dataSchema,
+ tuningConfigBuilder().build()
+ );
+ final Set<KafkaTopicPartition> expected = singlePartitionMap(topic, 0, 0L,
2, 0L).keySet();
+ selected.getPartitionGroups().put(0, expected);
+
+ Assertions.assertTrue(selected.isTaskPartitionSetCurrent(0, expected));
+ // A task reading fewer, more, or a different group's partitions is not
current.
+ Assertions.assertFalse(selected.isTaskPartitionSetCurrent(0,
singlePartitionMap(topic, 0, 0L).keySet()));
+ Assertions.assertFalse(selected.isTaskPartitionSetCurrent(0,
singlePartitionMap(topic, 0, 0L, 2, 0L, 4, 0L).keySet()));
+ Assertions.assertFalse(selected.isTaskPartitionSetCurrent(1, expected));
+ }
+
+ @Test
+ public void testSelectedPartitionResetValidation()
+ {
+ selectedPartitionIds = Set.of(0, 2);
+ final KafkaSupervisor selected = createSupervisor(
+ 1,
+ 2,
+ true,
+ "PT1H",
+ null,
+ null,
+ false,
+ kafkaHost,
+ dataSchema,
+ tuningConfigBuilder().build()
+ );
+ final Set<KafkaTopicPartition> expected = singlePartitionMap(topic, 0, 0L,
2, 0L).keySet();
+ selected.getPartitionGroups().put(0, expected);
+ final KafkaDataSourceMetadata excluded = new KafkaDataSourceMetadata(
+ new SeekableStreamEndSequenceNumbers<>(topic,
singlePartitionMap(topic, 0, 0L, 4, 0L))
+ );
+
+ assertResetRejected(selected, excluded);
+ Assertions.assertEquals(expected, selected.getPartitionGroups().get(0));
+
+ // Offsets for selected partitions only, and a full reset (null), are
accepted.
+ selected.validatePartitionReset(new KafkaDataSourceMetadata(
+ new SeekableStreamEndSequenceNumbers<>(topic,
singlePartitionMap(topic, 0, 0L))
+ ));
+ selected.validatePartitionReset(null);
+ }
+
+ @Test
+ public void testSelectedPartitionResetRejectsTopicQualifiedJsonKeys() throws
Exception
+ {
+ selectedPartitionIds = Set.of(0);
+ final KafkaSupervisor selected = createSupervisor(
+ 1,
+ 1,
+ true,
+ "PT1H",
+ null,
+ null,
+ false,
+ kafkaHost,
+ dataSchema,
+ tuningConfigBuilder().build()
+ );
+ final Set<KafkaTopicPartition> expected = singlePartitionMap(topic, 0,
0L).keySet();
+ selected.getPartitionGroups().put(0, expected);
+ final ObjectMapper resetMapper = OBJECT_MAPPER.copy();
+ resetMapper.registerSubtypes(new NamedType(KafkaDataSourceMetadata.class,
"kafka"));
+ EasyMock.replay(indexerMetadataStorageCoordinator, taskMaster, taskQueue);
+ for (final String key : List.of("other-topic:0", topic + ":0", ":0")) {
+ final KafkaDataSourceMetadata metadata = resetMapper.readValue(
+ StringUtils.format(
+
"{\"type\":\"kafka\",\"partitions\":{\"type\":\"end\",\"stream\":\"%s\","
+ + "\"partitionSequenceNumberMap\":{\"0\":5,\"%s\":10}}}",
+ topic,
+ key
+ ),
+ KafkaDataSourceMetadata.class
+ );
+
+ assertResetRejected(selected, metadata);
+ Assertions.assertEquals(expected, selected.getPartitionGroups().get(0));
+ }
+ // No interactions with storage or tasks are permitted, including for
mixed-key requests.
+ EasyMock.verify(indexerMetadataStorageCoordinator, taskMaster, taskQueue);
+ }
+
+ /** Every reset entry point must reject the request before touching metadata
storage or the task queue. */
+ private static void assertResetRejected(KafkaSupervisor selected,
KafkaDataSourceMetadata metadata)
+ {
+ Assertions.assertThrows(DruidException.class, () ->
selected.reset(metadata));
+ Assertions.assertThrows(DruidException.class, () ->
selected.resetOffsets(metadata));
+ Assertions.assertThrows(DruidException.class, () ->
selected.resetInternal(metadata));
+ Assertions.assertThrows(DruidException.class, () ->
selected.resetOffsetsInternal(metadata));
+ }
+
@Test
public void testIsTaskCurrent()
{
@@ -5826,6 +6105,7 @@ public class KafkaSupervisorTest extends EasyMockSupport
consumerProperties.put("bootstrap.servers", kafkaHost);
KafkaSupervisorIOConfig kafkaSupervisorIOConfig = new
KafkaIOConfigBuilder()
.withTopic(multiTopic ? null : topic)
+ .withPartitionIds(selectedPartitionIds)
.withTopicPattern(multiTopic ? topicPattern : null)
.withInputFormat(INPUT_FORMAT)
.withReplicas(replicas)
@@ -6000,6 +6280,7 @@ public class KafkaSupervisorTest extends EasyMockSupport
consumerProperties.put("isolation.level", "read_committed");
KafkaSupervisorIOConfig kafkaSupervisorIOConfig = new
KafkaIOConfigBuilder()
.withTopic(topic)
+ .withPartitionIds(selectedPartitionIds)
.withInputFormat(INPUT_FORMAT)
.withReplicas(replicas)
.withTaskCount(taskCount)
@@ -6305,7 +6586,9 @@ public class KafkaSupervisorTest extends EasyMockSupport
return new KafkaRecordSupplier(
new KafkaConsumer<>(props, keyDeserializerObject,
valueDeserializerObject),
getIoConfig().isMultiTopic(),
- null
+ null,
+ null,
+ getIoConfig().getPartitionIds()
);
}
@@ -6421,6 +6704,31 @@ public class KafkaSupervisorTest extends EasyMockSupport
.get(new
KafkaTopicPartition(false, topic, 0)));
}
+ @Test
+ public void testBoundedSupervisorReportsLagOnlyForBoundedPartitions() throws
Exception
+ {
+ // The topic has more partitions than the bounded config covers, as in a
backfill of a partition selection.
+ final BoundedStreamConfig boundedConfig = new
BoundedStreamConfig(Map.of("0", 0L), Map.of("0", 100L));
+ supervisor = getTestableSupervisorWithBoundedConfig(1, 1, "PT1H",
boundedConfig);
+ addSomeEvents(100);
+
EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.of(taskQueue)).anyTimes();
+
EasyMock.expect(taskMaster.getTaskRunner()).andReturn(Optional.of(taskRunner)).anyTimes();
+
EasyMock.expect(taskQueue.getActiveTasksForDatasource(DATASOURCE)).andReturn(Map.of()).anyTimes();
+
EasyMock.expect(indexerMetadataStorageCoordinator.retrieveDataSourceMetadata(DATASOURCE)).andReturn(
+ new KafkaDataSourceMetadata(null)
+ ).anyTimes();
+
EasyMock.expect(taskQueue.add(EasyMock.anyObject())).andReturn(true).anyTimes();
+ taskRunner.registerListener(EasyMock.anyObject(TaskRunnerListener.class),
EasyMock.anyObject(Executor.class));
+ replayAll();
+ supervisor.start();
+ supervisor.runInternal();
+
+ supervisor.updatePartitionLagFromStream();
+
+ // Partitions the bounded run never assigned must not be queried, as
reading their positions throws.
+ Assertions.assertEquals(Set.of(new KafkaTopicPartition(false, topic, 0)),
supervisor.getPartitionRecordLag().keySet());
+ }
+
@Test
public void testBoundedStreamConfig_withCheckpoint_resumesFromCheckpoint()
throws Exception
{
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java
index 3377090a29c..727a9e3a1bd 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/supervisor/SeekableStreamSupervisor.java
@@ -1364,9 +1364,15 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
});
}
+ protected void validatePartitionReset(@Nullable DataSourceMetadata metadata)
+ {
+ // Stream implementations may restrict which partitions can be reset.
+ }
+
@Override
public void reset(@Nullable final DataSourceMetadata dataSourceMetadata)
{
+ validatePartitionReset(dataSourceMetadata);
log.info("Posting ResetNotice with datasource metadata [%s]",
dataSourceMetadata);
addNotice(new ResetNotice(dataSourceMetadata));
}
@@ -1381,6 +1387,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
@Override
public void resetOffsets(@Nonnull DataSourceMetadata resetDataSourceMetadata)
{
+ validatePartitionReset(resetDataSourceMetadata);
if (resetDataSourceMetadata == null) {
throw InvalidInput.exception("Reset dataSourceMetadata is required for
resetOffsets.");
}
@@ -2135,6 +2142,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
@VisibleForTesting
public void resetInternal(DataSourceMetadata dataSourceMetadata)
{
+ validatePartitionReset(dataSourceMetadata);
if (dataSourceMetadata == null) {
// Reset everything
boolean result =
indexerMetadataStorageCoordinator.deleteDataSourceMetadata(supervisorId);
@@ -2260,6 +2268,7 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
*/
public void resetOffsetsInternal(@Nonnull final DataSourceMetadata
dataSourceMetadata)
{
+ validatePartitionReset(dataSourceMetadata);
log.info("Reset offsets for supervisor[%s] for dataSource[%s] with
metadata[%s]", supervisorId, dataSource, dataSourceMetadata);
@SuppressWarnings("unchecked")
@@ -3087,6 +3096,12 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
);
}
+ /** Whether a task's partitions match what its group now owns; if not, the
task is replaced. Default: always true. */
+ protected boolean isTaskPartitionSetCurrent(int taskGroupId,
Set<PartitionIdType> taskPartitions)
+ {
+ return true;
+ }
+
/**
* Determines whether a given task was created by the current version of the
supervisor.
* Uses the Task object mapped to this taskId in the {@code activeTaskMap}.
@@ -3116,6 +3131,13 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
SeekableStreamIndexTask<PartitionIdType, SequenceOffsetType, RecordType>
task =
(SeekableStreamIndexTask<PartitionIdType, SequenceOffsetType,
RecordType>) genericTask;
+ if (!isTaskPartitionSetCurrent(
+ taskGroupId,
+
task.getIOConfig().getStartSequenceNumbers().getPartitionSequenceNumberMap().keySet()
+ )) {
+ return false;
+ }
+
// We recompute the sequence name hash for the supervisor's own
configuration and compare this to the hash created
// by rehashing the task's sequence name using the most up-to-date class
definitions of tuning config and
// data schema. Recomputing both ensures that forwards-compatible tasks
won't be killed (which would occur
@@ -4223,6 +4245,12 @@ public abstract class
SeekableStreamSupervisor<PartitionIdType, SequenceOffsetTy
Integer groupId = taskGroupEntry.getKey();
TaskGroup taskGroup = taskGroupEntry.getValue();
+ if (!isTaskPartitionSetCurrent(groupId,
taskGroup.startingSequences.keySet())) {
+ futures.add(stopTasksInGroup(taskGroup, "Task group partition
selection has changed"));
+ iTaskGroups.remove();
+ continue;
+ }
+
// Iterate the list of known tasks in this group and:
// 1) Kill any tasks which are not "current" (have the partitions,
starting sequences, and minimumMessageTime
// & maximumMessageTime (if applicable) in
[activelyReadingTaskGroups])
diff --git
a/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java
b/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java
new file mode 100644
index 00000000000..f469ee94041
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/jackson/StrictIntegerDeserializer.java
@@ -0,0 +1,59 @@
+/*
+ * 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.druid.jackson;
+
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+
+import java.io.IOException;
+
+/**
+ * Deserializes an {@link Integer} from a JSON integer or a string holding an
integer, such as {@code 1} or
+ * {@code "1"}. Unlike Jackson's default, it rejects values that would be
silently truncated, such as {@code 0.5} and
+ * {@code 1.0}.
+ * <p>
+ * Use it with {@code @JsonDeserialize(using = ...)} on a field, or {@code
contentUsing = ...} on a collection.
+ */
+public class StrictIntegerDeserializer extends JsonDeserializer<Integer>
+{
+ @Override
+ public Integer deserialize(final JsonParser parser, final
DeserializationContext context) throws IOException
+ {
+ if (parser.hasToken(JsonToken.VALUE_NUMBER_INT)) {
+ return parser.getIntValue();
+ }
+ if (parser.hasToken(JsonToken.VALUE_STRING)) {
+ try {
+ return Integer.parseInt(parser.getText());
+ }
+ catch (NumberFormatException ignored) {
+ // Fall through to report the offending value.
+ }
+ }
+ return context.reportInputMismatch(
+ Integer.class,
+ "Expected a JSON integer but got %s [%s]",
+ parser.currentToken(),
+ parser.getText()
+ );
+ }
+}
diff --git
a/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java
b/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java
new file mode 100644
index 00000000000..9af5adbbebe
--- /dev/null
+++
b/processing/src/test/java/org/apache/druid/jackson/StrictIntegerDeserializerTest.java
@@ -0,0 +1,89 @@
+/*
+ * 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.druid.jackson;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.JsonMappingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class StrictIntegerDeserializerTest
+{
+ private final ObjectMapper mapper = new DefaultObjectMapper();
+
+ @Test
+ public void testAcceptsIntegers() throws Exception
+ {
+ final Holder holder = mapper.readValue("{\"ids\":[0, 7, 2147483647]}",
Holder.class);
+
+ Assertions.assertEquals(List.of(0, 7, Integer.MAX_VALUE), holder.ids);
+ }
+
+ @Test
+ public void testAcceptsStringsHoldingIntegers() throws Exception
+ {
+ final Holder holder = mapper.readValue("{\"ids\":[\"1\", \"-2\"]}",
Holder.class);
+
+ Assertions.assertEquals(List.of(1, -2), holder.ids);
+ }
+
+ @Test
+ public void testRejectsNonIntegersAndReportsTheValue()
+ {
+ assertRejected("0.5", "VALUE_NUMBER_FLOAT [0.5]");
+ assertRejected("1.0", "VALUE_NUMBER_FLOAT [1.0]");
+ assertRejected("true", "VALUE_TRUE [true]");
+ assertRejected("\"abc\"", "VALUE_STRING [abc]");
+ assertRejected("\"0.5\"", "VALUE_STRING [0.5]");
+ assertRejected("\"\"", "VALUE_STRING []");
+ }
+
+ @Test
+ public void testRejectsOutOfRangeIntegers()
+ {
+ Assertions.assertThrows(JsonMappingException.class, () ->
mapper.readValue("{\"ids\":[2147483648]}", Holder.class));
+ assertRejected("\"2147483648\"", "VALUE_STRING [2147483648]");
+ }
+
+ private void assertRejected(String json, String expectedValueDescription)
+ {
+ final JsonMappingException e = Assertions.assertThrows(
+ JsonMappingException.class,
+ () -> mapper.readValue("{\"ids\":[" + json + "]}", Holder.class)
+ );
+ Assertions.assertTrue(e.getMessage().contains("Expected a JSON integer but
got " + expectedValueDescription), e.getMessage());
+ }
+
+ private static class Holder
+ {
+ private final List<Integer> ids;
+
+ @JsonCreator
+ Holder(@JsonProperty("ids") @JsonDeserialize(contentUsing =
StrictIntegerDeserializer.class) List<Integer> ids)
+ {
+ this.ids = ids;
+ }
+ }
+}
diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts
b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts
index cac347cb31f..b8653045147 100644
--- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts
+++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.spec.ts
@@ -23,15 +23,77 @@ import {
adjustId,
cleanSpec,
DEFAULT_FORCE_SEGMENT_SORT_BY_TIME,
+ getIoConfigFormFields,
guessColumnTypeFromInput,
guessColumnTypeFromSampleResponse,
guessKafkaInputFormat,
guessSimpleInputFormat,
+ issueWithIoConfig,
updateSchemaWithSample,
upgradeSpec,
} from './ingestion-spec';
describe('ingestion-spec', () => {
+ it('validates Kafka partition selection without contacting Kafka', () => {
+ const ioConfig = { type: 'kafka', topic: 'events' };
+ expect(issueWithIoConfig(ioConfig, true)).toBeUndefined();
+ for (const partitionIds of [[0, 2], ['1', '2'], [2147483647],
['2147483647']]) {
+ expect(issueWithIoConfig({ ...ioConfig, partitionIds },
true)).toBeUndefined();
+ }
+ for (const partitionIds of [
+ [],
+ [-1],
+ [0.5],
+ [NaN],
+ [2147483648],
+ ['2147483648'],
+ ['abc'],
+ ['-1'],
+ ['1.5'],
+ ]) {
+ expect(issueWithIoConfig({ ...ioConfig, partitionIds },
true)).toBeDefined();
+ }
+ expect(
+ issueWithIoConfig({ ...ioConfig, topicPattern: 'events.*', partitionIds:
[0] }, true),
+ ).toBeDefined();
+ });
+
+ it('accepts comma-separated Kafka partition IDs in the form', () => {
+ const partitionIdsField = getIoConfigFormFields('kafka').find(
+ field => field.name === 'partitionIds',
+ );
+
+ expect(partitionIdsField?.type).toBe('string-array');
+ expect(partitionIdsField?.placeholder).toBe('Optional; comma-separated,
e.g. 0, 2');
+ expect(partitionIdsField?.valueAdjustment?.(['1', '2', '3'])).toEqual([1,
2, 3]);
+ // Values the server rejects are kept as strings, not coerced to numbers.
+ expect(partitionIdsField?.valueAdjustment?.(['1e3', '0x10', '-1',
'abc'])).toEqual([
+ '1e3',
+ '0x10',
+ '-1',
+ 'abc',
+ ]);
+ expect(partitionIdsField?.valueAdjustment?.([])).toBeUndefined();
+ expect(partitionIdsField?.valueAdjustment?.(undefined)).toBeUndefined();
+
+ // The field is shown for a topic, and stays shown while it holds a value
so it can be cleared.
+ const defined = partitionIdsField?.defined;
+ expect(typeof defined).toBe('function');
+ if (typeof defined === 'function') {
+ expect(defined({ type: 'kafka', topic: 'events' })).toBe(true);
+ expect(defined({ type: 'kafka', topicPattern: 'events.*' })).toBe(false);
+ expect(defined({ type: 'kafka', topicPattern: 'events.*', partitionIds:
[0] })).toBe(true);
+ expect(defined({ type: 'kinesis', partitionIds: [0] })).toBe(false);
+ }
+
+ const hideInMore = partitionIdsField?.hideInMore;
+ expect(typeof hideInMore).toBe('function');
+ if (typeof hideInMore === 'function') {
+ expect(hideInMore({})).toBe(true);
+ expect(hideInMore({ partitionIds: [1] })).toBe(false);
+ }
+ });
+
it('upgrades / downgrades task spec 1', () => {
const oldTaskSpec = {
type: 'index_parallel',
diff --git a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
index 778818f4680..2ada7a29a77 100644
--- a/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
+++ b/web-console/src/druid-models/ingestion-spec/ingestion-spec.tsx
@@ -590,6 +590,7 @@ export interface IoConfig {
appendToExisting?: boolean;
topic?: string;
topicPattern?: string;
+ partitionIds?: (number | string)[];
consumerProperties?: any;
replicas?: number;
taskCount?: number;
@@ -1156,6 +1157,23 @@ export function
getIoConfigFormFields(ingestionComboType: IngestionComboType): F
placeholder: 'your_kafka_topic',
info: 'The name of the Kafka topic to ingest from.',
},
+ {
+ name: 'partitionIds',
+ label: 'Partition IDs',
+ type: 'string-array',
+ // Keep the field visible while it holds a value, so it can still be
cleared after switching to topicPattern.
+ defined: ioConfig =>
+ ioConfig.type === 'kafka' && (Boolean(ioConfig.topic) ||
ioConfig.partitionIds != null),
+ placeholder: 'Optional; comma-separated, e.g. 0, 2',
+ hideInMore: ioConfig => ioConfig.partitionIds == null,
+ // Only plain decimal digits become numbers. Anything else stays a
string so that validation rejects it,
+ // as the server does, instead of being coerced (e.g. "1e3" or
"0x10").
+ valueAdjustment: partitionIds =>
+ partitionIds?.length
+ ? partitionIds.map((id: string) => (/^\d+$/.test(id) ?
Number(id) : id))
+ : undefined,
+ info: 'Optional comma-separated partition IDs, such as 0,2. Omit to
read all partitions. IDs must exist when Kafka metadata is discovered.',
+ },
{
name: 'topicPattern',
type: 'string',
@@ -1319,6 +1337,16 @@ export function
getIoConfigFormFields(ingestionComboType: IngestionComboType): F
throw new Error(`unknown input type ${ingestionComboType}`);
}
+// The server reads partition IDs as Java ints, accepting JSON integers and
integer strings such as "1".
+const MAX_KAFKA_PARTITION_ID = 2147483647;
+
+function isValidKafkaPartitionId(id: unknown): boolean {
+ if (typeof id === 'number')
+ return Number.isInteger(id) && id >= 0 && id <= MAX_KAFKA_PARTITION_ID;
+ if (typeof id === 'string') return /^\+?\d+$/.test(id) && Number(id) <=
MAX_KAFKA_PARTITION_ID;
+ return false;
+}
+
export function issueWithIoConfig(
ioConfig: IoConfig | undefined,
ignoreInputFormat = false,
@@ -1335,6 +1363,16 @@ export function issueWithIoConfig(
case 'kafka':
if (!ioConfig.topic && !ioConfig.topicPattern) return 'must have a topic
or topicPattern';
+ if (ioConfig.partitionIds != null) {
+ if (ioConfig.topicPattern) return 'partitionIds requires a single
topic';
+ if (
+ !Array.isArray(ioConfig.partitionIds) ||
+ !ioConfig.partitionIds.length ||
+ !ioConfig.partitionIds.every(isValidKafkaPartitionId)
+ ) {
+ return `partitionIds must be a nonempty array of integers between 0
and ${MAX_KAFKA_PARTITION_ID}`;
+ }
+ }
break;
case 'kinesis':
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
