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..e53b16d3a0b1 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 @@ -19,10 +19,14 @@ package org.apache.druid.query.filter; +import com.google.common.collect.Range; import com.google.common.collect.RangeSet; import org.apache.druid.error.InvalidInput; import org.apache.druid.segment.VirtualColumn; import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.timeline.ClusterGroupTuples; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.partition.ShardSpec; @@ -47,7 +51,7 @@ public class FilterSegmentPruner implements SegmentPruner private final Set filterFields; private final VirtualColumns virtualColumns; private final Map>> rangeCache; - private final Map> shardEquivalenceCache; + private final Map> virtualColumnEquivalenceCache; public FilterSegmentPruner( DimFilter filter, @@ -59,7 +63,7 @@ public FilterSegmentPruner( this.filterFields = filterFields == null ? filter.getRequiredColumns() : filterFields; this.virtualColumns = virtualColumns == null ? VirtualColumns.EMPTY : virtualColumns; this.rangeCache = new HashMap<>(); - this.shardEquivalenceCache = new HashMap<>(); + this.virtualColumnEquivalenceCache = new HashMap<>(); } @@ -73,39 +77,72 @@ public FilterSegmentPruner( public boolean include(DataSegment segment) { final ShardSpec shard = segment.getShardSpec(); - boolean include = true; if (shard != null) { final Map> filterDomain = new HashMap<>(); - final List dimensions = shard.getDomainDimensions(); - for (String dimension : dimensions) { - final VirtualColumns.Node shardNode = shard.getDomainVirtualColumns().getNode(dimension); - if (shardNode != null) { - final VirtualColumn queryEquivalent = getQueryEquivalent(shardNode); - if (queryEquivalent != null) { - if (filterFields == null || filterFields.contains(queryEquivalent.getOutputName())) { - final Optional> optFilterRangeSet = rangeCache - .computeIfAbsent( - queryEquivalent.getOutputName(), - d -> Optional.ofNullable(filter.getDimensionRangeSet(d)) - ); - optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put( - shardNode.getVirtualColumn().getOutputName(), - stringRangeSet - )); - } - } - } 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)); - } + for (String dimension : shard.getDomainDimensions()) { + addToFilterDomain(dimension, shard.getDomainVirtualColumns(), filterDomain); } if (!filterDomain.isEmpty() && !shard.possibleInDomain(filterDomain)) { - include = false; + return false; } } - return include; + + final ClusterGroupTuples clusterGroups = segment.getClusterGroups(); + if (clusterGroups != null && !possibleInClusterGroups(clusterGroups)) { + return false; + } + + return true; + } + + private boolean possibleInClusterGroups(ClusterGroupTuples clusterGroups) + { + final RowSignature clusteringColumns = clusterGroups.clusteringColumns(); + final int numColumns = clusteringColumns.size(); + + final Map> filterDomain = new HashMap<>(); + for (int i = 0; i < numColumns; i++) { + final String column = clusteringColumns.getColumnName(i); + if (!ColumnType.STRING.equals(clusteringColumns.getColumnType(i).orElse(null))) { + continue; + } + addToFilterDomain(column, clusterGroups.virtualColumns(), filterDomain); + } + + if (filterDomain.isEmpty()) { + // Filter doesn't constrain any string clustering column. + return true; + } + + for (final List tuple : clusterGroups.tuples()) { + if (tupleMatchesDomain(clusteringColumns, tuple, filterDomain)) { + return true; + } + } + + return false; + } + + private static boolean tupleMatchesDomain( + RowSignature clusteringColumns, + List tuple, + Map> filterDomain + ) + { + for (int i = 0; i < clusteringColumns.size(); i++) { + final RangeSet domainRangeSet = filterDomain.get(clusteringColumns.getColumnName(i)); + if (domainRangeSet == null) { + continue; + } + final Object rawValue = tuple.get(i); + // Nulls are less than empty String in segments + final Range valueRange = rawValue == null ? Range.lessThan("") : Range.singleton((String) rawValue); + if (domainRangeSet.subRangeSet(valueRange).isEmpty()) { + return false; + } + } + return true; } @Override @@ -164,10 +201,46 @@ public String toString() '}'; } + /** + * Adds the filter's {@link RangeSet} for {@code column} to {@code filterDomain}, if the filter constrains it. + *

+ * If {@code domainVirtualColumns} considers {@code column} virtual, only a query virtual column with an + * equivalent expression can be matched against it, if none exists, nothing is added and this column is never pruned on. + *

+ * Otherwise, {@code column} is a plain physical column, it can only be used for pruning if it's a non-virtual column in the query. + */ + private void addToFilterDomain( + String column, + VirtualColumns domainVirtualColumns, + Map> filterDomain + ) + { + final VirtualColumns.Node domainNode = domainVirtualColumns.getNode(column); + if (domainNode != null) { + final VirtualColumn queryEquivalent = getQueryEquivalent(domainNode); + if (queryEquivalent != null) { + addRangeSetIfPresent(queryEquivalent.getOutputName(), column, filterDomain); + } + } else if (virtualColumns.getNode(column) == null) { + // Query doesn't shadow the materialized column with its own virtual column of the same name. + addRangeSetIfPresent(column, column, filterDomain); + } + } + + private void addRangeSetIfPresent(String filterField, String domainColumn, Map> filterDomain) + { + if (!filterFields.contains(filterField)) { + return; + } + final Optional> optFilterRangeSet = + rangeCache.computeIfAbsent(filterField, d -> Optional.ofNullable(filter.getDimensionRangeSet(d))); + optFilterRangeSet.ifPresent(rangeSet -> filterDomain.put(domainColumn, rangeSet)); + } + @Nullable private VirtualColumn getQueryEquivalent(VirtualColumns.Node node) { - final Optional cached = shardEquivalenceCache.computeIfAbsent( + final Optional cached = virtualColumnEquivalenceCache.computeIfAbsent( node, n -> Optional.ofNullable(virtualColumns.findEquivalent(n)) ); 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..1585eef5b7df 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 @@ -26,8 +26,10 @@ import org.apache.druid.query.expression.TestExprMacroTable; import org.apache.druid.segment.VirtualColumns; import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.virtual.ExpressionVirtualColumn; import org.apache.druid.segment.virtual.NestedFieldVirtualColumn; +import org.apache.druid.timeline.ClusterGroupTuples; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.timeline.partition.DimensionRangeShardSpec; @@ -282,7 +284,7 @@ void testEqualsAndHashcode() { EqualsVerifier.forClass(FilterSegmentPruner.class) .usingGetClass() - .withIgnoredFields("rangeCache", "shardEquivalenceCache") + .withIgnoredFields("rangeCache", "virtualColumnEquivalenceCache") .verify(); } @@ -329,6 +331,197 @@ void testPruneNumericIn() Assertions.assertTrue(pruner.include(seg)); } + @Test + void testPruneClusterGroupTuples() + { + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final RowSignature clusteringColumns = RowSignature.builder().add("dim1", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + List.of(List.of("abc"), List.of("xyz")) + ); + + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + final DimFilter matchingFilter = new EqualityFilter("dim1", ColumnType.STRING, "abc", null); + final DimFilter nonMatchingFilter = new EqualityFilter("dim1", ColumnType.STRING, "foo", null); + + Assertions.assertTrue(new FilterSegmentPruner(matchingFilter, null, null).include(seg)); + Assertions.assertFalse(new FilterSegmentPruner(nonMatchingFilter, null, null).include(seg)); + } + + @Test + void testPruneClusterGroupTuplesMultipleColumns() + { + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final RowSignature clusteringColumns = RowSignature.builder() + .add("dim1", ColumnType.STRING) + .add("dim2", ColumnType.STRING) + .build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + List.of(List.of("abc", "xyz"), List.of("def", "uvw")) + ); + + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + // matches the first tuple on both columns + final DimFilter matchingFilter = new AndDimFilter( + new EqualityFilter("dim1", ColumnType.STRING, "abc", null), + new EqualityFilter("dim2", ColumnType.STRING, "xyz", null) + ); + Assertions.assertTrue(new FilterSegmentPruner(matchingFilter, null, null).include(seg)); + + // each value individually matches a tuple, but not the same tuple, so the combination must prune + final DimFilter mismatchedCombinationFilter = new AndDimFilter( + new EqualityFilter("dim1", ColumnType.STRING, "abc", null), + new EqualityFilter("dim2", ColumnType.STRING, "uvw", null) + ); + Assertions.assertFalse(new FilterSegmentPruner(mismatchedCombinationFilter, null, null).include(seg)); + + // constraining only one of the two clustering columns still matches via the second tuple + final DimFilter singleColumnFilter = new EqualityFilter("dim2", ColumnType.STRING, "uvw", null); + Assertions.assertTrue(new FilterSegmentPruner(singleColumnFilter, null, null).include(seg)); + + // neither tuple matches + final DimFilter nonMatchingFilter = new AndDimFilter( + new EqualityFilter("dim1", ColumnType.STRING, "abc", null), + new EqualityFilter("dim2", ColumnType.STRING, "foo", null) + ); + Assertions.assertFalse(new FilterSegmentPruner(nonMatchingFilter, null, null).include(seg)); + } + + @Test + void testClusterGroupTuplesSkipsNonStringColumns() + { + // Numeric columns are skipped for pruning (see druid issue #19408), so this must not prune. + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final RowSignature clusteringColumns = RowSignature.builder().add("id", ColumnType.LONG).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples(clusteringColumns, List.of(List.of(100L), List.of(200L))); + + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + final DimFilter filter = new EqualityFilter("id", ColumnType.LONG, 999L, null); + + Assertions.assertTrue(new FilterSegmentPruner(filter, null, null).include(seg)); + } + + @Test + void testPruneClusterGroupTuplesVirtualColumn() + { + final VirtualColumns clusterVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("vdim1", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final RowSignature clusteringColumns = RowSignature.builder().add("vdim1", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + clusterVirtualColumns, + List.of(List.of("abcfoo"), List.of("xyzfoo")) + ); + + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + // same expression, same name + VirtualColumns queryVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("vdim1", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final DimFilter matchingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "abcfoo", null); + final DimFilter nonMatchingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "deffoo", null); + Assertions.assertTrue(new FilterSegmentPruner(matchingFilter, null, queryVirtualColumns).include(seg)); + Assertions.assertFalse(new FilterSegmentPruner(nonMatchingFilter, null, queryVirtualColumns).include(seg)); + + // same expression, different name: still resolved via virtual column equivalence + queryVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final DimFilter matchingFilterDifferentName = new EqualityFilter("v0", ColumnType.STRING, "abcfoo", null); + final DimFilter nonMatchingFilterDifferentName = new EqualityFilter("v0", ColumnType.STRING, "deffoo", null); + Assertions.assertTrue(new FilterSegmentPruner(matchingFilterDifferentName, null, queryVirtualColumns).include(seg)); + Assertions.assertFalse(new FilterSegmentPruner(nonMatchingFilterDifferentName, null, queryVirtualColumns).include(seg)); + } + + @Test + void testPruneClusterGroupTuplesVirtualColumnSameNameDifferentExpressionNeverPrunes() + { + final VirtualColumns clusterVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("vdim1", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final RowSignature clusteringColumns = RowSignature.builder().add("vdim1", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + clusterVirtualColumns, + List.of(List.of("abcfoo"), List.of("xyzfoo")) + ); + + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + // query's vdim1 is a different expression, so it has no equivalent on the segment side and must never prune, + // even though the filter value would not match any tuple if it were (incorrectly) compared directly + final VirtualColumns queryVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("vdim1", "concat(dim1, 'bar')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final DimFilter nonMatchingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "nomatch", null); + Assertions.assertTrue(new FilterSegmentPruner(nonMatchingFilter, null, queryVirtualColumns).include(seg)); + + final DimFilter matchingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "abcfoo", null); + Assertions.assertTrue(new FilterSegmentPruner(matchingFilter, null, queryVirtualColumns).include(seg)); + } + + @Test + void testPruneClusterGroupTuplesVirtualColumnNoQueryVirtualColumnNeverPrunes() + { + // The segment's cluster groups record "vdim1" as derived from an expression. A query with no virtual column + // of its own named "vdim1" has no way to prove it means the same expression, so pruning must not assume the + // filter value can be compared against the tuple's virtual-column-derived value: never prune. + final VirtualColumns clusterVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("vdim1", "concat(dim1, 'foo')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final RowSignature clusteringColumns = RowSignature.builder().add("vdim1", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + clusterVirtualColumns, + List.of(List.of("abcfoo"), List.of("xyzfoo")) + ); + + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + final DimFilter matchingLookingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "abcfoo", null); + final DimFilter nonMatchingLookingFilter = new EqualityFilter("vdim1", ColumnType.STRING, "deffoo", null); + + // no query virtual columns at all: neither filter value can be resolved against the domain's virtual column + Assertions.assertTrue(new FilterSegmentPruner(matchingLookingFilter, null, null).include(seg)); + Assertions.assertTrue(new FilterSegmentPruner(nonMatchingLookingFilter, null, null).include(seg)); + } + + @Test + void testPruneClusterGroupTuplesShadowedByQueryVirtualColumnNeverPrunes() + { + // "dim1" is a plain materialized clustering column with no virtual columns on the segment side. If the query + // defines its own virtual column named "dim1", it shadows the real column, so the query is no longer + // referring to the materialized clustering values and the segment must never prune, even for a filter value + // that looks like it wouldn't match any of the real tuple values. + final String interval = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z"; + final RowSignature clusteringColumns = RowSignature.builder().add("dim1", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + List.of(List.of("abc"), List.of("xyz")) + ); + + final DataSegment seg = makeDataSegment(interval, makeRange("dim1", 0, null, null), tuples); + + final VirtualColumns queryVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("dim1", "concat(dim2, 'zzz')", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + + final DimFilter nonMatchingLookingFilter = new EqualityFilter("dim1", ColumnType.STRING, "nomatch", null); + Assertions.assertTrue(new FilterSegmentPruner(nonMatchingLookingFilter, null, queryVirtualColumns).include(seg)); + + final DimFilter matchingLookingFilter = new EqualityFilter("dim1", ColumnType.STRING, "abc", null); + Assertions.assertTrue(new FilterSegmentPruner(matchingLookingFilter, null, queryVirtualColumns).include(seg)); + } + private ShardSpec makeRange( String column, int partitionNumber, @@ -385,4 +578,13 @@ private DataSegment makeDataSegment(String intervalString, ShardSpec shardSpec) .shardSpec(shardSpec) .build(); } + + private DataSegment makeDataSegment(String intervalString, ShardSpec shardSpec, ClusterGroupTuples clusterGroups) + { + Interval interval = Intervals.of(intervalString); + return DataSegment.builder(SegmentId.of("prune-test", interval, "0", shardSpec)) + .shardSpec(shardSpec) + .clusterGroups(clusterGroups) + .build(); + } }