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..570dcb3c5e0c 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,75 @@ 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, false); } 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; + } + // Cluster group virtual columns are wired into real query execution (they're merged into the segment's + // live VirtualColumns at cursor construction time), so "column" is directly queryable by that exact name + // as long as the query doesn't shadow it with a differently-defined virtual column of its own. + addToFilterDomain(column, clusterGroups.virtualColumns(), filterDomain, true); + } + + 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 +204,47 @@ public String toString() '}'; } + /** + * Adds the filter's {@link RangeSet} for {@code column} to {@code filterDomain}, if the filter constrains it and + * {@code column} can be safely matched: as a query virtual column with an equivalent expression, directly by name + * when {@code allowDirectAccess} (only valid for domain virtual columns that are themselves directly queryable, + * e.g. clustering groups), or as an unshadowed physical column. Otherwise, nothing is added. + */ + private void addToFilterDomain( + String column, + VirtualColumns domainVirtualColumns, + Map> filterDomain, + boolean allowDirectAccess + ) + { + 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 (allowDirectAccess && virtualColumns.getNode(column) == null) { + addRangeSetIfPresent(column, 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/main/java/org/apache/druid/segment/VirtualColumns.java b/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java index 48d1decb559a..3a5416e24770 100644 --- a/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java +++ b/processing/src/main/java/org/apache/druid/segment/VirtualColumns.java @@ -282,7 +282,14 @@ public VirtualColumn findEquivalent(Node otherNode) toCheckForEquivalence = otherVirtualColumn.rewriteRequiredColumns(equivalenceRewriteMap); } - return equivalence.get().get(toCheckForEquivalence.getEquivalanceKey()); + VirtualColumn matched = equivalence.get().get(toCheckForEquivalence.getEquivalanceKey()); + if (matched != null && + // guardrail check for expression collision when a virtual column shadows the physical column + // e.x. otherNode v0 = dim1 and VCs dim1 = dim2 plus q = dim1, q can be treated as equivalent to v0 even though it reads physical column dim2 + getNode(matched.getOutputName()).getDependencies().size() == otherNode.getDependencies().size()) { + return matched; + } + return null; } /** 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..ca85d801155c 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,235 @@ 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 testPruneClusterGroupTuplesVirtualColumnNoQueryVirtualColumnPrunesDirectly() + { + // The segment's cluster groups record "vdim1" as derived from an expression, but cluster group virtual + // columns are directly queryable. A query with no virtual column of its own named "vdim1" doesn't shadow it, + // so the filter value can be compared directly against the tuple's virtual-column-derived value. + 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, and "vdim1" isn't shadowed, so it's matched directly by name + Assertions.assertTrue(new FilterSegmentPruner(matchingLookingFilter, null, null).include(seg)); + Assertions.assertFalse(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)); + } + + @Test + void testPruneClusterGroupTuplesTransitivelyShadowedByQueryVirtualColumnNeverPrunes() + { + // The segment's cluster group is keyed by virtual column "v0" = "dim1", i.e. the real physical "dim1" column. + // The query defines its own virtual column "dim1" = "dim2", which shadows the physical column, plus a second + // virtual column "q" = "dim1" that therefore transitively reads "dim2", not the real "dim1". Matching "q" + // against "v0" by comparing raw, unresolved expression text ("dim1" == "dim1") would wrongly treat them as + // equivalent, even though "q" doesn't actually read the column "v0" reads. Equivalence must be resolved + // through the query's own virtual column dependency graph, so pruning on "q" must never occur here. + final VirtualColumns clusterVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("v0", "dim1", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + final RowSignature clusteringColumns = RowSignature.builder().add("v0", ColumnType.STRING).build(); + final ClusterGroupTuples tuples = new ClusterGroupTuples( + clusteringColumns, + clusterVirtualColumns, + List.of(List.of("abc"), List.of("xyz")) + ); + + 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 VirtualColumns queryVirtualColumns = VirtualColumns.create( + new ExpressionVirtualColumn("dim1", "dim2", ColumnType.STRING, TestExprMacroTable.INSTANCE), + new ExpressionVirtualColumn("q", "dim1", ColumnType.STRING, TestExprMacroTable.INSTANCE) + ); + + // A filter value that matches none of the tuple's "v0" values must still never prune, since "q" doesn't + // actually read what "v0" reads. + final DimFilter nonMatchingLookingFilter = new EqualityFilter("q", ColumnType.STRING, "nomatch", null); + Assertions.assertTrue(new FilterSegmentPruner(nonMatchingLookingFilter, null, queryVirtualColumns).include(seg)); + + // Nor should a filter value that happens to match one of the tuple's "v0" values: that match is coincidental + // and says nothing about "dim2", which is what "q" actually reads. + final DimFilter matchingLookingFilter = new EqualityFilter("q", ColumnType.STRING, "abc", null); + Assertions.assertTrue(new FilterSegmentPruner(matchingLookingFilter, null, queryVirtualColumns).include(seg)); + } + private ShardSpec makeRange( String column, int partitionNumber, @@ -385,4 +616,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(); + } }