This is an automated email from the ASF dual-hosted git repository.
clintropolis 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 669e6d1d42b minor: add AcquireMode.PARTIAL_RULE_HELD for testing and
experimentation (#20617)
669e6d1d42b is described below
commit 669e6d1d42bca494271af1e60b081c655af6240b
Author: Clint Wylie <[email protected]>
AuthorDate: Wed Oct 21 04:31:49 2026 -0700
minor: add AcquireMode.PARTIAL_RULE_HELD for testing and experimentation
(#20617)
---
.../query/PartialProjectionLoadRuleQueryTest.java | 40 ++-
.../java/org/apache/druid/query/QueryContext.java | 4 +-
.../query/context/QueryContextParameters.java | 1 +
.../druid/segment/PartialQueryableIndex.java | 57 +++-
.../PartialQueryableIndexCursorFactory.java | 45 ++-
.../segment/PartialQueryableIndexSegment.java | 21 +-
.../segment/file/PartialSegmentFileMapperV10.java | 32 +-
.../apache/druid/segment/loading/AcquireMode.java | 19 +-
.../org/apache/druid/query/QueryContextTest.java | 7 +
...eIndexCursorFactoryClusteredProjectionTest.java | 79 ++++-
...alQueryableIndexCursorFactoryClusteredTest.java | 59 ++++
.../PartialQueryableIndexServableBundlesTest.java | 353 +++++++++++++++++++++
.../loading/PartialSegmentBundleCacheEntry.java | 29 +-
.../loading/PartialSegmentMetadataCacheEntry.java | 33 +-
.../druid/segment/loading/SegmentCacheManager.java | 8 +-
.../segment/loading/SegmentLocalCacheManager.java | 51 ++-
.../segment/loading/VirtualPlaceholderSegment.java | 7 +-
...egmentLocalCacheManagerPartialRuleLoadTest.java | 317 +++++++++++++++++-
18 files changed, 1110 insertions(+), 52 deletions(-)
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
index 90486463e26..1c13cc77c9b 100644
---
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
@@ -42,6 +42,7 @@ import org.apache.druid.query.QueryContext;
import org.apache.druid.query.aggregation.LongMinAggregatorFactory;
import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
import org.apache.druid.query.context.QueryContextParameters;
+import org.apache.druid.segment.loading.AcquireMode;
import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
import org.apache.druid.server.coordinator.rules.ForeverLoadRule;
import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule;
@@ -308,6 +309,38 @@ class PartialProjectionLoadRuleQueryTest extends
EmbeddedClusterTestBase
);
}
+ @Test
+ void testRuleHeldAcquireModeReadsOnlyRuleLoadedBundles()
+ {
+ // With segmentAcquireMode=PARTIAL_RULE_HELD the MSQ leaves only read the
bundles the partial-load rule selected.
+ // The projection is selected, so a query it can serve returns its usual
results. The cluster groups are not, so a
+ // projection miss reads no rows at all, rather than loading them on
demand as in
+ // testProjectionMissLoadsBaseBundleOnDemand (or reading them if an
earlier test already did). Neither query loads
+ // anything.
+ final LatchableEmitter emitter = quiesceAndFlushStorageMonitor();
+ final Map<String, Object> context =
+ QueryContext.ofMap(QueryContextParameters.SEGMENT_ACQUIRE_MODE,
AcquireMode.PARTIAL_RULE_HELD);
+
+ runDartQuery(
+ "SELECT \"countryName\", SUM(\"delta\") FROM \"" + dataSource + "\"
GROUP BY 1 ORDER BY 1",
+ context,
+ "CA,3\nFR,7\nUS,17"
+ );
+ runDartQuery(
+ "SELECT \"channel\", SUM(\"delta\") FROM \"" + dataSource + "\" GROUP
BY 1 ORDER BY 1",
+ context,
+ ""
+ );
+
+ emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT,
MONITOR_QUIESCE_TIMEOUT_MILLIS);
+
+ Assertions.assertEquals(
+ 0L,
+ emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_COUNT),
+ "expected no on-demand file loads under PARTIAL_RULE_HELD"
+ );
+ }
+
@Test
void testRuleTransitionBetweenPartialAndFullLoad()
{
@@ -454,13 +487,18 @@ class PartialProjectionLoadRuleQueryTest extends
EmbeddedClusterTestBase
* only download the files the query actually reads) and asserts the CSV
result matches {@code expectedCsv}.
*/
private void runDartQuery(String sql, String expectedCsv)
+ {
+ runDartQuery(sql, Map.of(), expectedCsv);
+ }
+
+ private void runDartQuery(String sql, Map<String, Object> context, String
expectedCsv)
{
final String sqlQueryId = UUID.randomUUID().toString();
final String result;
try {
result = msqApis.submitDartSqlAsync(
sql,
- QueryContext.ofMap(QueryContextParameters.SQL_QUERY_ID, sqlQueryId),
+
QueryContext.builder().putAll(context).put(QueryContextParameters.SQL_QUERY_ID,
sqlQueryId).toMap(),
broker
).get();
}
diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java
b/processing/src/main/java/org/apache/druid/query/QueryContext.java
index 6fc005cc9bf..b5ca45c04ec 100644
--- a/processing/src/main/java/org/apache/druid/query/QueryContext.java
+++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java
@@ -906,7 +906,9 @@ public class QueryContext
* {@link AcquireMode#PARTIAL} degrades to {@link AcquireMode#FULL} where
partial downloads are unavailable, such as
* when they are disabled by configuration or the segment is not
range-readable. Forcing {@link AcquireMode#PARTIAL}
* where the engine defaults to {@link AcquireMode#FULL} fails at cursor
creation if the engine reads segments through
- * the synchronous cursor API.
+ * the synchronous cursor API. The same applies to {@link
AcquireMode#PARTIAL_RULE_HELD}, which never degrades: it
+ * reads only the bundles selected by each segment's partial-load rule, so
it is only meaningful on servers that apply
+ * those rules (Historicals) and reads nothing elsewhere.
*/
public AcquireMode getSegmentAcquireMode(AcquireMode defaultValue)
{
diff --git
a/processing/src/main/java/org/apache/druid/query/context/QueryContextParameters.java
b/processing/src/main/java/org/apache/druid/query/context/QueryContextParameters.java
index 62cdd7edab3..e5461adb978 100644
---
a/processing/src/main/java/org/apache/druid/query/context/QueryContextParameters.java
+++
b/processing/src/main/java/org/apache/druid/query/context/QueryContextParameters.java
@@ -1832,6 +1832,7 @@ public final class QueryContextParameters
Overrides how leaf stages acquire segments from the segment
cache.
`FULL` downloads each segment up front, while `PARTIAL`
downloads only the parts a query reads.
`PARTIAL` falls back to `FULL` where partial downloads are
unavailable.
+ `PARTIAL_RULE_HELD` reads only the parts a segment's
partial-load rule selected, and nothing from segments without one.
Intended for testing, debugging and performance comparison.
"""
)
diff --git
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java
index 703254f1434..f81d8767334 100644
---
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java
+++
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndex.java
@@ -65,6 +65,7 @@ import java.util.Set;
import java.util.SortedSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
+import java.util.function.Predicate;
import java.util.stream.Collectors;
/**
@@ -353,6 +354,31 @@ public class PartialQueryableIndex implements
QueryableIndex
);
}
+ /**
+ * {@link #getProjection} restricted to the projections whose bundle {@code
servableBundles} accepts.
+ */
+ @Nullable
+ private QueryableProjection<QueryableIndex> getServableProjection(
+ CursorBuildSpec cursorBuildSpec,
+ Predicate<String> servableBundles
+ )
+ {
+ final SortedSet<AggregateProjectionMetadata> servable =
+ new ObjectAVLTreeSet<>(AggregateProjectionMetadata.COMPARATOR);
+ for (AggregateProjectionMetadata projection : projections) {
+ if (servableBundles.test(projection.getSchema().getName())) {
+ servable.add(projection);
+ }
+ }
+ return Projections.findMatchingProjection(
+ cursorBuildSpec,
+ servable,
+ dataInterval,
+ this::checkProjectionHasColumn,
+ this::getProjectionQueryableIndex
+ );
+ }
+
private boolean checkProjectionHasColumn(String projectionName, String
columnName)
{
final ProjectionMetadata meta = projectionSpecs.get(projectionName);
@@ -524,14 +550,31 @@ public class PartialQueryableIndex implements
QueryableIndex
*/
CursorPrefetchPlan planCursorPrefetch(CursorBuildSpec spec)
{
- final QueryableProjection<QueryableIndex> matched = getProjection(spec);
+ return planCursorPrefetch(spec, null);
+ }
+
+ /**
+ * Like {@link #planCursorPrefetch(CursorBuildSpec)}, but when {@code
servableBundles} is non-null, plans only over
+ * the bundles it accepts, as if they were the whole segment. Only accepted
projections are matching candidates, and
+ * only accepted cluster groups are planned. If the plan would otherwise
need an unaccepted base table bundle, it has
+ * no bundles at all, meaning the spec reads no rows. A matched projection's
base table parent bundle is a dependency
+ * of the projection rather than a row selector of its own, so it is planned
regardless.
+ */
+ CursorPrefetchPlan planCursorPrefetch(CursorBuildSpec spec, @Nullable
Predicate<String> servableBundles)
+ {
+ final QueryableProjection<QueryableIndex> matched = servableBundles == null
+ ? getProjection(spec)
+ :
getServableProjection(spec, servableBundles);
if (matched == null && clusteredBaseSummary != null) {
// Clustered base table: metadata-only group resolution; only the groups
whose clustering tuples survive the
// query's filters plan any downloads, one bundle per surviving group.
- final ClusterGroupQueryPlan clusterGroupPlan =
Projections.planClusterGroupQuery(
- new ArrayList<>(getClusterGroupSchemas()),
- spec
- );
+ final List<TableClusterGroupSpec> candidateGroups = new
ArrayList<>(getClusterGroupSchemas());
+ if (servableBundles != null) {
+ candidateGroups.removeIf(
+ group ->
!servableBundles.test(Projections.getClusterGroupBundleName(group.getClusteringValueIds()))
+ );
+ }
+ final ClusterGroupQueryPlan clusterGroupPlan =
Projections.planClusterGroupQuery(candidateGroups, spec);
final List<PrefetchBundle> bundles = new
ArrayList<>(clusterGroupPlan.survivingGroups().size());
for (TableClusterGroupSpec group : clusterGroupPlan.survivingGroups()) {
final QueryableIndex groupIndex = getClusterGroupQueryableIndex(group,
true);
@@ -554,6 +597,10 @@ public class PartialQueryableIndex implements
QueryableIndex
// parents.
final QueryableIndex rowSelector = matched != null ?
matched.getRowSelector() : this;
final String bundleName = matched != null ? matched.getName() :
Projections.BASE_TABLE_PROJECTION_NAME;
+ if (servableBundles != null && !servableBundles.test(bundleName)) {
+ // only reachable for the base table, since a matched projection was
already accepted
+ return new CursorPrefetchPlan(null, null, List.of());
+ }
final Set<String> required = requiredColumns(rowSelector, matched, spec);
final List<PrefetchBundle> bundles = new ArrayList<>(2);
bundles.add(
diff --git
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java
index 2542d80424a..5e9e573d0e0 100644
---
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java
+++
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexCursorFactory.java
@@ -39,6 +39,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
import java.util.function.Supplier;
/**
@@ -70,22 +71,42 @@ import java.util.function.Supplier;
* <b>Parallelism.</b> Every planned range read is its own task on the
download executor, so a bundle's independent
* deep-storage requests proceed concurrently rather than paying their round
trips sequentially. Once every run is
* resident, the ready callback materializes each bundle's required columns
and the cursor holder is constructed.
+ * <p>
+ * <b>Restricted serving.</b> When constructed with a {@code servableBundles}
predicate (see
+ * {@link org.apache.druid.segment.loading.AcquireMode#PARTIAL_RULE_HELD}),
both paths plan only over the bundles it
+ * accepts, and a spec that nothing servable can answer gets an {@link
EmptyCursorHolder}. Accepted bundles that are not
+ * resident are still downloaded by the async path as usual.
*/
public class PartialQueryableIndexCursorFactory implements CursorFactory
{
private final PartialQueryableIndex index;
private final QueryableIndexCursorFactory delegate;
private final PartialBundleAcquirer bundleAcquirer;
+ // when non-null, only the bundles this accepts are served, as if they were
the whole segment; see
+ // PartialQueryableIndex#planCursorPrefetch(CursorBuildSpec, Predicate)
+ @Nullable
+ private final Predicate<String> servableBundles;
public PartialQueryableIndexCursorFactory(
PartialQueryableIndex index,
TimeBoundaryInspector timeBoundaryInspector,
PartialBundleAcquirer bundleAcquirer
)
+ {
+ this(index, timeBoundaryInspector, bundleAcquirer, null);
+ }
+
+ public PartialQueryableIndexCursorFactory(
+ PartialQueryableIndex index,
+ TimeBoundaryInspector timeBoundaryInspector,
+ PartialBundleAcquirer bundleAcquirer,
+ @Nullable Predicate<String> servableBundles
+ )
{
this.index = index;
this.delegate = new QueryableIndexCursorFactory(index,
timeBoundaryInspector);
this.bundleAcquirer = bundleAcquirer;
+ this.servableBundles = servableBundles;
}
@Override
@@ -98,7 +119,18 @@ public class PartialQueryableIndexCursorFactory implements
CursorFactory
+ "on-demand loading, or acquire the segment via the eager path so
all files are loaded up front."
);
}
- return delegate.makeCursorHolder(spec);
+ if (servableBundles == null) {
+ return delegate.makeCursorHolder(spec);
+ }
+ // restricted: dispatch from the restricted plan rather than letting the
delegate match over the whole segment
+ final PartialQueryableIndex.CursorPrefetchPlan plan =
index.planCursorPrefetch(spec, servableBundles);
+ if (plan.bundles().isEmpty()) {
+ return EmptyCursorHolder.forSpec(spec);
+ }
+ if (plan.clusterGroupPlan() != null) {
+ return delegate.makeClusteredCursorHolder(spec, plan.clusterGroupPlan());
+ }
+ return delegate.makeCursorHolderForProjection(spec,
plan.matchedProjection());
}
@Override
@@ -106,12 +138,13 @@ public class PartialQueryableIndexCursorFactory
implements CursorFactory
{
// The index resolves what the spec needs (projection match / surviving
cluster groups / base table, required
// columns, planned range reads); this factory purely orchestrates the
downloads and cursor wiring.
- final PartialQueryableIndex.CursorPrefetchPlan plan =
index.planCursorPrefetch(spec);
+ final PartialQueryableIndex.CursorPrefetchPlan plan =
index.planCursorPrefetch(spec, servableBundles);
+ if (plan.bundles().isEmpty()) {
+ // Filter rules out every cluster group, or nothing servable can answer
the spec: no bundle to acquire, nothing
+ // to download.
+ return AsyncCursorHolder.completed(EmptyCursorHolder.forSpec(spec));
+ }
if (plan.clusterGroupPlan() != null) {
- if (plan.bundles().isEmpty()) {
- // Filter rules out every cluster group: no bundle to acquire, nothing
to download.
- return AsyncCursorHolder.completed(EmptyCursorHolder.forSpec(spec));
- }
// reuse the plan's cluster resolution
return buildAsyncCursorHolder(
plan.bundles(),
diff --git
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java
index 44856c52385..9414f9d7fcd 100644
---
a/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java
+++
b/processing/src/main/java/org/apache/druid/segment/PartialQueryableIndexSegment.java
@@ -26,6 +26,7 @@ import org.joda.time.Interval;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.Closeable;
+import java.util.function.Predicate;
/**
* {@link Segment} wrapper around a {@link PartialQueryableIndex}. Mirrors
{@link QueryableIndexSegment} but wires up
@@ -52,6 +53,23 @@ public class PartialQueryableIndexSegment implements
ReferenceCountedSegmentProv
Closeable onClose,
PartialBundleAcquirer bundleAcquirer
)
+ {
+ this(index, segmentId, onClose, bundleAcquirer, null);
+ }
+
+ /**
+ * @param servableBundles when non-null, the cursor factory only serves the
bundles this accepts, as if they were the
+ * whole segment (see {@link
PartialQueryableIndexCursorFactory}). Other {@link #as} views,
+ * such as {@link QueryableIndex}, {@link Metadata},
and {@link RowCountInspector}, still
+ * describe the whole segment.
+ */
+ public PartialQueryableIndexSegment(
+ PartialQueryableIndex index,
+ SegmentId segmentId,
+ Closeable onClose,
+ PartialBundleAcquirer bundleAcquirer,
+ @Nullable Predicate<String> servableBundles
+ )
{
this.index = index;
this.timeBoundaryInspector = V10TimeBoundaryInspector.forBaseProjection(
@@ -61,7 +79,8 @@ public class PartialQueryableIndexSegment implements
ReferenceCountedSegmentProv
this.cursorFactory = new PartialQueryableIndexCursorFactory(
index,
timeBoundaryInspector,
- bundleAcquirer
+ bundleAcquirer,
+ servableBundles
);
this.segmentId = segmentId;
this.onClose = onClose;
diff --git
a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java
b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java
index d5a372d635f..5bf6992cc83 100644
---
a/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java
+++
b/processing/src/main/java/org/apache/druid/segment/file/PartialSegmentFileMapperV10.java
@@ -412,6 +412,27 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
return bundleToContainerIndices.keySet();
}
+ /**
+ * The bundle that physically holds {@code requestedBundleName}, across this
mapper and its externals. Normally the
+ * name itself, but when no container carries that name and {@link
SegmentFileBuilder#ROOT_BUNDLE_NAME} is the
+ * segment's only bundle, the root bundle: legacy V10 segments written
before bundle names were recorded keep
+ * everything there. Any other unknown name is returned unchanged, so
callers fail loudly rather than silently reading
+ * the wrong data.
+ */
+ public String resolveBundleName(String requestedBundleName)
+ {
+ final Set<String> present = new HashSet<>(getBundleNames());
+ for (PartialSegmentFileMapperV10 external : externalMappers.values()) {
+ present.addAll(external.getBundleNames());
+ }
+ if (!present.contains(requestedBundleName)
+ && present.size() == 1
+ && present.contains(SegmentFileBuilder.ROOT_BUNDLE_NAME)) {
+ return SegmentFileBuilder.ROOT_BUNDLE_NAME;
+ }
+ return requestedBundleName;
+ }
+
/**
* Indices, into this mapper's {@link SegmentFileMetadata#getContainers()}
list, of the containers belonging to
* {@code bundleName}; empty if this mapper has no containers in that bundle.
@@ -663,9 +684,15 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
/**
* Plan the parallel-capped range reads for every not-yet-resident file in
{@code bundleName}'s containers, across
* this mapper AND every attached external mapper so runs come back paired
with their owning mapper, unlike the
- * self-scoped {@link #planParallelFetch}. Empty for an unknown bundle.
+ * self-scoped {@link #planParallelFetch}. The name is first resolved by
{@link #resolveBundleName}, so on a legacy
+ * root-only segment any bundle name plans the root bundle. Empty for an
unknown bundle.
*/
public List<PlannedFetch> planParallelFetchBundle(String bundleName)
+ {
+ return planParallelFetchResolvedBundle(resolveBundleName(bundleName));
+ }
+
+ private List<PlannedFetch> planParallelFetchResolvedBundle(String bundleName)
{
checkClosed();
final List<PlannedFetch> fetches = new ArrayList<>();
@@ -690,7 +717,8 @@ public class PartialSegmentFileMapperV10 implements
SegmentFileMapper
}
}
for (PartialSegmentFileMapperV10 external : externalMappers.values()) {
- fetches.addAll(external.planParallelFetchBundle(bundleName));
+ // already resolved against the whole segment, so don't re-resolve
against the external's bundles alone
+ fetches.addAll(external.planParallelFetchResolvedBundle(bundleName));
}
return fetches;
}
diff --git
a/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
b/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
index 60c0199a671..0d4adc16d49 100644
--- a/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
+++ b/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
@@ -45,5 +45,22 @@ public enum AcquireMode
* Degrades to {@link #FULL} behavior when the segment's deep-storage layout
cannot supply range reads (e.g. a
* non-V10 or zipped segment) or when partial downloads are disabled by
configuration.
*/
- PARTIAL
+ PARTIAL,
+
+ /**
+ * Testing/experimental. Like {@link #PARTIAL}, but the returned segment
only serves the bundles selected by the
+ * segment's current partial-load rule, as if those bundles were the whole
segment. Selected bundles that are not yet
+ * resident are still fetched on demand at cursor-build, so callers must use
the async API such as
+ * {@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync}.
+ * <p>
+ * Projection matching only considers rule-selected projections, so a query
falls back to another selected
+ * projection, a selected base table, or the selected cluster groups.
Unselected cluster groups are left out of the
+ * results, and a query whose projection or base table is not selected reads
no rows. A matched projection may still
+ * read the base table columns it depends on, since it cannot be read
without them.
+ * <p>
+ * Never degrades to {@link #FULL}: a segment with no partial-load rule
applied (including segments that cannot be
+ * partially loaded, or any segment when partial downloads are disabled)
serves no rows, and nothing is downloaded for
+ * it. The bundle selection is snapshotted when the segment is acquired.
+ */
+ PARTIAL_RULE_HELD
}
diff --git
a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
index 8ac2eb6c990..30d9a447693 100644
--- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
+++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
@@ -599,6 +599,13 @@ public class QueryContextTest
QueryContext.builder().putRaw(QueryContextParameters.SEGMENT_ACQUIRE_MODE.getName(),
"PARTIAL").toContext()
.getSegmentAcquireMode(AcquireMode.FULL)
);
+ assertEquals(
+ AcquireMode.PARTIAL_RULE_HELD,
+ QueryContext.builder()
+
.putRaw(QueryContextParameters.SEGMENT_ACQUIRE_MODE.getName(),
"partial_rule_held")
+ .toContext()
+ .getSegmentAcquireMode(AcquireMode.FULL)
+ );
assertThrows(
BadQueryContextException.class,
() ->
QueryContext.builder().putRaw(QueryContextParameters.SEGMENT_ACQUIRE_MODE.getName(),
"dummy").toContext()
diff --git
a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredProjectionTest.java
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredProjectionTest.java
index 7a689999b87..5f0c198c870 100644
---
a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredProjectionTest.java
+++
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredProjectionTest.java
@@ -44,6 +44,7 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
+import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
@@ -52,6 +53,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Predicate;
/**
* Partial (on-demand) download coverage for a segment that is <em>both</em>
clustered <em>and</em> carries an aggregate
@@ -191,12 +193,87 @@ class
PartialQueryableIndexCursorFactoryClusteredProjectionTest extends PartialQ
}
}
+ @Test
+ void testUnservedProjectionFallsBackToServedGroups() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader,
"combo_unserved_projection")) {
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), ACME_BUNDLE::equals);
+ // this spec matches the projection when it's servable (see
testProjectionMatchDownloadsOnlyProjectionBundle)
+ final CursorBuildSpec aggSpec = CursorBuildSpec.builder()
+
.setGroupingColumns(List.of("tenant"))
+
.setAggregators(List.of(new LongSumAggregatorFactory("sum_x", "x")))
+
.setPhysicalColumns(Set.of("tenant", "x"))
+ .build();
+
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(aggSpec);
+ CursorHolder holder = asyncHolder.release()) {
+ // only the served acme group answers (2 of the 3 rows)
+ final Cursor cursor = holder.asCursor();
+ int rows = 0;
+ while (!cursor.isDone()) {
+ rows++;
+ cursor.advance();
+ }
+ Assertions.assertEquals(2, rows);
+ final Set<String> downloaded = opened.mapper().getDownloadedFiles();
+ Assertions.assertTrue(downloaded.contains(ACME_BUNDLE + "/x"), "got: "
+ downloaded);
+ Assertions.assertTrue(
+ downloaded.stream().noneMatch(f -> f.startsWith(PROJECTION_BUNDLE
+ "/")
+ || f.startsWith(GLOBEX_BUNDLE +
"/")),
+ "unserved bundles must not download; got: " + downloaded
+ );
+ }
+ }
+ }
+
+ @Test
+ void testServedProjectionWithUnservedGroups() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader,
"combo_served_projection")) {
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), PROJECTION_BUNDLE::equals);
+ final CursorBuildSpec aggSpec = CursorBuildSpec.builder()
+
.setGroupingColumns(List.of("tenant"))
+
.setAggregators(List.of(new LongSumAggregatorFactory("sum_x", "x")))
+
.setPhysicalColumns(Set.of("tenant", "x"))
+ .build();
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(aggSpec);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertNotNull(holder.asCursor(), "the served projection
answers");
+ }
+
+ // a spec the projection can't answer would need the unserved groups, so
it reads nothing
+ final CursorBuildSpec spec = CursorBuildSpec.builder()
+ .setFilter(new
EqualityFilter("tenant", ColumnType.STRING, "acme", null))
+ .build();
+ try (AsyncCursorHolder asyncHolder = factory.makeCursorHolderAsync(spec);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertTrue(scanRegion(holder).isEmpty());
+ }
+ final Set<String> downloaded = opened.mapper().getDownloadedFiles();
+ Assertions.assertTrue(
+ downloaded.stream().noneMatch(f -> f.startsWith(ACME_BUNDLE + "/")
|| f.startsWith(GLOBEX_BUNDLE + "/")),
+ "unserved groups must not download; got: " + downloaded
+ );
+ }
+ }
+
private PartialQueryableIndexCursorFactory factory(PartialQueryableIndex
index)
+ {
+ return factory(index, null);
+ }
+
+ private PartialQueryableIndexCursorFactory factory(
+ PartialQueryableIndex index,
+ @Nullable Predicate<String> servableBundles
+ )
{
return new PartialQueryableIndexCursorFactory(
index,
QueryableIndexTimeBoundaryInspector.create(index),
- noOpAcquirer(directExec())
+ noOpAcquirer(directExec()),
+ servableBundles
);
}
diff --git
a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredTest.java
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredTest.java
index af4ea18a9ce..87a71f9c9eb 100644
---
a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredTest.java
+++
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexCursorFactoryClusteredTest.java
@@ -249,6 +249,65 @@ class PartialQueryableIndexCursorFactoryClusteredTest
extends PartialQueryableIn
}
}
+ @Test
+ void testServableBundlesDropsUnservedGroups() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "servable_acme")) {
+ final PartialSegmentFileMapperV10 mapper = opened.mapper();
+ final PartialQueryableIndexCursorFactory factory = new
PartialQueryableIndexCursorFactory(
+ opened.index(),
+ QueryableIndexTimeBoundaryInspector.create(opened.index()),
+ noOpAcquirer(directExec()),
+ ACME_BUNDLE::equals
+ );
+
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN);
+ CursorHolder holder = asyncHolder.release()) {
+ // the unserved globex group is left out entirely, rather than nil-ed
+ Assertions.assertEquals(
+ List.of(List.of("acme", "us-east-1"), List.of("acme",
"us-west-2")),
+ scanTenantRegion(holder)
+ );
+ final Set<String> downloaded = mapper.getDownloadedFiles();
+ Assertions.assertTrue(
+ downloaded.stream().noneMatch(f -> f.startsWith(GLOBEX_BUNDLE +
"/")),
+ "unserved globex group must not be downloaded; got: " + downloaded
+ );
+ }
+
+ // sync path drops the unserved group too, even once it is resident
+ mapper.fetchFiles(mapper.getSegmentFileMetadata().getFiles().keySet());
+ try (CursorHolder holder =
factory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) {
+ Assertions.assertEquals(
+ List.of(List.of("acme", "us-east-1"), List.of("acme",
"us-west-2")),
+ scanTenantRegion(holder)
+ );
+ }
+ }
+ }
+
+ @Test
+ void testNoServableGroupsDownloadsNothing() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "servable_none")) {
+ final PartialQueryableIndexCursorFactory factory = new
PartialQueryableIndexCursorFactory(
+ opened.index(),
+ QueryableIndexTimeBoundaryInspector.create(opened.index()),
+ noOpAcquirer(directExec()),
+ bundle -> false
+ );
+ rangeReader.resetCount();
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertTrue(scanTenantRegion(holder).isEmpty());
+ }
+ Assertions.assertEquals(0, rangeReader.getReadCount());
+ Assertions.assertTrue(opened.mapper().getDownloadedFiles().isEmpty());
+ }
+ }
+
@Test
void testSyncMakeCursorHolderAfterFullDownload() throws IOException
{
diff --git
a/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexServableBundlesTest.java
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexServableBundlesTest.java
new file mode 100644
index 00000000000..dd9cbb110a2
--- /dev/null
+++
b/processing/src/test/java/org/apache/druid/segment/PartialQueryableIndexServableBundlesTest.java
@@ -0,0 +1,353 @@
+/*
+ * 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.segment;
+
+import com.google.common.util.concurrent.ListeningExecutorService;
+import org.apache.druid.common.asyncresource.AsyncResource;
+import org.apache.druid.common.asyncresource.AsyncResources;
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.data.input.ListBasedInputRow;
+import org.apache.druid.data.input.impl.AggregateProjectionSpec;
+import org.apache.druid.data.input.impl.DimensionsSpec;
+import org.apache.druid.data.input.impl.LongDimensionSchema;
+import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.data.CompressionStrategy;
+import org.apache.druid.segment.file.CountingRangeReader;
+import org.apache.druid.segment.file.PartialSegmentFileMapperV10;
+import org.apache.druid.segment.incremental.IncrementalIndexSchema;
+import org.apache.druid.segment.projections.Projections;
+import
org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
+import org.apache.druid.testing.TemporaryFolderExtension;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import java.io.Closeable;
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * Coverage for planning and serving a {@link PartialQueryableIndex}
restricted to a set of servable bundles (what
+ * {@link org.apache.druid.segment.loading.AcquireMode#PARTIAL_RULE_HELD}
uses), over a plain (non-clustered) segment
+ * with two aggregate projections that can both answer the same query. The
clustered cases live in
+ * {@link PartialQueryableIndexCursorFactoryClusteredTest} and
+ * {@link PartialQueryableIndexCursorFactoryClusteredProjectionTest}.
+ */
+class PartialQueryableIndexServableBundlesTest extends
PartialQueryableIndexCursorFactoryTestBase
+{
+ private static final DateTime TIME = DateTimes.of("2025-01-01");
+ private static final String BASE = Projections.BASE_TABLE_PROJECTION_NAME;
+ // 2 rows (one per dim1 value), so it sorts ahead of BIG_PROJECTION when
both match
+ private static final String SMALL_PROJECTION = "dim1_count";
+ // 4 rows (one per dim1, metric1 pair)
+ private static final String BIG_PROJECTION = "dim1_metric1_count";
+
+ private static final RowSignature ROW_SIGNATURE = RowSignature.builder()
+ .add("dim1",
ColumnType.STRING)
+
.add("metric1", ColumnType.LONG)
+ .build();
+
+ private static final List<AggregateProjectionSpec> PROJECTIONS = List.of(
+ AggregateProjectionSpec.builder(SMALL_PROJECTION)
+ .groupingColumns(new
StringDimensionSchema("dim1"))
+ .aggregators(new CountAggregatorFactory("_count"))
+ .build(),
+ AggregateProjectionSpec.builder(BIG_PROJECTION)
+ .groupingColumns(new
StringDimensionSchema("dim1"), new LongDimensionSchema("metric1"))
+ .aggregators(new CountAggregatorFactory("_count"))
+ .build()
+ );
+
+ private static final List<InputRow> ROWS = Arrays.asList(
+ new ListBasedInputRow(ROW_SIGNATURE, TIME,
ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 1L)),
+ new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(1),
ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 2L)),
+ new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(2),
ROW_SIGNATURE.getColumnNames(), Arrays.asList("b", 3L)),
+ new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(3),
ROW_SIGNATURE.getColumnNames(), Arrays.asList("b", 4L))
+ );
+
+ // group by dim1 + count, which both projections can answer
+ private static final CursorBuildSpec COUNT_BY_DIM1 =
+ CursorBuildSpec.builder()
+ .setGroupingColumns(List.of("dim1"))
+ .setAggregators(List.of(new
CountAggregatorFactory("_count")))
+ .setPhysicalColumns(Set.of("dim1"))
+ .build();
+
+ @RegisterExtension
+ public static final TemporaryFolderExtension SHARED_TEMPORARY_FOLDER =
TemporaryFolderExtension.classScoped();
+
+ private static File segmentDir;
+
+ @BeforeAll
+ static void buildSegment() throws IOException
+ {
+ final File tmpDir = SHARED_TEMPORARY_FOLDER.newFolder("build_" +
ThreadLocalRandom.current().nextInt());
+ segmentDir = IndexBuilder.create()
+ .useV10()
+ .tmpDir(tmpDir)
+
.segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+ .schema(
+ IncrementalIndexSchema.builder()
+ .withDimensionsSpec(
+
DimensionsSpec.builder()
+
.setDimensions(
+
List.of(
+
new StringDimensionSchema("dim1"),
+
new LongDimensionSchema("metric1")
+ )
+ )
+
.build()
+ )
+ .withRollup(false)
+
.withMinTimestamp(TIME.getMillis())
+
.withProjections(PROJECTIONS)
+ .build()
+ )
+
.indexSpec(IndexSpec.builder().withMetadataCompression(CompressionStrategy.NONE).build())
+ .rows(ROWS)
+ .buildMMappedIndexFile();
+ }
+
+ @Test
+ void testUnrestrictedMatchesSmallestProjection() throws IOException
+ {
+ try (IndexAndMapper opened = openIndex(new
CountingRangeReader(segmentDir), "unrestricted")) {
+ final PartialQueryableIndex.CursorPrefetchPlan plan =
opened.index().planCursorPrefetch(COUNT_BY_DIM1, null);
+ Assertions.assertEquals(SMALL_PROJECTION,
plan.matchedProjection().getName());
+ }
+ }
+
+ @Test
+ void testUnservedBestProjectionFallsBackToServedProjection() throws
IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader,
"fallback_projection")) {
+ final PartialQueryableIndex.CursorPrefetchPlan plan =
+ opened.index().planCursorPrefetch(COUNT_BY_DIM1,
BIG_PROJECTION::equals);
+ Assertions.assertEquals(BIG_PROJECTION,
plan.matchedProjection().getName());
+ // the base parent of the dim1 projection column is a dependency of the
projection, so it's planned even though
+ // the base table itself is not servable
+ Assertions.assertEquals(List.of(BIG_PROJECTION, BASE),
bundleNames(plan));
+
+ final RecordingAcquirer acquirer = new RecordingAcquirer(directExec());
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), acquirer, BIG_PROJECTION::equals);
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(COUNT_BY_DIM1);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertEquals(4, countRows(holder), "the big projection has
one row per (dim1, metric1)");
+ }
+ Assertions.assertEquals(new TreeSet<>(Set.of(BIG_PROJECTION, BASE)),
acquirer.acquired);
+ final Set<String> downloaded = opened.mapper().getDownloadedFiles();
+ Assertions.assertTrue(
+ downloaded.stream().noneMatch(f -> f.startsWith(SMALL_PROJECTION +
"/")),
+ "the unserved projection must not download; got: " + downloaded
+ );
+ }
+ }
+
+ @Test
+ void testNoServedProjectionFallsBackToServedBaseTable() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "fallback_base")) {
+ final RecordingAcquirer acquirer = new RecordingAcquirer(directExec());
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), acquirer, BASE::equals);
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(COUNT_BY_DIM1);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertEquals(4, countRows(holder), "the base table has
every ingested row");
+ }
+ Assertions.assertEquals(new TreeSet<>(Set.of(BASE)), acquirer.acquired);
+ final Set<String> downloaded = opened.mapper().getDownloadedFiles();
+ Assertions.assertTrue(
+ downloaded.stream().allMatch(f -> f.startsWith(BASE + "/")),
+ "only base table files may download; got: " + downloaded
+ );
+ }
+ }
+
+ @Test
+ void testUnservedBaseTableReadsNothing() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "unserved_base")) {
+ // a projection is servable, but a non-aggregate scan can only be
answered by the base table, which is not
+ Assertions.assertTrue(
+ opened.index().planCursorPrefetch(CursorBuildSpec.FULL_SCAN,
SMALL_PROJECTION::equals).bundles().isEmpty()
+ );
+ final RecordingAcquirer acquirer = new RecordingAcquirer(directExec());
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), acquirer, SMALL_PROJECTION::equals);
+ rangeReader.resetCount();
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(CursorBuildSpec.FULL_SCAN)) {
+ Assertions.assertTrue(asyncHolder.isReady(), "nothing to download, so
the holder is ready immediately");
+ try (CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertEquals(0, countRows(holder));
+ }
+ }
+ Assertions.assertTrue(acquirer.acquired.isEmpty(), "got: " +
acquirer.acquired);
+ Assertions.assertEquals(0, rangeReader.getReadCount());
+ Assertions.assertTrue(opened.mapper().getDownloadedFiles().isEmpty());
+ }
+ }
+
+ @Test
+ void testNothingServableReadsNothing() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "nothing_servable")) {
+ final RecordingAcquirer acquirer = new RecordingAcquirer(directExec());
+ final PartialQueryableIndexCursorFactory factory =
factory(opened.index(), acquirer, bundle -> false);
+ rangeReader.resetCount();
+ for (CursorBuildSpec spec : List.of(COUNT_BY_DIM1,
CursorBuildSpec.FULL_SCAN)) {
+ try (AsyncCursorHolder asyncHolder =
factory.makeCursorHolderAsync(spec);
+ CursorHolder holder = asyncHolder.release()) {
+ Assertions.assertEquals(0, countRows(holder));
+ }
+ }
+ Assertions.assertTrue(acquirer.acquired.isEmpty(), "got: " +
acquirer.acquired);
+ Assertions.assertEquals(0, rangeReader.getReadCount());
+ }
+ }
+
+ @Test
+ void testSyncPathHonorsServableBundles() throws IOException
+ {
+ final CountingRangeReader rangeReader = new
CountingRangeReader(segmentDir);
+ try (IndexAndMapper opened = openIndex(rangeReader, "sync_restricted")) {
+ final PartialSegmentFileMapperV10 mapper = opened.mapper();
+ final PartialQueryableIndexCursorFactory factory =
+ factory(opened.index(), new RecordingAcquirer(directExec()),
BIG_PROJECTION::equals);
+
+ // still refuses to download on the sync path
+ Assertions.assertThrows(DruidException.class, () ->
factory.makeCursorHolder(COUNT_BY_DIM1));
+
+ mapper.fetchFiles(mapper.getSegmentFileMetadata().getFiles().keySet());
+ Assertions.assertTrue(mapper.isFullyDownloaded());
+ try (CursorHolder holder = factory.makeCursorHolder(COUNT_BY_DIM1)) {
+ Assertions.assertEquals(4, countRows(holder), "served from the big
projection, not the unserved small one");
+ }
+ try (CursorHolder holder =
factory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) {
+ Assertions.assertEquals(0, countRows(holder), "the base table is not
servable, even though it is resident");
+ }
+ }
+ }
+
+ @Test
+ void testUseProjectionNamingUnservedProjectionFails() throws IOException
+ {
+ try (IndexAndMapper opened = openIndex(new
CountingRangeReader(segmentDir), "use_projection")) {
+ final CursorBuildSpec spec = CursorBuildSpec.builder(COUNT_BY_DIM1)
+ .setQueryContext(
+ QueryContext.of(
+
Map.of(QueryContexts.USE_PROJECTION, SMALL_PROJECTION)
+ )
+ )
+ .build();
+ Assertions.assertEquals(
+ SMALL_PROJECTION,
+ opened.index().planCursorPrefetch(spec,
null).matchedProjection().getName()
+ );
+ final DruidException e = Assertions.assertThrows(
+ DruidException.class,
+ () -> opened.index().planCursorPrefetch(spec, BIG_PROJECTION::equals)
+ );
+ Assertions.assertTrue(e.getMessage().contains("does not satisfy query"),
e.getMessage());
+ }
+ }
+
+ private static PartialQueryableIndexCursorFactory factory(
+ PartialQueryableIndex index,
+ PartialBundleAcquirer acquirer,
+ Predicate<String> servableBundles
+ )
+ {
+ return new PartialQueryableIndexCursorFactory(
+ index,
+ QueryableIndexTimeBoundaryInspector.create(index),
+ acquirer,
+ servableBundles
+ );
+ }
+
+ private static List<String>
bundleNames(PartialQueryableIndex.CursorPrefetchPlan plan)
+ {
+ return plan.bundles()
+ .stream()
+ .map(PartialQueryableIndex.PrefetchBundle::bundleName)
+ .collect(Collectors.toList());
+ }
+
+ private static int countRows(CursorHolder holder)
+ {
+ final Cursor cursor = holder.asCursor();
+ if (cursor == null) {
+ return 0;
+ }
+ int rows = 0;
+ while (!cursor.isDone()) {
+ rows++;
+ cursor.advance();
+ }
+ return rows;
+ }
+
+ /**
+ * Records the bundles the cursor factory acquires; holds are no-ops.
+ */
+ private static final class RecordingAcquirer implements PartialBundleAcquirer
+ {
+ private final ListeningExecutorService downloadExec;
+ private final Set<String> acquired = new TreeSet<>();
+
+ private RecordingAcquirer(ListeningExecutorService downloadExec)
+ {
+ this.downloadExec = downloadExec;
+ }
+
+ @Override
+ public Closeable acquire(String bundleName)
+ {
+ acquired.add(bundleName);
+ return () -> {};
+ }
+
+ @Override
+ public <T> AsyncResource<T> submitDownload(Callable<T> task)
+ {
+ return AsyncResources.fromFutureUnmanaged(downloadExec.submit(task));
+ }
+ }
+}
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java
b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java
index ccf2acab3f9..3614ff6f5d4 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentBundleCacheEntry.java
@@ -651,20 +651,29 @@ public class PartialSegmentBundleCacheEntry implements
CacheEntry
* <li>A future writer that legitimately writes some files to the root
bundle <em>alongside</em> named bundles is
* unaffected: its reader asks for {@code __root__} by name, which
exists, so the fallback never fires.</li>
* </ul>
+ * Delegates to {@link PartialSegmentFileMapperV10#resolveBundleName}, which
whole-bundle prefetch planning also
+ * uses, so the cache layer and the planner always agree on the physical
bundle.
*/
public static String resolveBundleName(PartialSegmentFileMapperV10
fileMapper, String requestedBundleName)
{
- if (!findContainersForBundle(fileMapper, requestedBundleName).isEmpty()) {
- return requestedBundleName;
- }
- final Set<String> present = bundleNames(fileMapper);
- // legacy v10 segments from before bundle field was persisted in the
segment
- if (present.size() == 1 &&
present.contains(SegmentFileBuilder.ROOT_BUNDLE_NAME)) {
- return SegmentFileBuilder.ROOT_BUNDLE_NAME;
+ return fileMapper.resolveBundleName(requestedBundleName);
+ }
+
+ /**
+ * {@link #resolveBundleName} applied to each of {@code
requestedBundleNames}, giving the bundles that physically hold
+ * them. On a legacy root-only segment every name resolves to the root
bundle, so a partial-load rule that selects a
+ * projection by name selects the root bundle.
+ */
+ public static Set<String> resolveBundleNames(
+ PartialSegmentFileMapperV10 fileMapper,
+ Collection<String> requestedBundleNames
+ )
+ {
+ final Set<String> resolved = new HashSet<>();
+ for (String requestedBundleName : requestedBundleNames) {
+ resolved.add(resolveBundleName(fileMapper, requestedBundleName));
}
- // Requested bundle is absent and there are named bundles present: leave
the name as-is so the caller fails
- // loudly (forBundle throws "no containers") rather than silently serving
the wrong data.
- return requestedBundleName;
+ return Set.copyOf(resolved);
}
private static void awaitMount(SettableFuture<Void> future) throws
IOException
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java
b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java
index fa0f7bda6d3..c30391ad0e2 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntry.java
@@ -63,6 +63,7 @@ import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Predicate;
/**
* Cache entry for the metadata header of a V10 segment loaded via partial
download. Mounting this entry range-reads
@@ -1162,7 +1163,7 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
@Override
public Optional<Segment> acquireReference()
{
- return buildSegmentWithExtraClose(null);
+ return buildSegmentWithExtraClose(null, false);
}
/**
@@ -1172,12 +1173,23 @@ public class PartialSegmentMetadataCacheEntry
implements SegmentCacheEntry, Resi
@Override
public Optional<Segment> acquireReference(Closeable extraOnClose)
{
- return buildSegmentWithExtraClose(extraOnClose);
+ return buildSegmentWithExtraClose(extraOnClose, false);
}
- private Optional<Segment> buildSegmentWithExtraClose(@Nullable Closeable
extraOnClose)
+ /**
+ * Like {@link #acquireReference(Closeable)}, but for {@link
AcquireMode#PARTIAL_RULE_HELD}: the returned segment only
+ * serves the bundles selected by the partial-load rule applied when this is
called, and serves no rows if no rule is
+ * applied. Selected bundles that are not resident are still downloaded on
demand by the async cursor factory.
+ */
+ public Optional<Segment> acquireRuleHeldReference(@Nullable Closeable
extraOnClose)
+ {
+ return buildSegmentWithExtraClose(extraOnClose, true);
+ }
+
+ private Optional<Segment> buildSegmentWithExtraClose(@Nullable Closeable
extraOnClose, boolean ruleHeldOnly)
{
final PartialQueryableIndex index;
+ final Predicate<String> servableBundles;
entryLock.lock();
try {
if (fileMapper == null) {
@@ -1192,6 +1204,18 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
);
}
index = queryableIndex;
+ if (ruleHeldOnly) {
+ // Snapshot the selection (empty when no rule is applied). It names
physical bundles, since the cache manager
+ // resolves it when applying the rule, while the planner asks by
logical name, so resolve each requested name
+ // the same way the bundle acquirer does: on a legacy root-only
segment every name maps to the root bundle.
+ final Set<String> selected = ruleSelectedBundleNames;
+ final PartialSegmentFileMapperV10 mapper = fileMapper;
+ servableBundles = bundleName -> selected.contains(
+ PartialSegmentBundleCacheEntry.resolveBundleName(mapper,
bundleName)
+ );
+ } else {
+ servableBundles = null;
+ }
}
catch (Throwable t) {
// We took ownership of extraOnClose; a failure building the queryable
index must release it rather than leak
@@ -1233,7 +1257,8 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
index,
segmentId,
onClose,
- bundleAcquirer
+ bundleAcquirer,
+ servableBundles
)
);
}
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
index 58f05c099c9..119942576c4 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
@@ -116,7 +116,9 @@ public interface SegmentCacheManager
* {@link AcquireMode#FULL} a cached entry is returned only if it is
behaviorally fully-materialized. With
* {@link AcquireMode#PARTIAL} a mounted (but possibly
not-yet-fully-downloaded) entry is sufficient, and callers
* acknowledge that synchronous methods may refuse if the segment is not
fully downloaded; they must use asynchronous
- * methods like {@link CursorFactory#makeCursorHolderAsync} instead. See
{@link AcquireMode} for details.
+ * methods like {@link CursorFactory#makeCursorHolderAsync} instead. {@link
AcquireMode#PARTIAL_RULE_HELD} is like
+ * {@link AcquireMode#PARTIAL}, but only returns an entry that a
partial-load rule could apply to, and the returned
+ * segment only serves the rule-selected bundles. See {@link AcquireMode}
for details.
*/
Optional<Segment> acquireCachedSegment(SegmentId segmentId, AcquireMode
acquireMode);
@@ -134,7 +136,9 @@ public interface SegmentCacheManager
* {@link AcquireMode#FULL} the segment is fully materialized up front. With
{@link AcquireMode#PARTIAL}, when the
* segment's {@link LoadSpec} supports range reads, only a minimal amount of
metadata is downloaded and additional
* loading is deferred to the async methods callers use to interact with the
returned segment; partial-ineligible
- * segments fall back to {@link AcquireMode#FULL} behavior. See {@link
AcquireMode} for details.
+ * segments fall back to {@link AcquireMode#FULL} behavior. With {@link
AcquireMode#PARTIAL_RULE_HELD} nothing is
+ * loaded up front: the segment serves only the bundles selected by its
partial-load rule, and if there is no rule it
+ * serves no rows. See {@link AcquireMode} for details.
*/
AcquireSegmentAction acquireSegment(DataSegment dataSegment, AcquireMode
acquireMode);
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
b/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
index a310672439c..95bb7991711 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/SegmentLocalCacheManager.java
@@ -462,16 +462,20 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
@Override
public Optional<Segment> acquireCachedSegment(final SegmentId segmentId,
final AcquireMode acquireMode)
{
- // PARTIAL accepts a mounted-but-not-fully-downloaded entry; FULL (and
PARTIAL when partial downloads are disabled)
- // requires a behaviorally fully-materialized entry.
- final boolean requireFullyDownloaded =
- acquireMode == AcquireMode.FULL ||
!config.isVirtualStoragePartialDownloadsEnabled();
- return acquireCachedInternal(segmentId, requireFullyDownloaded);
+ return acquireCachedInternal(segmentId, acquireMode);
}
@Override
public AcquireSegmentAction acquireSegment(final DataSegment dataSegment,
final AcquireMode acquireMode)
{
+ if (acquireMode == AcquireMode.PARTIAL_RULE_HELD) {
+ // Only an existing partial entry can have a rule applied, so never
reserve, mount, or download anything here.
+ // Anything else gets an empty placeholder that serves no rows.
+ final Segment segment = acquireCachedInternal(dataSegment.getId(),
acquireMode).orElseGet(
+ () -> new VirtualPlaceholderSegment(dataSegment)
+ );
+ return
AcquireSegmentAction.completed(AcquireSegmentResult.of(Optional.of(segment)));
+ }
// Partial-eligible segments route through the partial machinery: PARTIAL
mounts lazily (header only, bundles load
// on demand), FULL force-downloads everything up front.
Partial-ineligible segments fall through to the eager path
// below for both modes.
@@ -540,20 +544,25 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
/**
* Shared implementation for {@link #acquireCachedSegment}. Walks the
storage locations checking for existing entries.
* <p>
- * When {@code requireFullyDownloaded} is {@code true} the entry must also
report
- * {@link SegmentCacheEntry#isFullyDownloaded()} ({@link AcquireMode#FULL}),
which never hands back an entry that
- * isn't behaviorally fully-materialized. When {@code false} a mounted entry
is sufficient
- * ({@link AcquireMode#PARTIAL}'s lazy-mount contract), even if some bundles
are still not yet on disk. In either case
- * {@link SegmentCacheEntry#acquireReference(Closeable)} composes the
weak-entry hold into the returned segment's close
- * lifecycle.
+ * {@link AcquireMode#FULL} (and {@link AcquireMode#PARTIAL} when partial
downloads are disabled) requires the entry to
+ * also report {@link SegmentCacheEntry#isFullyDownloaded()}, which never
hands back an entry that isn't behaviorally
+ * fully-materialized. Otherwise a mounted entry is sufficient ({@link
AcquireMode#PARTIAL}'s lazy-mount contract),
+ * even if some bundles are still not yet on disk. In either case {@link
SegmentCacheEntry#acquireReference(Closeable)}
+ * composes the weak-entry hold into the returned segment's close lifecycle.
+ * <p>
+ * {@link AcquireMode#PARTIAL_RULE_HELD} only accepts a mounted {@link
PartialSegmentMetadataCacheEntry}, since no
+ * other entry can have a partial-load rule applied, and returns empty for
anything else.
*/
- private Optional<Segment> acquireCachedInternal(SegmentId segmentId, boolean
requireFullyDownloaded)
+ private Optional<Segment> acquireCachedInternal(SegmentId segmentId,
AcquireMode acquireMode)
{
+ final boolean ruleHeldOnly = acquireMode == AcquireMode.PARTIAL_RULE_HELD;
+ final boolean requireFullyDownloaded =
+ acquireMode == AcquireMode.FULL ||
!config.isVirtualStoragePartialDownloadsEnabled();
final SegmentCacheEntryIdentifier id = new
SegmentCacheEntryIdentifier(segmentId);
for (StorageLocation location : locations) {
final SegmentCacheEntry staticEntry = location.getStaticCacheEntry(id);
if (staticEntry != null) {
- if (staticEntry.isMounted() && (!requireFullyDownloaded ||
staticEntry.isFullyDownloaded())) {
+ if (!ruleHeldOnly && staticEntry.isMounted() &&
(!requireFullyDownloaded || staticEntry.isFullyDownloaded())) {
return staticEntry.acquireReference();
}
return Optional.empty();
@@ -564,6 +573,13 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
final StorageLocation.ReservationHold<SegmentCacheEntry> hold =
location.addWeakReservationHoldIfExists(id);
if (hold != null) {
final SegmentCacheEntry entry = hold.getEntry();
+ if (ruleHeldOnly) {
+ if (entry instanceof PartialSegmentMetadataCacheEntry partial &&
partial.isMounted()) {
+ return partial.acquireRuleHeldReference(hold);
+ }
+ hold.close();
+ return Optional.empty();
+ }
if (entry.isMounted() && (!requireFullyDownloaded ||
entry.isFullyDownloaded())) {
if (requireFullyDownloaded && entry instanceof
PartialSegmentMetadataCacheEntry partial) {
// A FULL reference on a partial segment must pin every bundle for
the segment's lifetime (the sync cursor
@@ -1059,7 +1075,10 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
// awaitEagerDownloadsOrClearRule leaves the metadata entry
weak-reserved with the hook attached; when
// cache eventually reclaims, the info file gets deleted along with
the entry.
metadata.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
- final Set<String> selected = Set.copyOf(
+ // Resolve the selection to the bundles that physically hold it
(they differ on legacy root-only segments), so
+ // rule holds, eager downloads, and AcquireMode.PARTIAL_RULE_HELD
all work with real bundle names.
+ final Set<String> selected =
PartialSegmentBundleCacheEntry.resolveBundleNames(
+ mapper,
wrapper.getSelectedBundleNames(dataSegment,
mapper.getSegmentFileMetadata())
);
// Snapshot the rule this call is about to replace, then pin the
UNION of it and the new selection for the
@@ -1804,7 +1823,9 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
"Bootstrap-restored partial metadata for segment[%s] has no file
mapper", dataSegment.getId()
);
}
- final Set<String> selected = Set.copyOf(
+ // resolved to physical bundles, as in loadPartial
+ final Set<String> selected =
PartialSegmentBundleCacheEntry.resolveBundleNames(
+ mapper,
wrapper.getSelectedBundleNames(dataSegment,
mapper.getSegmentFileMetadata())
);
final String priorFingerprint = partial.getRuleFingerprint();
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/VirtualPlaceholderSegment.java
b/server/src/main/java/org/apache/druid/segment/loading/VirtualPlaceholderSegment.java
index b86b1fd7e09..0800806638f 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/VirtualPlaceholderSegment.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/VirtualPlaceholderSegment.java
@@ -23,6 +23,7 @@ import org.apache.druid.data.input.Row;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.segment.DimensionDictionarySelector;
import org.apache.druid.segment.PhysicalSegmentColumnInspector;
+import org.apache.druid.segment.ReferenceCountedSegmentProvider;
import org.apache.druid.segment.RowAdapters;
import org.apache.druid.segment.RowBasedSegment;
import org.apache.druid.segment.RowCountInspector;
@@ -41,9 +42,13 @@ import javax.annotation.Nullable;
/**
* Empty segment with nothing but a time column, used when {@link
SegmentLoaderConfig#isVirtualStorage()} is true
* to create synthetic segments for {@link
org.apache.druid.query.metadata.metadata.SegmentMetadataQuery} so to not
- * force download of segments from deep storage when computing the SQL schema.
+ * force download of segments from deep storage when computing the SQL schema.
Also what
+ * {@link AcquireMode#PARTIAL_RULE_HELD} acquires for a segment that has no
partial-load rule, so it serves no rows.
+ * <p>
+ * Holds nothing in the cache, so it is a {@link
ReferenceCountedSegmentProvider.LeafReference} whose close does nothing.
*/
public class VirtualPlaceholderSegment extends RowBasedSegment<Row>
+ implements ReferenceCountedSegmentProvider.LeafReference
{
private final DataSegment dataSegment;
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
index 587f5280a38..5ae5f4115ef 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
@@ -40,9 +40,16 @@ import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.query.aggregation.CountAggregatorFactory;
import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
import org.apache.druid.query.expression.TestExprMacroTable;
+import org.apache.druid.segment.AsyncCursorHolder;
+import org.apache.druid.segment.Cursor;
+import org.apache.druid.segment.CursorBuildSpec;
+import org.apache.druid.segment.CursorFactory;
+import org.apache.druid.segment.CursorHolder;
import org.apache.druid.segment.IndexBuilder;
import org.apache.druid.segment.IndexIO;
import org.apache.druid.segment.IndexSpec;
+import org.apache.druid.segment.PartialQueryableIndexSegment;
+import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentLazyLoadFailCallback;
import org.apache.druid.segment.TestHelper;
import org.apache.druid.segment.column.ColumnConfig;
@@ -50,6 +57,9 @@ import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.data.CompressionStrategy;
import org.apache.druid.segment.file.PartialSegmentFileMapperV10;
+import org.apache.druid.segment.file.SegmentFileBuilder;
+import org.apache.druid.segment.file.SegmentFileContainerMetadata;
+import org.apache.druid.segment.file.SegmentFileMetadata;
import org.apache.druid.segment.incremental.IncrementalIndexSchema;
import org.apache.druid.segment.projections.Projections;
import
org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
@@ -69,12 +79,20 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import javax.annotation.Nullable;
import java.io.File;
+import java.io.FileOutputStream;
import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.file.Files;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
/**
* Manager-level tests for the partial-load-rule {@link
SegmentLocalCacheManager#load} path: a segment whose
@@ -120,6 +138,22 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
.build()
);
+ // group by dim1 + count, which both projections can answer
+ private static final CursorBuildSpec COUNT_BY_DIM1 =
+ CursorBuildSpec.builder()
+ .setGroupingColumns(List.of("dim1"))
+ .setAggregators(List.of(new
CountAggregatorFactory("_count")))
+ .setPhysicalColumns(Set.of("dim1"))
+ .build();
+
+ // group by dim1 + sum(metric1), which only AGG_BUNDLE (or the base table)
can answer
+ private static final CursorBuildSpec SUM_BY_DIM1 =
+ CursorBuildSpec.builder()
+ .setGroupingColumns(List.of("dim1"))
+ .setAggregators(List.of(new
LongSumAggregatorFactory("_metric1_sum", "metric1")))
+ .setPhysicalColumns(Set.of("dim1", "metric1"))
+ .build();
+
private static final List<InputRow> ROWS = Arrays.asList(
new ListBasedInputRow(ROW_SIGNATURE, TIME,
ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 1L)),
new ListBasedInputRow(ROW_SIGNATURE, TIME.plusMinutes(1),
ROW_SIGNATURE.getColumnNames(), Arrays.asList("a", 2L)),
@@ -131,6 +165,10 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
public static final TemporaryFolderExtension SHARED_TEMPORARY_FOLDER =
TemporaryFolderExtension.classScoped();
private static File DEEP_STORAGE_DIR;
+ // the same segment, laid out as a legacy V10 segment whose only bundle is
the root bundle
+ private static File LEGACY_ROOT_ONLY_DEEP_STORAGE_DIR;
+ // as above, but also without per-column file lists
+ private static File LEGACY_ROOT_ONLY_NO_COLUMN_FILES_DEEP_STORAGE_DIR;
@RegisterExtension
public final TemporaryFolderExtension temporaryFolder =
TemporaryFolderExtension.testCaseScoped();
@@ -169,9 +207,59 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
.build())
.rows(ROWS)
.buildMMappedIndexFile();
+ LEGACY_ROOT_ONLY_DEEP_STORAGE_DIR = writeRootOnlyCopy(
+ DEEP_STORAGE_DIR,
+ SHARED_TEMPORARY_FOLDER.newFolder("legacy_" +
ThreadLocalRandom.current().nextInt()),
+ false
+ );
+ LEGACY_ROOT_ONLY_NO_COLUMN_FILES_DEEP_STORAGE_DIR = writeRootOnlyCopy(
+ DEEP_STORAGE_DIR,
+ SHARED_TEMPORARY_FOLDER.newFolder("legacy_no_column_files_" +
ThreadLocalRandom.current().nextInt()),
+ true
+ );
EmittingLogger.registerEmitter(new NoopServiceEmitter());
}
+ /**
+ * Writes a copy of the V10 segment in {@code sourceDir} to {@code
targetDir} with every container re-tagged to the
+ * root bundle, the layout of a legacy segment written before bundle names
were recorded. Container offsets are
+ * relative to the end of the header, so only the metadata changes.
+ *
+ * @param dropColumnFiles also drop the per-column file lists, which legacy
segments predate, so prefetches fall back
+ * to downloading whole bundles
+ */
+ private static File writeRootOnlyCopy(File sourceDir, File targetDir,
boolean dropColumnFiles) throws IOException
+ {
+ final ObjectMapper mapper = TestHelper.makeJsonMapper();
+ final byte[] source = Files.readAllBytes(new File(sourceDir,
IndexIO.V10_FILE_NAME).toPath());
+ Assertions.assertEquals(IndexIO.V10_VERSION, source[0]);
+ Assertions.assertEquals(CompressionStrategy.NONE.getId(), source[1],
"fixture metadata must be uncompressed");
+ final int metadataLength = ByteBuffer.wrap(source, 2,
Integer.BYTES).order(ByteOrder.LITTLE_ENDIAN).getInt();
+ final int headerLength = 2 + Integer.BYTES + metadataLength;
+ final SegmentFileMetadata metadata =
+ mapper.readValue(source, 2 + Integer.BYTES, metadataLength,
SegmentFileMetadata.class);
+ final SegmentFileMetadata rootOnly = new SegmentFileMetadata(
+ metadata.getContainers()
+ .stream()
+ .map(container -> new
SegmentFileContainerMetadata(container.getStartOffset(), container.getSize(),
null))
+ .toList(),
+ metadata.getFiles(),
+ metadata.getInterval(),
+ metadata.getColumnDescriptors(),
+ dropColumnFiles ? null : metadata.getColumnFiles(),
+ metadata.getProjections(),
+ metadata.getBitmapEncoding()
+ );
+ final byte[] rootOnlyMetadata = mapper.writeValueAsBytes(rootOnly);
+ try (OutputStream out = new FileOutputStream(new File(targetDir,
IndexIO.V10_FILE_NAME))) {
+ out.write(new byte[]{IndexIO.V10_VERSION,
CompressionStrategy.NONE.getId()});
+
out.write(ByteBuffer.allocate(Integer.BYTES).order(ByteOrder.LITTLE_ENDIAN).putInt(rootOnlyMetadata.length).array());
+ out.write(rootOnlyMetadata);
+ out.write(source, headerLength, source.length - headerLength);
+ }
+ return targetDir;
+ }
+
@BeforeEach
void setup() throws IOException
{
@@ -1064,6 +1152,221 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
}
}
+ @Test
+ void testRuleHeldAcquireServesOnlySelectedBundles() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+ final DataSegment segment =
partialWrapperSegment(List.of(OTHER_AGG_BUNDLE));
+ manager.load(segment);
+ final long loadsAfterRule = location.getWeakStats().getLoadCount();
+
+ try (AcquireSegmentAction action = manager.acquireSegment(segment,
AcquireMode.PARTIAL_RULE_HELD)) {
+ action.await();
+ try (Segment acquired = action.release().getSegment().orElseThrow()) {
+ Assertions.assertInstanceOf(PartialQueryableIndexSegment.class,
acquired);
+ Assertions.assertEquals(2, countRows(acquired, COUNT_BY_DIM1), "served
from the selected projection");
+ // only the unselected projection has the sum, and the base table
fallback is resident (as a dependency of the
+ // selected projection) but was not selected itself
+ Assertions.assertEquals(0, countRows(acquired, SUM_BY_DIM1));
+ Assertions.assertEquals(0, countRows(acquired,
CursorBuildSpec.FULL_SCAN));
+ }
+ }
+
+ // the cached fast path is restricted the same way
+ final Segment cached = manager.acquireCachedSegment(SEGMENT_ID,
AcquireMode.PARTIAL_RULE_HELD).orElseThrow();
+ try (cached) {
+ Assertions.assertEquals(2, countRows(cached, COUNT_BY_DIM1));
+ Assertions.assertEquals(0, countRows(cached, SUM_BY_DIM1));
+ }
+
+ Assertions.assertNull(
+ location.getCacheEntry(new
PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE)),
+ "the unselected projection must never be loaded"
+ );
+ Assertions.assertEquals(loadsAfterRule,
location.getWeakStats().getLoadCount(), "nothing new was loaded");
+ }
+
+ @Test
+ void testRuleHeldAcquireDownloadsSelectedBundleOnDemand() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+ final PartialSegmentBundleCacheEntryIdentifier aggId =
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE);
+
+ // mount the metadata the way an on-demand query would, then select a
bundle that has not been loaded, as the rule
+ // state looks while a rule's eager downloads are still running
+ try (AcquireSegmentAction action = manager.acquireSegment(plainSegment(),
AcquireMode.PARTIAL)) {
+ action.await();
+ action.release().getSegment().orElseThrow().close();
+ }
+ final PartialSegmentMetadataCacheEntry metadata =
weakReservedMetadata(location, SEGMENT_ID);
+ metadata.applyRule(FINGERPRINT, Set.of(AGG_BUNDLE));
+ Assertions.assertNull(location.getCacheEntry(aggId), "precondition: the
selected bundle is not loaded yet");
+
+ final Segment cached = manager.acquireCachedSegment(SEGMENT_ID,
AcquireMode.PARTIAL_RULE_HELD).orElseThrow();
+ try (cached) {
+ Assertions.assertEquals(2, countRows(cached, SUM_BY_DIM1));
+ Assertions.assertEquals(0, countRows(cached, CursorBuildSpec.FULL_SCAN),
"the base table was not selected");
+ }
+ Assertions.assertTrue(location.isWeakReserved(aggId), "the selected bundle
was loaded on demand");
+ }
+
+ @Test
+ void testRuleHeldAcquireOfEntryWithoutRuleServesNothing() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ // an entry that exists only because a query loaded it has no rule applied
+ try (AcquireSegmentAction action = manager.acquireSegment(plainSegment(),
AcquireMode.PARTIAL)) {
+ action.await();
+ action.release().getSegment().orElseThrow().close();
+ }
+ final long loadsBefore = location.getWeakStats().getLoadCount();
+
+ final Segment cached = manager.acquireCachedSegment(SEGMENT_ID,
AcquireMode.PARTIAL_RULE_HELD).orElseThrow();
+ try (cached) {
+ Assertions.assertEquals(0, countRows(cached, COUNT_BY_DIM1));
+ Assertions.assertEquals(0, countRows(cached, CursorBuildSpec.FULL_SCAN));
+ }
+ Assertions.assertEquals(loadsBefore,
location.getWeakStats().getLoadCount(), "nothing was loaded");
+ }
+
+ @Test
+ void testRuleHeldAcquireOfUncachedSegmentLoadsNothing() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ Assertions.assertTrue(manager.acquireCachedSegment(SEGMENT_ID,
AcquireMode.PARTIAL_RULE_HELD).isEmpty());
+ try (AcquireSegmentAction action = manager.acquireSegment(plainSegment(),
AcquireMode.PARTIAL_RULE_HELD)) {
+ Assertions.assertTrue(action.isReady(), "nothing to load, so the action
is ready immediately");
+ try (Segment acquired = action.release().getSegment().orElseThrow()) {
+ Assertions.assertInstanceOf(VirtualPlaceholderSegment.class, acquired);
+ Assertions.assertEquals(SEGMENT_ID, acquired.getId());
+ Assertions.assertEquals(0, countRows(acquired, COUNT_BY_DIM1));
+ Assertions.assertEquals(0, countRows(acquired,
CursorBuildSpec.FULL_SCAN));
+ }
+ }
+ Assertions.assertNull(
+ location.getCacheEntry(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+ "no cache entry may be created"
+ );
+ Assertions.assertEquals(0, location.getWeakStats().getLoadCount());
+ }
+
+ @Test
+ void testProjectionRuleOnLegacyRootOnlySegmentHoldsRootBundle() throws
Exception
+ {
+ // On a legacy segment the root bundle physically holds everything,
including the selected projection, so the rule
+ // must resolve its selection to the root bundle: that's what gets
rule-held, eagerly downloaded, and served under
+ // PARTIAL_RULE_HELD. Without resolving, the rule holds and downloads
nothing, and PARTIAL_RULE_HELD reads nothing.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+ final DataSegment segment =
+ partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), FINGERPRINT,
LEGACY_ROOT_ONLY_DEEP_STORAGE_DIR);
+ manager.load(segment);
+
+ final PartialSegmentMetadataCacheEntry metadata =
weakReservedMetadata(location, SEGMENT_ID);
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ Assertions.assertNotNull(mapper);
+ Assertions.assertEquals(
+ Set.of(SegmentFileBuilder.ROOT_BUNDLE_NAME),
+ PartialSegmentBundleCacheEntry.bundleNames(mapper),
+ "test setup: the segment must have only the root bundle"
+ );
+ Assertions.assertEquals(Set.of(SegmentFileBuilder.ROOT_BUNDLE_NAME),
metadata.getRuleSelectedBundleNames());
+
Assertions.assertTrue(metadata.isBundleRuleHeld(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+
Assertions.assertTrue(mapper.isBundleFullyDownloaded(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+ final long loadsAfterRule = location.getWeakStats().getLoadCount();
+
+ try (AcquireSegmentAction action = manager.acquireSegment(segment,
AcquireMode.PARTIAL_RULE_HELD)) {
+ action.await();
+ try (Segment acquired = action.release().getSegment().orElseThrow()) {
+ Assertions.assertEquals(2, countRows(acquired, COUNT_BY_DIM1), "served
from the selected projection");
+ // the root bundle is the only bundle, and it holds the other
projection and the base table too
+ Assertions.assertEquals(2, countRows(acquired, SUM_BY_DIM1));
+ Assertions.assertEquals(4, countRows(acquired,
CursorBuildSpec.FULL_SCAN));
+ }
+ }
+ Assertions.assertEquals(loadsAfterRule,
location.getWeakStats().getLoadCount(), "nothing new was loaded");
+ }
+
+ @Test
+ void testOnDemandAcquireOfLegacyRootOnlySegmentWithoutColumnFiles() throws
Exception
+ {
+ // Without column file lists the prefetch downloads whole bundles, and it
must resolve the projection it matched
+ // to the root bundle that physically holds it, or it plans nothing and
reading the columns fails.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ try (AcquireSegmentAction action =
+
manager.acquireSegment(plainSegment(LEGACY_ROOT_ONLY_NO_COLUMN_FILES_DEEP_STORAGE_DIR),
AcquireMode.PARTIAL)) {
+ action.await();
+ try (Segment acquired = action.release().getSegment().orElseThrow()) {
+ final PartialSegmentFileMapperV10 mapper =
weakReservedMetadata(location, SEGMENT_ID).getFileMapper();
+ Assertions.assertNotNull(mapper);
+
Assertions.assertNull(mapper.getSegmentFileMetadata().getColumnFiles(), "test
setup: no column file lists");
+
Assertions.assertFalse(mapper.isBundleFullyDownloaded(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+
+ Assertions.assertEquals(2, countRows(acquired, COUNT_BY_DIM1));
+
Assertions.assertTrue(mapper.isBundleFullyDownloaded(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+ }
+ }
+ }
+
+ @Test
+ void
testRuleHeldAcquireDownloadsSelectedRootBundleOfLegacySegmentWithoutColumnFiles()
throws Exception
+ {
+ // The selected root bundle is not resident yet, so PARTIAL_RULE_HELD must
download it on demand through the same
+ // whole-bundle prefetch.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ try (AcquireSegmentAction action =
+
manager.acquireSegment(plainSegment(LEGACY_ROOT_ONLY_NO_COLUMN_FILES_DEEP_STORAGE_DIR),
AcquireMode.PARTIAL)) {
+ action.await();
+ action.release().getSegment().orElseThrow().close();
+ }
+ final PartialSegmentMetadataCacheEntry metadata =
weakReservedMetadata(location, SEGMENT_ID);
+ metadata.applyRule(FINGERPRINT,
Set.of(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ Assertions.assertNotNull(mapper);
+
Assertions.assertFalse(mapper.isBundleFullyDownloaded(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+
+ final Segment cached = manager.acquireCachedSegment(SEGMENT_ID,
AcquireMode.PARTIAL_RULE_HELD).orElseThrow();
+ try (cached) {
+ Assertions.assertEquals(2, countRows(cached, COUNT_BY_DIM1));
+ }
+
Assertions.assertTrue(mapper.isBundleFullyDownloaded(SegmentFileBuilder.ROOT_BUNDLE_NAME));
+ }
+
+ /**
+ * Builds an async cursor over {@code segment} for {@code spec}, waits for
it, and counts the rows it reads.
+ */
+ private static int countRows(Segment segment, CursorBuildSpec spec) throws
InterruptedException
+ {
+ try (AsyncCursorHolder asyncHolder =
segment.as(CursorFactory.class).makeCursorHolderAsync(spec)) {
+ final CountDownLatch ready = new CountDownLatch(1);
+ asyncHolder.addReadyCallback(ready::countDown);
+ Assertions.assertTrue(ready.await(15, TimeUnit.SECONDS), "async cursor
must materialize within the test timeout");
+ try (CursorHolder holder = asyncHolder.release()) {
+ final Cursor cursor = holder.asCursor();
+ if (cursor == null) {
+ return 0;
+ }
+ int rows = 0;
+ while (!cursor.isDone()) {
+ rows++;
+ cursor.advance();
+ }
+ return rows;
+ }
+ }
+ }
+
private static CacheEntry bundleEntry(StorageLocation location, String
bundleName)
{
final CacheEntry entry =
@@ -1122,10 +1425,15 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
}
private DataSegment partialWrapperSegment(List<String> selectedProjections,
String fingerprint)
+ {
+ return partialWrapperSegment(selectedProjections, fingerprint,
DEEP_STORAGE_DIR);
+ }
+
+ private DataSegment partialWrapperSegment(List<String> selectedProjections,
String fingerprint, File deepStorageDir)
{
final Map<String, Object> delegate = Map.of(
"type", "local",
- "path", DEEP_STORAGE_DIR.getAbsolutePath()
+ "path", deepStorageDir.getAbsolutePath()
);
final Map<String, Object> wrapperWire =
PartialProjectionLoadSpec.wireForm(delegate, selectedProjections,
fingerprint);
@@ -1205,10 +1513,15 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
* for a regular full load.
*/
private DataSegment plainSegment()
+ {
+ return plainSegment(DEEP_STORAGE_DIR);
+ }
+
+ private DataSegment plainSegment(File deepStorageDir)
{
return DataSegment.builder(SEGMENT_ID)
.shardSpec(NoneShardSpec.instance())
- .loadSpec(Map.of("type", "local", "path",
DEEP_STORAGE_DIR.getAbsolutePath()))
+ .loadSpec(Map.of("type", "local", "path",
deepStorageDir.getAbsolutePath()))
.size(0)
.build();
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
