diff --git a/docs/ingestion/kafka-ingestion.md b/docs/ingestion/kafka-ingestion.md index 79232c9f5a28..0441dde69d16 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 scenarios, ingestion can stall or pause frequently, causing it to fall behind. You can use additional threads to parallelize the segment creation without blocking ingestion as long as there are sufficient CPU resources available.|No|1| -|`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](#streaming-partitions-spec) for details.|No|null| ## 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 4bf668bbc867..b19f4d7a3f31 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 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 void testSerdeWithStreamingPartitionsSpec() throws Exception ); 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 void testSerdeWithoutStreamingPartitionsSpecIsNull() throws Exception 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 void testSerdeWithEmptyStringPartitionDimension() throws Exception { // 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 void testSerdeWithNumericLookingPartitionDimension() throws Exception { // 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 @@ private KafkaIndexTaskTuningConfig roundTripWithStreamingPartitionsSpec(String p ); } + private static List 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 74ad0c6ecdb9..781f1e2cced7 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.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 void test_multiDimensionAndMultiValuePartitionDimensionValues() 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 void test_numericDimension_isNotPruned() .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 void test_incorrectValuesInTopic_inMemory_handledGracefully() .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 void test_maxValuesPerDimensionCap_overCapDimensionOmitted() .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 @@ private KafkaSupervisorSpecBuilder newSupervisorBuilder() 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 @@ private KafkaSupervisorSpec dayGranularitySupervisor(String topic, List .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 000000000000..46156298cbd6 --- /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. + * + *

A {@code null} element denotes an observed null/missing value (kept distinct from {@code ""}) so that + * {@code IS NULL} queries are not pruned. + * + *

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 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>> 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 restartSpannedSegments = Sets.newConcurrentHashSet(); + + public DimensionValueSetCollector(List partitionDimensions, @Nullable Integer maxValuesPerDimension) + { + this.partitionDimensions = partitionDimensions; + this.maxValuesPerDimension = maxValuesPerDimension; + } + + @Override + public void collect(SegmentId segmentId, InputRow row) + { + if (partitionDimensions.isEmpty()) { + return; + } + final Map> segValues = + observedPartitionDimValuesBySegment.computeIfAbsent(segmentId, k -> new ConcurrentHashMap<>()); + for (String dim : partitionDimensions) { + final Set 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 dimValues = row.getDimension(dim); + if (dimValues == null || dimValues.isEmpty()) { + dimSet.add(null); + } else { + dimSet.addAll(dimValues); + } + } + } + + @Override + public void onSegmentsRestored(Collection 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> snapshotFilters = new HashMap<>(); + final SegmentId lookupKey = s.getId(); + final Map> 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 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 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/DimensionValueSetPartitionsSpec.java b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetPartitionsSpec.java new file mode 100644 index 000000000000..60b02e181ce8 --- /dev/null +++ b/indexing-service/src/main/java/org/apache/druid/indexing/seekablestream/DimensionValueSetPartitionsSpec.java @@ -0,0 +1,127 @@ +/* + * 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.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; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +/** + * {@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}. + * + *

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 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 DimensionValueSetPartitionsSpec implements StreamingPartitionsSpec +{ + public static final String TYPE = "dim_value_set"; + + private final List partitionDimensions; + @Nullable + private final Integer maxValuesPerDimension; + + @JsonCreator + public DimensionValueSetPartitionsSpec( + @JsonProperty("partitionDimensions") @Nullable List partitionDimensions, + @JsonProperty("maxValuesPerDimension") @Nullable Integer maxValuesPerDimension + ) + { + this.partitionDimensions = Configs.valueOrDefault(partitionDimensions, Collections.emptyList()); + 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 DimensionValueSetPartitionsSpec(@Nullable List partitionDimensions) + { + this(partitionDimensions, null); + } + + @JsonProperty + @JsonInclude(JsonInclude.Include.NON_EMPTY) + public List getPartitionDimensions() + { + return partitionDimensions; + } + + @Nullable + @JsonProperty + @JsonInclude(JsonInclude.Include.NON_NULL) + public Integer getMaxValuesPerDimension() + { + 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) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DimensionValueSetPartitionsSpec that = (DimensionValueSetPartitionsSpec) 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 "DimensionValueSetPartitionsSpec{" + + "partitionDimensions=" + partitionDimensions + + ", maxValuesPerDimension=" + maxValuesPerDimension + + '}'; + } +} 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 d2054550fd63..b127ff7f5aad 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.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.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 enum Status private final Map 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>> 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 restartSpannedSegments = Sets.newConcurrentHashSet(); + @Nullable + private final StreamingShardSpecCollector shardSpecCollector; private volatile DateTime minMessageTime; private volatile DateTime maxMessageTime; @@ -288,6 +277,8 @@ public SeekableStreamIndexTaskRunner( 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 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception //milliseconds waited for created segments to be handed off long handoffWaitMs = 0L; - final List partitionDimensions = - StreamingPartitionsSpec.getPartitionDimensionsOrEmpty(tuningConfig.getStreamingPartitionsSpec()); - try (final RecordSupplier recordSupplier = task.newTaskRecordSupplier(toolbox)) { this.recordSupplier = recordSupplier; @@ -523,20 +511,15 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception } ); - // 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 void run() ); 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> segValues = observedPartitionDimValuesBySegment - .computeIfAbsent(segmentId, k -> new ConcurrentHashMap<>()); - for (String dim : partitionDimensions) { - final Set 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 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 @@ private void checkPublishAndHandoffFailure() throws ExecutionException, Interrup 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 partitionDimensions = StreamingPartitionsSpec.getPartitionDimensionsOrEmpty(partitionsSpec); - if (CollectionUtils.isNullOrEmpty(partitionDimensions)) { - return s; - } - final Integer maxValuesPerDimension = StreamingPartitionsSpec.getMaxValuesPerDimensionOrNull(partitionsSpec); - final Map> snapshotFilters = new HashMap<>(); - final SegmentId lookupKey = s.getId(); - final Map> 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 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 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 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 publishFuture = Futures.transform( driver.publish( @@ -1189,10 +1101,10 @@ public void onSuccess(SegmentsAndCommitMetadata publishedSegmentsAndCommitMetada ); 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 79638885fe2c..01a00b21486b 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 what 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. * - *

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. + *

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. + * + *

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 partitionDimensions; - @Nullable - private final Integer maxValuesPerDimension; - - @JsonCreator - public StreamingPartitionsSpec( - @JsonProperty("partitionDimensions") @Nullable List 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 partitionDimensions) - { - this(partitionDimensions, null); - } - - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_EMPTY) - public List 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 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 000000000000..32720bd38009 --- /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. + * + *

One instance is created per task run (via {@link StreamingPartitionsSpec#createCollector()}) and shared across all + * of that task's segments. Implementations must 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 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. + * + *

Implementations must 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 cbc1ca7c7a61..01c9f2564f19 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 Authorizer getAuthorizer(String name) 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 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 bbc39c156cae..d4cf8bb8961d 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.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.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 void testAnnotateSegmentStampsDimensionValueSetShardSpecForObservedValues { 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 void testRestartSpannedSegmentGetsEmptyFilterDimensionValueSetShardSpec() { 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 void testRestartBatchMixingFallbackAndObservedSegmentsPublishesWithDimens { 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 sameIntervalPartitions = CreateDataSegments @@ -477,17 +477,16 @@ public void testNullValuedDimensionDeclaresNullInPartitionDimensionValues() thro { 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 void testOnlyNullValuedDimensionDeclaresNull() throws Exception { 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 void testSegmentWithNoObservedValuesGetsEmptyFilterDimensionValueSetShard { 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 void testSegmentWithNoObservedValuesGetsEmptyFilterDimensionValueSetShard @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 void testCapAtBoundaryStampsValuesNormally() throws Exception { 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 void testCapExceededOmitsDimensionFromFilterMap() throws Exception { 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 void testCapEnforcedPerDimensionIndependently() throws Exception { 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 void testNullCountsTowardCap() throws Exception { 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 @@ private static DataSegment createSingleSegment() .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 @@ private static void observe( ) { 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 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 @@ private TestSeekableStreamIndexTaskRunner createRunner( Map endOffsets ) { - return createRunner(createDataSchema(), null, null, null, startOffsets, endOffsets); + return createRunner(startOffsets, endOffsets, null); + } + + private TestSeekableStreamIndexTaskRunner createRunner( + Map startOffsets, + Map endOffsets, + @Nullable StreamingPartitionsSpec streamingPartitionsSpec + ) + { + return createRunner(createDataSchema(), null, null, null, startOffsets, endOffsets, streamingPartitionsSpec); } private TestSeekableStreamIndexTaskRunner createRunnerWithMessageTimeBounds( @@ -765,7 +799,8 @@ private TestSeekableStreamIndexTaskRunner createRunnerWithMessageTimeBounds( minMessageTime, maxMessageTime, ImmutableMap.of(), - ImmutableMap.of() + ImmutableMap.of(), + null ); } @@ -775,7 +810,8 @@ private TestSeekableStreamIndexTaskRunner createRunner( @Nullable DateTime minMessageTime, @Nullable DateTime maxMessageTime, Map startOffsets, - Map endOffsets + Map endOffsets, + @Nullable StreamingPartitionsSpec streamingPartitionsSpec ) { final SeekableStreamIndexTaskTuningConfig tuningConfig = Mockito.mock(SeekableStreamIndexTaskTuningConfig.class); @@ -791,6 +827,7 @@ private TestSeekableStreamIndexTaskRunner createRunner( ); 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);