This is an automated email from the ASF dual-hosted git repository.
cwylie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 3d8b81ceba9 feat: add SegmentPruner support for datasources/policies
(#19228)
3d8b81ceba9 is described below
commit 3d8b81ceba96ea4f34bb205ef499ef5b5d60152a
Author: Clint Wylie <[email protected]>
AuthorDate: Mon Mar 30 14:21:25 2026 -0700
feat: add SegmentPruner support for datasources/policies (#19228)
* Add SegmentPruner support for RestrictedDataSource policy filters
changes:
* adds new `include` method to `SegmentPruner` for checking individual
segments for whether or not to prune
* adds default implementation of `prune` method which calls `include`
* adds new `combine` method to `SegmentPruner` for merging pruners
* adds new `CompositeSegmentPruner` for cases where pruners cannot be
naturally combined
* adds new `createSegmentPruner` method to `DataSource` and `Policy` so
that they can participate in pruning
* updates `ExecutionVertex` to combine the new datasource pruner with the
pruner of the filter
---
.../querykit/RestrictedInputNumberDataSource.java | 9 +
.../RestrictedInputNumberDataSourceTest.java | 11 +
.../java/org/apache/druid/query/DataSource.java | 14 ++
.../apache/druid/query/RestrictedDataSource.java | 9 +
.../druid/query/filter/CompositeSegmentPruner.java | 141 +++++++++++
.../druid/query/filter/FilterSegmentPruner.java | 121 ++++-----
.../apache/druid/query/filter/SegmentPruner.java | 39 ++-
.../druid/query/planning/ExecutionVertex.java | 14 +-
.../java/org/apache/druid/query/policy/Policy.java | 12 +
.../apache/druid/query/policy/RowFilterPolicy.java | 11 +
.../druid/query/RestrictedDataSourceTest.java | 19 ++
.../query/filter/CompositeSegmentPrunerTest.java | 280 +++++++++++++++++++++
.../query/filter/FilterSegmentPrunerTest.java | 71 ++++++
.../query/policy/NoRestrictionPolicyTest.java | 6 +
.../druid/query/policy/RowFilterPolicyTest.java | 14 ++
15 files changed, 710 insertions(+), 61 deletions(-)
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSource.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSource.java
index 280ba5fc10c..4a936a65f83 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSource.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSource.java
@@ -25,10 +25,12 @@ import com.fasterxml.jackson.annotation.JsonTypeName;
import com.google.common.base.Preconditions;
import org.apache.druid.query.LeafDataSource;
import org.apache.druid.query.Query;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.query.policy.Policy;
import org.apache.druid.segment.RestrictedSegment;
import org.apache.druid.segment.SegmentMapFunction;
+import javax.annotation.Nullable;
import java.util.Collections;
import java.util.Objects;
import java.util.Set;
@@ -101,6 +103,13 @@ public class RestrictedInputNumberDataSource extends
LeafDataSource
return SegmentMapFunction.IDENTITY.thenMap(segment -> new
RestrictedSegment(segment, policy));
}
+ @Nullable
+ @Override
+ public SegmentPruner createSegmentPruner()
+ {
+ return policy.createSegmentPruner();
+ }
+
@Override
public byte[] getCacheKey()
{
diff --git
a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSourceTest.java
b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSourceTest.java
index 848e2b0eb48..0682ab62efb 100644
---
a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSourceTest.java
+++
b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/RestrictedInputNumberDataSourceTest.java
@@ -25,10 +25,12 @@ import nl.jqno.equalsverifier.EqualsVerifier;
import org.apache.druid.msq.guice.MSQIndexingModule;
import org.apache.druid.query.DataSource;
import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.filter.FilterSegmentPruner;
import org.apache.druid.query.filter.TrueDimFilter;
import org.apache.druid.query.policy.NoRestrictionPolicy;
import org.apache.druid.query.policy.RowFilterPolicy;
import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.VirtualColumns;
import org.junit.Assert;
import org.junit.Test;
@@ -154,4 +156,13 @@ public class RestrictedInputNumberDataSourceTest
mapper.readValue(mapper.writeValueAsString(dataSource),
DataSource.class)
);
}
+
+ @Test
+ public void test_createSegmentPruner_withRowFilterPolicy()
+ {
+ Assert.assertEquals(
+ new FilterSegmentPruner(TrueDimFilter.instance(), null,
VirtualColumns.EMPTY),
+ restrictedFooDataSource.createSegmentPruner()
+ );
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/DataSource.java
b/processing/src/main/java/org/apache/druid/query/DataSource.java
index f978a0abfa1..9b0d65f2d83 100644
--- a/processing/src/main/java/org/apache/druid/query/DataSource.java
+++ b/processing/src/main/java/org/apache/druid/query/DataSource.java
@@ -22,6 +22,7 @@ package org.apache.druid.query;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import org.apache.druid.java.util.common.Cacheable;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.query.planning.PreJoinableClause;
import org.apache.druid.query.policy.Policy;
import org.apache.druid.query.policy.PolicyEnforcer;
@@ -105,6 +106,19 @@ public interface DataSource extends Cacheable
*/
SegmentMapFunction createSegmentMapFunction(Query query);
+ /**
+ * Returns a {@link SegmentPruner} if this datasource embeds in any
information which can be used to determine if a
+ * segment needs processed or not. Note that callers of this method will
always only be processing segments for the
+ * datasource, so there is no need for a 'default' pruner that ensures the
segment has the proper datasource. A return
+ * value of null indicates that no pruning can be performed from this
datasource, though other sources of pruning,
+ * such as filters may still be used.
+ */
+ @Nullable
+ default SegmentPruner createSegmentPruner()
+ {
+ return null;
+ }
+
/**
* Returns an updated datasource based on the policy restrictions on tables.
* <p>
diff --git
a/processing/src/main/java/org/apache/druid/query/RestrictedDataSource.java
b/processing/src/main/java/org/apache/druid/query/RestrictedDataSource.java
index 5fb212b81af..98b82695060 100644
--- a/processing/src/main/java/org/apache/druid/query/RestrictedDataSource.java
+++ b/processing/src/main/java/org/apache/druid/query/RestrictedDataSource.java
@@ -24,12 +24,14 @@ import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.collect.ImmutableList;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.query.policy.NoRestrictionPolicy;
import org.apache.druid.query.policy.Policy;
import org.apache.druid.query.policy.PolicyEnforcer;
import org.apache.druid.segment.RestrictedSegment;
import org.apache.druid.segment.SegmentMapFunction;
+import javax.annotation.Nullable;
import java.util.List;
import java.util.Map;
import java.util.Objects;
@@ -127,6 +129,13 @@ public class RestrictedDataSource implements DataSource
return base.createSegmentMapFunction(query).thenMap(segment -> new
RestrictedSegment(segment, policy));
}
+ @Nullable
+ @Override
+ public SegmentPruner createSegmentPruner()
+ {
+ return policy.createSegmentPruner();
+ }
+
@Override
public DataSource withPolicies(Map<String, Optional<Policy>> policyMap,
PolicyEnforcer policyEnforcer)
{
diff --git
a/processing/src/main/java/org/apache/druid/query/filter/CompositeSegmentPruner.java
b/processing/src/main/java/org/apache/druid/query/filter/CompositeSegmentPruner.java
new file mode 100644
index 00000000000..5ca1343b38a
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/query/filter/CompositeSegmentPruner.java
@@ -0,0 +1,141 @@
+/*
+ * 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.query.filter;
+
+import org.apache.druid.error.DruidException;
+import org.apache.druid.timeline.DataSegment;
+
+import java.util.LinkedHashSet;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * {@link SegmentPruner} implementation that applies a set of {@link
SegmentPruner} against each {@link DataSegment}
+ * and will return false for {@link #include(DataSegment)} if ANY pruner
indicates that it should not be included.
+ */
+public class CompositeSegmentPruner implements SegmentPruner
+{
+ private final Set<SegmentPruner> pruners;
+
+ public CompositeSegmentPruner(Set<SegmentPruner> pruners)
+ {
+ this.pruners = pruners;
+ validate(pruners);
+ }
+
+ @Override
+ public boolean include(DataSegment segment)
+ {
+ for (SegmentPruner pruner : pruners) {
+ if (!pruner.include(segment)) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ @Override
+ public SegmentPruner combine(SegmentPruner other)
+ {
+ final Set<SegmentPruner> combinedPruners;
+ if (other instanceof CompositeSegmentPruner composite) {
+ // combine both sets, folding filter pruners in or adding if not
+ Set<SegmentPruner> combining = new LinkedHashSet<>(pruners);
+ for (SegmentPruner pruner : composite.pruners) {
+ if (pruner instanceof FilterSegmentPruner filter) {
+ combining = foldFilterPruner(combining, filter);
+ } else {
+ combining.add(pruner);
+ }
+ }
+ combinedPruners = combining;
+ } else if (other instanceof FilterSegmentPruner filter) {
+ // fold the filter pruner into our set
+ combinedPruners = foldFilterPruner(pruners, filter);
+ } else {
+ // default, add the other one to our set
+ combinedPruners = new LinkedHashSet<>(pruners);
+ combinedPruners.add(other);
+ }
+ return new CompositeSegmentPruner(combinedPruners);
+ }
+
+
+ @Override
+ public boolean equals(Object o)
+ {
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ CompositeSegmentPruner that = (CompositeSegmentPruner) o;
+ return Objects.equals(pruners, that.pruners);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return Objects.hashCode(pruners);
+ }
+
+ @Override
+ public String toString()
+ {
+ return "CompositeSegmentPruner{" +
+ "pruners=" + pruners +
+ '}';
+ }
+
+ private static void validate(Set<SegmentPruner> pruners)
+ {
+ FilterSegmentPruner filter = null;
+ for (SegmentPruner pruner : pruners) {
+ if (pruner instanceof FilterSegmentPruner filterPruner) {
+ if (filter != null) {
+ throw DruidException.defensive(
+ "Combine multiple filter pruners prior to creating composite,
found[%s] and [%s]",
+ filter,
+ filterPruner
+ );
+ }
+ filter = filterPruner;
+ }
+ }
+ }
+
+ private static Set<SegmentPruner> foldFilterPruner(Set<SegmentPruner>
pruners, FilterSegmentPruner filter)
+ {
+ Set<SegmentPruner> combinedPruners = new LinkedHashSet<>();
+ // if other is a filterPruner, check to see if we contain any filter
pruners to combine with it
+ // a composite cannot have more than 1 filter pruner
+ boolean notCombined = true;
+ for (SegmentPruner pruner : pruners) {
+ if (pruner instanceof FilterSegmentPruner) {
+ combinedPruners.add(filter.combine(pruner));
+ notCombined = false;
+ } else {
+ combinedPruners.add(pruner);
+ }
+ }
+ if (notCombined) {
+ combinedPruners.add(filter);
+ }
+ return combinedPruners;
+ }
+}
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 aef51e2c05d..0df393661cb 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
@@ -27,7 +27,7 @@ import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.partition.ShardSpec;
import javax.annotation.Nullable;
-import java.util.Collection;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
@@ -35,7 +35,6 @@ import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
-import java.util.function.Function;
/**
* Uses a {@link DimFilter} to check the {@link
DimFilter#getDimensionRangeSet(String)} against
@@ -61,70 +60,78 @@ public class FilterSegmentPruner implements SegmentPruner
this.rangeCache = new HashMap<>();
}
- /**
- * Filter the given iterable of objects by removing any object whose {@link
DataSegment}, obtained from the converter
- * function, does not fit in the RangeSet of the dimFilter {@link
DimFilter#getDimensionRangeSet(String)}. The
- * returned set contains the filtered objects in the same order as they
appear in input.
- *
- * {@link #rangeCache} stores the RangeSets of different dimensions for the
filter, so it can be re-used between
- * calls to save redundant evaluation of {@link
DimFilter#getDimensionRangeSet(String)} on the same columns.
- *
- * @param input The iterable of objects to be filtered
- * @param converter The function to convert T to {@link DataSegment} that
can be filtered by
- * @param <T> This can be any type, as long as transform function is
provided to extract a {@link DataSegment}
- *
- * @return The set of pruned object, in the same order as input
- */
+
+ /**
+ * Returns false if the {@link DataSegment} does not fit in {@link
DimFilter#getDimensionRangeSet(String)}.
+ * <p>
+ * {@link #rangeCache} stores the RangeSets of different dimensions for the
filter, so it can be re-used between
+ * calls to save redundant evaluation of {@link
DimFilter#getDimensionRangeSet(String)} on the same columns.
+ */
@Override
- public <T> Collection<T> prune(Iterable<T> input, Function<T, DataSegment>
converter)
+ public boolean include(DataSegment segment)
{
- // LinkedHashSet retains order from "input".
- final Set<T> retSet = new LinkedHashSet<>();
-
- for (T obj : input) {
- final DataSegment segment = converter.apply(obj);
- if (segment == null) {
- continue;
- }
- final ShardSpec shard = segment.getShardSpec();
- boolean include = true;
+ final ShardSpec shard = segment.getShardSpec();
+ boolean include = true;
- if (shard != null) {
- Map<String, RangeSet<String>> filterDomain = new HashMap<>();
- List<String> dimensions = shard.getDomainDimensions();
- for (String dimension : dimensions) {
- final VirtualColumn shardVirtualColumn =
shard.getDomainVirtualColumns().getVirtualColumn(dimension);
- if (shardVirtualColumn != null) {
- final VirtualColumn queryEquivalent =
virtualColumns.findEquivalent(shardVirtualColumn);
- if (queryEquivalent != null) {
- if (filterFields == null ||
filterFields.contains(queryEquivalent.getOutputName())) {
- Optional<RangeSet<String>> optFilterRangeSet = rangeCache
- .computeIfAbsent(
- queryEquivalent.getOutputName(),
- d ->
Optional.ofNullable(filter.getDimensionRangeSet(d))
- );
- optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put(
- shardVirtualColumn.getOutputName(),
- stringRangeSet
- ));
- }
+ if (shard != null) {
+ final Map<String, RangeSet<String>> filterDomain = new HashMap<>();
+ final List<String> dimensions = shard.getDomainDimensions();
+ for (String dimension : dimensions) {
+ final VirtualColumn shardVirtualColumn =
shard.getDomainVirtualColumns().getVirtualColumn(dimension);
+ if (shardVirtualColumn != null) {
+ final VirtualColumn queryEquivalent =
virtualColumns.findEquivalent(shardVirtualColumn);
+ if (queryEquivalent != null) {
+ if (filterFields == null ||
filterFields.contains(queryEquivalent.getOutputName())) {
+ final Optional<RangeSet<String>> optFilterRangeSet = rangeCache
+ .computeIfAbsent(
+ queryEquivalent.getOutputName(),
+ d -> Optional.ofNullable(filter.getDimensionRangeSet(d))
+ );
+ optFilterRangeSet.ifPresent(stringRangeSet -> filterDomain.put(
+ shardVirtualColumn.getOutputName(),
+ stringRangeSet
+ ));
}
- } else if (filterFields == null || filterFields.contains(dimension))
{
- Optional<RangeSet<String>> optFilterRangeSet =
- rangeCache.computeIfAbsent(dimension, d ->
Optional.ofNullable(filter.getDimensionRangeSet(d)));
- optFilterRangeSet.ifPresent(stringRangeSet ->
filterDomain.put(dimension, stringRangeSet));
}
- }
- if (!filterDomain.isEmpty() && !shard.possibleInDomain(filterDomain)) {
- include = false;
+ } else if (filterFields == null || filterFields.contains(dimension)) {
+ final Optional<RangeSet<String>> optFilterRangeSet =
+ rangeCache.computeIfAbsent(dimension, d ->
Optional.ofNullable(filter.getDimensionRangeSet(d)));
+ optFilterRangeSet.ifPresent(stringRangeSet ->
filterDomain.put(dimension, stringRangeSet));
}
}
-
- if (include) {
- retSet.add(obj);
+ if (!filterDomain.isEmpty() && !shard.possibleInDomain(filterDomain)) {
+ include = false;
}
}
- return retSet;
+ return include;
+ }
+
+ @Override
+ public SegmentPruner combine(SegmentPruner other)
+ {
+ if (other instanceof FilterSegmentPruner pruner) {
+ final List<VirtualColumn> combinedVirtualColumns = new ArrayList<>();
+
combinedVirtualColumns.addAll(List.of(virtualColumns.getVirtualColumns()));
+
combinedVirtualColumns.addAll(List.of(pruner.virtualColumns.getVirtualColumns()));
+
+ final Set<String> combinedFields = new LinkedHashSet<>();
+ combinedFields.addAll(filterFields);
+ combinedFields.addAll(pruner.filterFields);
+
+ final DimFilter combinedFilter = new AndDimFilter(filter, pruner.filter);
+
+ return new FilterSegmentPruner(
+ combinedFilter,
+ combinedFields,
+ VirtualColumns.create(combinedVirtualColumns)
+ );
+ } else if (other instanceof CompositeSegmentPruner composite) {
+ // composite pruner can combine a filter pruner with any filter pruners
it already has, so call it
+ return composite.combine(this);
+ }
+ return new CompositeSegmentPruner(
+ Set.of(this, other)
+ );
}
@Override
diff --git
a/processing/src/main/java/org/apache/druid/query/filter/SegmentPruner.java
b/processing/src/main/java/org/apache/druid/query/filter/SegmentPruner.java
index 4b372271659..fdb5b500e9d 100644
--- a/processing/src/main/java/org/apache/druid/query/filter/SegmentPruner.java
+++ b/processing/src/main/java/org/apache/druid/query/filter/SegmentPruner.java
@@ -22,13 +22,50 @@ package org.apache.druid.query.filter;
import org.apache.druid.timeline.DataSegment;
import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.Set;
import java.util.function.Function;
public interface SegmentPruner
{
+ /**
+ * Check if a {@link DataSegment} should be included or not. If this method
returns false, the segment can be skipped
+ * during processing.
+ */
+ boolean include(DataSegment segment);
+
/**
* Filter the given {@link Iterable} of objects containing a {@link
DataSegment} (obtained from the converter
* function), to reduce the overall working set which need to be processed.
+ *
+ * @param input The iterable of objects to be filtered
+ * @param converter The function to convert T to {@link DataSegment} to
{@link #include(DataSegment)}
+ * @param <T> This can be any type, as long as transform function is
provided to extract a {@link DataSegment}
+ *
+ * @return The set of pruned object, in the same order as input
+ */
+ default <T> Collection<T> prune(Iterable<T> input, Function<T, DataSegment>
converter)
+ {
+ // LinkedHashSet retains order from "input".
+ final Set<T> retSet = new LinkedHashSet<>();
+
+ for (T obj : input) {
+ final DataSegment segment = converter.apply(obj);
+ if (segment == null) {
+ continue;
+ }
+
+ if (include(segment)) {
+ retSet.add(obj);
+ }
+ }
+ return retSet;
+ }
+
+ /**
+ * @return combined {@link SegmentPruner}, which if both are of the same
type may be merged into a new pruner of the
+ * same type that contains the information of both, or if not directly
combinable, implementors should create a
+ * {@link CompositeSegmentPruner}
*/
- <T> Collection<T> prune(Iterable<T> input, Function<T, DataSegment>
converter);
+ SegmentPruner combine(SegmentPruner other);
}
diff --git
a/processing/src/main/java/org/apache/druid/query/planning/ExecutionVertex.java
b/processing/src/main/java/org/apache/druid/query/planning/ExecutionVertex.java
index 64a974771a2..462a73bc0a1 100644
---
a/processing/src/main/java/org/apache/druid/query/planning/ExecutionVertex.java
+++
b/processing/src/main/java/org/apache/druid/query/planning/ExecutionVertex.java
@@ -220,8 +220,11 @@ public class ExecutionVertex
if (!topQuery.context().isSecondaryPartitionPruningEnabled()) {
return null;
}
+
+ final SegmentPruner forDataSource =
topQuery.getDataSource().createSegmentPruner();
+
if (topQuery.getFilter() == null) {
- return null;
+ return forDataSource;
}
final Set<String> baseFields = new HashSet<>();
for (final String field : topQuery.getFilter().getRequiredColumns()) {
@@ -229,12 +232,17 @@ public class ExecutionVertex
baseFields.add(field);
}
}
-
- return new FilterSegmentPruner(
+ final FilterSegmentPruner forFilters = new FilterSegmentPruner(
topQuery.getFilter(),
baseFields,
topQuery.getVirtualColumns()
);
+
+ if (forDataSource == null) {
+ return forFilters;
+ }
+
+ return forDataSource.combine(forFilters);
}
/**
diff --git a/processing/src/main/java/org/apache/druid/query/policy/Policy.java
b/processing/src/main/java/org/apache/druid/query/policy/Policy.java
index 16d54f753af..74940476ed1 100644
--- a/processing/src/main/java/org/apache/druid/query/policy/Policy.java
+++ b/processing/src/main/java/org/apache/druid/query/policy/Policy.java
@@ -22,8 +22,11 @@ package org.apache.druid.query.policy;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import org.apache.druid.guice.annotations.UnstableApi;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.segment.CursorBuildSpec;
+import javax.annotation.Nullable;
+
/**
* Extensible interface for a granular-level (e.x. row filter) restriction on
read-table access. Implementations must be
* Jackson-serializable.
@@ -47,4 +50,13 @@ public interface Policy
*/
CursorBuildSpec visit(CursorBuildSpec spec);
+ /**
+ * @return a {@link SegmentPruner} which can be used to skip processing
entire segments if it is possible to determine
+ * if the policy restricts access.
+ */
+ @Nullable
+ default SegmentPruner createSegmentPruner()
+ {
+ return null;
+ }
}
diff --git
a/processing/src/main/java/org/apache/druid/query/policy/RowFilterPolicy.java
b/processing/src/main/java/org/apache/druid/query/policy/RowFilterPolicy.java
index a83c89d4268..4dc9a86a2a7 100644
---
a/processing/src/main/java/org/apache/druid/query/policy/RowFilterPolicy.java
+++
b/processing/src/main/java/org/apache/druid/query/policy/RowFilterPolicy.java
@@ -23,9 +23,13 @@ import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Preconditions;
import org.apache.druid.query.filter.DimFilter;
+import org.apache.druid.query.filter.FilterSegmentPruner;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.segment.CursorBuildSpec;
+import org.apache.druid.segment.VirtualColumns;
import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
import java.util.Objects;
/**
@@ -58,6 +62,13 @@ public class RowFilterPolicy implements Policy
return
CursorBuildSpec.builder(spec).andFilter(rowFilter.toFilter()).build();
}
+ @Nullable
+ @Override
+ public SegmentPruner createSegmentPruner()
+ {
+ return new FilterSegmentPruner(rowFilter, rowFilter.getRequiredColumns(),
VirtualColumns.EMPTY);
+ }
+
@Override
public String toString()
{
diff --git
a/processing/src/test/java/org/apache/druid/query/RestrictedDataSourceTest.java
b/processing/src/test/java/org/apache/druid/query/RestrictedDataSourceTest.java
index f5b5c1e71d6..4401f7d1d2f 100644
---
a/processing/src/test/java/org/apache/druid/query/RestrictedDataSourceTest.java
+++
b/processing/src/test/java/org/apache/druid/query/RestrictedDataSourceTest.java
@@ -23,10 +23,13 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import nl.jqno.equalsverifier.EqualsVerifier;
import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.query.filter.FilterSegmentPruner;
+import org.apache.druid.query.filter.SegmentPruner;
import org.apache.druid.query.filter.TrueDimFilter;
import org.apache.druid.query.policy.NoRestrictionPolicy;
import org.apache.druid.query.policy.RowFilterPolicy;
import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.VirtualColumns;
import org.junit.Assert;
import org.junit.Test;
@@ -150,6 +153,22 @@ public class RestrictedDataSourceTest
);
}
+ @Test
+ public void test_createSegmentPruner_withRowFilterPolicy()
+ {
+ Assert.assertEquals(
+ new FilterSegmentPruner(TrueDimFilter.instance(), null,
VirtualColumns.EMPTY),
+ restrictedFooDataSource.createSegmentPruner()
+ );
+ }
+
+ @Test
+ public void test_createSegmentPruner_withNoRestrictionPolicy()
+ {
+ SegmentPruner pruner = restrictedBarDataSource.createSegmentPruner();
+ Assert.assertNull(pruner);
+ }
+
@Test
public void testStringRep()
{
diff --git
a/processing/src/test/java/org/apache/druid/query/filter/CompositeSegmentPrunerTest.java
b/processing/src/test/java/org/apache/druid/query/filter/CompositeSegmentPrunerTest.java
new file mode 100644
index 00000000000..00068320e0f
--- /dev/null
+++
b/processing/src/test/java/org/apache/druid/query/filter/CompositeSegmentPrunerTest.java
@@ -0,0 +1,280 @@
+/*
+ * 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.query.filter;
+
+import nl.jqno.equalsverifier.EqualsVerifier;
+import org.apache.druid.data.input.StringTuple;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.ColumnType;
+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.ShardSpec;
+import org.joda.time.Interval;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.function.Function;
+
+class CompositeSegmentPrunerTest
+{
+ private static final DimFilter FILTER_DIM1 = new RangeFilter(
+ "dim1",
+ ColumnType.STRING,
+ null,
+ "aaa",
+ null,
+ null,
+ null
+ );
+ private static final DimFilter FILTER_DIM2 = new RangeFilter(
+ "dim2",
+ ColumnType.STRING,
+ null,
+ "bbb",
+ null,
+ null,
+ null
+ );
+
+ @Test
+ void testIncludeAllPrunersPass()
+ {
+ DataSegment seg = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, null, "zzz"));
+ FilterSegmentPruner filter = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ // a pruner that always includes
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filter);
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ Assertions.assertTrue(composite.include(seg));
+ }
+
+ @Test
+ void testIncludeOnePrunerExcludes()
+ {
+ // seg has dim1 range [lmn, null) which is outside filter range (null, aaa)
+ DataSegment seg = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, "lmn", null));
+ FilterSegmentPruner filter = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filter);
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ Assertions.assertFalse(composite.include(seg));
+ }
+
+ @Test
+ void testPrune()
+ {
+ DataSegment seg1 = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, null, "abc"));
+ DataSegment seg2 = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 1, "lmn", null));
+ FilterSegmentPruner filter = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filter);
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ Assertions.assertEquals(
+ Set.of(seg1),
+ composite.prune(List.of(seg1, seg2), Function.identity())
+ );
+ }
+
+ @Test
+ void testValidateRejectsMultipleFilterPruners()
+ {
+ FilterSegmentPruner filter1 = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ FilterSegmentPruner filter2 = new FilterSegmentPruner(FILTER_DIM2, null,
VirtualColumns.EMPTY);
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filter1);
+ pruners.add(filter2);
+
+ Assertions.assertThrows(DruidException.class, () -> new
CompositeSegmentPruner(pruners));
+ }
+
+ @Test
+ void testCombineWithFilterPrunerFoldsIntoExisting()
+ {
+ FilterSegmentPruner filter1 = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filter1);
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ FilterSegmentPruner filter2 = new FilterSegmentPruner(FILTER_DIM2, null,
VirtualColumns.EMPTY);
+ SegmentPruner combined = composite.combine(filter2);
+
+ // result should be a composite with the filter pruners merged into one
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+
+ // seg has dim1 < aaa (included by filter1) but dim2 starts at "ccc"
(excluded by filter2)
+ DataSegment includedBoth = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, null, "abc"));
+ DataSegment excludedByDim2 = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim2", 0, "ccc", null));
+
+ Assertions.assertTrue(combined.include(includedBoth));
+ Assertions.assertFalse(combined.include(excludedByDim2));
+ }
+
+ @Test
+ void testCombineWithFilterPrunerAddsWhenNoExistingFilter()
+ {
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ FilterSegmentPruner filter = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ SegmentPruner combined = composite.combine(filter);
+
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+
+ // the filter should be active in the result
+ DataSegment excluded = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, "lmn", null));
+ Assertions.assertFalse(combined.include(excluded));
+ }
+
+ @Test
+ void testCombineWithComposite()
+ {
+ FilterSegmentPruner filter1 = new FilterSegmentPruner(FILTER_DIM1, null,
VirtualColumns.EMPTY);
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners1 = new LinkedHashSet<>();
+ pruners1.add(filter1);
+ pruners1.add(alwaysInclude);
+ CompositeSegmentPruner composite1 = new CompositeSegmentPruner(pruners1);
+
+ FilterSegmentPruner filter2 = new FilterSegmentPruner(FILTER_DIM2, null,
VirtualColumns.EMPTY);
+
+ Set<SegmentPruner> pruners2 = new LinkedHashSet<>();
+ pruners2.add(filter2);
+ CompositeSegmentPruner composite2 = new CompositeSegmentPruner(pruners2);
+
+ SegmentPruner combined = composite1.combine(composite2);
+
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+
+ // filter pruners from both should be merged, not cause a validation error
+ DataSegment excludedByDim1 = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, "lmn", null));
+ DataSegment excludedByDim2 = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim2", 0, "ccc", null));
+ DataSegment included = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, null, "abc"));
+
+ Assertions.assertFalse(combined.include(excludedByDim1));
+ Assertions.assertFalse(combined.include(excludedByDim2));
+ Assertions.assertTrue(combined.include(included));
+ }
+
+ @Test
+ void testCombineWithOtherPrunerType()
+ {
+ SegmentPruner alwaysInclude = new AlwaysIncludePruner();
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(alwaysInclude);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ NeverIncludePruner neverInclude = new NeverIncludePruner();
+ SegmentPruner combined = composite.combine(neverInclude);
+
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+
+ DataSegment seg = makeDataSegment("2026-01-01/2026-01-02",
makeRange("dim1", 0, null, "abc"));
+ Assertions.assertFalse(combined.include(seg));
+ }
+
+ @Test
+ void testEqualsAndHashcode()
+ {
+
EqualsVerifier.forClass(CompositeSegmentPruner.class).usingGetClass().verify();
+ }
+
+ private ShardSpec makeRange(String column, int partitionNumber, String
start, String end)
+ {
+ return new DimensionRangeShardSpec(
+ List.of(column),
+ null,
+ start == null ? null : StringTuple.create(start),
+ end == null ? null : StringTuple.create(end),
+ partitionNumber,
+ 0
+ );
+ }
+
+ private DataSegment makeDataSegment(String intervalString, ShardSpec
shardSpec)
+ {
+ Interval interval = Intervals.of(intervalString);
+ return DataSegment.builder(SegmentId.of("prune-test", interval, "0",
shardSpec))
+ .shardSpec(shardSpec)
+ .build();
+ }
+
+ /**
+ * Simple test pruner that always includes segments.
+ */
+ private static class AlwaysIncludePruner implements SegmentPruner
+ {
+ @Override
+ public boolean include(DataSegment segment)
+ {
+ return true;
+ }
+
+ @Override
+ public SegmentPruner combine(SegmentPruner other)
+ {
+ return other;
+ }
+ }
+
+ /**
+ * Simple test pruner that never includes segments.
+ */
+ private static class NeverIncludePruner implements SegmentPruner
+ {
+ @Override
+ public boolean include(DataSegment segment)
+ {
+ return false;
+ }
+
+ @Override
+ public SegmentPruner combine(SegmentPruner other)
+ {
+ return this;
+ }
+ }
+}
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 a1a1dba0c28..8d7628ad7e6 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
@@ -37,6 +37,7 @@ import org.junit.jupiter.api.Test;
import javax.annotation.Nullable;
import java.util.Collections;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Function;
@@ -141,6 +142,76 @@ class FilterSegmentPrunerTest
Assertions.assertEquals(Set.copyOf(segs), prunerEmptyFields.prune(segs,
Function.identity()));
}
+ @Test
+ void testCombineWithFilterPruner()
+ {
+ DimFilter filterA = new RangeFilter("dim1", ColumnType.STRING, null,
"aaa", null, null, null);
+ DimFilter filterB = new RangeFilter("dim2", ColumnType.STRING, null,
"bbb", null, null, null);
+
+ FilterSegmentPruner prunerA = new FilterSegmentPruner(filterA, null,
VirtualColumns.EMPTY);
+ FilterSegmentPruner prunerB = new FilterSegmentPruner(filterB, null,
VirtualColumns.EMPTY);
+
+ SegmentPruner combined = prunerA.combine(prunerB);
+ Assertions.assertInstanceOf(FilterSegmentPruner.class, combined);
+
+ // combined pruner should prune based on both filters
+ String interval1 = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z";
+ DataSegment includedByBoth = makeDataSegment(interval1, makeRange("dim1",
0, null, "abc"));
+ DataSegment excludedByDim1 = makeDataSegment(interval1, makeRange("dim1",
1, "lmn", null));
+ DataSegment excludedByDim2 = makeDataSegment(interval1, makeRange("dim2",
0, "ccc", null));
+
+ Assertions.assertTrue(combined.include(includedByBoth));
+ Assertions.assertFalse(combined.include(excludedByDim1));
+ Assertions.assertFalse(combined.include(excludedByDim2));
+ }
+
+ @Test
+ void testCombineWithCompositePruner()
+ {
+ DimFilter filterA = new RangeFilter("dim1", ColumnType.STRING, null,
"aaa", null, null, null);
+ DimFilter filterB = new RangeFilter("dim2", ColumnType.STRING, null,
"bbb", null, null, null);
+
+ FilterSegmentPruner filterPrunerA = new FilterSegmentPruner(filterA, null,
VirtualColumns.EMPTY);
+ FilterSegmentPruner filterPrunerB = new FilterSegmentPruner(filterB, null,
VirtualColumns.EMPTY);
+
+ Set<SegmentPruner> pruners = new LinkedHashSet<>();
+ pruners.add(filterPrunerB);
+ CompositeSegmentPruner composite = new CompositeSegmentPruner(pruners);
+
+ // FilterSegmentPruner.combine(CompositeSegmentPruner) should delegate to
composite
+ SegmentPruner combined = filterPrunerA.combine(composite);
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+
+ String interval1 = "2026-01-01T00:00:00Z/2026-01-02T00:00:00Z";
+ DataSegment excludedByDim1 = makeDataSegment(interval1, makeRange("dim1",
0, "lmn", null));
+ Assertions.assertFalse(combined.include(excludedByDim1));
+ }
+
+ @Test
+ void testCombineWithUnknownPrunerType()
+ {
+ DimFilter filterA = new RangeFilter("dim1", ColumnType.STRING, null,
"aaa", null, null, null);
+ FilterSegmentPruner filterPruner = new FilterSegmentPruner(filterA, null,
VirtualColumns.EMPTY);
+
+ SegmentPruner other = new SegmentPruner()
+ {
+ @Override
+ public boolean include(DataSegment segment)
+ {
+ return false;
+ }
+
+ @Override
+ public SegmentPruner combine(SegmentPruner other)
+ {
+ return this;
+ }
+ };
+
+ SegmentPruner combined = filterPruner.combine(other);
+ Assertions.assertInstanceOf(CompositeSegmentPruner.class, combined);
+ }
+
@Test
void testEqualsAndHashcode()
{
diff --git
a/processing/src/test/java/org/apache/druid/query/policy/NoRestrictionPolicyTest.java
b/processing/src/test/java/org/apache/druid/query/policy/NoRestrictionPolicyTest.java
index 0f20e11b416..89d0c749a78 100644
---
a/processing/src/test/java/org/apache/druid/query/policy/NoRestrictionPolicyTest.java
+++
b/processing/src/test/java/org/apache/druid/query/policy/NoRestrictionPolicyTest.java
@@ -57,4 +57,10 @@ public class NoRestrictionPolicyTest
final NoRestrictionPolicy policy = NoRestrictionPolicy.instance();
Assert.assertEquals(CursorBuildSpec.FULL_SCAN,
policy.visit(CursorBuildSpec.FULL_SCAN));
}
+
+ @Test
+ public void testCreateSegmentPruner_noRestriction()
+ {
+ Assert.assertNull(NoRestrictionPolicy.instance().createSegmentPruner());
+ }
}
diff --git
a/processing/src/test/java/org/apache/druid/query/policy/RowFilterPolicyTest.java
b/processing/src/test/java/org/apache/druid/query/policy/RowFilterPolicyTest.java
index 265b7322092..c10dd1f508c 100644
---
a/processing/src/test/java/org/apache/druid/query/policy/RowFilterPolicyTest.java
+++
b/processing/src/test/java/org/apache/druid/query/policy/RowFilterPolicyTest.java
@@ -26,8 +26,10 @@ import nl.jqno.equalsverifier.EqualsVerifier;
import org.apache.druid.query.filter.DimFilter;
import org.apache.druid.query.filter.EqualityFilter;
import org.apache.druid.query.filter.Filter;
+import org.apache.druid.query.filter.FilterSegmentPruner;
import org.apache.druid.segment.CursorBuildSpec;
import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.filter.AndFilter;
import org.junit.Assert;
@@ -79,6 +81,18 @@ public class RowFilterPolicyTest
Assert.assertEquals(policyFilter,
policy.visit(CursorBuildSpec.FULL_SCAN).getFilter());
}
+ @Test
+ public void testCreateSegmentPruner()
+ {
+ DimFilter policyFilter = new EqualityFilter("col", ColumnType.STRING,
"val", null);
+ final RowFilterPolicy policy = RowFilterPolicy.from(policyFilter);
+
+ Assert.assertEquals(
+ new FilterSegmentPruner(policyFilter, null, VirtualColumns.EMPTY),
+ policy.createSegmentPruner()
+ );
+ }
+
@Test
public void testVisit_combineFilters()
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
