Skip to content

feat: numeric pruning for dim_value_set segments #19704

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jaykanakiya wants to merge 3 commits into
base: master
Choose a base branch
from
Open

feat: numeric pruning for dim_value_set segments #19704

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f75ef78
feat: LONG numeric pruning for dim_value_set segments
jaykanakiya Jul 20, 2026
933d504
Merge branch 'master' into dim_value_set
jaykanakiya Jul 20, 2026
4925275
skip legacy string pruning for type-stamped LONG dimensions
jaykanakiya Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 43 additions & 22 deletions ...java/org/apache/druid/indexing/kafka/simulate/EmbeddedDimensionValueSetShardSpecTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,8 +244,8 @@ public void test_multiDimensionAndMultiValuePartitionDimensionValues()
Assertions.assertEquals("5", cluster.runSql(
"SELECT COUNT(*) FROM %s WHERE %s = 'tenant_b'", dataSource, COL_TENANT));

// Numeric dimension equality filter still returns correct counts. Note this does NOT prune (numeric filters
// opt out of segment pruning); pruning behavior is asserted in test_numericDimension_isNotPruned.
// Numeric equality filter correctness. Its LongDimensionSchema makes it eligible for type-gated pruning, which is
// asserted in test_numericLongDimension_isPruned.
Assertions.assertEquals("5", cluster.runSql(
"SELECT COUNT(*) FROM %s WHERE %s = 1", dataSource, colRegionCode));
Assertions.assertEquals("5", cluster.runSql(
Expand Down Expand Up @@ -277,26 +277,24 @@ public void test_multiDimensionAndMultiValuePartitionDimensionValues()
}

/**
* Numeric (Long) tracked dimensions are recorded in the shard spec but are NOT used for pruning: numeric query
* filters opt out of segment pruning (their getDimensionRangeSet returns null, since pruning compares values
* lexicographically), so a numeric equality filter scans every segment even though each segment declares exactly
* one numeric value. Queries stay correct; there is simply no pruning benefit.
*
* <p>This is intentional for now. We could either consider extending pruning to numeric types with type-aware
* (non-lexicographic) comparison, or (b) reject numeric dimensions outright when they're declared.
* A numeric (LONG) tracked dimension with an explicit {@link LongDimensionSchema} IS pruned: the broker's
* equality/IN/IS NULL channel is type-gated on {@code dimensionColumnTypes}, which the task stamps only for such a
* schema (a LONG stringifies identically on ingest and query, so set membership is exact). This is exact
* set-membership, not a range comparison, so it works even though the dimension is numeric.
*/
@Test
public void test_numericDimension_isNotPruned()
public void test_numericLongDimension_isPruned()
{
final String colCode = "code"; // numeric (Long), tracked
final String colCode = "code"; // numeric (Long), tracked, explicit LongDimensionSchema
final String topic = dataSource + "_topic";
kafkaServer.createTopicWithPartitions(topic, 1);

// One distinct numeric code per DAY segment: Day1=1, Day2=2, Day3=3. If numeric pruning worked, "code = 1" would
// scan only Day1; because it does NOT, all three segments are scanned.
// One code per day: Day1=1, Day2=2, Day3=3; stamped LONG, so "code = 1" prunes to Day1 only. Day1's raw tokens are
// deliberately non-canonical ("00001", "+1") to exercise the capture-time coercion to the canonical "1" (the
// broker's query value); without it Day1 would stamp ["00001","+1"], miss the query set {"1"}, and be wrongly pruned.
final List<ProducerRecord<byte[], byte[]>> records = new ArrayList<>();
records.add(record(topic, "%s,1,val_0", DateTimes.of("2025-01-01T01:00:00")));
records.add(record(topic, "%s,1,val_1", DateTimes.of("2025-01-01T02:00:00")));
records.add(record(topic, "%s,00001,val_0", DateTimes.of("2025-01-01T01:00:00")));
records.add(record(topic, "%s,+1,val_1", DateTimes.of("2025-01-01T02:00:00")));
records.add(record(topic, "%s,2,val_2", DateTimes.of("2025-01-02T01:00:00")));
records.add(record(topic, "%s,2,val_3", DateTimes.of("2025-01-02T02:00:00")));
records.add(record(topic, "%s,3,val_4", DateTimes.of("2025-01-03T01:00:00")));
Expand Down Expand Up @@ -341,20 +339,43 @@ public void test_numericDimension_isNotPruned()
awaitRowsProcessed(6);
suspendAndAwaitHandoff(spec, 1);

// The numeric value IS recorded in the shard spec (stringified), even though it is never used for pruning.
verifyAllSegmentsHaveDimensionValueSetShardSpec(dataSource);

// The published shard spec stamps colCode's type as LONG, gating the broker's numeric pruning channel.
final List<Map<String, Object>> shardSpecs = getShardSpecs(dataSource);
for (Map<String, Object> shardSpec : shardSpecs) {
@SuppressWarnings("unchecked")
final Map<String, Object> columnTypes = (Map<String, Object>) shardSpec.get("dimensionColumnTypes");
Assertions.assertNotNull(columnTypes, "Expected dimensionColumnTypes on shard spec: " + shardSpec);
Assertions.assertEquals(
"LONG",
columnTypes.get(colCode),
"Expected " + colCode + " stamped as LONG in dimensionColumnTypes: " + shardSpec
);
}

final Map<String, String> startToSegmentId = getStartToSegmentId(dataSource);
final String day1 = startToSegmentId.get("2025-01-01T00:00:00.000Z");
final String day2 = startToSegmentId.get("2025-01-02T00:00:00.000Z");
final String day3 = startToSegmentId.get("2025-01-03T00:00:00.000Z");
final Set<String> allDays = Set.of(day1, day2, day3);
Assertions.assertNotNull(day1, "Missing Day1 segment id in: " + startToSegmentId);
Assertions.assertNotNull(day2, "Missing Day2 segment id in: " + startToSegmentId);
Assertions.assertNotNull(day3, "Missing Day3 segment id in: " + startToSegmentId);

// Correct count, but ALL segments scanned: numeric equality does not prune even though each segment holds one code.
assertScan("2", allDays, "SELECT COUNT(*) FROM %s WHERE %s = 1", dataSource, colCode);
assertScan("2", allDays, "SELECT COUNT(*) FROM %s WHERE %s = 3", dataSource, colCode);
// A non-existent code returns 0 rows but is still NOT pruned (would be a full prune if numeric pruning worked).
assertScan("0", allDays, "SELECT COUNT(*) FROM %s WHERE %s = 999", dataSource, colCode);
// Day1's non-canonical tokens ("00001", "+1") are stamped as the canonical Long string "1", not verbatim.
@SuppressWarnings("unchecked")
final Set<String> allCodes = shardSpecs.stream()
.map(shardSpec -> ((Map<String, List<String>>) shardSpec.get("partitionDimensionValues")).get(colCode))
.flatMap(List::stream)
.collect(Collectors.toSet());
Assertions.assertEquals(Set.of("1", "2", "3"), allCodes, "Expected canonical LONG codes in partitionDimensionValues");

// Numeric equality now prunes exactly like the string case: each query scans only the one segment whose
// stamped value set contains the queried code.
assertScan("2", Set.of(day1), "SELECT COUNT(*) FROM %s WHERE %s = 1", dataSource, colCode);
assertScan("2", Set.of(day3), "SELECT COUNT(*) FROM %s WHERE %s = 3", dataSource, colCode);
// A non-existent code is a full prune: zero rows and zero segments scanned.
assertScan("0", Set.of(), "SELECT COUNT(*) FROM %s WHERE %s = 999", dataSource, colCode);

// Strict correctness: the RIGHT rows survive, not just the right count.
assertValues(Set.of("val_0", "val_1"), "SELECT \"%s\" FROM %s WHERE %s = 1", COL_VALUE, dataSource, colCode);
Expand Down
81 changes: 74 additions & 7 deletions ...src/main/java/org/apache/druid/indexing/seekablestream/SeekableStreamIndexTaskRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@
import org.apache.druid.data.input.InputRow;
import org.apache.druid.data.input.InputRowSchema;
import org.apache.druid.data.input.impl.ByteEntity;
import org.apache.druid.data.input.impl.DimensionsSpec;
import org.apache.druid.data.input.impl.LongDimensionSchema;
import org.apache.druid.discovery.DiscoveryDruidNode;
import org.apache.druid.discovery.LookupNodeService;
import org.apache.druid.discovery.NodeRole;
Expand Down Expand Up @@ -79,9 +81,12 @@
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.parsers.ParseException;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.metadata.PendingSegmentRecord;
import org.apache.druid.query.DruidMetrics;
import org.apache.druid.segment.DimensionHandlerUtils;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.incremental.InputRowFilterResult;
import org.apache.druid.segment.incremental.ParseExceptionHandler;
import org.apache.druid.segment.incremental.ParseExceptionReport;
Expand Down Expand Up @@ -478,6 +483,9 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception

final List<String> partitionDimensions =
StreamingPartitionsSpec.getPartitionDimensionsOrEmpty(tuningConfig.getStreamingPartitionsSpec());
// Partition dimensions declared as LONG. Their observed values are coerced to a canonical Long string at capture
// time so the stamped set matches the broker's canonical query value exactly (see the capture loop below).
final Set<String> longPartitionDimensions = getLongPartitionDimensions(partitionDimensions);

try (final RecordSupplier<PartitionIdType, SequenceOffsetType, RecordType> recordSupplier =
task.newTaskRecordSupplier(toolbox)) {
Expand Down Expand Up @@ -747,13 +755,17 @@ public void run()
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);
if (longPartitionDimensions.contains(dim)) {
dimSet.add(canonicalLongValue(row.getRaw(dim), dim));

@FrankChen021 FrankChen021 Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Disable legacy string pruning for stamped LONG dimensions

Canonicalizing raw "00001" or "+1" to "1" is correct for the new typed domain, but the dimension still passes through possibleInDomain first. A native SelectorDimFilter("code", "00001", ...) matches indexed LONG 1 through runtime numeric coercion, while its range set remains literal "00001". The range check therefore sees no overlap with stamped "1" and prunes the segment before typed pruning runs, dropping valid rows. Exclude type-stamped LONG dimensions from the legacy string/range channel, or safely canonicalize that channel, and add a regression test for noncanonical string selectors.

@jaykanakiya jaykanakiya Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed so that long dims now skip the string/range channel and prune only via the type-gated possibleInValueDomain.

} else {
dimSet.addAll(dimValues);
// 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);
}
}
}
}
Expand Down Expand Up @@ -1079,6 +1091,50 @@ void markSegmentRestartSpannedForTest(SegmentId segmentId)
restartSpannedSegments.add(segmentId);
}

/**
* The subset of the given partition dimensions declared with an explicit {@link LongDimensionSchema}. These are the
* only dimensions whose observed values are canonicalized to a Long string at capture and stamped
* {@link ColumnType#LONG} at publish, enabling type-gated numeric pruning at the broker. Schemaless/auto or non-LONG
* dims are excluded (no numeric pruning), avoiding canonicalization hazards.
*/
private Set<String> getLongPartitionDimensions(List<String> partitionDimensions)
{
if (partitionDimensions.isEmpty()) {
return Collections.emptySet();
}
final DimensionsSpec dimensionsSpec = task.getDataSchema() == null ? null : task.getDataSchema().getDimensionsSpec();
if (dimensionsSpec == null) {
return Collections.emptySet();
}
final Set<String> longDimensions = new HashSet<>();
for (String dim : partitionDimensions) {
if (dimensionsSpec.getSchema(dim) instanceof LongDimensionSchema) {
longDimensions.add(dim);
}
}
return longDimensions;
}

/**
* Canonical Long string for a LONG partition dimension's raw value, matching the string the indexer stores and the
* broker queries with, so set membership is exact (e.g. "00001" and 1.0 yield "1", "+5" yields "5"). Returns null for
* a null/missing, unparseable, or multi-value value (which the indexer would reject as a long) so {@code IS NULL} is
* not pruned; mirrors LongDimensionIndexer's convertObjectToLong on getRaw.
*/
@Nullable
@VisibleForTesting
static String canonicalLongValue(@Nullable Object rawValue, String dimension)
{
Long coerced = null;
try {
coerced = DimensionHandlerUtils.convertObjectToLong(rawValue, false, dimension);
}
catch (ParseException ignored) {
// Multi-value or invalid type: treated as null (no numeric pruning for this value).
}
return coerced == null ? null : Long.toString(coerced);
}

/**
* 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
Expand All @@ -1096,6 +1152,11 @@ DataSegment annotateSegmentWithPartitionDimensionValues(DataSegment s)
}
final Integer maxValuesPerDimension = StreamingPartitionsSpec.getMaxValuesPerDimensionOrNull(partitionsSpec);
final Map<String, List<String>> snapshotFilters = new HashMap<>();
// Per-dimension type stamp, gating the broker's typed numeric-value pruning (possibleInValueDomain). Only
// populated for explicit LONG dimensions: a LONG stringifies identically at ingest and query, so set membership
// is exact. Others are left unstamped (no numeric pruning).
final Map<String, ColumnType> dimensionColumnTypes = new HashMap<>();
final Set<String> longPartitionDimensions = getLongPartitionDimensions(partitionDimensions);
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.
Expand Down Expand Up @@ -1130,13 +1191,19 @@ DataSegment annotateSegmentWithPartitionDimensionValues(DataSegment s)
// Sort for deterministic published metadata; null (missing value) sorts first.
snapshot.sort(Comparator.nullsFirst(Comparator.naturalOrder()));
snapshotFilters.put(dim, snapshot);

// Stamp the type only for an explicitly-declared LONG dimension; any other case is left unstamped.
if (longPartitionDimensions.contains(dim)) {
dimensionColumnTypes.put(dim, ColumnType.LONG);
}
}
}
return s.withShardSpec(
new DimensionValueSetShardSpec(
s.getShardSpec().getPartitionNum(),
s.getShardSpec().getNumCorePartitions(),
snapshotFilters
snapshotFilters,
dimensionColumnTypes
)
);
}
Expand Down
Loading
Loading