refactor: Make shard spec collector an interface for extensibility

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
 

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,37 +166,80 @@ 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
   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
   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<String> partitionDimensionsOf(KafkaIndexTaskTuningConfig config)
+  {
+    return ((DimensionValueSetPartitionsSpec) config.getStreamingPartitionsSpec()).getPartitionDimensions();
+  }
+
   @Test
   public void testConvert()
   {

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<String>
         .withTuningConfig(t -> {
           t.withMaxRowsPerSegment(1000).withReleaseLocksOnHandoff(true);
           if (!partitionFilterDims.isEmpty()) {
-            t.withStreamingPartitionsSpec(new StreamingPartitionsSpec(partitionFilterDims));
+            t.withStreamingPartitionsSpec(new DimensionValueSetPartitionsSpec(partitionFilterDims));
           }
         })
         .withIoConfig(