This is an automated email from the ASF dual-hosted git repository.
abhishekrb19 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 a85d64c4cbd refactor: Make shard spec collector an interface for
extensibility (#19689)
a85d64c4cbd is described below
commit a85d64c4cbd8d76614c0fa64c97ba1159a4fc851
Author: Abhishek Radhakrishnan <[email protected]>
AuthorDate: Thu Jul 23 13:26:58 2026 +0530
refactor: Make shard spec collector an interface for extensibility (#19689)
refactor: Make shard spec collector an interface for extensibility
This is a follow-up to #19571. This patch addresses this review comment
#19571 (comment) by refactoring the logic into an interface so it's extensible
and maintainable. No other functional change.
---
docs/ingestion/kafka-ingestion.md | 4 +-
.../kafka/KafkaIndexTaskTuningConfigTest.java | 64 ++++++-
.../EmbeddedDimensionValueSetShardSpecTest.java | 14 +-
.../seekablestream/DimensionValueSetCollector.java | 191 +++++++++++++++++++++
...c.java => DimensionValueSetPartitionsSpec.java} | 51 +++---
.../SeekableStreamIndexTaskRunner.java | 160 ++++-------------
.../seekablestream/StreamingPartitionsSpec.java | 119 ++++---------
.../StreamingShardSpecCollector.java | 75 ++++++++
.../SeekableStreamIndexTaskRunnerAuthTest.java | 1 +
.../SeekableStreamIndexTaskRunnerTest.java | 119 ++++++++-----
10 files changed, 504 insertions(+), 294 deletions(-)
diff --git a/docs/ingestion/kafka-ingestion.md
b/docs/ingestion/kafka-ingestion.md
index 79232c9f5a2..0441dde69d1 100644
--- a/docs/ingestion/kafka-ingestion.md
+++ b/docs/ingestion/kafka-ingestion.md
@@ -267,6 +267,8 @@ When you set `streamingPartitionsSpec.partitionDimensions`
in the tuning config,
This enables segment pruning for streaming-ingested data without waiting for
compaction to produce hash or range-partitioned segments. The
`partitionDimensions` should be kept in sync with the compaction config's
`partitionDimensions` for the same datasource.
+`streamingPartitionsSpec` is a typed object. The default and currently only
type is `dim_value_set`, which records the distinct values observed for each
tracked dimension. The `type` may be omitted, in which case it defaults to
`dim_value_set`; you can set it explicitly (`"streamingPartitionsSpec":
{"type": "dim_value_set", "partitionDimensions": [...]}`) to be
forward-compatible with future strategies.
+
**Usage guidelines:**
- Only string-typed dimensions are currently supported.
@@ -484,7 +486,7 @@ For configuration properties shared across all streaming
ingestion methods, refe
|Property|Type|Description|Required|Default|
|--------|----|-----------|--------|-------|
|`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. Contains `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](#streaming-partitions-spec) for details.|No|null|
+|`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 [...]
## Deployment notes on Kafka partitions and Druid segments
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java
index 4bf668bbc86..b19f4d7a3f3 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTuningConfigTest.java
@@ -24,7 +24,7 @@ import nl.jqno.equalsverifier.EqualsVerifier;
import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorTuningConfig;
import org.apache.druid.indexing.kafka.supervisor.KafkaTuningConfigBuilder;
import
org.apache.druid.indexing.kafka.test.TestModifiedKafkaIndexTaskTuningConfig;
-import org.apache.druid.indexing.seekablestream.StreamingPartitionsSpec;
+import
org.apache.druid.indexing.seekablestream.DimensionValueSetPartitionsSpec;
import org.apache.druid.jackson.DefaultObjectMapper;
import org.apache.druid.segment.IndexSpec;
import org.apache.druid.segment.data.CompressionStrategy;
@@ -146,10 +146,10 @@ public class KafkaIndexTaskTuningConfigTest
);
Assert.assertEquals(
- new StreamingPartitionsSpec(List.of("tenant", "region")),
+ new DimensionValueSetPartitionsSpec(List.of("tenant", "region")),
config.getStreamingPartitionsSpec()
);
- Assert.assertEquals(List.of("tenant", "region"),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(List.of("tenant", "region"),
partitionDimensionsOf(config));
}
@Test
@@ -166,14 +166,14 @@ public class KafkaIndexTaskTuningConfigTest
public void testSerdeWithEmptyPartitionDimensions() throws Exception
{
final KafkaIndexTaskTuningConfig config =
roundTripWithStreamingPartitionsSpec("[]");
- Assert.assertEquals(Collections.emptyList(),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(Collections.emptyList(),
partitionDimensionsOf(config));
}
@Test
public void testSerdeWithNullPartitionDimensionsCoalescesToEmpty() throws
Exception
{
final KafkaIndexTaskTuningConfig config =
roundTripWithStreamingPartitionsSpec("null");
- Assert.assertEquals(Collections.emptyList(),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(Collections.emptyList(),
partitionDimensionsOf(config));
}
@Test
@@ -181,7 +181,7 @@ public class KafkaIndexTaskTuningConfigTest
{
// An empty-string dimension name is preserved verbatim (it simply never
matches an ingested value).
final KafkaIndexTaskTuningConfig config =
roundTripWithStreamingPartitionsSpec("[\"\"]");
- Assert.assertEquals(List.of(""),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(List.of(""), partitionDimensionsOf(config));
}
@Test
@@ -189,14 +189,57 @@ public class KafkaIndexTaskTuningConfigTest
{
// Dimension names are plain strings; a numeric-looking name is just a
string.
final KafkaIndexTaskTuningConfig config =
roundTripWithStreamingPartitionsSpec("[\"123\"]");
- Assert.assertEquals(List.of("123"),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(List.of("123"), partitionDimensionsOf(config));
}
@Test
public void testSerdeWithNullElementInPartitionDimensions() throws Exception
{
final KafkaIndexTaskTuningConfig config =
roundTripWithStreamingPartitionsSpec("[\"tenant\", null]");
- Assert.assertEquals(Arrays.asList("tenant", null),
config.getStreamingPartitionsSpec().getPartitionDimensions());
+ Assert.assertEquals(Arrays.asList("tenant", null),
partitionDimensionsOf(config));
+ }
+
+ @Test
+ public void testSerdeWithExplicitDimValueSetType() throws Exception
+ {
+ // An explicit "type": "dim_value_set" round-trips to the same spec as the
untyped (default) form.
+ final String jsonStr = "{\n"
+ + " \"type\": \"kafka\",\n"
+ + " \"streamingPartitionsSpec\": "
+ + "{\"type\": \"dim_value_set\",
\"partitionDimensions\": [\"tenant\", \"region\"]}\n"
+ + "}";
+
+ final KafkaIndexTaskTuningConfig config = (KafkaIndexTaskTuningConfig)
mapper.readValue(
+ mapper.writeValueAsString(mapper.readValue(jsonStr,
TuningConfig.class)),
+ TuningConfig.class
+ );
+
+ Assert.assertEquals(
+ new DimensionValueSetPartitionsSpec(List.of("tenant", "region")),
+ config.getStreamingPartitionsSpec()
+ );
+ Assert.assertEquals(List.of("tenant", "region"),
partitionDimensionsOf(config));
+ }
+
+ @Test
+ public void testSerdeWithUnknownStreamingPartitionsSpecTypeIsRejected()
+ {
+ // An explicit but unknown type (e.g. a typo, or a subtype whose extension
isn't loaded on this peon) must fail
+ // rather than silently falling back to the default
DimensionValueSetPartitionsSpec.
+ final String jsonStr = "{\n"
+ + " \"type\": \"kafka\",\n"
+ + " \"streamingPartitionsSpec\": "
+ + "{\"type\": \"dim_value_sets\",
\"partitionDimensions\": [\"tenant\"]}\n"
+ + "}";
+
+ final Exception e = Assert.assertThrows(
+ Exception.class,
+ () -> mapper.readValue(jsonStr, TuningConfig.class)
+ );
+ Assert.assertTrue(
+ "Expected the unknown type id to be surfaced, got: " + e.getMessage(),
+ e.getMessage().contains("dim_value_sets")
+ );
}
private KafkaIndexTaskTuningConfig
roundTripWithStreamingPartitionsSpec(String partitionDimensionsJson)
@@ -212,6 +255,11 @@ public class KafkaIndexTaskTuningConfigTest
);
}
+ private static List<String> partitionDimensionsOf(KafkaIndexTaskTuningConfig
config)
+ {
+ return ((DimensionValueSetPartitionsSpec)
config.getStreamingPartitionsSpec()).getPartitionDimensions();
+ }
+
@Test
public void testConvert()
{
diff --git
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedDimensionValueSetShardSpecTest.java
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedDimensionValueSetShardSpecTest.java
index 74ad0c6ecdb..781f1e2cced 100644
---
a/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedDimensionValueSetShardSpecTest.java
+++
b/extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/simulate/EmbeddedDimensionValueSetShardSpecTest.java
@@ -31,7 +31,7 @@ import org.apache.druid.indexing.common.task.Tasks;
import org.apache.druid.indexing.kafka.KafkaIndexTaskModule;
import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpec;
import org.apache.druid.indexing.kafka.supervisor.KafkaSupervisorSpecBuilder;
-import org.apache.druid.indexing.seekablestream.StreamingPartitionsSpec;
+import
org.apache.druid.indexing.seekablestream.DimensionValueSetPartitionsSpec;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.granularity.Granularities;
@@ -212,7 +212,7 @@ public class EmbeddedDimensionValueSetShardSpecTest extends
EmbeddedClusterTestB
t -> t.withMaxRowsPerSegment(1)
.withReleaseLocksOnHandoff(true)
// Track both a string dimension and a numeric dimension
- .withStreamingPartitionsSpec(new
StreamingPartitionsSpec(List.of(COL_TENANT, colRegionCode)))
+ .withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(List.of(COL_TENANT, colRegionCode)))
)
.withIoConfig(
ioConfig -> ioConfig
@@ -321,7 +321,7 @@ public class EmbeddedDimensionValueSetShardSpecTest extends
EmbeddedClusterTestB
.withTuningConfig(
t -> t.withMaxRowsPerSegment(1000)
.withReleaseLocksOnHandoff(true)
- .withStreamingPartitionsSpec(new
StreamingPartitionsSpec(List.of(colCode)))
+ .withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(List.of(colCode)))
)
.withIoConfig(
ioConfig -> ioConfig
@@ -400,7 +400,7 @@ public class EmbeddedDimensionValueSetShardSpecTest extends
EmbeddedClusterTestB
.withMaxRowsPerSegment(100_000)
// Long enough that time-based persist/push doesn't fire
during the test
.withIntermediatePersistPeriod(Period.hours(1))
- .withStreamingPartitionsSpec(new
StreamingPartitionsSpec(List.of(COL_TENANT)))
+ .withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(List.of(COL_TENANT)))
)
.withIoConfig(
ioConfig -> ioConfig
@@ -552,7 +552,7 @@ public class EmbeddedDimensionValueSetShardSpecTest extends
EmbeddedClusterTestB
.withTuningConfig(
t -> t.withMaxRowsPerSegment(1000)
.withReleaseLocksOnHandoff(true)
- .withStreamingPartitionsSpec(new
StreamingPartitionsSpec(List.of(COL_TENANT, colRegion), 2))
+ .withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(List.of(COL_TENANT, colRegion), 2))
)
.withIoConfig(
ioConfig -> ioConfig
@@ -1048,7 +1048,7 @@ public class EmbeddedDimensionValueSetShardSpecTest
extends EmbeddedClusterTestB
tuningConfig -> tuningConfig
.withMaxRowsPerSegment(1)
.withReleaseLocksOnHandoff(true)
- .withStreamingPartitionsSpec(new
StreamingPartitionsSpec(List.of(COL_TENANT)))
+ .withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(List.of(COL_TENANT)))
)
.withIoConfig(
ioConfig -> ioConfig
@@ -1085,7 +1085,7 @@ public class EmbeddedDimensionValueSetShardSpecTest
extends EmbeddedClusterTestB
.withTuningConfig(t -> {
t.withMaxRowsPerSegment(1000).withReleaseLocksOnHandoff(true);
if (!partitionFilterDims.isEmpty()) {
- t.withStreamingPartitionsSpec(new
StreamingPartitionsSpec(partitionFilterDims));
+ t.withStreamingPartitionsSpec(new
DimensionValueSetPartitionsSpec(partitionFilterDims));
}
})
.withIoConfig(
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetCollector.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetCollector.java
new file mode 100644
index 00000000000..46156298cbd
--- /dev/null
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetCollector.java
@@ -0,0 +1,191 @@
+/*
+ * 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.indexing.seekablestream;
+
+import com.google.common.collect.Sets;
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.DimensionValueSetShardSpec;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * {@link StreamingShardSpecCollector} for {@link
DimensionValueSetPartitionsSpec}: it records, per segment, the
+ * distinct values observed for each configured dimension and stamps each
segment with a
+ * {@link DimensionValueSetShardSpec} at publish time so the broker can prune
it.
+ *
+ * <p>A {@code null} element denotes an observed null/missing value (kept
distinct from {@code ""}) so that
+ * {@code IS NULL} queries are not pruned.
+ *
+ * <p>Thread-safety follows the {@link StreamingShardSpecCollector} contract:
+ * {@link #observedPartitionDimValuesBySegment} is a
+ * {@link ConcurrentHashMap} keyed by {@link SegmentId} whose value sets are
{@link Collections#synchronizedSet} written
+ * by the run loop ({@link #collect}) and snapshotted under their own monitor
by the publish path ({@link #annotate}).
+ */
+public class DimensionValueSetCollector implements StreamingShardSpecCollector
+{
+ private static final Logger log = new
Logger(DimensionValueSetCollector.class);
+
+ private final List<String> partitionDimensions;
+
+ /**
+ * If non-null, the maximum number of distinct values to record for a single
dimension on a single segment. A segment
+ * whose observed values for a dimension exceed this cap omits that
dimension from its stamped filter map (pruning
+ * disabled for it on that segment), guarding against shard-spec bloat.
+ */
+ @Nullable
+ private final Integer maxValuesPerDimension;
+
+ /**
+ * Observed values per tracked dimension, keyed by segment identifier. Inner
sets permit null and are written by the
+ * run loop / read by the publish path under their own monitor. Entries are
removed on successful publish via
+ * {@link #onSegmentPublished}; a publish failure is terminal for the task,
so any remaining entries are reclaimed at
+ * task teardown.
+ */
+ private final ConcurrentHashMap<SegmentId, Map<String, Set<String>>>
observedPartitionDimValuesBySegment =
+ new ConcurrentHashMap<>();
+
+ /**
+ * Segment identifiers restored from disk at startup (i.e. spanning a task
restart). Their pre-restart rows are not
+ * re-read, so {@link #observedPartitionDimValuesBySegment} would
under-include values; to avoid wrongly pruning them,
+ * such segments are published with an empty-filter (non-pruning) {@link
DimensionValueSetShardSpec} instead of one
+ * declaring observed values.
+ */
+ private final Set<SegmentId> restartSpannedSegments =
Sets.newConcurrentHashSet();
+
+ public DimensionValueSetCollector(List<String> partitionDimensions,
@Nullable Integer maxValuesPerDimension)
+ {
+ this.partitionDimensions = partitionDimensions;
+ this.maxValuesPerDimension = maxValuesPerDimension;
+ }
+
+ @Override
+ public void collect(SegmentId segmentId, InputRow row)
+ {
+ if (partitionDimensions.isEmpty()) {
+ return;
+ }
+ final Map<String, Set<String>> segValues =
+ observedPartitionDimValuesBySegment.computeIfAbsent(segmentId, k ->
new ConcurrentHashMap<>());
+ for (String dim : partitionDimensions) {
+ final Set<String> dimSet = segValues.computeIfAbsent(
+ dim,
+ k -> Collections.synchronizedSet(new HashSet<>())
+ );
+ // Empty getDimension result means a null/missing value; record null so
IS NULL is not pruned
+ // (distinct from "", which getDimension returns as [""]).
+ final List<String> dimValues = row.getDimension(dim);
+ if (dimValues == null || dimValues.isEmpty()) {
+ dimSet.add(null);
+ } else {
+ dimSet.addAll(dimValues);
+ }
+ }
+ }
+
+ @Override
+ public void onSegmentsRestored(Collection<SegmentId> segmentIds)
+ {
+ if (segmentIds.isEmpty()) {
+ return;
+ }
+ restartSpannedSegments.addAll(segmentIds);
+ log.warn(
+ "Disabling partition-filter pruning for %d segment(s) restored across
a task restart: %s",
+ segmentIds.size(),
+ segmentIds
+ );
+ }
+
+ /**
+ * Stamps a segment with a {@link DimensionValueSetShardSpec} declaring its
observed dimension values so the broker
+ * can prune it. We always return a {@link DimensionValueSetShardSpec},
falling back to an empty (non-pruning) filter
+ * map when values can't be safely declared, so segments in an interval stay
class-uniform for
+ * {@link
org.apache.druid.segment.realtime.appenderator.SegmentPublisherHelper}. A null
observed value is carried
+ * through (distinct from {@code ""}) so {@code IS NULL} queries are not
pruned.
+ */
+ @Override
+ public DataSegment annotate(DataSegment s)
+ {
+ final Map<String, List<String>> snapshotFilters = new HashMap<>();
+ final SegmentId lookupKey = s.getId();
+ final Map<String, Set<String>> segObserved =
observedPartitionDimValuesBySegment.get(lookupKey);
+ // Leave filters empty for restart-spanned segments: their pre-restart
values can't be re-observed.
+ if (!restartSpannedSegments.contains(lookupKey) && segObserved != null) {
+ for (String dim : partitionDimensions) {
+ final Set<String> vals = segObserved.get(dim);
+ if (vals == null) {
+ continue;
+ }
+ // vals is a synchronized set written by the run loop; copy it under
its monitor to iterate safely.
+ final List<String> snapshot;
+ synchronized (vals) {
+ if (vals.isEmpty()) {
+ continue;
+ }
+ // Over-cap: omit this dim from the stamped filter map (still a
DimensionValueSetShardSpec for
+ // class-uniformity; possibleInDomain treats an absent dim as
unconstrained, so pruning is disabled
+ // for it on this segment).
+ if (maxValuesPerDimension != null && vals.size() >
maxValuesPerDimension) {
+ log.warn(
+ "Segment[%s] dimension[%s] observed [%d] distinct values,
exceeds maxValuesPerDimension[%d]; "
+ + "pruning disabled for this dimension on this segment.",
+ lookupKey,
+ dim,
+ vals.size(),
+ maxValuesPerDimension
+ );
+ continue;
+ }
+ snapshot = new ArrayList<>(vals);
+ }
+ // Sort for deterministic published metadata; null (missing value)
sorts first.
+ snapshot.sort(Comparator.nullsFirst(Comparator.naturalOrder()));
+ snapshotFilters.put(dim, snapshot);
+ }
+ }
+ return s.withShardSpec(
+ new DimensionValueSetShardSpec(
+ s.getShardSpec().getPartitionNum(),
+ s.getShardSpec().getNumCorePartitions(),
+ snapshotFilters
+ )
+ );
+ }
+
+ @Override
+ public void onSegmentPublished(SegmentId segmentId)
+ {
+ observedPartitionDimValuesBySegment.remove(segmentId);
+ restartSpannedSegments.remove(segmentId);
+ }
+}
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetPartitionsSpec.java
similarity index 65%
copy from
indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
copy to
indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetPartitionsSpec.java
index 79638885fe2..60b02e181ce 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetPartitionsSpec.java
@@ -22,6 +22,7 @@ package org.apache.druid.indexing.seekablestream;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.common.config.Configs;
import org.apache.druid.error.DruidException;
import javax.annotation.Nullable;
@@ -30,28 +31,32 @@ import java.util.List;
import java.util.Objects;
/**
- * Streaming analog of the batch/compaction {@code partitionsSpec}. Configured
in the streaming tuning config, it
- * declares the dimensions whose observed values each published segment should
record in a
- * {@link org.apache.druid.timeline.partition.DimensionValueSetShardSpec} so
the broker can prune segments at query time
- * without waiting for compaction. Unlike batch partitioning this does not
route rows into shards; it only annotates
- * segments with the values they happened to ingest.
+ * {@link StreamingPartitionsSpec} that records, per published segment, the
distinct values observed for a configured set
+ * of dimensions and stamps them onto a {@link
org.apache.druid.timeline.partition.DimensionValueSetShardSpec} so the
+ * broker can prune segments at query time without waiting for compaction. The
collection and stamping are performed by
+ * {@link DimensionValueSetCollector}.
*
* <p>Use low-to-medium cardinality dimensions; the {@code
partitionDimensions} here should be kept in sync with the
- * {@code partitionDimensions} of the compaction config for the same
datasource.
+ * {@code partitionDimensions} of the compaction config for the same
datasource. {@code maxValuesPerDimension}, if set,
+ * caps the number of distinct values recorded for a dimension on a single
segment: a segment whose observed values for a
+ * dimension exceed the cap omits that dimension from its filter map (pruning
is disabled for it on that segment),
+ * guarding against shard-spec bloat from an unexpectedly high-cardinality
dimension.
*/
-public class StreamingPartitionsSpec
+public class DimensionValueSetPartitionsSpec implements StreamingPartitionsSpec
{
+ public static final String TYPE = "dim_value_set";
+
private final List<String> partitionDimensions;
@Nullable
private final Integer maxValuesPerDimension;
@JsonCreator
- public StreamingPartitionsSpec(
+ public DimensionValueSetPartitionsSpec(
@JsonProperty("partitionDimensions") @Nullable List<String>
partitionDimensions,
@JsonProperty("maxValuesPerDimension") @Nullable Integer
maxValuesPerDimension
)
{
- this.partitionDimensions = partitionDimensions == null ?
Collections.emptyList() : partitionDimensions;
+ this.partitionDimensions = Configs.valueOrDefault(partitionDimensions,
Collections.emptyList());
if (maxValuesPerDimension != null && maxValuesPerDimension <= 0) {
throw DruidException.forPersona(DruidException.Persona.USER)
.ofCategory(DruidException.Category.INVALID_INPUT)
@@ -60,7 +65,7 @@ public class StreamingPartitionsSpec
this.maxValuesPerDimension = maxValuesPerDimension;
}
- public StreamingPartitionsSpec(@Nullable List<String> partitionDimensions)
+ public DimensionValueSetPartitionsSpec(@Nullable List<String>
partitionDimensions)
{
this(partitionDimensions, null);
}
@@ -80,6 +85,17 @@ public class StreamingPartitionsSpec
return maxValuesPerDimension;
}
+ @Nullable
+ @Override
+ public StreamingShardSpecCollector createCollector()
+ {
+ // Nothing to collect without dimensions; treat as feature-off so segments
are published unchanged.
+ if (partitionDimensions.isEmpty()) {
+ return null;
+ }
+ return new DimensionValueSetCollector(partitionDimensions,
maxValuesPerDimension);
+ }
+
@Override
public boolean equals(Object o)
{
@@ -89,7 +105,7 @@ public class StreamingPartitionsSpec
if (o == null || getClass() != o.getClass()) {
return false;
}
- StreamingPartitionsSpec that = (StreamingPartitionsSpec) o;
+ DimensionValueSetPartitionsSpec that = (DimensionValueSetPartitionsSpec) o;
return Objects.equals(partitionDimensions, that.partitionDimensions)
&& Objects.equals(maxValuesPerDimension,
that.maxValuesPerDimension);
}
@@ -103,20 +119,9 @@ public class StreamingPartitionsSpec
@Override
public String toString()
{
- return "StreamingPartitionsSpec{"
+ return "DimensionValueSetPartitionsSpec{"
+ "partitionDimensions=" + partitionDimensions
+ ", maxValuesPerDimension=" + maxValuesPerDimension
+ '}';
}
-
- public static List<String> getPartitionDimensionsOrEmpty(@Nullable
StreamingPartitionsSpec spec)
- {
- return spec == null ? List.of() : spec.getPartitionDimensions();
- }
-
- @Nullable
- public static Integer getMaxValuesPerDimensionOrNull(@Nullable
StreamingPartitionsSpec spec)
- {
- return spec == null ? null : spec.getMaxValuesPerDimension();
- }
}
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
index d2054550fd6..b127ff7f5aa 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
@@ -97,8 +97,6 @@ import
org.apache.druid.segment.realtime.appenderator.StreamAppenderatorDriver;
import org.apache.druid.server.security.AuthorizationUtils;
import org.apache.druid.server.security.AuthorizerMapper;
import org.apache.druid.timeline.DataSegment;
-import org.apache.druid.timeline.SegmentId;
-import org.apache.druid.timeline.partition.DimensionValueSetShardSpec;
import org.apache.druid.utils.CollectionUtils;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.joda.time.DateTime;
@@ -120,7 +118,6 @@ import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
-import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
@@ -259,21 +256,13 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
private final Map<PartitionIdType, Long> partitionsThroughput = new
HashMap<>();
/**
- * Observed values per tracked dimension, keyed by segment identifier, used
to stamp the {@link DimensionValueSetShardSpec}
- * at publish time. A {@code null} element denotes an observed null/missing
value (distinct from {@code ""}) so that
- * {@code IS NULL} queries are not pruned. Inner sets permit null and are
written by the run loop / read by the
- * publish thread under their own monitor. Entries are cleared on successful
publish; a publish failure is terminal
- * for the task, so any remaining entries are reclaimed at task teardown
rather than removed individually.
+ * Per-task collector that accumulates information from ingested rows and
stamps each published segment with a prunable
+ * shard spec, or {@code null} when {@link
SeekableStreamIndexTaskTuningConfig#getStreamingPartitionsSpec()} is unset
+ * or has nothing to collect. Created once in the constructor and shared
across the run loop and the
+ * publish path (which runs on the future-completing thread via {@code
MoreExecutors.directExecutor()}).
*/
- private final ConcurrentHashMap<SegmentId, Map<String, Set<String>>>
observedPartitionDimValuesBySegment = new ConcurrentHashMap<>();
-
- /**
- * Segment identifiers restored from disk at startup (i.e. spanning a task
restart). Their pre-restart rows are not
- * re-read, so {@link #observedPartitionDimValuesBySegment} would
under-include values; to avoid wrongly pruning them,
- * such segments are published with an empty-filter (non-pruning) {@link
DimensionValueSetShardSpec} instead of one
- * declaring observed values.
- */
- private final Set<SegmentId> restartSpannedSegments =
Sets.newConcurrentHashSet();
+ @Nullable
+ private final StreamingShardSpecCollector shardSpecCollector;
private volatile DateTime minMessageTime;
private volatile DateTime maxMessageTime;
@@ -288,6 +277,8 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
this.task = task;
this.ioConfig = task.getIOConfig();
this.tuningConfig = task.getTuningConfig();
+ final StreamingPartitionsSpec streamingPartitionsSpec =
tuningConfig.getStreamingPartitionsSpec();
+ this.shardSpecCollector = streamingPartitionsSpec == null ? null :
streamingPartitionsSpec.createCollector();
this.inputRowSchema = InputRowSchemas.fromDataSchema(task.getDataSchema());
this.inputFormat = ioConfig.getInputFormat();
this.stream = ioConfig.getStartSequenceNumbers().getStream();
@@ -476,9 +467,6 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
//milliseconds waited for created segments to be handed off
long handoffWaitMs = 0L;
- final List<String> partitionDimensions =
-
StreamingPartitionsSpec.getPartitionDimensionsOrEmpty(tuningConfig.getStreamingPartitionsSpec());
-
try (final RecordSupplier<PartitionIdType, SequenceOffsetType, RecordType>
recordSupplier =
task.newTaskRecordSupplier(toolbox)) {
this.recordSupplier = recordSupplier;
@@ -523,20 +511,15 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
}
);
- // Segments restored from disk span a task restart; their pre-restart
values can't be re-observed, so record them
- // to fall back to an empty-filter (non-pruning)
DimensionValueSetShardSpec at publish rather than stamping an
- // incomplete filter.
- if (!partitionDimensions.isEmpty()) {
- for (SegmentIdWithShardSpec restored : appenderator.getSegments()) {
- restartSpannedSegments.add(restored.asSegmentId());
- }
- if (!restartSpannedSegments.isEmpty()) {
- log.warn(
- "Disabling partition-filter pruning for %d segment(s) restored
across a task restart: %s",
- restartSpannedSegments.size(),
- restartSpannedSegments
- );
- }
+ // Segments already in the appenderator at startup were restored from
disk across a task restart; tell the
+ // collector.
+ if (shardSpecCollector != null) {
+ shardSpecCollector.onSegmentsRestored(
+ appenderator.getSegments()
+ .stream()
+ .map(SegmentIdWithShardSpec::asSegmentId)
+ .collect(Collectors.toList())
+ );
}
if (restoredMetadata == null) {
@@ -737,25 +720,9 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
);
if (addResult.isOk()) {
- // Accumulate observed dimension values per segment for
DimensionValueSetShardSpec at publish time.
- if (!partitionDimensions.isEmpty()) {
- final SegmentId segmentId =
addResult.getSegmentIdentifier().asSegmentId();
- final Map<String, Set<String>> segValues =
observedPartitionDimValuesBySegment
- .computeIfAbsent(segmentId, k -> new
ConcurrentHashMap<>());
- for (String dim : partitionDimensions) {
- final Set<String> dimSet = segValues.computeIfAbsent(
- dim,
- k -> Collections.synchronizedSet(new HashSet<>())
- );
- // Empty getDimension result means a null/missing value;
record null so IS NULL is not pruned
- // (distinct from "", which getDimension returns as [""
]).
- final List<String> dimValues = row.getDimension(dim);
- if (dimValues == null || dimValues.isEmpty()) {
- dimSet.add(null);
- } else {
- dimSet.addAll(dimValues);
- }
- }
+ // Accumulate per-row info for the segment's prunable shard
spec, stamped at publish time.
+ if (shardSpecCollector != null) {
+
shardSpecCollector.collect(addResult.getSegmentIdentifier().asSegmentId(), row);
}
// If the number of rows in the segment exceeds the
threshold after adding a row,
@@ -1064,88 +1031,33 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
handOffWaitList.removeAll(handoffFinished);
}
+ /**
+ * Returns the per-task {@link StreamingShardSpecCollector} created in the
constructor, or {@code null} when no
+ * {@link StreamingPartitionsSpec} is configured (or the configured one has
nothing to collect).
+ */
+ @Nullable
@VisibleForTesting
- void recordObservedDimensionValueForTest(SegmentId segmentId, String
dimension, @Nullable String value)
- {
- observedPartitionDimValuesBySegment
- .computeIfAbsent(segmentId, k -> new ConcurrentHashMap<>())
- .computeIfAbsent(dimension, k -> Collections.synchronizedSet(new
HashSet<>()))
- .add(value);
- }
-
- @VisibleForTesting
- void markSegmentRestartSpannedForTest(SegmentId segmentId)
+ StreamingShardSpecCollector getShardSpecCollector()
{
- restartSpannedSegments.add(segmentId);
+ return shardSpecCollector;
}
/**
- * Stamps a segment with a {@link DimensionValueSetShardSpec} declaring its
observed dimension values so the broker can
- * prune it. When the feature is on we always return a {@link
DimensionValueSetShardSpec}, falling back to an empty
- * (non-pruning) filter map when values can't be safely declared, so
segments in an interval stay class-uniform for
- * {@link
org.apache.druid.segment.realtime.appenderator.SegmentPublisherHelper}. A null
observed value is carried
- * through (distinct from {@code ""}) so {@code IS NULL} queries are not
pruned.
+ * Delegates to {@link StreamingShardSpecCollector#annotate} to stamp a
segment's shard spec at publish time,
+ * returning the segment unchanged when no {@link StreamingPartitionsSpec}
is configured. Safe to apply
+ * unconditionally on the publish path.
*/
@VisibleForTesting
DataSegment annotateSegmentWithPartitionDimensionValues(DataSegment s)
{
- final StreamingPartitionsSpec partitionsSpec =
tuningConfig.getStreamingPartitionsSpec();
- final List<String> partitionDimensions =
StreamingPartitionsSpec.getPartitionDimensionsOrEmpty(partitionsSpec);
- if (CollectionUtils.isNullOrEmpty(partitionDimensions)) {
- return s;
- }
- final Integer maxValuesPerDimension =
StreamingPartitionsSpec.getMaxValuesPerDimensionOrNull(partitionsSpec);
- final Map<String, List<String>> snapshotFilters = new HashMap<>();
- final SegmentId lookupKey = s.getId();
- final Map<String, Set<String>> segObserved =
observedPartitionDimValuesBySegment.get(lookupKey);
- // Leave filters empty for restart-spanned segments: their pre-restart
values can't be re-observed.
- if (!restartSpannedSegments.contains(lookupKey) && segObserved != null) {
- for (String dim : partitionDimensions) {
- final Set<String> vals = segObserved.get(dim);
- if (vals == null) {
- continue;
- }
- // vals is a synchronized set written by the run loop; copy it under
its monitor to iterate safely.
- final List<String> snapshot;
- synchronized (vals) {
- if (vals.isEmpty()) {
- continue;
- }
- // Over-cap: omit this dim from the stamped filter map (still a
DimensionValueSetShardSpec for
- // class-uniformity; possibleInDomain treats an absent dim as
unconstrained, so pruning is disabled
- // for it on this segment).
- if (maxValuesPerDimension != null && vals.size() >
maxValuesPerDimension) {
- log.warn(
- "Segment[%s] dimension[%s] observed [%d] distinct values,
exceeds maxValuesPerDimension[%d]; "
- + "pruning disabled for this dimension on this segment.",
- lookupKey,
- dim,
- vals.size(),
- maxValuesPerDimension
- );
- continue;
- }
- snapshot = new ArrayList<>(vals);
- }
- // Sort for deterministic published metadata; null (missing value)
sorts first.
- snapshot.sort(Comparator.nullsFirst(Comparator.naturalOrder()));
- snapshotFilters.put(dim, snapshot);
- }
- }
- return s.withShardSpec(
- new DimensionValueSetShardSpec(
- s.getShardSpec().getPartitionNum(),
- s.getShardSpec().getNumCorePartitions(),
- snapshotFilters
- )
- );
+ return shardSpecCollector == null ? s : shardSpecCollector.annotate(s);
}
private void publishAndRegisterHandoff(SequenceMetadata<PartitionIdType,
SequenceOffsetType> sequenceMetadata)
{
log.debug("Publishing segments for sequence [%s].", sequenceMetadata);
- // annotateSegmentWithPartitionDimensionValues returns the segment
unchanged when partition filters are not configured,
+ // annotateSegmentWithPartitionDimensionValues returns the segment
unchanged when there is no shardSpecCollector,
// so it is always safe to apply here.
final ListenableFuture<SegmentsAndCommitMetadata> publishFuture =
Futures.transform(
driver.publish(
@@ -1189,10 +1101,10 @@ public abstract class
SeekableStreamIndexTaskRunner<PartitionIdType, SequenceOff
);
log.infoSegments(publishedSegmentsAndCommitMetadata.getSegments(),
"Published segments");
- for (DataSegment segment :
publishedSegmentsAndCommitMetadata.getSegments()) {
- final SegmentId segmentId = segment.getId();
- observedPartitionDimValuesBySegment.remove(segmentId);
- restartSpannedSegments.remove(segmentId);
+ if (shardSpecCollector != null) {
+ for (DataSegment segment :
publishedSegmentsAndCommitMetadata.getSegments()) {
+ shardSpecCollector.onSegmentPublished(segment.getId());
+ }
}
publishedSequences.add(sequenceMetadata.getSequenceName());
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
index 79638885fe2..01a00b21486 100644
---
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingPartitionsSpec.java
@@ -19,104 +19,43 @@
package org.apache.druid.indexing.seekablestream;
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonInclude;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import org.apache.druid.error.DruidException;
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.databind.annotation.JsonTypeResolver;
+import org.apache.druid.jackson.StrictTypeIdResolver;
import javax.annotation.Nullable;
-import java.util.Collections;
-import java.util.List;
-import java.util.Objects;
/**
* Streaming analog of the batch/compaction {@code partitionsSpec}. Configured
in the streaming tuning config, it
- * declares the dimensions whose observed values each published segment should
record in a
- * {@link org.apache.druid.timeline.partition.DimensionValueSetShardSpec} so
the broker can prune segments at query time
+ * declares <em>what</em> a streaming task should collect from its rows so
each published segment can be stamped with a
+ * prunable {@link org.apache.druid.timeline.partition.ShardSpec}, letting the
broker prune segments at query time
* without waiting for compaction. Unlike batch partitioning this does not
route rows into shards; it only annotates
- * segments with the values they happened to ingest.
+ * segments based on the values they happened to ingest.
*
- * <p>Use low-to-medium cardinality dimensions; the {@code
partitionDimensions} here should be kept in sync with the
- * {@code partitionDimensions} of the compaction config for the same
datasource.
+ * <p>This is a polymorphic, pluggable type (mirroring {@link
org.apache.druid.indexer.partitions.PartitionsSpec}): each
+ * implementation pairs a per-task {@link StreamingShardSpecCollector} (the
"collect something per row, then build a
+ * shard spec" operation) with the shard-spec type it produces. The built-in
{@link DimensionValueSetPartitionsSpec}
+ * collects the distinct values of a set of dimensions; future strategies
(e.g. min/max ranges or bloom filters) can be
+ * added by registering a new subtype below with its own collector, without
changing the task runner.
+ *
+ * <p>An omitted {@code type} defaults to {@link
DimensionValueSetPartitionsSpec} for backward compatibility, but an
+ * explicit unknown {@code type} (a typo, or a subtype whose extension isn't
loaded here) is rejected by
+ * {@link StrictTypeIdResolver} rather than silently falling back.
*/
-public class StreamingPartitionsSpec
+@JsonTypeResolver(StrictTypeIdResolver.Builder.class)
+@JsonTypeInfo(use = JsonTypeInfo.Id.CUSTOM, property = "type", defaultImpl =
DimensionValueSetPartitionsSpec.class)
+@JsonSubTypes({
+ @JsonSubTypes.Type(name = DimensionValueSetPartitionsSpec.TYPE, value =
DimensionValueSetPartitionsSpec.class)
+})
+public interface StreamingPartitionsSpec
{
- private final List<String> partitionDimensions;
- @Nullable
- private final Integer maxValuesPerDimension;
-
- @JsonCreator
- public StreamingPartitionsSpec(
- @JsonProperty("partitionDimensions") @Nullable List<String>
partitionDimensions,
- @JsonProperty("maxValuesPerDimension") @Nullable Integer
maxValuesPerDimension
- )
- {
- this.partitionDimensions = partitionDimensions == null ?
Collections.emptyList() : partitionDimensions;
- if (maxValuesPerDimension != null && maxValuesPerDimension <= 0) {
- throw DruidException.forPersona(DruidException.Persona.USER)
- .ofCategory(DruidException.Category.INVALID_INPUT)
- .build("maxValuesPerDimension must be > 0, got
[%d]", maxValuesPerDimension);
- }
- this.maxValuesPerDimension = maxValuesPerDimension;
- }
-
- public StreamingPartitionsSpec(@Nullable List<String> partitionDimensions)
- {
- this(partitionDimensions, null);
- }
-
- @JsonProperty
- @JsonInclude(JsonInclude.Include.NON_EMPTY)
- public List<String> getPartitionDimensions()
- {
- return partitionDimensions;
- }
-
- @Nullable
- @JsonProperty
- @JsonInclude(JsonInclude.Include.NON_NULL)
- public Integer getMaxValuesPerDimension()
- {
- return maxValuesPerDimension;
- }
-
- @Override
- public boolean equals(Object o)
- {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
- StreamingPartitionsSpec that = (StreamingPartitionsSpec) o;
- return Objects.equals(partitionDimensions, that.partitionDimensions)
- && Objects.equals(maxValuesPerDimension,
that.maxValuesPerDimension);
- }
-
- @Override
- public int hashCode()
- {
- return Objects.hash(partitionDimensions, maxValuesPerDimension);
- }
-
- @Override
- public String toString()
- {
- return "StreamingPartitionsSpec{"
- + "partitionDimensions=" + partitionDimensions
- + ", maxValuesPerDimension=" + maxValuesPerDimension
- + '}';
- }
-
- public static List<String> getPartitionDimensionsOrEmpty(@Nullable
StreamingPartitionsSpec spec)
- {
- return spec == null ? List.of() : spec.getPartitionDimensions();
- }
-
+ /**
+ * Creates a task-scoped {@link StreamingShardSpecCollector} that
accumulates per-row information and stamps each
+ * segment's shard spec with it at publish time. Returns {@code null} when
this spec is configured such that there is
+ * nothing to collect (e.g. no dimensions), in which case no collector runs
and segments are published unchanged.
+ * One collector is created per task run.
+ */
@Nullable
- public static Integer getMaxValuesPerDimensionOrNull(@Nullable
StreamingPartitionsSpec spec)
- {
- return spec == null ? null : spec.getMaxValuesPerDimension();
- }
+ StreamingShardSpecCollector createCollector();
}
diff --git
a/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingShardSpecCollector.java
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingShardSpecCollector.java
new file mode 100644
index 00000000000..32720bd3800
--- /dev/null
+++
b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/StreamingShardSpecCollector.java
@@ -0,0 +1,75 @@
+/*
+ * 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.indexing.seekablestream;
+
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.ShardSpec;
+
+import java.util.Collection;
+
+/**
+ * Accumulates information from the rows a streaming task ingests and, at
publish time, stamps each segment with a
+ * {@link ShardSpec} derived from that information. This is the pluggable
strategy behind {@link StreamingPartitionsSpec}:
+ * a new strategy is added by writing a {@link StreamingPartitionsSpec}
subtype plus a matching collector, without
+ * touching the task runner. The built-in {@code DimensionValueSetCollector}
records each segment's observed dimension
+ * values so the broker can prune it at query time without waiting for
compaction.
+ *
+ * <p>One instance is created per task run (via {@link
StreamingPartitionsSpec#createCollector()}) and shared across all
+ * of that task's segments. Implementations <b>must</b> be thread-safe: {@link
#collect} and {@link #annotate} can run
+ * concurrently for the same {@link SegmentId} (see {@link #annotate}).
+ */
+public interface StreamingShardSpecCollector
+{
+ /**
+ * Records whatever information this strategy needs from {@code row} for
{@code segmentId}. Called on the run-loop
+ * thread, once per row successfully added to {@code segmentId}.
+ */
+ void collect(SegmentId segmentId, InputRow row);
+
+ /**
+ * Notifies the collector that {@code segmentIds} were restored from disk
across a task restart. Such segments'
+ * pre-restart rows are not re-read, so the collected information is
incomplete; to avoid wrongly pruning those rows,
+ * {@link #annotate} must return a non-pruning shard spec for them. The full
restored set is known at once, so it is
+ * passed in a single call rather than one segment at a time. Called only at
task startup, before the run loop
+ * begins; idempotent.
+ */
+ void onSegmentsRestored(Collection<SegmentId> segmentIds);
+
+ /**
+ * Returns {@code segment} stamped with a shard spec derived from the
information collected for it. When the segment
+ * was restored across a restart (see {@link #onSegmentsRestored}) or
nothing was collected for it, this returns a
+ * non-pruning fallback shard spec.
+ *
+ * <p>Implementations <b>must</b> keep the shard-spec class uniform within a
time interval: all segments handed to a
+ * single publish must share one shard-spec class, or
+ * {@link
org.apache.druid.segment.realtime.appenderator.SegmentPublisherHelper} rejects
the publish. In practice this
+ * means returning the same shard-spec type for both the pruning and the
non-pruning (fallback) case, differing only
+ * in the declared values. Must be safe to call concurrently with {@link
#collect}.
+ */
+ DataSegment annotate(DataSegment segment);
+
+ /**
+ * Notifies the collector that {@code segmentId} has been successfully
published and handed off, so any per-segment
+ * state accumulated for it can be released. Called only from the
publish-success callback.
+ */
+ void onSegmentPublished(SegmentId segmentId);
+}
diff --git
a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerAuthTest.java
b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerAuthTest.java
index cbc1ca7c7a6..01c9f2564f1 100644
---
a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerAuthTest.java
+++
b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerAuthTest.java
@@ -125,6 +125,7 @@ public class SeekableStreamIndexTaskRunnerAuthTest
EasyMock.expect(tuningConfig.isLogParseExceptions()).andReturn(false).anyTimes();
EasyMock.expect(tuningConfig.getMaxParseExceptions()).andReturn(10).anyTimes();
EasyMock.expect(tuningConfig.getMaxSavedParseExceptions()).andReturn(10).anyTimes();
+
EasyMock.expect(tuningConfig.getStreamingPartitionsSpec()).andReturn(null).anyTimes();
replay(tuningConfig);
SeekableStreamIndexTaskIOConfig<String, String> ioConfig = new
TestSeekableStreamIndexTaskIOConfig();
diff --git
a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerTest.java
b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerTest.java
index bbc39c156ca..d4cf8bb8961 100644
---
a/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerTest.java
+++
b/indexing-service/src/test/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunnerTest.java
@@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures;
import org.apache.druid.client.DruidServer;
import org.apache.druid.data.input.InputRow;
+import org.apache.druid.data.input.MapBasedInputRow;
import org.apache.druid.data.input.impl.DimensionsSpec;
import org.apache.druid.data.input.impl.JsonInputFormat;
import org.apache.druid.data.input.impl.StringDimensionSchema;
@@ -82,10 +83,12 @@ import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.CopyOnWriteArrayList;
@@ -362,10 +365,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
Map.of("partition", "0"),
- Map.of("partition", "100")
+ Map.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant")));
final DataSegment segment = createSingleSegment();
final SegmentId lookupKey = segment.getId();
@@ -397,10 +399,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant")));
final DataSegment segment = createSingleSegment();
final SegmentId lookupKey = segment.getId();
@@ -432,10 +433,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant")));
// Two partitions in one interval: partition 0 was restored from disk
across a restart, partition 1 created after.
final List<DataSegment> sameIntervalPartitions = CreateDataSegments
@@ -477,17 +477,16 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant", "region"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant",
"region")));
final DataSegment segment = createSingleSegment();
final SegmentId lookupKey = segment.getId();
- // tenant saw a non-null value and (in another row) a null/missing value;
region only saw non-null values.
- observe(runner, lookupKey, "tenant", "tenant_a", null);
- observe(runner, lookupKey, "region", "us-west");
+ // Row 1: tenant=tenant_a, region=us-west. Row 2: region=us-west but
tenant missing (a null/missing tenant value).
+ collectRow(runner, lookupKey, Map.of("tenant", "tenant_a", "region",
"us-west"));
+ collectRow(runner, lookupKey, Map.of("region", "us-west"));
final DataSegment annotated =
runner.annotateSegmentWithPartitionDimensionValues(segment);
@@ -515,10 +514,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant")));
final DataSegment segment = createSingleSegment();
final SegmentId lookupKey = segment.getId();
@@ -546,10 +544,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"))
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant")));
// No observe(...) call: nothing was recorded for this segment.
final DataSegment annotated =
runner.annotateSegmentWithPartitionDimensionValues(createSingleSegment());
@@ -568,11 +565,11 @@ public class SeekableStreamIndexTaskRunnerTest
@Test
public void testFeatureOffReturnsSegmentUnchanged() throws Exception
{
+ // No streamingPartitionsSpec passed: the feature is off.
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
ImmutableMap.of("partition", "100")
);
-
Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec()).thenReturn(null);
final DataSegment segment = createSingleSegment();
final DataSegment annotated =
runner.annotateSegmentWithPartitionDimensionValues(segment);
@@ -586,10 +583,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"), 3)
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant"), 3));
final DataSegment segment = createSingleSegment();
observe(runner, segment.getId(), "tenant", "tenant_a", "tenant_b",
"tenant_c");
@@ -609,10 +605,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"), 2)
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant"), 2));
final DataSegment segment = createSingleSegment();
observe(runner, segment.getId(), "tenant", "tenant_a", "tenant_b",
"tenant_c");
@@ -632,14 +627,16 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant", "region"), 2)
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant",
"region"), 2));
final DataSegment segment = createSingleSegment();
- observe(runner, segment.getId(), "tenant", "tenant_a", "tenant_b",
"tenant_c");
- observe(runner, segment.getId(), "region", "us-west", "us-east");
+ // Each row sets both tracked dims (collect evaluates all configured dims
per row). tenant sees 3 distinct values
+ // (over cap), region sees 2 (at cap).
+ collectRow(runner, segment.getId(), Map.of("tenant", "tenant_a", "region",
"us-west"));
+ collectRow(runner, segment.getId(), Map.of("tenant", "tenant_b", "region",
"us-east"));
+ collectRow(runner, segment.getId(), Map.of("tenant", "tenant_c", "region",
"us-west"));
final DataSegment annotated =
runner.annotateSegmentWithPartitionDimensionValues(segment);
@@ -661,10 +658,9 @@ public class SeekableStreamIndexTaskRunnerTest
{
final TestSeekableStreamIndexTaskRunner runner = createRunner(
ImmutableMap.of("partition", "0"),
- ImmutableMap.of("partition", "100")
+ ImmutableMap.of("partition", "100"),
+ new DimensionValueSetPartitionsSpec(List.of("tenant"), 2)
);
- Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec())
- .thenReturn(new StreamingPartitionsSpec(List.of("tenant"), 2));
final DataSegment segment = createSingleSegment();
observe(runner, segment.getId(), "tenant", "tenant_a", "tenant_b", null);
@@ -688,6 +684,11 @@ public class SeekableStreamIndexTaskRunnerTest
.get(0);
}
+ /**
+ * Feeds the collector one row per value through its real {@link
StreamingShardSpecCollector#collect} API. A
+ * {@code null} value is sent as a row missing {@code dimension} (so {@code
getDimension} returns empty and the
+ * collector records a null), matching how a null/missing ingested value is
observed in production.
+ */
private static void observe(
SeekableStreamIndexTaskRunner runner,
SegmentId segmentId,
@@ -696,13 +697,37 @@ public class SeekableStreamIndexTaskRunnerTest
)
{
for (String value : values) {
- runner.recordObservedDimensionValueForTest(segmentId, dimension, value);
+ collectRow(runner, segmentId, value == null ? Map.of() :
Map.of(dimension, value));
}
}
+ /**
+ * Feeds the collector a single row built from {@code event} through its real
+ * {@link StreamingShardSpecCollector#collect} API. A dimension absent from
{@code event} is observed as a
+ * null/missing value.
+ */
+ private static void collectRow(
+ SeekableStreamIndexTaskRunner runner,
+ SegmentId segmentId,
+ Map<String, Object> event
+ )
+ {
+ final StreamingShardSpecCollector collector = Objects.requireNonNull(
+ runner.getShardSpecCollector(),
+ "streamingPartitionsSpec must be configured before collecting rows"
+ );
+ collector.collect(
+ segmentId,
+ new MapBasedInputRow(DateTimes.nowUtc(), new
ArrayList<>(event.keySet()), event)
+ );
+ }
+
private static void markRestartSpanned(SeekableStreamIndexTaskRunner runner,
SegmentId segmentId)
{
- runner.markSegmentRestartSpannedForTest(segmentId);
+ Objects.requireNonNull(
+ runner.getShardSpecCollector(),
+ "streamingPartitionsSpec must be configured before marking
restart-spanned segments"
+ ).onSegmentsRestored(Collections.singletonList(segmentId));
}
private TaskToolbox createTaskToolbox()
@@ -750,7 +775,16 @@ public class SeekableStreamIndexTaskRunnerTest
Map<String, String> endOffsets
)
{
- return createRunner(createDataSchema(), null, null, null, startOffsets,
endOffsets);
+ return createRunner(startOffsets, endOffsets, null);
+ }
+
+ private TestSeekableStreamIndexTaskRunner createRunner(
+ Map<String, String> startOffsets,
+ Map<String, String> endOffsets,
+ @Nullable StreamingPartitionsSpec streamingPartitionsSpec
+ )
+ {
+ return createRunner(createDataSchema(), null, null, null, startOffsets,
endOffsets, streamingPartitionsSpec);
}
private TestSeekableStreamIndexTaskRunner createRunnerWithMessageTimeBounds(
@@ -765,7 +799,8 @@ public class SeekableStreamIndexTaskRunnerTest
minMessageTime,
maxMessageTime,
ImmutableMap.of(),
- ImmutableMap.of()
+ ImmutableMap.of(),
+ null
);
}
@@ -775,7 +810,8 @@ public class SeekableStreamIndexTaskRunnerTest
@Nullable DateTime minMessageTime,
@Nullable DateTime maxMessageTime,
Map<String, String> startOffsets,
- Map<String, String> endOffsets
+ Map<String, String> endOffsets,
+ @Nullable StreamingPartitionsSpec streamingPartitionsSpec
)
{
final SeekableStreamIndexTaskTuningConfig tuningConfig =
Mockito.mock(SeekableStreamIndexTaskTuningConfig.class);
@@ -791,6 +827,7 @@ public class SeekableStreamIndexTaskRunnerTest
);
Mockito.when(tuningConfig.getIntermediateHandoffPeriod()).thenReturn(Period.minutes(1));
+
Mockito.when(tuningConfig.getStreamingPartitionsSpec()).thenReturn(streamingPartitionsSpec);
Mockito.when(ioConfig.getRefreshRejectionPeriodsInMinutes()).thenReturn(refreshRejectionPeriodsInMinutes);
Mockito.when(ioConfig.getMaximumMessageTime()).thenReturn(maxMessageTime);
Mockito.when(ioConfig.getMinimumMessageTime()).thenReturn(minMessageTime);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
