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..d332e4f99873 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 @@ -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( @@ -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. - * - *

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> 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"))); @@ -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> shardSpecs = getShardSpecs(dataSource); + for (Map shardSpec : shardSpecs) { + @SuppressWarnings("unchecked") + final Map columnTypes = (Map) 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 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 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 allCodes = shardSpecs.stream() + .map(shardSpec -> ((Map>) 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); 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..309318b9426b 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 @@ -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; @@ -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; @@ -478,6 +483,9 @@ private TaskStatus runInternal(TaskToolbox toolbox) throws Exception final List 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 longPartitionDimensions = getLongPartitionDimensions(partitionDimensions); try (final RecordSupplier recordSupplier = task.newTaskRecordSupplier(toolbox)) { @@ -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 dimValues = row.getDimension(dim); - if (dimValues == null || dimValues.isEmpty()) { - dimSet.add(null); + if (longPartitionDimensions.contains(dim)) { + dimSet.add(canonicalLongValue(row.getRaw(dim), dim)); } 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 dimValues = row.getDimension(dim); + if (dimValues == null || dimValues.isEmpty()) { + dimSet.add(null); + } else { + dimSet.addAll(dimValues); + } } } } @@ -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 getLongPartitionDimensions(List 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 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 @@ -1096,6 +1152,11 @@ DataSegment annotateSegmentWithPartitionDimensionValues(DataSegment s) } final Integer maxValuesPerDimension = StreamingPartitionsSpec.getMaxValuesPerDimensionOrNull(partitionsSpec); final Map> 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 dimensionColumnTypes = new HashMap<>(); + final Set longPartitionDimensions = getLongPartitionDimensions(partitionDimensions); 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. @@ -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 ) ); } 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..662ec3656b48 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 @@ -28,6 +28,7 @@ import org.apache.druid.data.input.InputRow; import org.apache.druid.data.input.impl.DimensionsSpec; import org.apache.druid.data.input.impl.JsonInputFormat; +import org.apache.druid.data.input.impl.LongDimensionSchema; import org.apache.druid.data.input.impl.StringDimensionSchema; import org.apache.druid.data.input.impl.TimestampSpec; import org.apache.druid.discovery.DataNodeService; @@ -49,6 +50,7 @@ import org.apache.druid.java.util.emitter.service.ServiceMetricEvent; import org.apache.druid.java.util.metrics.StubServiceEmitter; import org.apache.druid.segment.TestHelper; +import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.incremental.InputRowFilterResult; import org.apache.druid.segment.incremental.NoopRowIngestionMeters; import org.apache.druid.segment.indexing.DataSchema; @@ -385,6 +387,110 @@ public void testAnnotateSegmentStampsDimensionValueSetShardSpecForObservedValues ); } + /** + * A partition dimension explicitly declared as a {@link LongDimensionSchema} gets a {@link ColumnType#LONG} stamp in + * {@link DimensionValueSetShardSpec#getDimensionColumnTypes()}, gating the broker's typed numeric-value pruning. + */ + @Test + public void testAnnotateSegmentStampsLongColumnTypeForLongDimensionSchema() throws Exception + { + final TestSeekableStreamIndexTaskRunner runner = createRunner( + createDataSchemaWithPartitionDimensionSchemas(), + null, + null, + null, + Map.of("partition", "0"), + Map.of("partition", "100") + ); + Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec()) + .thenReturn(new StreamingPartitionsSpec(List.of("tenant"))); + + final DataSegment segment = createSingleSegment(); + final SegmentId lookupKey = segment.getId(); + observe(runner, lookupKey, "tenant", "100", "200"); + + final DataSegment annotated = runner.annotateSegmentWithPartitionDimensionValues(segment); + + Assert.assertTrue(annotated.getShardSpec() instanceof DimensionValueSetShardSpec); + final DimensionValueSetShardSpec shardSpec = (DimensionValueSetShardSpec) annotated.getShardSpec(); + Assert.assertEquals(ColumnType.LONG, shardSpec.getDimensionColumnTypes().get("tenant")); + // The new type stamp does not disturb value-stamping. + Assert.assertEquals( + Arrays.asList("100", "200"), + shardSpec.getPartitionDimensionValues().get("tenant") + ); + } + + /** + * A LONG partition dimension's raw value canonicalizes to the same string the indexer stores and the broker queries + * with, so a non-canonical token (leading zeros, leading '+', trailing ".0", a boxed number) is not wrongly pruned. + * Unparseable, multi-value, and null/missing values coerce to null so {@code IS NULL} is not pruned. + */ + @Test + public void testCanonicalLongValueMatchesIndexerCoercion() + { + // Non-canonical tokens collapse to the canonical Long string. + Assert.assertEquals("1", SeekableStreamIndexTaskRunner.canonicalLongValue("00001", "tenant")); + Assert.assertEquals("5", SeekableStreamIndexTaskRunner.canonicalLongValue("+5", "tenant")); + Assert.assertEquals("1", SeekableStreamIndexTaskRunner.canonicalLongValue("1.0", "tenant")); + Assert.assertEquals("1", SeekableStreamIndexTaskRunner.canonicalLongValue(1L, "tenant")); + Assert.assertEquals("5", SeekableStreamIndexTaskRunner.canonicalLongValue(5, "tenant")); + Assert.assertEquals("1", SeekableStreamIndexTaskRunner.canonicalLongValue(1.0, "tenant")); + + // Null, unparseable, whitespace-padded, and multi-value all become null (so IS NULL is not falsely pruned). + Assert.assertNull(SeekableStreamIndexTaskRunner.canonicalLongValue(null, "tenant")); + Assert.assertNull(SeekableStreamIndexTaskRunner.canonicalLongValue("abc", "tenant")); + Assert.assertNull(SeekableStreamIndexTaskRunner.canonicalLongValue(" 5", "tenant")); + Assert.assertNull(SeekableStreamIndexTaskRunner.canonicalLongValue(Arrays.asList(1, 2), "tenant")); + } + + /** + * A partition dimension declared as a {@link StringDimensionSchema} (or not declared at all, i.e. schemaless) is left + * out of {@link DimensionValueSetShardSpec#getDimensionColumnTypes()}, but is still stamped in + * {@link DimensionValueSetShardSpec#getPartitionDimensionValues()}. + */ + @Test + public void testAnnotateSegmentOmitsColumnTypeForStringOrSchemalessDimension() throws Exception + { + final TestSeekableStreamIndexTaskRunner runner = createRunner( + createDataSchemaWithPartitionDimensionSchemas(), + null, + null, + null, + Map.of("partition", "0"), + Map.of("partition", "100") + ); + Mockito.when(task.getTuningConfig().getStreamingPartitionsSpec()) + .thenReturn(new StreamingPartitionsSpec(List.of("region", "unknown"))); + + final DataSegment segment = createSingleSegment(); + final SegmentId lookupKey = segment.getId(); + observe(runner, lookupKey, "region", "us-west", "us-east"); + observe(runner, lookupKey, "unknown", "1", "2"); + + final DataSegment annotated = runner.annotateSegmentWithPartitionDimensionValues(segment); + + Assert.assertTrue(annotated.getShardSpec() instanceof DimensionValueSetShardSpec); + final DimensionValueSetShardSpec shardSpec = (DimensionValueSetShardSpec) annotated.getShardSpec(); + Assert.assertFalse( + "A StringDimensionSchema dimension must not get a column type stamp", + shardSpec.getDimensionColumnTypes().containsKey("region") + ); + Assert.assertFalse( + "A dimension absent from the DimensionsSpec (schemaless) must not get a column type stamp", + shardSpec.getDimensionColumnTypes().containsKey("unknown") + ); + // Both are still stamped with their observed values. + Assert.assertEquals( + ImmutableSet.of("us-west", "us-east"), + ImmutableSet.copyOf(shardSpec.getPartitionDimensionValues().get("region")) + ); + Assert.assertEquals( + Arrays.asList("1", "2"), + shardSpec.getPartitionDimensionValues().get("unknown") + ); + } + /** * A segment that spans a task restart has incomplete observed values, so it must NOT declare any partition filters * (no pruning), to avoid wrongly pruning pre-restart rows. It is still stamped with an empty-filter @@ -469,8 +575,8 @@ public void testRestartBatchMixingFallbackAndObservedSegmentsPublishesWithDimens } /** - * A dimension that ingested a null/missing value declares null (as a null list element) alongside its non-null - * values, so {@code IS NULL} queries are not pruned. Here tenant saw tenant_a and a null; region saw only us-west. + * A dimension that ingested a null/missing value declares null (a null list element) alongside its non-null values, + * so {@code IS NULL} queries are not pruned. */ @Test public void testNullValuedDimensionDeclaresNullInPartitionDimensionValues() throws Exception @@ -485,7 +591,6 @@ public void testNullValuedDimensionDeclaresNullInPartitionDimensionValues() thro 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"); @@ -495,7 +600,7 @@ public void testNullValuedDimensionDeclaresNullInPartitionDimensionValues() thro annotated.getShardSpec() instanceof DimensionValueSetShardSpec ); final DimensionValueSetShardSpec shardSpec = (DimensionValueSetShardSpec) annotated.getShardSpec(); - // tenant declares both its non-null value AND null, so IS NULL queries are not pruned. + // tenant declares both null and its non-null value. Assert.assertEquals( Arrays.asList(null, "tenant_a"), shardSpec.getPartitionDimensionValues().get("tenant") @@ -828,6 +933,29 @@ private static DataSchema createDataSchema() .build(); } + /** + * A DataSchema declaring "tenant" as a {@link LongDimensionSchema} and "region" as a {@link StringDimensionSchema}, + * for exercising the dimensionColumnTypes stamping in {@code annotateSegmentWithPartitionDimensionValues}. Note + * that "unknown" is deliberately absent, standing in for a schemaless/not-explicitly-declared dimension. + */ + private static DataSchema createDataSchemaWithPartitionDimensionSchemas() + { + final DimensionsSpec dimensionsSpec = new DimensionsSpec( + Arrays.asList( + new LongDimensionSchema("tenant"), + new StringDimensionSchema("region") + ) + ); + return DataSchema.builder() + .withDataSource(DATA_SOURCE) + .withTimestamp(TimestampSpec.DEFAULT) + .withDimensions(dimensionsSpec) + .withGranularity( + new UniformGranularitySpec(Granularities.MINUTE, Granularities.NONE, null) + ) + .build(); + } + private TestSeekableStreamIndexTaskRunner createInitializedRunner( Map startOffsets, Map endOffsets diff --git a/processing/src/main/java/org/apache/druid/query/filter/AndDimFilter.java b/processing/src/main/java/org/apache/druid/query/filter/AndDimFilter.java index 43e7e147ec04..ebca313e5b33 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/AndDimFilter.java +++ b/processing/src/main/java/org/apache/druid/query/filter/AndDimFilter.java @@ -26,8 +26,11 @@ import com.google.common.collect.RangeSet; import com.google.common.collect.TreeRangeSet; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.filter.Filters; +import org.apache.druid.timeline.partition.TypedValueSet; +import javax.annotation.Nullable; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; @@ -114,6 +117,33 @@ public RangeSet getDimensionRangeSet(String dimension) return retSet; } + @Nullable + @Override + public TypedValueSet getDimensionValueSet(String dimension) + { + // A row matches AND iff it matches every branch, so the allowed value set is the intersection of the constraining + // branches' value sets. Unconstrained (null) branches are ignored, mirroring getDimensionRangeSet above. + Set retValues = null; + ColumnType retType = null; + for (DimFilter field : fields) { + final TypedValueSet valueSet = field.getDimensionValueSet(dimension); + if (valueSet == null) { + continue; + } + if (retValues == null) { + retValues = new HashSet<>(valueSet.getValues()); + retType = valueSet.getType(); + } else { + // Bail out on a type mismatch: intersecting across incompatible canonicalizations would risk unsound pruning. + if (!retType.equals(valueSet.getType())) { + return null; + } + retValues.retainAll(valueSet.getValues()); + } + } + return retValues == null ? null : new TypedValueSet(retValues, retType); + } + @Override public Set getRequiredColumns() { diff --git a/processing/src/main/java/org/apache/druid/query/filter/DimFilter.java b/processing/src/main/java/org/apache/druid/query/filter/DimFilter.java index db1f0ecd7fa9..a051a8461e9b 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/DimFilter.java +++ b/processing/src/main/java/org/apache/druid/query/filter/DimFilter.java @@ -25,8 +25,10 @@ import org.apache.druid.java.util.common.Cacheable; import org.apache.druid.query.extraction.ExtractionFn; import org.apache.druid.segment.CursorBuildSpec; +import org.apache.druid.timeline.partition.TypedValueSet; import javax.annotation.Nullable; +import java.util.Map; import java.util.Set; /** @@ -106,6 +108,30 @@ public interface DimFilter extends Cacheable @Nullable RangeSet getDimensionRangeSet(String dimension); + /** + * A type-aware, set-membership parallel to {@link #getDimensionRangeSet(String)}, consumed only by the + * {@link org.apache.druid.timeline.partition.DimensionValueSetShardSpec} numeric pruning channel + * ({@link org.apache.druid.timeline.partition.ShardSpec#possibleInValueDomain(Map)}). + * + *

Returns the stringified match values (carrying the filter's match {@link org.apache.druid.segment.column.ColumnType}), + * or {@code null} if this filter does not constrain {@code dimension} to a finite, type-safe value set. Only + * equality/IN filters on {@code LONG} return non-null; {@code null} means no numeric-value pruning is possible. + * + *

Separate from {@link #getDimensionRangeSet(String)} on purpose: that feeds the lexicographic + * {@link org.apache.druid.timeline.partition.ShardSpec#possibleInDomain(Map)} path used by + * {@link org.apache.druid.timeline.partition.DimensionRangeShardSpec}, which must stay string-only to avoid the + * numeric-ordering bug of [#19408]/[#19415]. The type carried here lets a shard spec gate pruning on an exact type + * match, which a typeless {@code RangeSet} cannot express. + * + * @param dimension name of the dimension to get the typed value set for + * @return the typed set of match values, or {@code null} if not determinable / not type-safe for this filter + */ + @Nullable + default TypedValueSet getDimensionValueSet(String dimension) + { + return null; + } + /** * @return a HashSet that represents all columns' name which the DimFilter required to do filter. */ diff --git a/processing/src/main/java/org/apache/druid/query/filter/EqualityFilter.java b/processing/src/main/java/org/apache/druid/query/filter/EqualityFilter.java index 482a96c9906c..4266c908c9a6 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/EqualityFilter.java +++ b/processing/src/main/java/org/apache/druid/query/filter/EqualityFilter.java @@ -63,10 +63,12 @@ import org.apache.druid.segment.index.semantic.ValueIndexes; import org.apache.druid.segment.nested.StructuredData; import org.apache.druid.segment.vector.VectorColumnSelectorFactory; +import org.apache.druid.timeline.partition.TypedValueSet; import javax.annotation.Nullable; import java.nio.ByteBuffer; import java.util.Arrays; +import java.util.Collections; import java.util.Comparator; import java.util.Map; import java.util.Objects; @@ -228,6 +230,23 @@ public RangeSet getDimensionRangeSet(String dimension) return retSet; } + @Nullable + @Override + public TypedValueSet getDimensionValueSet(String dimension) + { + if (!Objects.equals(getColumn(), dimension)) { + return null; + } + // Only LONG is safe: a LONG stringifies identically on ingest and query, so set membership is exact. DOUBLE/FLOAT + // ("1.0" vs "1") is not safe and is excluded; STRING is handled by getDimensionRangeSet. + if (!matchValueType.is(ValueType.LONG)) { + return null; + } + // Use the coerced matchValueEval (not the raw getMatchValue()) so an EqualityFilter(col, LONG, "1.0") yields "1", + // matching how a LONG is stamped at ingest. + return new TypedValueSet(Collections.singleton(matchValueEval.asString()), matchValueType); + } + @Nullable @Override public BitmapColumnIndex getBitmapColumnIndex(ColumnIndexSelector selector) diff --git a/processing/src/main/java/org/apache/druid/query/filter/FilterSegmentPruner.java b/processing/src/main/java/org/apache/druid/query/filter/FilterSegmentPruner.java index 6ef7565302b0..b36d2bc74dfa 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/FilterSegmentPruner.java +++ b/processing/src/main/java/org/apache/druid/query/filter/FilterSegmentPruner.java @@ -25,6 +25,7 @@ import org.apache.druid.segment.VirtualColumns; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.partition.ShardSpec; +import org.apache.druid.timeline.partition.TypedValueSet; import javax.annotation.Nullable; import java.util.ArrayList; @@ -47,6 +48,7 @@ public class FilterSegmentPruner implements SegmentPruner private final Set filterFields; private final VirtualColumns virtualColumns; private final Map>> rangeCache; + private final Map> valueSetCache; private final Map> shardEquivalenceCache; public FilterSegmentPruner( @@ -59,6 +61,7 @@ public FilterSegmentPruner( this.filterFields = filterFields == null ? filter.getRequiredColumns() : filterFields; this.virtualColumns = virtualColumns == null ? VirtualColumns.EMPTY : virtualColumns; this.rangeCache = new HashMap<>(); + this.valueSetCache = new HashMap<>(); this.shardEquivalenceCache = new HashMap<>(); } @@ -77,6 +80,10 @@ public boolean include(DataSegment segment) if (shard != null) { final Map> filterDomain = new HashMap<>(); + // Parallel, type-carrying domain consumed only by shard specs that override possibleInValueDomain + // (DimensionValueSetShardSpec). Built from the same dimensions and virtual-column resolution as the string + // domain, so numeric equality/IN filters can prune without going through the lexicographic range path. + final Map filterValueDomain = new HashMap<>(); final List dimensions = shard.getDomainDimensions(); for (String dimension : dimensions) { final VirtualColumns.Node shardNode = shard.getDomainVirtualColumns().getNode(dimension); @@ -84,25 +91,36 @@ public boolean include(DataSegment segment) final VirtualColumn queryEquivalent = getQueryEquivalent(shardNode); if (queryEquivalent != null) { if (filterFields == null || filterFields.contains(queryEquivalent.getOutputName())) { + final String shardOutputName = shardNode.getVirtualColumn().getOutputName(); final Optional> optFilterRangeSet = rangeCache .computeIfAbsent( queryEquivalent.getOutputName(), d -> Optional.ofNullable(filter.getDimensionRangeSet(d)) ); - optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put( - shardNode.getVirtualColumn().getOutputName(), - stringRangeSet - )); + optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put(shardOutputName, stringRangeSet)); + + final Optional optFilterValueSet = valueSetCache + .computeIfAbsent( + queryEquivalent.getOutputName(), + d -> Optional.ofNullable(filter.getDimensionValueSet(d)) + ); + optFilterValueSet.ifPresent(valueSet -> filterValueDomain.put(shardOutputName, valueSet)); } } } else if (filterFields == null || filterFields.contains(dimension)) { final Optional> optFilterRangeSet = rangeCache.computeIfAbsent(dimension, d -> Optional.ofNullable(filter.getDimensionRangeSet(d))); optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put(dimension, stringRangeSet)); + + final Optional optFilterValueSet = + valueSetCache.computeIfAbsent(dimension, d -> Optional.ofNullable(filter.getDimensionValueSet(d))); + optFilterValueSet.ifPresent(valueSet -> filterValueDomain.put(dimension, valueSet)); } } if (!filterDomain.isEmpty() && !shard.possibleInDomain(filterDomain)) { include = false; + } else if (!filterValueDomain.isEmpty() && !shard.possibleInValueDomain(filterValueDomain)) { + include = false; } } return include; diff --git a/processing/src/main/java/org/apache/druid/query/filter/OrDimFilter.java b/processing/src/main/java/org/apache/druid/query/filter/OrDimFilter.java index 3f7ce37c7b03..925308324475 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/OrDimFilter.java +++ b/processing/src/main/java/org/apache/druid/query/filter/OrDimFilter.java @@ -26,8 +26,11 @@ import com.google.common.collect.RangeSet; import com.google.common.collect.TreeRangeSet; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.filter.Filters; +import org.apache.druid.timeline.partition.TypedValueSet; +import javax.annotation.Nullable; import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; @@ -119,6 +122,31 @@ public RangeSet getDimensionRangeSet(String dimension) return retSet; } + @Nullable + @Override + public TypedValueSet getDimensionValueSet(String dimension) + { + // A row matches OR iff it matches any branch, so the allowed value set is the union of the branches' value sets. + // If any branch is unconstrained (null) it could match any value, so the OR cannot prune — return null. Mirrors + // getDimensionRangeSet. + final Set retValues = new HashSet<>(); + ColumnType retType = null; + for (DimFilter field : fields) { + final TypedValueSet valueSet = field.getDimensionValueSet(dimension); + if (valueSet == null) { + return null; + } + if (retType == null) { + retType = valueSet.getType(); + } else if (!retType.equals(valueSet.getType())) { + // Mixed types across branches cannot be unioned into one type-gated set; don't prune. + return null; + } + retValues.addAll(valueSet.getValues()); + } + return retType == null ? null : new TypedValueSet(retValues, retType); + } + @Override public Set getRequiredColumns() { diff --git a/processing/src/main/java/org/apache/druid/query/filter/TypedInFilter.java b/processing/src/main/java/org/apache/druid/query/filter/TypedInFilter.java index 8bcd3e7075f5..fe7c759e6c29 100644 --- a/processing/src/main/java/org/apache/druid/query/filter/TypedInFilter.java +++ b/processing/src/main/java/org/apache/druid/query/filter/TypedInFilter.java @@ -66,12 +66,14 @@ import org.apache.druid.segment.index.semantic.Utf8ValueSetIndexes; import org.apache.druid.segment.index.semantic.ValueSetIndexes; import org.apache.druid.segment.vector.VectorColumnSelectorFactory; +import org.apache.druid.timeline.partition.TypedValueSet; import javax.annotation.Nullable; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.Comparator; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -277,6 +279,27 @@ public RangeSet getDimensionRangeSet(String dimension) return retSet; } + @Nullable + @Override + public TypedValueSet getDimensionValueSet(String dimension) + { + if (!Objects.equals(getColumn(), dimension)) { + return null; + } + // Only LONG is safe (identical stringification on ingest and query); DOUBLE/FLOAT are excluded and STRING is + // handled by getDimensionRangeSet. + if (!matchValueType.is(ValueType.LONG)) { + return null; + } + // sortedMatchValues is coerced to matchValueType, so Evals.asString yields the same canonical strings as ingest. + // A null element is kept as a null set element so it matches a shard's stamped null (missing value). + final Set values = new HashSet<>(); + for (Object value : sortedMatchValues.get()) { + values.add(Evals.asString(value)); + } + return new TypedValueSet(values, matchValueType); + } + @Override public Set getRequiredColumns() { diff --git a/processing/src/main/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpec.java b/processing/src/main/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpec.java index 6be48e370266..93b692aa766f 100644 --- a/processing/src/main/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpec.java +++ b/processing/src/main/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpec.java @@ -24,6 +24,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Range; import com.google.common.collect.RangeSet; +import org.apache.druid.segment.column.ColumnType; import javax.annotation.Nullable; import java.util.Collections; @@ -44,15 +45,44 @@ public class DimensionValueSetShardSpec extends NumberedShardSpec */ private final Map> partitionDimensionValues; + /** + * Maps dimension name → the {@link ColumnType} of the values stamped in {@link #partitionDimensionValues} for that + * dimension. Only populated for types safe for typed value-set pruning (currently {@link ColumnType#LONG}); a + * dimension absent from this map is pruned only via the string/range channel ({@link #possibleInDomain(Map)}), + * never via {@link #possibleInValueDomain(Map)}. + * + *

This is the safety gate for numeric pruning: {@link #possibleInValueDomain(Map)} may only prune a dimension + * when the query filter's match type equals the type recorded here, because value stringification only agrees + * across ingest and query for identical types. + */ + private final Map dimensionColumnTypes; + @JsonCreator public DimensionValueSetShardSpec( @JsonProperty("partitionNum") int partitionNum, @JsonProperty("partitions") int partitions, - @JsonProperty("partitionDimensionValues") @Nullable Map> partitionDimensionValues + @JsonProperty("partitionDimensionValues") @Nullable Map> partitionDimensionValues, + @JsonProperty("dimensionColumnTypes") @Nullable Map dimensionColumnTypes ) { super(partitionNum, partitions); this.partitionDimensionValues = partitionDimensionValues == null ? Collections.emptyMap() : partitionDimensionValues; + this.dimensionColumnTypes = dimensionColumnTypes == null ? Collections.emptyMap() : dimensionColumnTypes; + } + + /** + * Convenience constructor for callers that carry no per-dimension type information (e.g. non-numeric streaming, + * the append→replace upgrade fallback, or tests). A null {@code dimensionColumnTypes} disables the typed + * {@link #possibleInValueDomain(Map)} numeric-pruning channel but leaves the string/range channel + * ({@link #possibleInDomain(Map)}) fully functional. + */ + public DimensionValueSetShardSpec( + int partitionNum, + int partitions, + @Nullable Map> partitionDimensionValues + ) + { + this(partitionNum, partitions, partitionDimensionValues, null); } @JsonProperty("partitionDimensionValues") @@ -61,6 +91,12 @@ public Map> getPartitionDimensionValues() return partitionDimensionValues; } + @JsonProperty("dimensionColumnTypes") + public Map getDimensionColumnTypes() + { + return dimensionColumnTypes; + } + @Override public List getDomainDimensions() { @@ -89,6 +125,11 @@ public boolean possibleInDomain(Map> domain) final String dimension = entry.getKey(); final List allowedValues = entry.getValue(); + if (dimensionColumnTypes.containsKey(dimension)) { + // Canonicalized numeric dimension: pruned only via possibleInValueDomain, not the literal string range. + continue; + } + final RangeSet domainRangeSet = domain.get(dimension); if (domainRangeSet == null || domainRangeSet.isEmpty()) { // Query doesn't constrain this dimension — cannot prune on it. @@ -112,6 +153,60 @@ public boolean possibleInDomain(Map> domain) return true; } + /** + * Type-aware, set-membership pruning for numeric ({@code LONG}) equality/IN filters. For each dimension the query + * constrains to a {@link TypedValueSet}, this shard can be pruned only when: + *

+ * + *

The type-equality check is the data-safety gate: if the stamped type is missing or differs from the filter's + * type (e.g. a LONG filter against a DOUBLE-stamped dimension), stringified values may not agree ({@code "1"} vs + * {@code "1.0"}), so this method must not prune and keeps the segment. Filters only produce LONG value sets today. + * + * @return true if the segment must be considered for the query, false if it can be pruned + */ + @Override + public boolean possibleInValueDomain(Map valueDomain) + { + if (partitionDimensionValues.isEmpty()) { + return true; + } + + for (Map.Entry entry : valueDomain.entrySet()) { + final String dimension = entry.getKey(); + final TypedValueSet filterValues = entry.getValue(); + + final List allowedValues = partitionDimensionValues.get(dimension); + if (allowedValues == null) { + // This shard declares no values for this dimension, so it cannot be pruned on it. + continue; + } + + final ColumnType stampedType = dimensionColumnTypes.get(dimension); + if (stampedType == null || !stampedType.equals(filterValues.getType())) { + // No stamped type or a type mismatch: stringified values may not agree, so it is not safe to prune. + continue; + } + + boolean anyMatch = false; + for (String value : allowedValues) { + if (filterValues.contains(value)) { + anyMatch = true; + break; + } + } + if (!anyMatch) { + return false; + } + } + + return true; + } + @Override public String getType() { @@ -121,13 +216,13 @@ public String getType() @Override public ShardSpec withPartitionNum(int partitionNum) { - return new DimensionValueSetShardSpec(partitionNum, getNumCorePartitions(), partitionDimensionValues); + return new DimensionValueSetShardSpec(partitionNum, getNumCorePartitions(), partitionDimensionValues, dimensionColumnTypes); } @Override public ShardSpec withCorePartitions(int partitions) { - return new DimensionValueSetShardSpec(getPartitionNum(), partitions, partitionDimensionValues); + return new DimensionValueSetShardSpec(getPartitionNum(), partitions, partitionDimensionValues, dimensionColumnTypes); } @Override @@ -143,13 +238,14 @@ public boolean equals(Object o) return false; } DimensionValueSetShardSpec that = (DimensionValueSetShardSpec) o; - return Objects.equals(partitionDimensionValues, that.partitionDimensionValues); + return Objects.equals(partitionDimensionValues, that.partitionDimensionValues) && + Objects.equals(dimensionColumnTypes, that.dimensionColumnTypes); } @Override public int hashCode() { - return Objects.hash(super.hashCode(), partitionDimensionValues); + return Objects.hash(super.hashCode(), partitionDimensionValues, dimensionColumnTypes); } @Override @@ -159,6 +255,7 @@ public String toString() "partitionNum=" + getPartitionNum() + ", partitions=" + getNumCorePartitions() + ", partitionDimensionValues=" + partitionDimensionValues + + ", dimensionColumnTypes=" + dimensionColumnTypes + '}'; } } diff --git a/processing/src/main/java/org/apache/druid/timeline/partition/ShardSpec.java b/processing/src/main/java/org/apache/druid/timeline/partition/ShardSpec.java index f1098d2282d1..5515745be29e 100644 --- a/processing/src/main/java/org/apache/druid/timeline/partition/ShardSpec.java +++ b/processing/src/main/java/org/apache/druid/timeline/partition/ShardSpec.java @@ -169,6 +169,25 @@ default VirtualColumns getDomainVirtualColumns() @JsonIgnore boolean possibleInDomain(Map> domain); + /** + * A type-aware, set-membership parallel to {@link #possibleInDomain(Map)}. Given a domain of typed value sets + * (produced by {@link org.apache.druid.query.filter.DimFilter#getDimensionValueSet(String)} for equality/IN + * filters), return false only if this shard provably cannot contain any of the requested values; otherwise return + * true. + * + *

The default returns {@code true} (no pruning); only {@link DimensionValueSetShardSpec} overrides it. + * {@link DimensionRangeShardSpec} deliberately does not participate: its pruning is lexicographic-ordering based + * (see [#19408]) and is not safe for numeric equality, so this no-op default keeps numeric filters from reaching it + * and reintroducing that bug. + * + * @return possibility of the requested values being present in this shard + */ + @JsonIgnore + default boolean possibleInValueDomain(Map valueDomain) + { + return true; + } + /** * Get the type name of this ShardSpec. */ diff --git a/processing/src/main/java/org/apache/druid/timeline/partition/TypedValueSet.java b/processing/src/main/java/org/apache/druid/timeline/partition/TypedValueSet.java new file mode 100644 index 000000000000..bc6ebc919f11 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/timeline/partition/TypedValueSet.java @@ -0,0 +1,97 @@ +/* + * 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.timeline.partition; + +import org.apache.druid.segment.column.ColumnType; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Objects; +import java.util.Set; + +/** + * A type-carrying set of stringified match values, produced by + * {@link org.apache.druid.query.filter.DimFilter#getDimensionValueSet(String)} and consumed by + * {@link ShardSpec#possibleInValueDomain(java.util.Map)}. + * + *

The value-domain analog of the {@code RangeSet} used by + * {@link org.apache.druid.query.filter.DimFilter#getDimensionRangeSet(String)} / + * {@link ShardSpec#possibleInDomain(java.util.Map)}, but it also carries the filter's match {@link ColumnType}. That + * type is the safety gate for numeric segment pruning: a {@link DimensionValueSetShardSpec} may only prune when the + * filter's match type equals the type stamped for the dimension, because value stringification only agrees across + * ingest and query for identical types (LONG {@code 1} → {@code "1"}, but DOUBLE {@code 1.0} → {@code "1.0"}). A + * typeless {@code RangeSet} cannot express this gate, which is why this is a separate channel. + * + *

The contained {@link #values} are the stringified match values (via {@code Evals.asString}), matching the + * ingest side. The set may contain a {@code null} element to denote a null/missing match value (e.g. an {@code IN} + * list containing {@code NULL}). + */ +public class TypedValueSet +{ + private final Set values; + private final ColumnType type; + + public TypedValueSet(Set values, ColumnType type) + { + // Defensive copy into a set that tolerates a null element (denoting a null match value). + this.values = Collections.unmodifiableSet(new HashSet<>(values)); + this.type = type; + } + + public Set getValues() + { + return values; + } + + public ColumnType getType() + { + return type; + } + + public boolean contains(String value) + { + return values.contains(value); + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TypedValueSet that = (TypedValueSet) o; + return Objects.equals(values, that.values) && Objects.equals(type, that.type); + } + + @Override + public int hashCode() + { + return Objects.hash(values, type); + } + + @Override + public String toString() + { + return "TypedValueSet{values=" + values + ", type=" + type + '}'; + } +} diff --git a/processing/src/test/java/org/apache/druid/query/filter/AndDimFilterTest.java b/processing/src/test/java/org/apache/druid/query/filter/AndDimFilterTest.java index 4ad8c8904a3b..28f3c2ead583 100644 --- a/processing/src/test/java/org/apache/druid/query/filter/AndDimFilterTest.java +++ b/processing/src/test/java/org/apache/druid/query/filter/AndDimFilterTest.java @@ -21,13 +21,18 @@ import com.google.common.collect.Lists; import com.google.common.collect.Sets; +import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.filter.FalseFilter; import org.apache.druid.segment.filter.FilterTestUtils; import org.apache.druid.segment.filter.TrueFilter; import org.apache.druid.testing.InitializedNullHandlingTest; +import org.apache.druid.timeline.partition.TypedValueSet; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.Arrays; +import java.util.HashSet; + public class AndDimFilterTest extends InitializedNullHandlingTest { @Test @@ -102,4 +107,56 @@ public void testToFilterAndOfMultipleFiltersReturningAsItIs() ); Assertions.assertEquals(filter.toFilter(), filter.toFilter()); } + + @Test + public void testGetDimensionValueSet_singleConstrainingBranch() + { + // AND with one constraining branch → value set is that branch. + final DimFilter filter = new AndDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new SelectorDimFilter("region", "us", null) + ); + final TypedValueSet valueSet = filter.getDimensionValueSet("code"); + Assertions.assertNotNull(valueSet); + Assertions.assertEquals(ColumnType.LONG, valueSet.getType()); + Assertions.assertEquals(new HashSet<>(java.util.List.of("5")), valueSet.getValues()); + // A dimension no branch constrains → null. + Assertions.assertNull(filter.getDimensionValueSet("region")); + } + + @Test + public void testGetDimensionValueSet_intersectsBranchesOnSameDim() + { + // code = 5 AND code IN (5,6) → intersection {5}. + final DimFilter filter = new AndDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new TypedInFilter("code", ColumnType.LONG, Arrays.asList(5L, 6L), null, null) + ); + final TypedValueSet valueSet = filter.getDimensionValueSet("code"); + Assertions.assertNotNull(valueSet); + Assertions.assertEquals(new HashSet<>(java.util.List.of("5")), valueSet.getValues()); + } + + @Test + public void testGetDimensionValueSet_disjointBranches_emptyIntersection() + { + // code = 5 AND code = 6 → empty intersection, which prunes everything. + final DimFilter filter = new AndDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new EqualityFilter("code", ColumnType.LONG, 6L, null) + ); + final TypedValueSet valueSet = filter.getDimensionValueSet("code"); + Assertions.assertNotNull(valueSet); + Assertions.assertTrue(valueSet.getValues().isEmpty()); + } + + @Test + public void testGetDimensionValueSet_noConstrainingBranch_returnsNull() + { + final DimFilter filter = new AndDimFilter( + new SelectorDimFilter("region", "us", null), + new SelectorDimFilter("tenant", "a", null) + ); + Assertions.assertNull(filter.getDimensionValueSet("code")); + } } diff --git a/processing/src/test/java/org/apache/druid/query/filter/FilterSegmentPrunerTest.java b/processing/src/test/java/org/apache/druid/query/filter/FilterSegmentPrunerTest.java index d5c0d38bebaa..196284002bcb 100644 --- a/processing/src/test/java/org/apache/druid/query/filter/FilterSegmentPrunerTest.java +++ b/processing/src/test/java/org/apache/druid/query/filter/FilterSegmentPrunerTest.java @@ -31,6 +31,7 @@ import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.timeline.partition.DimensionRangeShardSpec; +import org.apache.druid.timeline.partition.DimensionValueSetShardSpec; import org.apache.druid.timeline.partition.ShardSpec; import org.joda.time.Interval; import org.junit.jupiter.api.Assertions; @@ -40,6 +41,7 @@ import java.util.Collections; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.function.Function; @@ -282,7 +284,7 @@ void testEqualsAndHashcode() { EqualsVerifier.forClass(FilterSegmentPruner.class) .usingGetClass() - .withIgnoredFields("rangeCache", "shardEquivalenceCache") + .withIgnoredFields("rangeCache", "valueSetCache", "shardEquivalenceCache") .verify(); } @@ -329,6 +331,75 @@ void testPruneNumericIn() Assertions.assertTrue(pruner.include(seg)); } + @Test + void testPruneDimValueSetNumericEquality() + { + // LONG equality prunes on an absent value, keeps on a present one. Prune twice to exercise the value-set cache. + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment( + interval, + new DimensionValueSetShardSpec( + 0, + 1, + Map.of("id", List.of("1", "2")), + Map.of("id", ColumnType.LONG) + ) + ); + + final FilterSegmentPruner keeps = + new FilterSegmentPruner(new EqualityFilter("id", ColumnType.LONG, 1L, null), null, null); + Assertions.assertTrue(keeps.include(seg)); + Assertions.assertTrue(keeps.include(seg)); + + final FilterSegmentPruner prunes = + new FilterSegmentPruner(new EqualityFilter("id", ColumnType.LONG, 3L, null), null, null); + Assertions.assertFalse(prunes.include(seg)); + Assertions.assertFalse(prunes.include(seg)); + + final FilterSegmentPruner prunesIn = + new FilterSegmentPruner(new TypedInFilter("id", ColumnType.LONG, List.of(7L, 8L), null, null), null, null); + Assertions.assertFalse(prunesIn.include(seg)); + } + + @Test + void testPruneDimValueSetNumericEquality_compoundAndFilter() + { + // WHERE id = 3 AND note = 'x' reaches the pruner as an AndDimFilter; the LONG constraint still prunes when 3 is absent. + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment( + interval, + new DimensionValueSetShardSpec( + 0, + 1, + Map.of("id", List.of("1", "2")), + Map.of("id", ColumnType.LONG) + ) + ); + + final DimFilter compound = new AndDimFilter( + new EqualityFilter("id", ColumnType.LONG, 3L, null), + new SelectorDimFilter("note", "x", null) + ); + final FilterSegmentPruner pruner = new FilterSegmentPruner(compound, null, null); + Assertions.assertFalse(pruner.include(seg)); + } + + @Test + void testDimValueSetNoStampedType_isNotPruned() + { + // A dim_value_set segment WITHOUT a stamped type (legacy) must NOT be pruned by a numeric filter: the typed + // channel returns true (cannot prune) and there is no string domain for a LONG equality. + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment( + interval, + new DimensionValueSetShardSpec(0, 1, Map.of("id", List.of("1", "2"))) + ); + + final FilterSegmentPruner pruner = + new FilterSegmentPruner(new EqualityFilter("id", ColumnType.LONG, 3L, null), null, null); + Assertions.assertTrue(pruner.include(seg)); + } + private ShardSpec makeRange( String column, int partitionNumber, diff --git a/processing/src/test/java/org/apache/druid/query/filter/OrDimFilterTest.java b/processing/src/test/java/org/apache/druid/query/filter/OrDimFilterTest.java index 66103e1e06ac..6a24f171790f 100644 --- a/processing/src/test/java/org/apache/druid/query/filter/OrDimFilterTest.java +++ b/processing/src/test/java/org/apache/druid/query/filter/OrDimFilterTest.java @@ -19,13 +19,18 @@ package org.apache.druid.query.filter; +import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.filter.FalseFilter; import org.apache.druid.segment.filter.FilterTestUtils; import org.apache.druid.segment.filter.TrueFilter; import org.apache.druid.testing.InitializedNullHandlingTest; +import org.apache.druid.timeline.partition.TypedValueSet; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.Arrays; +import java.util.HashSet; + public class OrDimFilterTest extends InitializedNullHandlingTest { @Test @@ -87,4 +92,40 @@ public void testOptimizeOrOfMultipleFiltersReturningAsItIs() ); Assertions.assertEquals(filter.toFilter(), filter.toFilter()); } + + @Test + public void testGetDimensionValueSet_unionsBranchesOnSameDim() + { + // code = 5 OR code = 6 → union {5,6}. + final DimFilter filter = new OrDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new EqualityFilter("code", ColumnType.LONG, 6L, null) + ); + final TypedValueSet valueSet = filter.getDimensionValueSet("code"); + Assertions.assertNotNull(valueSet); + Assertions.assertEquals(ColumnType.LONG, valueSet.getType()); + Assertions.assertEquals(new HashSet<>(Arrays.asList("5", "6")), valueSet.getValues()); + } + + @Test + public void testGetDimensionValueSet_anyUnconstrainedBranch_returnsNull() + { + // code = 5 OR region = 'us': the region branch matches any code, so code cannot be pruned. + final DimFilter filter = new OrDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new SelectorDimFilter("region", "us", null) + ); + Assertions.assertNull(filter.getDimensionValueSet("code")); + } + + @Test + public void testGetDimensionValueSet_branchOnDifferentDim_returnsNull() + { + // Querying a dimension the OR does not constrain → null. + final DimFilter filter = new OrDimFilter( + new EqualityFilter("code", ColumnType.LONG, 5L, null), + new EqualityFilter("code", ColumnType.LONG, 6L, null) + ); + Assertions.assertNull(filter.getDimensionValueSet("region")); + } } diff --git a/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java b/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java index 375e16e7fc54..1d3c99c3412f 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/EqualityFilterTests.java @@ -43,6 +43,7 @@ import org.apache.druid.segment.IndexBuilder; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.testing.InitializedNullHandlingTest; +import org.apache.druid.timeline.partition.TypedValueSet; import org.junit.AfterClass; import org.junit.Assert; import org.junit.Assume; @@ -52,6 +53,7 @@ import java.io.Closeable; import java.util.Arrays; +import java.util.Collections; public class EqualityFilterTests { @@ -1622,6 +1624,43 @@ public void testGetDimensionRangeSet() ); } + @Test + public void testGetDimensionValueSet() + { + // LONG equality yields a single-element, LONG-typed value set for the matching column, null for another. + final EqualityFilter longFilter = new EqualityFilter("x", ColumnType.LONG, 1L, null); + final TypedValueSet valueSet = longFilter.getDimensionValueSet("x"); + Assert.assertNotNull(valueSet); + Assert.assertEquals(ColumnType.LONG, valueSet.getType()); + Assert.assertEquals(Collections.singleton("1"), valueSet.getValues()); + Assert.assertNull(longFilter.getDimensionValueSet("y")); + + // Non-LONG match value types must NOT return a typed value set (STRING is served by the range channel; + // DOUBLE/FLOAT/ARRAY are excluded for canonicalization safety). + Assert.assertNull(new EqualityFilter("x", ColumnType.STRING, "hello", null).getDimensionValueSet("x")); + Assert.assertNull(new EqualityFilter("x", ColumnType.DOUBLE, 1.5, null).getDimensionValueSet("x")); + Assert.assertNull(new EqualityFilter("x", ColumnType.FLOAT, 1.5f, null).getDimensionValueSet("x")); + Assert.assertNull( + new EqualityFilter( + "x", + ColumnType.STRING_ARRAY, + ExprEval.ofType(ExpressionType.STRING_ARRAY, new Object[]{"abc", "def"}).value(), + null + ).getDimensionValueSet("x") + ); + } + + @Test + public void testGetDimensionValueSet_usesCoercedValue() + { + // The coerced value is stringified, not the raw one: a LONG filter built from 1.0 coerces to 1, so the value + // set holds "1" (matching ingest), not "1.0". + final EqualityFilter filter = new EqualityFilter("x", ColumnType.LONG, 1.0, null); + final TypedValueSet valueSet = filter.getDimensionValueSet("x"); + Assert.assertNotNull(valueSet); + Assert.assertEquals(Collections.singleton("1"), valueSet.getValues()); + } + @Test public void test_equals() { diff --git a/processing/src/test/java/org/apache/druid/segment/filter/InFilterTests.java b/processing/src/test/java/org/apache/druid/segment/filter/InFilterTests.java index 31d875023630..d769072e0851 100644 --- a/processing/src/test/java/org/apache/druid/segment/filter/InFilterTests.java +++ b/processing/src/test/java/org/apache/druid/segment/filter/InFilterTests.java @@ -51,6 +51,7 @@ import org.apache.druid.segment.IndexBuilder; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.testing.InitializedNullHandlingTest; +import org.apache.druid.timeline.partition.TypedValueSet; import org.junit.AfterClass; import org.junit.Assert; import org.junit.Assume; @@ -62,8 +63,10 @@ import java.io.Closeable; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.stream.Collectors; @@ -698,6 +701,36 @@ public void testGetDimensionRangeSet() ); } + @Test + public void testGetDimensionValueSet() + { + // LONG IN yields a LONG-typed set of stringified values for the matching column, null for another. + final TypedInFilter longFilter = inFilter("x", ColumnType.LONG, Arrays.asList(1L, 2L, 3L)); + final TypedValueSet valueSet = longFilter.getDimensionValueSet("x"); + Assert.assertNotNull(valueSet); + Assert.assertEquals(ColumnType.LONG, valueSet.getType()); + Assert.assertEquals(new HashSet<>(Arrays.asList("1", "2", "3")), valueSet.getValues()); + Assert.assertNull(longFilter.getDimensionValueSet("y")); + + // A null IN element is kept as a null set element (a null match value). + final TypedValueSet withNull = inFilter("x", ColumnType.LONG, Arrays.asList(null, 1L, 2L)).getDimensionValueSet("x"); + Assert.assertNotNull(withNull); + final Set expectedWithNull = new HashSet<>(Arrays.asList("1", "2")); + expectedWithNull.add(null); + Assert.assertEquals(expectedWithNull, withNull.getValues()); + + // Non-LONG match value types must NOT return a typed value set. + Assert.assertNull( + inFilter("x", ColumnType.STRING, Arrays.asList("a", "b")).getDimensionValueSet("x") + ); + Assert.assertNull( + inFilter("x", ColumnType.DOUBLE, Arrays.asList(1.1, 2.2)).getDimensionValueSet("x") + ); + Assert.assertNull( + inFilter("x", ColumnType.FLOAT, Arrays.asList(1.1f, 2.2f)).getDimensionValueSet("x") + ); + } + @Test public void testRequiredColumnRewrite() { diff --git a/processing/src/test/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpecTest.java b/processing/src/test/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpecTest.java index 4d1e6285854e..dbc01a5c5afe 100644 --- a/processing/src/test/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpecTest.java +++ b/processing/src/test/java/org/apache/druid/timeline/partition/DimensionValueSetShardSpecTest.java @@ -25,23 +25,46 @@ import com.google.common.collect.Range; import com.google.common.collect.RangeSet; import com.google.common.collect.TreeRangeSet; +import org.apache.druid.segment.column.ColumnType; import org.junit.Assert; import org.junit.Test; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; public class DimensionValueSetShardSpecTest { private static final String TENANT = "tenant"; + private static final String CODE = "code"; private static DimensionValueSetShardSpec spec(Map> filters) { return new DimensionValueSetShardSpec(0, 1, filters); } + private static DimensionValueSetShardSpec spec( + Map> filters, + Map types + ) + { + return new DimensionValueSetShardSpec(0, 1, filters, types); + } + + private static TypedValueSet longValues(String... values) + { + final Set set = new HashSet<>(Arrays.asList(values)); + return new TypedValueSet(set, ColumnType.LONG); + } + + private static Map longValueDomain(String dimension, String... values) + { + return ImmutableMap.of(dimension, longValues(values)); + } + private static RangeSet points(String... values) { final RangeSet rangeSet = TreeRangeSet.create(); @@ -310,4 +333,207 @@ public void testEmptyAllowedList_prunesEverything() final DimensionValueSetShardSpec s = spec(ImmutableMap.of(TENANT, List.of())); Assert.assertFalse(s.possibleInDomain(domain(TENANT, "tenant_a"))); } + + @Test + public void testStringDomain_skipsTypeStampedDimension() + { + // A LONG-stamped dim holds canonicalized values ("1"); a non-canonical selector like code = "00001" matches + // the indexed LONG via coercion, so the literal string range must not prune this segment. + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + Assert.assertTrue(s.possibleInDomain(domain(CODE, "00001"))); + } + + @Test + public void testValueDomain_noFilters_alwaysTrue() + { + final DimensionValueSetShardSpec s = spec(Collections.emptyMap(), Collections.emptyMap()); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "1"))); + Assert.assertTrue(s.possibleInValueDomain(Collections.emptyMap())); + } + + @Test + public void testValueDomain_longMatch_returnsTrue() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1", "2")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "1"))); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "2"))); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "3", "2"))); + } + + @Test + public void testValueDomain_longNoMatch_returnsFalse() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1", "2")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + Assert.assertFalse(s.possibleInValueDomain(longValueDomain(CODE, "3"))); + Assert.assertFalse(s.possibleInValueDomain(longValueDomain(CODE, "999", "1000"))); + } + + @Test + public void testValueDomain_noStampedType_cannotPrune() + { + // Values stamped but no type recorded (legacy/schemaless): must NOT prune without a matching type. + final DimensionValueSetShardSpec s = spec(ImmutableMap.of(CODE, List.of("1", "2"))); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "3"))); + } + + @Test + public void testValueDomain_typeMismatch_cannotPrune() + { + // A LONG filter must not prune a DOUBLE-stamped dim: cross-type stringification is unsafe (data-loss guard). + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1.0", "2.0")), + ImmutableMap.of(CODE, ColumnType.DOUBLE) + ); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "3"))); + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "1"))); + } + + @Test + public void testValueDomain_dimensionNotDeclared_cannotPrune() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(TENANT, List.of("tenant_a")), + ImmutableMap.of(TENANT, ColumnType.STRING) + ); + // Shard declares no values for the queried dimension → cannot prune. + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "3"))); + } + + @Test + public void testValueDomain_multipleDimensions_oneMismatches_returnsFalse() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1"), "other", List.of("10")), + ImmutableMap.of(CODE, ColumnType.LONG, "other", ColumnType.LONG) + ); + Assert.assertTrue(s.possibleInValueDomain(ImmutableMap.of( + CODE, longValues("1"), + "other", longValues("10") + ))); + Assert.assertFalse(s.possibleInValueDomain(ImmutableMap.of( + CODE, longValues("1"), + "other", longValues("999") + ))); + } + + @Test + public void testValueDomain_nullValueMembership() + { + // A null match value (e.g. IN (NULL)) must match a shard that declares a null value. + final DimensionValueSetShardSpec withNull = spec( + ImmutableMap.of(CODE, Arrays.asList("1", null)), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + final Set nullSet = new HashSet<>(); + nullSet.add(null); + Assert.assertTrue(withNull.possibleInValueDomain(ImmutableMap.of(CODE, new TypedValueSet(nullSet, ColumnType.LONG)))); + + // A shard with only concrete values is pruned for a null-only match. + final DimensionValueSetShardSpec concreteOnly = spec( + ImmutableMap.of(CODE, List.of("1")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + Assert.assertFalse(concreteOnly.possibleInValueDomain(ImmutableMap.of(CODE, new TypedValueSet(nullSet, ColumnType.LONG)))); + } + + @Test + public void testValueDomain_emptyAllowedList_prunes() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of()), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + Assert.assertFalse(s.possibleInValueDomain(longValueDomain(CODE, "1"))); + } + + @Test + public void testWithPartitionNum_preservesTypes() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + final DimensionValueSetShardSpec renumbered = (DimensionValueSetShardSpec) s.withPartitionNum(5); + Assert.assertEquals(5, renumbered.getPartitionNum()); + Assert.assertEquals(ImmutableMap.of(CODE, ColumnType.LONG), renumbered.getDimensionColumnTypes()); + Assert.assertFalse(renumbered.possibleInValueDomain(longValueDomain(CODE, "999"))); + } + + @Test + public void testWithCorePartitions_preservesTypes() + { + final DimensionValueSetShardSpec s = spec( + ImmutableMap.of(CODE, List.of("1")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + final DimensionValueSetShardSpec recored = (DimensionValueSetShardSpec) s.withCorePartitions(9); + Assert.assertEquals(9, recored.getNumCorePartitions()); + Assert.assertEquals(ImmutableMap.of(CODE, ColumnType.LONG), recored.getDimensionColumnTypes()); + } + + @Test + public void testEquals_distinguishesByType() + { + final DimensionValueSetShardSpec withLong = spec( + ImmutableMap.of(CODE, List.of("1")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + final DimensionValueSetShardSpec withoutType = spec(ImmutableMap.of(CODE, List.of("1"))); + Assert.assertNotEquals(withLong, withoutType); + Assert.assertEquals(withLong, spec(ImmutableMap.of(CODE, List.of("1")), ImmutableMap.of(CODE, ColumnType.LONG))); + } + + @Test + public void testThreeArgConstructor_hasEmptyTypes() + { + final DimensionValueSetShardSpec s = spec(ImmutableMap.of(CODE, List.of("1"))); + Assert.assertTrue(s.getDimensionColumnTypes().isEmpty()); + // With no stamped type, the value channel cannot prune. + Assert.assertTrue(s.possibleInValueDomain(longValueDomain(CODE, "999"))); + } + + @Test + public void testJsonSerde_withColumnTypes_roundTrips() throws Exception + { + final ObjectMapper mapper = newMapper(); + final DimensionValueSetShardSpec original = new DimensionValueSetShardSpec( + 3, + 8, + ImmutableMap.of(CODE, List.of("1", "2")), + ImmutableMap.of(CODE, ColumnType.LONG) + ); + final String json = mapper.writeValueAsString(original); + Assert.assertTrue(json.contains("\"dimensionColumnTypes\"")); + + final DimensionValueSetShardSpec deserialized = + mapper.readValue(json, DimensionValueSetShardSpec.class); + Assert.assertEquals(original, deserialized); + Assert.assertEquals(ImmutableMap.of(CODE, ColumnType.LONG), deserialized.getDimensionColumnTypes()); + Assert.assertFalse(deserialized.possibleInValueDomain(longValueDomain(CODE, "3"))); + } + + @Test + public void testJsonSerde_legacyWithoutColumnTypes_deserializes() throws Exception + { + // Back-compat: JSON written before dimensionColumnTypes existed deserializes to an empty map and does not prune + // via the typed channel. + final ObjectMapper mapper = newMapper(); + final String legacyJson = + "{\"type\":\"dim_value_set\",\"partitionNum\":0,\"partitions\":1," + + "\"partitionDimensionValues\":{\"code\":[\"1\",\"2\"]}}"; + final DimensionValueSetShardSpec deserialized = + mapper.readValue(legacyJson, DimensionValueSetShardSpec.class); + Assert.assertTrue(deserialized.getDimensionColumnTypes().isEmpty()); + Assert.assertEquals(ImmutableMap.of(CODE, List.of("1", "2")), deserialized.getPartitionDimensionValues()); + Assert.assertTrue(deserialized.possibleInValueDomain(longValueDomain(CODE, "3"))); + } } diff --git a/processing/src/test/java/org/apache/druid/timeline/partition/TypedValueSetTest.java b/processing/src/test/java/org/apache/druid/timeline/partition/TypedValueSetTest.java new file mode 100644 index 000000000000..0863c033a527 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/timeline/partition/TypedValueSetTest.java @@ -0,0 +1,84 @@ +/* + * 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.timeline.partition; + +import com.google.common.collect.ImmutableSet; +import org.apache.druid.segment.column.ColumnType; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; + +public class TypedValueSetTest +{ + @Test + public void testContainsAndGetters() + { + final TypedValueSet set = new TypedValueSet(ImmutableSet.of("1", "2"), ColumnType.LONG); + Assert.assertEquals(ColumnType.LONG, set.getType()); + Assert.assertEquals(ImmutableSet.of("1", "2"), set.getValues()); + Assert.assertTrue(set.contains("1")); + Assert.assertTrue(set.contains("2")); + Assert.assertFalse(set.contains("3")); + Assert.assertFalse(set.contains(null)); + } + + @Test + public void testNullElementIsRetained() + { + final Set withNull = new HashSet<>(Arrays.asList("1", null)); + final TypedValueSet set = new TypedValueSet(withNull, ColumnType.LONG); + Assert.assertTrue(set.contains(null)); + Assert.assertTrue(set.contains("1")); + } + + @Test + public void testDefensiveCopy() + { + final Set mutable = new HashSet<>(Arrays.asList("1")); + final TypedValueSet set = new TypedValueSet(mutable, ColumnType.LONG); + mutable.add("2"); + // Later mutation of the source must not leak in. + Assert.assertFalse(set.contains("2")); + } + + @Test + public void testValuesUnmodifiable() + { + final TypedValueSet set = new TypedValueSet(ImmutableSet.of("1"), ColumnType.LONG); + Assert.assertThrows(UnsupportedOperationException.class, () -> set.getValues().add("2")); + } + + @Test + public void testEqualsAndHashCode() + { + final TypedValueSet a = new TypedValueSet(ImmutableSet.of("1", "2"), ColumnType.LONG); + final TypedValueSet b = new TypedValueSet(ImmutableSet.of("2", "1"), ColumnType.LONG); + final TypedValueSet differentValues = new TypedValueSet(ImmutableSet.of("1", "3"), ColumnType.LONG); + final TypedValueSet differentType = new TypedValueSet(ImmutableSet.of("1", "2"), ColumnType.DOUBLE); + + Assert.assertEquals(a, b); + Assert.assertEquals(a.hashCode(), b.hashCode()); + Assert.assertNotEquals(a, differentValues); + Assert.assertNotEquals(a, differentType); + } +} diff --git a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java index 1317fd339a6a..00446a633404 100644 --- a/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java +++ b/server/src/main/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinator.java @@ -1299,10 +1299,12 @@ private static ShardSpec getUpgradedSegmentShardSpec(DataSegment oldSegment, Pen final ShardSpec pendingShardSpec = pendingSegment.getId().getShardSpec(); final ShardSpec oldShardSpec = oldSegment.getShardSpec(); if (oldShardSpec instanceof DimensionValueSetShardSpec) { + final DimensionValueSetShardSpec oldValueSetSpec = (DimensionValueSetShardSpec) oldShardSpec; return new DimensionValueSetShardSpec( pendingShardSpec.getPartitionNum(), pendingShardSpec.getNumCorePartitions(), - ((DimensionValueSetShardSpec) oldShardSpec).getPartitionDimensionValues() + oldValueSetSpec.getPartitionDimensionValues(), + oldValueSetSpec.getDimensionColumnTypes() ); } return pendingShardSpec;