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
new file mode 100644
index 000000000000..a3efef184dcc
--- /dev/null
+++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
@@ -0,0 +1,372 @@
+/*
+ * 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.testing.embedded.query;
+
+import org.apache.druid.catalog.guice.CatalogClientModule;
+import org.apache.druid.catalog.guice.CatalogCoordinatorModule;
+import org.apache.druid.catalog.model.Columns;
+import org.apache.druid.catalog.model.TableMetadata;
+import org.apache.druid.catalog.model.table.TableBuilder;
+import org.apache.druid.common.utils.IdUtils;
+import org.apache.druid.data.input.impl.AggregateProjectionSpec;
+import org.apache.druid.data.input.impl.ClusteredValueGroupsBaseTableProjectionSpec;
+import org.apache.druid.data.input.impl.LongDimensionSchema;
+import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.data.input.impl.TimestampSpec;
+import org.apache.druid.indexer.granularity.SegmentGranularitySpec;
+import org.apache.druid.indexing.common.task.TaskBuilder;
+import org.apache.druid.indexing.common.task.batch.parallel.ParallelIndexSupervisorTask;
+import org.apache.druid.java.util.common.HumanReadableBytes;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
+import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
+import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule;
+import org.apache.druid.server.coordinator.rules.WildcardProjectionPartialLoadMatcher;
+import org.apache.druid.server.metrics.LatchableEmitter;
+import org.apache.druid.server.metrics.StorageMonitor;
+import org.apache.druid.testing.embedded.EmbeddedBroker;
+import org.apache.druid.testing.embedded.EmbeddedCoordinator;
+import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
+import org.apache.druid.testing.embedded.EmbeddedHistorical;
+import org.apache.druid.testing.embedded.EmbeddedIndexer;
+import org.apache.druid.testing.embedded.EmbeddedOverlord;
+import org.apache.druid.testing.embedded.EmbeddedRouter;
+import org.apache.druid.testing.embedded.catalog.TestCatalogClient;
+import org.apache.druid.testing.embedded.junit5.EmbeddedClusterTestBase;
+import org.apache.druid.testing.embedded.minio.MinIOStorageResource;
+import org.apache.druid.testing.embedded.msq.EmbeddedMSQApis;
+import org.hamcrest.MatcherAssert;
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ * End-to-end coverage for the partial-load-rule wiring: a segment with a clustered base + aggregate projection
+ * loaded under a {@link ForeverPartialLoadRule} whose matcher selects only the projection. The coordinator emits
+ * a {@code PartialProjectionLoadSpec} wrapper on the load request; the historical resolves the wrapper into the
+ * projection's bundle name, eagerly downloads that one bundle, and pins it via a rule-hold. Queries that the
+ * projection can serve are proven to run entirely off the pre-loaded bundle (zero on-demand loads); queries that
+ * miss the projection are proven to trigger on-demand loads of the base bundles.
+ *
+ * Exercises the full happy-path flow end-to-end: coordinator matcher → wire form → historical
+ * {@code loadPartial} → {@code applyRule} + eager download → announced fingerprint → query with
+ * projection hit vs miss → on-demand metric emission.
+ */
+class PartialProjectionLoadRuleQueryTest extends EmbeddedClusterTestBase
+{
+ private static final String PROJECTION_NAME = "country_delta";
+ private static final long CACHE_SIZE = HumanReadableBytes.parse("1MiB");
+ private static final long MAX_SIZE = HumanReadableBytes.parse("100MiB");
+ private static final long ESTIMATE_SIZE = HumanReadableBytes.parse("2KiB");
+ // Quiescence wait for the storage monitor: a few times its PT1s emission period so a single missed tick
+ // doesn't falsely read as "idle" while still bounding how long we wait once activity has actually stopped.
+ private static final long MONITOR_QUIESCE_TIMEOUT_MILLIS = 3_000L;
+
+ private final EmbeddedBroker broker = new EmbeddedBroker();
+ private final EmbeddedIndexer indexer = new EmbeddedIndexer();
+ private final EmbeddedOverlord overlord = new EmbeddedOverlord();
+ private final EmbeddedHistorical historical = new EmbeddedHistorical();
+ private final EmbeddedCoordinator coordinator = new EmbeddedCoordinator();
+ private final EmbeddedRouter router = new EmbeddedRouter();
+ private final MinIOStorageResource storageResource = new MinIOStorageResource();
+
+ private EmbeddedMSQApis msqApis;
+
+ @Override
+ public EmbeddedDruidCluster createCluster()
+ {
+ historical.setServerMemory(500_000_000)
+ .addProperty("druid.segmentCache.virtualStorage", "true")
+ .addProperty("druid.segmentCache.virtualStoragePartialDownloadsEnabled", "true")
+ .addProperty(
+ "druid.segmentCache.virtualStorageMetadataReservationEstimate",
+ String.valueOf(ESTIMATE_SIZE)
+ )
+ .addProperty(
+ "druid.segmentCache.virtualStorageLoadThreads",
+ String.valueOf(Runtime.getRuntime().availableProcessors())
+ )
+ .addBeforeStartHook(
+ (cluster, self) -> self.addProperty(
+ "druid.segmentCache.locations",
+ StringUtils.format(
+ "[{\"path\":\"%s\",\"maxSize\":\"%s\"}]",
+ cluster.getTestFolder().newFolder().getAbsolutePath(),
+ CACHE_SIZE
+ )
+ )
+ )
+ .addProperty("druid.server.maxSize", String.valueOf(MAX_SIZE));
+
+ broker.setServerMemory(200_000_000)
+ .addProperty("druid.sql.planner.enableSysQueriesTable", "true")
+ // Poll the catalog often so the broker's DruidSchema picks up the table definition quickly.
+ // Under virtualStorage the historical exposes only __time via VirtualPlaceholderSegment, so
+ // BrokerSegmentMetadataCache alone can't infer the real column schema — the catalog fills that gap.
+ .addProperty("druid.catalog.client.pollingPeriod", "100")
+ // Retain Dart reports long enough to fetch them after the query completes.
+ .addProperty("druid.msq.dart.controller.maxRetainedReportCount", "10");
+
+ coordinator.addProperty("druid.manager.segments.useIncrementalCache", "always");
+
+ overlord.addProperty("druid.manager.segments.useIncrementalCache", "always")
+ .addProperty("druid.manager.segments.pollDuration", "PT0.1s")
+ .addProperty("druid.catalog.client.maxSyncRetries", "0");
+
+ indexer.setServerMemory(300_000_000)
+ .addProperty("druid.worker.capacity", "2")
+ .addProperty("druid.processing.numThreads", "2")
+ .addProperty("druid.segment.handoff.pollDuration", "PT0.1s");
+
+ return EmbeddedDruidCluster
+ .withEmbeddedDerbyAndZookeeper()
+ .useLatchableEmitter()
+ .useDefaultTimeoutForLatchableEmitter(60)
+ .addResource(storageResource)
+ // Catalog: table schema must be published explicitly since historical exposes only __time under
+ // virtualStorage=true (see broker.druid.catalog.client.pollingPeriod above).
+ .addExtensions(CatalogClientModule.class, CatalogCoordinatorModule.class)
+ // Run queries via Dart so MSQ leaf processors take the PARTIAL acquire path (async cursor factory);
+ // native SQL takes the FULL acquire path which pins every bundle for the reference lifetime, defeating
+ // the memory savings of partial loads.
+ .addCommonProperty("druid.msq.dart.enabled", "true")
+ // Clustered base-table + projection segments require the V10 segment format.
+ .addCommonProperty("druid.indexer.task.buildV10", "true")
+ // Range reads over MinIO — required for partial loading to be functional.
+ .addCommonProperty("druid.storage.zip", "false")
+ .addCommonProperty("druid.monitoring.emissionPeriod", "PT1s")
+ .addServer(coordinator)
+ .addServer(overlord)
+ .addServer(indexer)
+ .addServer(historical)
+ .addServer(broker)
+ .addServer(router);
+ }
+
+ @BeforeAll
+ void loadDataAndConfigureRule() throws IOException
+ {
+ msqApis = new EmbeddedMSQApis(cluster, overlord);
+ dataSource = "partial-projection-" + IdUtils.getRandomId();
+
+ // Under virtualStorage=true the historical announces segments via VirtualPlaceholderSegment, which only exposes
+ // __time. That means BrokerSegmentMetadataCache alone can't populate a full DruidSchema for SQL validation, so
+ // the queries below would fail at plan-time with "column not found". Publishing an explicit catalog table
+ // definition gives DruidSchema the real column list. Future work: derive a RowSignature from the partial-load
+ // header so this catalog registration becomes unnecessary.
+ final TestCatalogClient catalog = new TestCatalogClient(cluster);
+ final TableMetadata table = TableBuilder.datasource(dataSource, "HOUR")
+ .column(Columns.TIME_COLUMN, Columns.LONG)
+ .column("channel", Columns.SQL_VARCHAR)
+ .column("countryName", Columns.SQL_VARCHAR)
+ .column("delta", Columns.SQL_BIGINT)
+ .build();
+ catalog.createTable(table, true);
+
+ // Configure the partial-load rule BEFORE ingestion so the initial load applies the rule directly, rather than
+ // needing a rule change afterward. The matcher selects only the country_delta projection; base-table bundles
+ // (cluster groups) are NOT rule-loaded.
+ cluster.callApi().onLeaderCoordinator(
+ c -> c.updateRulesForDatasource(
+ dataSource,
+ List.of(
+ new ForeverPartialLoadRule(
+ Map.of("_default_tier", 1),
+ null,
+ new WildcardProjectionPartialLoadMatcher(List.of(PROJECTION_NAME), null),
+ CannotMatchBehavior.FALL_THROUGH
+ )
+ )
+ )
+ );
+
+ ingestClusteredSegmentWithProjection();
+ }
+
+ @Override
+ protected void refreshDatasourceName()
+ {
+ // Fixed datasource across tests — rule + ingest are one-time setup.
+ }
+
+ @Test
+ void testProjectionHitLoadsNothingOnDemand()
+ {
+ // The country_delta projection can serve group-by-countryName / sum-delta. The projection bundle was
+ // rule-loaded on ingest, so an MSQ leaf PARTIAL acquire should read it without triggering any downloads.
+ // MSQ input channel counters don't (yet) observe partial-load on-demand downloads (those happen inside the
+ // async cursor factory, downstream of the input channel), so we observe the historical's StorageMonitor
+ // directly — VSF_LOAD_* are incremented as each file is actually pulled from deep storage.
+ final LatchableEmitter emitter = quiesceAndFlushStorageMonitor();
+
+ runDartQuery(
+ "SELECT \"countryName\", SUM(\"delta\") FROM \"" + dataSource + "\" GROUP BY 1 ORDER BY 1",
+ "CA,3\nFR,7\nUS,17"
+ );
+
+ // Wait for post-query storage activity (if any) to settle before we sum the counters.
+ 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 for a projection-served Dart query, but the projection bundle was fetched at query time"
+ );
+ Assertions.assertEquals(
+ 0L,
+ emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_BYTES),
+ "expected zero on-demand load bytes for a projection-served Dart query"
+ );
+ }
+
+ @Test
+ void testProjectionMissLoadsBaseBundleOnDemand()
+ {
+ // Group by the clustering column — the projection can't serve this, so MSQ must fall back to the clustered
+ // base table, whose cluster-group bundles were NOT rule-loaded. Under PARTIAL acquire the async cursor
+ // factory triggers per-file downloads on demand, which show up on StorageMonitor's VSF_LOAD_* counters.
+ final LatchableEmitter emitter = quiesceAndFlushStorageMonitor();
+
+ runDartQuery(
+ "SELECT \"channel\", SUM(\"delta\") FROM \"" + dataSource + "\" GROUP BY 1 ORDER BY 1",
+ "#en,18\n#fr,9"
+ );
+
+ emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT, MONITOR_QUIESCE_TIMEOUT_MILLIS);
+
+ MatcherAssert.assertThat(
+ "expected on-demand file loads for a projection-miss Dart query that needs the base-table bundles",
+ emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_COUNT),
+ Matchers.greaterThan(0L)
+ );
+ MatcherAssert.assertThat(
+ "expected non-zero on-demand load bytes for a projection-miss Dart query",
+ emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_BYTES),
+ Matchers.greaterThan(0L)
+ );
+ }
+
+ /**
+ * Waits for any lingering storage-monitor activity from ingest / rule-application to settle, then flushes the
+ * emitter's event queue so subsequent {@code getMetricEventLongSum(...)} calls only reflect activity that
+ * happens after this point.
+ */
+ private LatchableEmitter quiesceAndFlushStorageMonitor()
+ {
+ final LatchableEmitter emitter = historical.latchableEmitter();
+ emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BEGIN_COUNT, MONITOR_QUIESCE_TIMEOUT_MILLIS);
+ emitter.flush();
+ return emitter;
+ }
+
+ /**
+ * Submits {@code sql} via the Dart engine (MSQ leaves take the PARTIAL acquire path, so partial-load bundles
+ * only download the files the query actually reads) and asserts the CSV result matches {@code expectedCsv}.
+ */
+ private void runDartQuery(String sql, String expectedCsv)
+ {
+ final String sqlQueryId = UUID.randomUUID().toString();
+ final String result;
+ try {
+ result = msqApis.submitDartSqlAsync(
+ sql,
+ Map.of(QueryContexts.CTX_SQL_QUERY_ID, sqlQueryId),
+ broker
+ ).get();
+ }
+ catch (Exception e) {
+ throw new RuntimeException("failed to run Dart query [" + sqlQueryId + "]: " + sql, e);
+ }
+ Assertions.assertEquals(expectedCsv, result.trim(), "unexpected result for Dart query [" + sqlQueryId + "]");
+ }
+
+ /**
+ * Ingests a single clustered base-table segment (clustered by {@code channel}) with a {@code country_delta}
+ * aggregate projection (group by {@code countryName}, sum {@code delta}). Copied from
+ * {@link ClusteredSegmentProjectionQueryTest#ingestClusteredSegmentWithProjection}; kept local so this test
+ * doesn't cross-depend on that class.
+ */
+ private void ingestClusteredSegmentWithProjection() throws IOException
+ {
+ final File tmpDir = cluster.getTestFolder().newFolder();
+ final File inputFile = new File(tmpDir, "clustered-input.json");
+ final String inputData =
+ "{\"time\":\"2024-01-01T00:10:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":10}\n"
+ + "{\"time\":\"2024-01-01T00:20:00Z\",\"channel\":\"#en\",\"countryName\":\"US\",\"delta\":5}\n"
+ + "{\"time\":\"2024-01-01T00:30:00Z\",\"channel\":\"#en\",\"countryName\":\"CA\",\"delta\":3}\n"
+ + "{\"time\":\"2024-01-01T00:40:00Z\",\"channel\":\"#fr\",\"countryName\":\"FR\",\"delta\":7}\n"
+ + "{\"time\":\"2024-01-01T00:50:00Z\",\"channel\":\"#fr\",\"countryName\":\"US\",\"delta\":2}\n";
+ Files.write(inputFile.toPath(), inputData.getBytes(StandardCharsets.UTF_8));
+
+ final ClusteredValueGroupsBaseTableProjectionSpec clusterSpec =
+ ClusteredValueGroupsBaseTableProjectionSpec.builder()
+ .columns(
+ new StringDimensionSchema("channel"),
+ new StringDimensionSchema("countryName"),
+ new LongDimensionSchema("delta"),
+ new LongDimensionSchema("__time")
+ )
+ .clusteringColumns("channel")
+ .build();
+
+ final AggregateProjectionSpec projection =
+ AggregateProjectionSpec.builder(PROJECTION_NAME)
+ .groupingColumns(new StringDimensionSchema("countryName"))
+ .aggregators(new LongSumAggregatorFactory("sumDelta", "delta"))
+ .build();
+
+ final SegmentGranularitySpec segmentGranularitySpec = new SegmentGranularitySpec(
+ Granularities.HOUR,
+ List.of(Intervals.of("2024-01-01/2024-01-02"))
+ );
+
+ final String taskId = IdUtils.getRandomId();
+ final ParallelIndexSupervisorTask task = TaskBuilder
+ .ofTypeIndexParallel()
+ .jsonInputFormat()
+ .localInputSourceWithFiles(inputFile)
+ .dataSchema(
+ builder -> builder
+ .withDataSource(dataSource)
+ .withTimestamp(new TimestampSpec("time", "iso", null))
+ .withSegmentGranularity(segmentGranularitySpec)
+ .withBaseTable(clusterSpec)
+ .withProjections(List.of(projection))
+ )
+ .tuningConfig(t -> t.withMaxNumConcurrentSubTasks(1))
+ .withId(taskId);
+
+ cluster.callApi().onLeaderOverlord(o -> o.runTask(taskId, task));
+ cluster.callApi().waitForTaskToSucceed(taskId, overlord);
+ cluster.callApi().waitForAllSegmentsToBeAvailable(dataSource, coordinator, broker);
+ }
+}
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 ea4895a7f2e4..78582bd7b6b9 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
@@ -448,7 +448,8 @@ public void ensureAllDownloaded() throws IOException
}
/**
- * Download every container belonging to {@code bundleName} in this mapper, each in a single range read (see
+ * Download every container belonging to {@code bundleName}, recursing into external mappers to cover bundles that
+ * span the main file plus one or more externals. Each container is downloaded in a single range read (see
* {@link #downloadContainer}). No-op for an unknown bundle.
*/
public void ensureBundleDownloaded(String bundleName) throws IOException
@@ -457,6 +458,32 @@ public void ensureBundleDownloaded(String bundleName) throws IOException
for (int containerIndex : getContainerIndicesForBundle(bundleName)) {
downloadContainer(containerIndex);
}
+ for (PartialSegmentFileMapperV10 external : externalMappers.values()) {
+ external.ensureBundleDownloaded(bundleName);
+ }
+ }
+
+ /**
+ * Whether every container belonging to {@code bundleName} has all its files present in {@link #downloadedFiles},
+ * across the main mapper AND every attached external mapper.
+ *
+ * Returns {@code true} for a bundle name unknown to any mapper (no containers to check).
+ */
+ public boolean isBundleFullyDownloaded(String bundleName)
+ {
+ checkClosed();
+ for (int containerIndex : getContainerIndicesForBundle(bundleName)) {
+ final List fileNames = containerFileNames.get(containerIndex);
+ if (!downloadedFiles.containsAll(fileNames)) {
+ return false;
+ }
+ }
+ for (PartialSegmentFileMapperV10 external : externalMappers.values()) {
+ if (!external.isBundleFullyDownloaded(bundleName)) {
+ return false;
+ }
+ }
+ return true;
}
/**
diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java
index 636de7bc776b..cbc27ad3cf99 100644
--- a/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java
+++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpec.java
@@ -25,7 +25,16 @@
import com.fasterxml.jackson.annotation.JsonTypeName;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Preconditions;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.segment.file.SegmentFileMetadata;
+import org.apache.druid.segment.projections.ClusteredValueGroupsBaseTableSchema;
+import org.apache.druid.segment.projections.ProjectionMetadata;
+import org.apache.druid.segment.projections.Projections;
+import org.apache.druid.segment.projections.TableClusterGroupSpec;
+import org.apache.druid.timeline.ClusterGroupTuples;
+import org.apache.druid.timeline.DataSegment;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
@@ -83,6 +92,64 @@ public List getClusterGroupIndices()
return clusterGroupIndices;
}
+ /**
+ * Resolve each {@code clusterGroupIndex} against the segment's {@link ClusteredValueGroupsBaseTableSchema} to a
+ * {@code __base$} bundle name via {@link Projections#getClusterGroupBundleName}. The base
+ * projection carries the authoritative group list; each index picks the group at that position and its
+ * {@code clusteringValueIds} disambiguate the bundle in the V10 layout.
+ *
+ * Defensive tripwires fire on any structural inconsistency: empty projections list, non-clustered base projection,
+ * segment's cluster-group tuple count vs. metadata's group count mismatch, or index out of range. These conditions
+ * are unreachable under a healthy writer/reader contract; a throw here indicates a coding bug.
+ *
+ * Returns an empty list when {@code clusterGroupIndices} is empty (the "sibling-empty" case where a matcher
+ * applied to a clustered segment but no configured pattern matched any group).
+ */
+ @Override
+ public List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata)
+ {
+ if (clusterGroupIndices.isEmpty()) {
+ return List.of();
+ }
+ final List projections = metadata.getProjections();
+ if (projections == null || projections.isEmpty()) {
+ throw DruidException.defensive(
+ "Cannot resolve cluster-group bundle names for segment[%s]: metadata has no projections",
+ segment.getId()
+ );
+ }
+ if (!(projections.getFirst().getSchema() instanceof ClusteredValueGroupsBaseTableSchema clusteredSummary)) {
+ throw DruidException.defensive(
+ "Cannot resolve cluster-group bundle names for segment[%s]: base projection is not clustered",
+ segment.getId()
+ );
+ }
+ final List metadataGroups = clusteredSummary.getClusterGroups();
+ final ClusterGroupTuples segmentTuples = segment.getClusterGroups();
+ final int segmentTupleCount = segmentTuples == null ? 0 : segmentTuples.tuples().size();
+ if (segmentTupleCount != metadataGroups.size()) {
+ throw DruidException.defensive(
+ "Cluster-group count mismatch for segment[%s]: DataSegment has [%s] tuples, metadata has [%s] groups",
+ segment.getId(),
+ segmentTupleCount,
+ metadataGroups.size()
+ );
+ }
+ final List bundleNames = new ArrayList<>(clusterGroupIndices.size());
+ for (int idx : clusterGroupIndices) {
+ if (idx < 0 || idx >= metadataGroups.size()) {
+ throw DruidException.defensive(
+ "Cluster-group index [%s] is out of range [0, %s) for segment[%s]",
+ idx,
+ metadataGroups.size(),
+ segment.getId()
+ );
+ }
+ bundleNames.add(Projections.getClusterGroupBundleName(metadataGroups.get(idx).getClusteringValueIds()));
+ }
+ return bundleNames;
+ }
+
@Override
public boolean equals(Object o)
{
diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java
index f1924dd6c4db..c9b0e82acdd2 100644
--- a/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java
+++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialLoadSpec.java
@@ -24,10 +24,13 @@
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
+import org.apache.druid.segment.file.SegmentFileMetadata;
+import org.apache.druid.timeline.DataSegment;
import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
+import java.util.List;
import java.util.Map;
/**
@@ -136,4 +139,17 @@ public SegmentRangeReader openRangeReader() throws IOException
{
return materializedDelegate().openRangeReader();
}
+
+ /**
+ * Returns the cache-layer bundle names that this partial-load spec selects for {@code segment}, resolving the
+ * scheme-specific identifiers carried in the load spec against the segment's parsed {@code metadata}. The
+ * historical's cache layer uses the returned names to acquire rule holds on the matching bundles as part of
+ * {@link SegmentCacheManager#load} for a partial-load rule; bundles not in the returned list are left to the
+ * ordinary on-demand acquire path.
+ *
+ * Returns an empty list when the load spec's selection is empty (the "sibling-empty" case — coordinator matcher
+ * applied but produced no positive selection; the historical still installs a rule hold on the metadata so
+ * bootstrap re-discovers the segment).
+ */
+ public abstract List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata);
}
diff --git a/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java b/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java
index 737a10655898..c063f33db8e1 100644
--- a/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java
+++ b/processing/src/main/java/org/apache/druid/segment/loading/PartialProjectionLoadSpec.java
@@ -25,11 +25,17 @@
import com.fasterxml.jackson.annotation.JsonTypeName;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Preconditions;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.segment.file.SegmentFileMetadata;
+import org.apache.druid.segment.projections.ProjectionMetadata;
+import org.apache.druid.timeline.DataSegment;
import org.apache.druid.utils.CollectionUtils;
import java.util.List;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
/**
* A {@link PartialLoadSpec} that requests partial loading of a segment's projections. The base class carries the
@@ -84,6 +90,42 @@ public List getProjections()
return projections;
}
+ /**
+ * Projection names double as bundle names in the V10 partial-segment layout (each projection's containers are
+ * tagged with its name as the bundle prefix), so the load spec selection maps to bundle names verbatim after
+ * validating that each requested name refers to a projection actually present on the segment.
+ *
+ * These are pure defensive tripwires: the coordinator-side matcher derives the wire form from the same segment
+ * metadata this method reads, so a mismatch here indicates a coding bug (matcher/reader drift, writer/reader
+ * contract violation, or serialization corruption).
+ */
+ @Override
+ public List getSelectedBundleNames(DataSegment segment, SegmentFileMetadata metadata)
+ {
+ final List segmentProjections = metadata.getProjections();
+ if (segmentProjections == null || segmentProjections.isEmpty()) {
+ throw DruidException.defensive(
+ "Cannot resolve projection bundles for segment[%s]: metadata has no projections",
+ segment.getId()
+ );
+ }
+ final Set known = segmentProjections.stream()
+ .map(pm -> pm.getSchema().getName())
+ .collect(Collectors.toSet());
+ for (String projection : projections) {
+ if (!known.contains(projection)) {
+ throw DruidException.defensive(
+ "Segment[%s] does not contain projection[%s]; matcher/reader drift or writer/reader contract violation."
+ + " Known projections: %s",
+ segment.getId(),
+ projection,
+ known
+ );
+ }
+ }
+ return projections;
+ }
+
@Override
public boolean equals(Object o)
{
diff --git a/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java b/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java
index 28906eba94d9..0c75eb31aed8 100644
--- a/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java
+++ b/processing/src/test/java/org/apache/druid/segment/loading/PartialClusterGroupLoadSpecTest.java
@@ -27,7 +27,26 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.google.common.collect.ImmutableMap;
+import org.apache.druid.error.DruidException;
import org.apache.druid.jackson.DefaultObjectMapper;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.query.OrderBy;
+import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.ColumnHolder;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.file.SegmentFileMetadata;
+import org.apache.druid.segment.projections.ClusteredValueGroupsBaseTableSchema;
+import org.apache.druid.segment.projections.ClusteringDictionaries;
+import org.apache.druid.segment.projections.ProjectionMetadata;
+import org.apache.druid.segment.projections.ProjectionSchema;
+import org.apache.druid.segment.projections.Projections;
+import org.apache.druid.segment.projections.TableClusterGroupSpec;
+import org.apache.druid.segment.projections.TableProjectionSchema;
+import org.apache.druid.timeline.ClusterGroupTuples;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NumberedShardSpec;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
@@ -163,6 +182,194 @@ void testRejectsNullFingerprint()
);
}
+ @Test
+ void testGetSelectedBundleNamesResolvesIndicesToBundleNames()
+ {
+ // Three cluster groups in the metadata; pick groups 0 and 2. Each index maps to the group's
+ // clusteringValueIds → bundle name.
+ final SegmentFileMetadata metadata = clusteredMetadata(
+ List.of(
+ new TableClusterGroupSpec(List.of(0), 10),
+ new TableClusterGroupSpec(List.of(1), 20),
+ new TableClusterGroupSpec(List.of(2), 30)
+ )
+ );
+ final DataSegment segment = clusteredSegment(
+ List.of(List.of("acme"), List.of("globex"), List.of("initech"))
+ );
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(0, 2),
+ FINGERPRINT,
+ jsonMapper
+ );
+ Assertions.assertEquals(
+ List.of(
+ Projections.getClusterGroupBundleName(List.of(0)),
+ Projections.getClusterGroupBundleName(List.of(2))
+ ),
+ spec.getSelectedBundleNames(segment, metadata)
+ );
+ }
+
+ @Test
+ void testGetSelectedBundleNamesEmptyIndicesReturnsEmpty()
+ {
+ // Sibling-empty: matcher applied but no positive selection. Caller pins only the metadata.
+ final SegmentFileMetadata metadata = clusteredMetadata(
+ List.of(new TableClusterGroupSpec(List.of(0), 1))
+ );
+ final DataSegment segment = clusteredSegment(List.of(List.of("acme")));
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(),
+ FINGERPRINT,
+ jsonMapper
+ );
+ Assertions.assertEquals(List.of(), spec.getSelectedBundleNames(segment, metadata));
+ }
+
+ @Test
+ void testGetSelectedBundleNamesOutOfRangeIndexThrows()
+ {
+ final SegmentFileMetadata metadata = clusteredMetadata(
+ List.of(new TableClusterGroupSpec(List.of(0), 1))
+ );
+ final DataSegment segment = clusteredSegment(List.of(List.of("acme")));
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(5),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(segment, metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("Cluster-group index [5] is out of range"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ @Test
+ void testGetSelectedBundleNamesSizeMismatchThrows()
+ {
+ // Metadata has 2 groups, segment has 1 tuple — writer/reader contract violation; tripwire fires.
+ final SegmentFileMetadata metadata = clusteredMetadata(
+ List.of(
+ new TableClusterGroupSpec(List.of(0), 1),
+ new TableClusterGroupSpec(List.of(1), 1)
+ )
+ );
+ final DataSegment segment = clusteredSegment(List.of(List.of("acme")));
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(0),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(segment, metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("Cluster-group count mismatch"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ @Test
+ void testGetSelectedBundleNamesNonClusteredBaseThrows()
+ {
+ // Base projection is a regular (non-clustered) table — partial cluster-group load is nonsensical.
+ final ProjectionSchema baseSchema = new TableProjectionSchema(
+ VirtualColumns.EMPTY,
+ List.of(ColumnHolder.TIME_COLUMN_NAME, "tenant"),
+ null,
+ List.of(OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME))
+ );
+ final SegmentFileMetadata metadata = new SegmentFileMetadata(
+ List.of(),
+ Map.of(),
+ null,
+ null,
+ List.of(new ProjectionMetadata(1, baseSchema)),
+ null
+ );
+ final DataSegment segment = clusteredSegment(List.of(List.of("acme")));
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(0),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(segment, metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("base projection is not clustered"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ @Test
+ void testGetSelectedBundleNamesNoProjectionsThrows()
+ {
+ final SegmentFileMetadata metadata = new SegmentFileMetadata(List.of(), Map.of(), null, null, null, null);
+ final DataSegment segment = clusteredSegment(List.of(List.of("acme")));
+ final PartialClusterGroupLoadSpec spec = new PartialClusterGroupLoadSpec(
+ DELEGATE,
+ List.of(0),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(segment, metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("metadata has no projections"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ private static final RowSignature CLUSTERING_TENANT = RowSignature.builder()
+ .add("tenant", ColumnType.STRING)
+ .build();
+
+ private static SegmentFileMetadata clusteredMetadata(List groups)
+ {
+ final ClusteredValueGroupsBaseTableSchema baseSchema = new ClusteredValueGroupsBaseTableSchema(
+ VirtualColumns.EMPTY,
+ List.of(ColumnHolder.TIME_COLUMN_NAME, "tenant", "metric"),
+ List.of(OrderBy.ascending("tenant"), OrderBy.ascending(ColumnHolder.TIME_COLUMN_NAME)),
+ CLUSTERING_TENANT,
+ null,
+ new ClusteringDictionaries(List.of("acme", "globex", "initech"), null, null, null),
+ groups
+ );
+ return new SegmentFileMetadata(
+ List.of(),
+ Map.of(),
+ null,
+ null,
+ List.of(new ProjectionMetadata(groups.stream().mapToInt(TableClusterGroupSpec::getNumRows).sum(), baseSchema)),
+ null
+ );
+ }
+
+ private static DataSegment clusteredSegment(List> tuples)
+ {
+ return DataSegment.builder(
+ SegmentId.of("ds", Intervals.ETERNITY, "v1", new NumberedShardSpec(0, 1))
+ )
+ .size(0)
+ .clusterGroups(new ClusterGroupTuples(CLUSTERING_TENANT, tuples))
+ .build();
+ }
+
/**
* Stub LoadSpec used to verify delegation. Uses the same JSON "type"=="stub" key as the test {@link #DELEGATE}.
*/
diff --git a/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java b/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java
index b64a0d0b7c1a..178f767d70be 100644
--- a/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java
+++ b/processing/src/test/java/org/apache/druid/segment/loading/PartialProjectionLoadSpecTest.java
@@ -27,13 +27,26 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.google.common.collect.ImmutableMap;
+import org.apache.druid.error.DruidException;
import org.apache.druid.jackson.DefaultObjectMapper;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.query.OrderBy;
+import org.apache.druid.query.aggregation.AggregatorFactory;
+import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.file.SegmentFileMetadata;
+import org.apache.druid.segment.projections.AggregateProjectionSchema;
+import org.apache.druid.segment.projections.ProjectionMetadata;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NumberedShardSpec;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import javax.annotation.Nullable;
import java.io.ByteArrayInputStream;
import java.io.File;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
@@ -167,6 +180,92 @@ void testRejectsNullFingerprint()
);
}
+ @Test
+ void testGetSelectedBundleNamesReturnsProjections()
+ {
+ // Projection names double as bundle names; each requested name must appear in the segment's projections.
+ PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec(
+ DELEGATE,
+ List.of("user_daily", "user_hourly"),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final SegmentFileMetadata metadata = metadataWithProjections("user_daily", "user_hourly", "some_other");
+ Assertions.assertEquals(List.of("user_daily", "user_hourly"), spec.getSelectedBundleNames(anySegment(), metadata));
+ }
+
+ @Test
+ void testGetSelectedBundleNamesThrowsWhenProjectionMissing()
+ {
+ // Defensive: a wire form referring to a projection this segment doesn't have. Should only happen on
+ // matcher/reader drift or a coding bug.
+ PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec(
+ DELEGATE,
+ List.of("user_daily", "vanished"),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final SegmentFileMetadata metadata = metadataWithProjections("user_daily");
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(anySegment(), metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("does not contain projection[vanished]"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ @Test
+ void testGetSelectedBundleNamesThrowsWhenMetadataHasNoProjections()
+ {
+ // Defensive: a segment with no projections in metadata can't satisfy any projection rule.
+ PartialProjectionLoadSpec spec = new PartialProjectionLoadSpec(
+ DELEGATE,
+ List.of("user_daily"),
+ FINGERPRINT,
+ jsonMapper
+ );
+ final SegmentFileMetadata metadata = new SegmentFileMetadata(List.of(), Map.of(), null, null, null, null);
+ final DruidException thrown = Assertions.assertThrows(
+ DruidException.class,
+ () -> spec.getSelectedBundleNames(anySegment(), metadata)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage().contains("metadata has no projections"),
+ "unexpected message: " + thrown.getMessage()
+ );
+ }
+
+ private static DataSegment anySegment()
+ {
+ return DataSegment.builder(SegmentId.of("ds", Intervals.ETERNITY, "v1", new NumberedShardSpec(0, 1)))
+ .size(0)
+ .build();
+ }
+
+ private static SegmentFileMetadata metadataWithProjections(String... names)
+ {
+ final List projections = new ArrayList<>(names.length);
+ for (String name : names) {
+ projections.add(new ProjectionMetadata(1, projectionSchemaNamed(name)));
+ }
+ return new SegmentFileMetadata(List.of(), Map.of(), null, null, projections, null);
+ }
+
+ private static AggregateProjectionSchema projectionSchemaNamed(String name)
+ {
+ return new AggregateProjectionSchema(
+ name,
+ null,
+ null,
+ VirtualColumns.EMPTY,
+ List.of("dim1"),
+ new AggregatorFactory[]{new CountAggregatorFactory("cnt")},
+ List.of(OrderBy.ascending("dim1"))
+ );
+ }
+
/**
* Stub LoadSpec used to verify delegation. Uses the same JSON "type"=="stub" key as the test {@link #DELEGATE}.
*/
diff --git a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java
index 8a53c1b5f139..b4d78ae3cfc3 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/PartialSegmentCacheBootstrap.java
@@ -56,11 +56,12 @@
* {@link #restoreBundlesFromDisk} to discover, reserve, and mount any bundles whose container files survived. The
* same call from the fresh acquire path is a no-op (no on-disk containers to restore).
*
- * Parent-set inference is delegated to {@link PartialSegmentMetadataCacheEntry#inferParentBundles}. A bundle whose
- * inferred parent isn't itself present on disk is treated as orphaned: its on-disk container files are deleted
- * (via {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the relevant bitmap bits) and the bundle
- * is not restored. The next access through the cache manager acquire path then triggers a clean cold re-fetch, the
- * same fall-back as when the cache manager finds a segment listed in the info directory but missing on disk.
+ * Dependency inference is delegated to {@link PartialSegmentMetadataCacheEntry#inferBundleDependencies}. A bundle
+ * whose inferred dependency isn't itself present on disk is treated as orphaned: its on-disk container files
+ * are deleted (via {@link PartialSegmentFileMapperV10#evictContainer}, which also clears the relevant bitmap bits) and
+ * the bundle is not restored. The next access through the cache manager acquire path then triggers a clean cold
+ * re-fetch, the same fall-back as when the cache manager finds a segment listed in the info directory but missing on
+ * disk.
*/
public final class PartialSegmentCacheBootstrap
{
@@ -173,8 +174,8 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St
final Set orphanedBundleNames = new HashSet<>();
for (String name : presentBundleNames) {
boolean orphaned = false;
- for (PartialSegmentBundleCacheEntryIdentifier parent : metadata.inferParentBundles(name)) {
- if (!presentBundleNames.contains(parent.bundleName())) {
+ for (PartialSegmentBundleCacheEntryIdentifier dep : metadata.inferBundleDependencies(name)) {
+ if (!presentBundleNames.contains(dep.bundleName())) {
orphaned = true;
break;
}
@@ -192,23 +193,23 @@ static void restoreBundlesFromDisk(PartialSegmentMetadataCacheEntry metadata, St
fileMapper.mapperForContainer(ref.externalFilename()).evictContainer(ref.containerIndex());
}
LOG.debug(
- "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (parent unrestorable); next access "
+ "Deleted on-disk state of orphaned bundle[%s] for segment[%s] (dependency unrestorable); next access "
+ "will trigger cold re-fetch",
orphanName,
segmentId
);
}
- // mount base bundle before any dependent bundle so its hold is available when dependents acquire parent holds
+ // Mount the base bundle before any dependent bundle so its hold is available when dependents acquire deps.
mountableBundleNames.sort(Comparator.comparing(name -> !Projections.BASE_TABLE_PROJECTION_NAME.equals(name)));
final List mountedBundles = new ArrayList<>();
boolean success = false;
try {
for (String bundleName : mountableBundleNames) {
- // Mountable bundles have all parents present by construction (orphans were filtered out above), so the
- // inferred parent set is exactly what we want, no further filtering needed.
- final List parentIds = metadata.inferParentBundles(bundleName);
+ // Mountable bundles have all dependencies present by construction (orphans were filtered out above), so the
+ // inferred dependency set is exactly what we want, no further filtering needed.
+ final List parentIds = metadata.inferBundleDependencies(bundleName);
final PartialSegmentBundleCacheEntry bundle = PartialSegmentBundleCacheEntry.forBundle(
metadata,
bundleName,
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 d597b656da57..94927bfb9efe 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
@@ -48,7 +48,11 @@
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
import java.util.List;
+import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;
@@ -132,8 +136,8 @@ public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, Resi
// bundle entries that are currently mounted against this segment, registered by PartialSegmentBundleCacheEntry on
// successful mount and removed on unmount. Lets the drop path enumerate bundles for cascade-close without scanning
- // the StorageLocation's entry maps.
- private final Set linkedBundles = ConcurrentHashMap.newKeySet();
+ // the StorageLocation's entry maps. Keyed by bundle name (one bundle per name per segment).
+ private final ConcurrentHashMap linkedBundles = new ConcurrentHashMap<>();
// Reference-counted gate over the actual cleanup work (close file mapper, delete header files). Set on
// successful mount; unmount() closes the wrapper which defers running cleanup until all outstanding references
@@ -147,6 +151,33 @@ public class PartialSegmentMetadataCacheEntry implements SegmentCacheEntry, Resi
// cleared so retries get a fresh attempt; on success the gate stays set until doActualUnmount clears it.
private final AtomicReference> mountFuture = new AtomicReference<>();
+ /**
+ * Rule-holds state machine (all guarded by {@link #entryLock}). A partial-load rule is applied via
+ * {@link #applyRule(String, Set)} after {@link #mount(StorageLocation)}. While a rule is applied
+ * ({@code ruleFingerprint != null}):
+ *
+ * - The entry holds a self-referential {@link StorageLocation.ReservationHold} on its own weak reservation,
+ * preventing cache from evicting the metadata entry.
+ * - For each bundle name in {@link #ruleSelectedBundleNames} that has registered with this entry via
+ * {@link #registerBundle}, the entry holds a {@link StorageLocation.ReservationHold} on the bundle's weak
+ * reservation, preventing cache from evicting the bundle.
+ *
+ * {@link #clearRule()} releases the holds and clears the state. {@link #applyRule(String, Set)} with a different
+ * fingerprint / selection set diffs the current state and closes / acquires only the deltas, so overlapping bundle
+ * holds stay live across a rule swap.
+ */
+ @GuardedBy("entryLock")
+ @Nullable
+ private String ruleFingerprint;
+ @GuardedBy("entryLock")
+ private Set ruleSelectedBundleNames = Set.of();
+ @GuardedBy("entryLock")
+ private final Map> ruleBundleHolds =
+ new HashMap<>();
+ @GuardedBy("entryLock")
+ @Nullable
+ private StorageLocation.ReservationHold metadataSelfHold;
+
public PartialSegmentMetadataCacheEntry(
SegmentId segmentId,
File localCacheDir,
@@ -199,6 +230,384 @@ File getLocalCacheDir()
return localCacheDir;
}
+ /**
+ * The fingerprint of the partial-load rule currently applied to this entry, or {@code null} if no rule has been
+ * applied. Set by {@link #applyRule(String, Set)} and cleared by {@link #clearRule()}. {@link #doActualUnmount()}
+ * asserts that this field is {@code null} at unmount time, the {@code metadataSelfHold} taken by applyRule pins
+ * the entry's phaser and prevents eviction reclaim / doActualUnmount from firing while a rule is applied. A fresh
+ * mount therefore always starts with no rule applied.
+ */
+ @Nullable
+ public String getRuleFingerprint()
+ {
+ entryLock.lock();
+ try {
+ return ruleFingerprint;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * The set of bundle names currently pinned by an applied partial-load rule, or an empty set if no rule is applied.
+ * Return value is immutable; safe to inspect without further synchronization.
+ */
+ public Set getRuleSelectedBundleNames()
+ {
+ entryLock.lock();
+ try {
+ return ruleSelectedBundleNames;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Whether a partial-load rule is currently applied to this entry. Equivalent to {@code getRuleFingerprint() != null}.
+ * While {@code true}, this entry holds a self-referential {@link StorageLocation.ReservationHold} that prevents cache
+ * from evicting it, so the metadata's V10 header stays resident until {@link #clearRule()} runs.
+ */
+ public boolean isRuleHeld()
+ {
+ entryLock.lock();
+ try {
+ return ruleFingerprint != null;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Apply (or replace) a partial-load rule on this entry. Must be called after {@link #mount(StorageLocation)}, the
+ * entry must be registered with its {@link StorageLocation} so it can acquire holds on itself and on its bundles.
+ *
+ * Concurrency contract. The caller MUST serialize {@code applyRule} and {@link #clearRule} for a given
+ * segment id through an external per-segment lock ({@code SegmentLocalCacheManager} uses its
+ * {@code ReferenceCountingLock segmentLocks}). This method releases {@link #entryLock} between Phase 2 (acquire
+ * holds outside the lock) and Phase 4 (release holds outside the lock); a concurrent invocation of
+ * {@code applyRule}/{@code clearRule} on the same entry can observe intermediate state during those windows and
+ * silently miss rule-hold installations or leak them. {@link #entryLock} alone is insufficient because it is
+ * released across the acquire/release phases by design (holding it across {@link StorageLocation} lock acquisition
+ * would invert the writeLock → entryLock ordering the storage layer already uses for eviction).
+ *
+ * On first application ({@code ruleFingerprint} transitions from {@code null}), a self-referential
+ * {@link StorageLocation.ReservationHold} is taken to pin the metadata entry in-cache. On repeat calls with the same
+ * {@code fingerprint} and matching {@code selectedBundleNames}, this is a no-op. On a genuine rule swap, only the
+ * delta between the previous and new selection is applied: bundle holds for names dropped from the selection are
+ * closed, and holds for names newly added are acquired for any bundle already registered with this entry (later
+ * arrivals get their hold via {@link #registerBundle}).
+ *
+ * Bundles whose {@link StorageLocation.ReservationHold} is currently zero-refcount (never mounted, or evicted but
+ * unregistered before this call) will not appear in {@link #linkedBundles} yet, so no hold is taken for them here;
+ * the caller is expected to drive an on-demand acquire per selected bundle name to trigger a fresh mount, which will
+ * call {@link #registerBundle} and pick up the rule-hold at that point.
+ *
+ * @throws DruidException if the entry is not currently mounted
+ */
+ public void applyRule(String fingerprint, Set selectedBundleNames)
+ {
+ Objects.requireNonNull(fingerprint, "fingerprint");
+ Objects.requireNonNull(selectedBundleNames, "selectedBundleNames");
+ final Set newSelection = Set.copyOf(selectedBundleNames);
+
+ // Phase 1: snapshot under entryLock and compute the diff. Do NOT call StorageLocation methods here — that would
+ // acquire the location's readLock while holding entryLock and could deadlock with a concurrent writeLock holder
+ // that needs entryLock
+ final StorageLocation loc;
+ final boolean needsSelfHold;
+ final List namesToRelease;
+ final List namesToAcquire;
+ entryLock.lock();
+ try {
+ if (location == null) {
+ throw DruidException.defensive(
+ "applyRule on partial metadata entry[%s] requires the entry to be mounted",
+ id
+ );
+ }
+ if (fingerprint.equals(ruleFingerprint) && newSelection.equals(ruleSelectedBundleNames)) {
+ return;
+ }
+ loc = location;
+ needsSelfHold = (metadataSelfHold == null);
+ namesToRelease = new ArrayList<>();
+ for (String name : ruleBundleHolds.keySet()) {
+ if (!newSelection.contains(name)) {
+ namesToRelease.add(name);
+ }
+ }
+ namesToAcquire = new ArrayList<>();
+ for (String name : newSelection) {
+ if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) {
+ namesToAcquire.add(name);
+ }
+ }
+ }
+ finally {
+ entryLock.unlock();
+ }
+
+ // Phase 2 + 3: acquire outside entryLock, then commit under entryLock. Track every hold this call acquires in
+ // `uncommittedHolds`; when a hold's ownership transfers into `metadataSelfHold` or `ruleBundleHolds`, it is
+ // removed from the list. Any throw between acquire and commit (or a race-lose at commit) leaves the hold in
+ // `uncommittedHolds`, and the finally at Phase 4 releases it. This is what makes applyRule all-or-nothing at
+ // the ReservationHold level: no leaks on partial failure, no stranded self-hold under a stale fingerprint.
+ final List> uncommittedHolds = new ArrayList<>();
+ final List> displacedHolds = new ArrayList<>();
+ try {
+ final StorageLocation.ReservationHold newSelfHold;
+ if (needsSelfHold) {
+ newSelfHold = loc.addWeakReservationHoldIfExists(id);
+ if (newSelfHold == null) {
+ throw DruidException.defensive(
+ "Failed to acquire self-referential rule-hold on partial metadata entry[%s]; entry is not weak-reserved",
+ id
+ );
+ }
+ uncommittedHolds.add(newSelfHold);
+ } else {
+ newSelfHold = null;
+ }
+ final Map> acquired = new HashMap<>();
+ for (String name : namesToAcquire) {
+ final StorageLocation.ReservationHold h =
+ loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
+ if (h != null) {
+ acquired.put(name, h);
+ uncommittedHolds.add(h);
+ }
+ }
+
+ // Phase 3: commit under entryLock. Ownership transfers are done by removing from `uncommittedHolds` after
+ // installing into the field. Anything left in `uncommittedHolds` at the end lost a race and gets released
+ // in Phase 4 alongside `displacedHolds` (the pre-existing holds diffed out of ruleBundleHolds).
+ entryLock.lock();
+ try {
+ if (newSelfHold != null) {
+ if (metadataSelfHold == null) {
+ metadataSelfHold = newSelfHold;
+ uncommittedHolds.remove(newSelfHold);
+ }
+ }
+ for (String name : namesToRelease) {
+ final StorageLocation.ReservationHold h = ruleBundleHolds.remove(name);
+ if (h != null) {
+ displacedHolds.add(h);
+ }
+ }
+ for (var e : acquired.entrySet()) {
+ if (!ruleBundleHolds.containsKey(e.getKey())) {
+ ruleBundleHolds.put(e.getKey(), e.getValue());
+ uncommittedHolds.remove(e.getValue());
+ }
+ }
+ ruleFingerprint = fingerprint;
+ ruleSelectedBundleNames = newSelection;
+ }
+ finally {
+ entryLock.unlock();
+ }
+
+ // Phase 3b: close the narrow race with concurrent registerBundle
+ reconcileLinkedBundlesWithSelection(loc, uncommittedHolds);
+ }
+ finally {
+ // Phase 4: release outside entryLock. `uncommittedHolds` contains every hold this call acquired but did not
+ // successfully transfer into a field (race-lose, mid-acquire failure, or commit failure). `displacedHolds`
+ // contains pre-existing holds that the diff removed from ruleBundleHolds after successful commit.
+ releaseHolds(uncommittedHolds);
+ releaseHolds(displacedHolds);
+ }
+ }
+
+ /**
+ * Post-{@link #applyRule} reconciliation: scan {@link #linkedBundles} against the freshly-committed
+ * {@link #ruleSelectedBundleNames} and install a rule-hold for any selected bundle that's present in
+ * {@code linkedBundles} but missing from {@link #ruleBundleHolds}.
+ *
+ * Closes a race that can occur when bundle mounts triggered by concurrent queries run on the storage-loading pool
+ * without the per-segment lock, and can call {@link #registerBundle} concurrently with
+ * {@link #applyRule}. If the bundle lands in {@link #linkedBundles} after {@link #applyRule}'s Phase 1
+ * snapshot but its {@code registerBundle} Phase 1 runs before Phase 3 has committed {@code newSelection},
+ * {@code registerBundle} observes the OLD selection and bails without installing a hold, and Phase 1's
+ * {@code namesToAcquire} also missed it (it wasn't in {@code linkedBundles} at snapshot time). Without this
+ * reconciliation, the entry would incorrectly be evictable despite the rule.
+ */
+ private void reconcileLinkedBundlesWithSelection(
+ StorageLocation loc,
+ List> uncommittedHolds
+ )
+ {
+ final List reconcileNames;
+ entryLock.lock();
+ try {
+ reconcileNames = new ArrayList<>();
+ for (String name : ruleSelectedBundleNames) {
+ if (!ruleBundleHolds.containsKey(name) && findLinkedBundleByName(name) != null) {
+ reconcileNames.add(name);
+ }
+ }
+ }
+ finally {
+ entryLock.unlock();
+ }
+ if (reconcileNames.isEmpty()) {
+ return;
+ }
+
+ final Map> acquired = new HashMap<>();
+ for (String name : reconcileNames) {
+ final StorageLocation.ReservationHold h =
+ loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, name));
+ if (h != null) {
+ acquired.put(name, h);
+ uncommittedHolds.add(h);
+ }
+ }
+
+ entryLock.lock();
+ try {
+ for (var e : acquired.entrySet()) {
+ if (ruleSelectedBundleNames.contains(e.getKey()) && !ruleBundleHolds.containsKey(e.getKey())) {
+ ruleBundleHolds.put(e.getKey(), e.getValue());
+ uncommittedHolds.remove(e.getValue());
+ }
+ }
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Release the partial-load rule applied to this entry. Closes every {@link StorageLocation.ReservationHold} taken by
+ * {@link #applyRule(String, Set)} (the metadata self-hold and every bundle hold) and clears the rule state. Bundle
+ * entries that were kept resident only by the rule become eviction eligible; the metadata entry itself becomes
+ * eviction eligible as well (unless a concurrent query holds a reference). Safe to call when no rule is applied, no-op.
+ *
+ * Concurrency contract. Must be serialized against {@link #applyRule} for a given segment id via an external
+ * per-segment lock (see {@code applyRule}'s Concurrency contract). Hold releases run OUTSIDE {@link #entryLock}, so
+ * a concurrent {@code applyRule}/{@code clearRule} on the same entry can observe intermediate state.
+ *
+ * Snapshot-under-lock, release-outside-lock. Closing a {@link StorageLocation.ReservationHold} eventually acquires
+ * the location's writeLock (via {@code createWeakEntryReleaseRunnable}), so it must not run under {@link #entryLock}
+ * lest it invert the writeLock → entryLock ordering the storage layer already relies on.
+ */
+ public void clearRule()
+ {
+ final List> toClose = new ArrayList<>();
+ entryLock.lock();
+ try {
+ toClose.addAll(ruleBundleHolds.values());
+ ruleBundleHolds.clear();
+ if (metadataSelfHold != null) {
+ toClose.add(metadataSelfHold);
+ metadataSelfHold = null;
+ }
+ ruleFingerprint = null;
+ ruleSelectedBundleNames = Set.of();
+ }
+ finally {
+ entryLock.unlock();
+ }
+ releaseHolds(toClose);
+ }
+
+ private void releaseHolds(List> holds)
+ {
+ for (StorageLocation.ReservationHold h : holds) {
+ CloseableUtils.closeAndSuppressExceptions(
+ h,
+ t -> LOG.warn(t, "Failed releasing rule-hold on partial segment[%s]", segmentId)
+ );
+ }
+ }
+
+ @Nullable
+ private PartialSegmentBundleCacheEntry findLinkedBundleByName(String bundleName)
+ {
+ return linkedBundles.get(bundleName);
+ }
+
+ /**
+ * Whether every container belonging to {@code bundleName} has been downloaded.
+ *
+ * Returns {@code false} when the entry is not mounted (no file mapper yet).
+ */
+ public boolean isBundleFullyDownloaded(String bundleName)
+ {
+ final PartialSegmentFileMapperV10 mapper = getFileMapper();
+ return mapper != null && mapper.isBundleFullyDownloaded(bundleName);
+ }
+
+ /**
+ * Whether {@code bundleName} is currently held by an applied partial-load rule on this entry (i.e. present in
+ * {@link #ruleBundleHolds}).
+ *
+ * NOT the same as "is the bundle linked": a bundle can be linked ({@link #linkedBundles}) without being rule-held,
+ * because {@link #linkedBundles} is a lock-free {@link ConcurrentHashMap} keyset that
+ * {@link #registerBundle}/{@link #unregisterBundle} mutate outside {@link #entryLock} (to preserve lock ordering
+ * against the storage layer's writeLock → entryLock chain).
+ */
+ public boolean isBundleRuleHeld(String bundleName)
+ {
+ entryLock.lock();
+ try {
+ return ruleBundleHolds.containsKey(bundleName);
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
+ /**
+ * Mount and eagerly download the named bundle. Drives an acquire through the internal
+ * {@link #bundleAcquirer} (which mounts the bundle if needed, causing {@link #registerBundle} to fire), downloads
+ * every container in the bundle via {@link PartialSegmentFileMapperV10#ensureBundleDownloaded}, then releases the
+ * transient acquire hold. When a partial-load rule selects this bundle, the metadata's own rule-hold (installed by
+ * {@link #registerBundle} on the mount that this call drives) keeps the bundle resident after the transient hold
+ * releases. When no rule selects this bundle, the transient hold's release leaves the bundle eviction eligible; the
+ * caller should already have driven {@link #applyRule} first to install the rule state before calling this.
+ */
+ public void ensureBundleResidentForRule(String bundleName) throws IOException
+ {
+ entryLock.lock();
+ try {
+ if (location == null || fileMapper == null) {
+ throw DruidException.defensive(
+ "ensureBundleResidentForRule on unmounted partial metadata entry[%s]",
+ id
+ );
+ }
+ }
+ finally {
+ entryLock.unlock();
+ }
+ final Closeable acquired = bundleAcquirer.acquire(bundleName);
+ try {
+ final PartialSegmentFileMapperV10 mapper = getFileMapper();
+ if (mapper == null) {
+ throw DruidException.defensive(
+ "Partial metadata entry[%s] lost its file mapper mid-acquire for bundle[%s]",
+ id,
+ bundleName
+ );
+ }
+ mapper.ensureBundleDownloaded(bundleName);
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(acquired, t -> LOG.warn(
+ t,
+ "Failed to release transient bundle-acquire hold for bundle[%s] on segment[%s]",
+ bundleName,
+ segmentId
+ ));
+ }
+ }
+
@Override
public long getSize()
{
@@ -273,9 +682,12 @@ public void mount(StorageLocation mountLocation) throws IOException
ours.set(null);
}
catch (Throwable t) {
- // clear the future so the next caller gets a fresh attempt
+ // Clear the future so the next caller gets a fresh attempt. Signal awaiters with the exception BEFORE firing
+ // the onUnmount hook so any thread observing the failure via awaitMount(future) sees the exception before
+ // observing the hook's file-system side effect (info-file deletion).
mountFuture.set(null);
ours.setException(t);
+ runOnUnmountHookOnce();
Throwables.propagateIfInstanceOf(t, IOException.class);
Throwables.propagateIfPossible(t);
throw DruidException.defensive(t, "Failed to mount metadata entry[%s]", id);
@@ -463,6 +875,17 @@ private static void awaitMount(SettableFuture future) throws IOException
* outstanding, the actual unmap-and-delete work is deferred until the last reference releases; in that case this
* method returns immediately and {@link #doActualUnmount} will fire later on the thread that closes the last
* reference. With no outstanding references, cleanup runs synchronously on the caller's thread.
+ *
+ * If this entry was reserved but never mounted, there is no reference-counted gate to close; instead run the
+ * {@link #setOnUnmount onUnmount} hook directly so external cleanup (e.g. info-file deletion) still fires.
+ *
+ * Mid-mount race guard. A concurrent reclaim can invoke {@code unmount()} while {@link #doMount} is still
+ * running its slow deep-storage header fetch and has not yet installed {@link #references}. Firing the hook here
+ * would race with {@code doMount}. Detect this via {@link #mountFuture} (set by {@link #mount} before
+ * {@code doMount} runs; cleared by {@code doMount}'s failure path and by {@link #doActualUnmount}) and defer
+ * hook-firing to the mount path: on success {@code verifyStillReservedOrRollback} calls {@code unmount()} again
+ * with {@link #references} installed and the hook fires via {@link #doActualUnmount}; on failure {@link #mount}'s
+ * catch block fires the hook directly (see below).
*/
@Override
public void unmount()
@@ -470,6 +893,30 @@ public void unmount()
final ReferenceCountingCloseableObject current = references.get();
if (current != null && !current.isClosed()) {
current.close();
+ return;
+ }
+ if (mountFuture.get() != null) {
+ // doMount is in flight and has not yet installed `references`; hook-firing is deferred to the mount path.
+ return;
+ }
+ // Never-mounted (or already-completed-cleanup) release: doActualUnmount won't run, so fire the hook here.
+ runOnUnmountHookOnce();
+ }
+
+ /**
+ * Atomically extract and run the {@link #setOnUnmount onUnmount} hook if one is set. Idempotent across concurrent
+ * callers, the {@code getAndSet(null)} ensures exactly one caller ever observes a non-null hook.
+ */
+ private void runOnUnmountHookOnce()
+ {
+ final Runnable hook = onUnmount.getAndSet(null);
+ if (hook != null) {
+ try {
+ hook.run();
+ }
+ catch (Throwable t) {
+ LOG.warn(t, "onUnmount hook failed for partial segment metadata entry[%s]", segmentId);
+ }
}
}
@@ -634,12 +1081,21 @@ public Optional acquireFullReference(@Nullable Closeable extraOnClose)
*/
private void doActualUnmount()
{
- final Runnable hook;
entryLock.lock();
try {
if (fileMapper == null) {
return;
}
+ // doActualUnmount must never fire while a partial-load rule is applied.
+ if (ruleFingerprint != null) {
+ throw DruidException.defensive(
+ "doActualUnmount fired for partial segment[%s] with rule state still applied (fingerprint=%s, "
+ + "selectedBundles=%s); the metadata self-hold invariant has been broken",
+ segmentId,
+ ruleFingerprint,
+ ruleSelectedBundleNames
+ );
+ }
try {
fileMapper.close();
}
@@ -654,21 +1110,13 @@ private void doActualUnmount()
// Clear the mount-dedup gate so a subsequent mount() on this same instance starts a fresh attempt.
mountFuture.set(null);
deleteHeaderFiles();
- hook = onUnmount.getAndSet(null);
}
finally {
entryLock.unlock();
}
// Run the hook outside entryLock so it can touch the file system / cache manager without contending with
// concurrent status reads, and so a slow or buggy hook can't deadlock against acquireReference paths.
- if (hook != null) {
- try {
- hook.run();
- }
- catch (Throwable t) {
- LOG.warn(t, "onUnmount hook failed for partial segment metadata entry[%s]", segmentId);
- }
- }
+ runOnUnmountHookOnce();
}
/**
@@ -804,7 +1252,7 @@ public Closeable acquire(String requestedBundleName)
() -> PartialSegmentBundleCacheEntry.forBundle(
PartialSegmentMetadataCacheEntry.this,
bundleName,
- inferParentBundles(bundleName)
+ inferBundleDependencies(bundleName)
)
);
if (hold == null) {
@@ -861,13 +1309,13 @@ private void mountBundleOrClose(
// holds + references on each parent and fails with a defensive error if a parent isn't registered+mounted at
// the location. The bootstrap restore path orders base-before-dependents; this is the equivalent ordering for
// the runtime acquire path, which would otherwise reach a projection bundle directly with no __base mounted.
- // The bundle keeps its own parent holds/refs for its lifetime, so we hold these transient ones only across the
- // mount and release them immediately after a successful mount (acquire() is acyclic: base/root have no
- // parents, so the recursion terminates).
+ // The bundle keeps its own dependency holds/refs for its lifetime, so we hold these transient ones only across
+ // the mount and release them immediately after a successful mount (acquire() is acyclic: base/root have no
+ // dependencies, so the recursion terminates).
final Closer parentHolds = Closer.create();
try {
- for (PartialSegmentBundleCacheEntryIdentifier parentId : inferParentBundles(bundleName)) {
- parentHolds.register(acquire(parentId.bundleName()));
+ for (PartialSegmentBundleCacheEntryIdentifier depId : inferBundleDependencies(bundleName)) {
+ parentHolds.register(acquire(depId.bundleName()));
}
bundle.mount(loc);
}
@@ -885,17 +1333,19 @@ private void mountBundleOrClose(
}
/**
- * Inference of the parent bundles that the given {@code bundleName} depends on within this segment.
+ * The bundles that {@code bundleName} depends on within this segment. Depending here means: at mount time, this
+ * bundle's file mapper must be able to acquire a hold on each dependency, and the dependency's containers must
+ * remain resident for as long as this bundle is mounted.
*
* The rule is uniform: the base bundle and the {@link SegmentFileBuilder#ROOT_BUNDLE_NAME root bundle} have no
- * parents (the root bundle owns everything written without an explicit {@code startFileBundle} call, for older
- * fileGroup-less segments, or any future shared internal metadata and is structurally a peer of the base);
+ * dependencies (the root bundle owns everything written without an explicit {@code startFileBundle} call, for
+ * older fileGroup-less segments, or any future shared internal metadata and is structurally a peer of the base);
* every other bundle depends on the base bundle, but only if this segment actually carries one.
*
* If future writers introduce richer dependency graphs, the rule will need to grow, likely by reading dependency
* metadata the writer records explicitly.
*/
- public List inferParentBundles(String bundleName)
+ public List inferBundleDependencies(String bundleName)
{
if (Projections.BASE_TABLE_PROJECTION_NAME.equals(bundleName)
|| SegmentFileBuilder.ROOT_BUNDLE_NAME.equals(bundleName)
@@ -910,6 +1360,34 @@ public List inferParentBundles(String
);
}
+ /**
+ * Return {@code roots} plus every transitive dependency (as inferred by {@link #inferBundleDependencies}) in
+ * dependency-first order.
+ *
+ * No cycle guard: {@link #inferBundleDependencies} only returns {@code []} or {@code [__base]}, and {@code __base}
+ * itself has no dependencies, so the graph is at most one level deep by construction. A richer future dependency
+ * graph should revisit this.
+ */
+ public List bundlesInMountOrder(Iterable roots)
+ {
+ final LinkedHashSet ordered = new LinkedHashSet<>();
+ for (String root : roots) {
+ visitDependenciesFirst(root, ordered);
+ }
+ return List.copyOf(ordered);
+ }
+
+ private void visitDependenciesFirst(String bundleName, LinkedHashSet ordered)
+ {
+ if (ordered.contains(bundleName)) {
+ return;
+ }
+ for (PartialSegmentBundleCacheEntryIdentifier dep : inferBundleDependencies(bundleName)) {
+ visitDependenciesFirst(dep.bundleName(), ordered);
+ }
+ ordered.add(bundleName);
+ }
+
/**
* Whether this segment carries a {@code __base} bundle (shared base-table column data). Probed from the mounted file
* mapper's actual bundle set; returns false when the entry is not mounted.
@@ -925,19 +1403,71 @@ private boolean hasBaseBundle()
* Register a bundle entry as a current dependent of this metadata entry. Called by
* {@link PartialSegmentBundleCacheEntry} after a successful mount; the drop path uses {@link #snapshotLinkedBundles}
* to enumerate dependents for cascade-close.
+ *
+ * If a partial-load rule is currently applied and {@code bundle}'s name is in the rule's selection set, take a
+ * {@link StorageLocation.ReservationHold} on the bundle so it stays resident until the rule is cleared or swapped
+ * away. The hold acquisition uses a snapshot-under-lock, acquire-outside-lock, commit-under-lock pattern to avoid
+ * holding {@link #entryLock} while calling into {@link StorageLocation}'s read/write locks.
*/
void registerBundle(PartialSegmentBundleCacheEntry bundle)
{
- linkedBundles.add(bundle);
+ final String bundleName = bundle.getBundleName();
+ linkedBundles.put(bundleName, bundle);
+
+ // Phase 1: snapshot under entryLock.
+ final StorageLocation loc;
+ entryLock.lock();
+ try {
+ if (location == null
+ || !ruleSelectedBundleNames.contains(bundleName)
+ || ruleBundleHolds.containsKey(bundleName)) {
+ return;
+ }
+ loc = location;
+ }
+ finally {
+ entryLock.unlock();
+ }
+
+ // Phase 2: acquire outside entryLock. Caller's mount path (via bundleAcquirer.acquire) already holds a transient
+ // ReservationHold on this bundle, so the weak entry is guaranteed present and this acquire cannot return null
+ // for a "just evicted" reason.
+ final StorageLocation.ReservationHold hold =
+ loc.addWeakReservationHoldIfExists(new PartialSegmentBundleCacheEntryIdentifier(segmentId, bundleName));
+ if (hold == null) {
+ return;
+ }
+
+ // Phase 3: commit under entryLock, or discard if the rule state moved under us.
+ boolean stored = false;
+ entryLock.lock();
+ try {
+ if (ruleSelectedBundleNames.contains(bundleName) && !ruleBundleHolds.containsKey(bundleName)) {
+ ruleBundleHolds.put(bundleName, hold);
+ stored = true;
+ }
+ }
+ finally {
+ entryLock.unlock();
+ }
+ if (!stored) {
+ CloseableUtils.closeAndSuppressExceptions(
+ hold,
+ t -> LOG.warn(t, "Failed releasing discarded rule-hold for bundle[%s] on segment[%s]", bundleName, segmentId)
+ );
+ }
}
/**
* Reverse of {@link #registerBundle}. Called by {@link PartialSegmentBundleCacheEntry#unmount} so the metadata's
- * view stays consistent with which bundles are actually mounted.
+ * view stays consistent with which bundles are actually mounted. Bookkeeping-only: never touches {@link
+ * #ruleBundleHolds}.
*/
void unregisterBundle(PartialSegmentBundleCacheEntry bundle)
{
- linkedBundles.remove(bundle);
+ // Remove only if the exact same instance is still registered under this name. A late-arriving unregister for a
+ // stale bundle instance (e.g. one that was replaced by a fresh remount) must not evict the fresh instance.
+ linkedBundles.remove(bundle.getBundleName(), bundle);
}
/**
@@ -947,7 +1477,7 @@ void unregisterBundle(PartialSegmentBundleCacheEntry bundle)
*/
public Collection snapshotLinkedBundles()
{
- return new ArrayList<>(linkedBundles);
+ return new ArrayList<>(linkedBundles.values());
}
/**
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 2c6e4505a5a8..3518cd7e7414 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
@@ -54,15 +54,21 @@
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
+import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Lock;
@@ -454,6 +460,23 @@ private void deleteSegmentInfoFile(DataSegment segment)
}
}
+ /**
+ * Write the info file for a partial-load segment, overwriting any existing content atomically. Distinct from
+ * {@link #storeInfoFile} which skips the write when the file already exists, for partial segments we must
+ * unconditionally rewrite so an incoming rule swap (new {@code fingerprint}/{@code delegate} inside the
+ * wrapped load spec) reaches disk. Otherwise bootstrap after a restart would restore the segment using the
+ * prior wrapper and re-announce the old rule until the coordinator resyncs.
+ */
+ private void writePartialInfoFile(DataSegment segment) throws IOException
+ {
+ final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), segment.getId().toString());
+ FileUtils.mkdirp(getEffectiveInfoDir());
+ FileUtils.writeAtomically(segmentInfoCacheFile, out -> {
+ jsonMapper.writeValue(out, segment);
+ return null;
+ });
+ }
+
@Override
public Optional acquireCachedSegment(final SegmentId segmentId, final AcquireMode acquireMode)
{
@@ -746,6 +769,48 @@ private ReservedPartial findExistingPartialWithHold(SegmentId segmentId)
return null;
}
+ /**
+ * If any storage location has a non-{@link PartialSegmentMetadataCacheEntry} cache entry at {@code segmentId},
+ * attempt to evict it via {@link StorageLocation#removeUnheldWeakEntry}. This handles the config-flip scenario
+ * where a segment with a {@link PartialLoadSpec} wrapper was previously queried while
+ * {@code virtualStoragePartialDownloadsEnabled=false} ({@link #acquireSegment} created a
+ * {@link CompleteSegmentCacheEntry} at this id, and a subsequent partial-load rule cannot reserve here without
+ * evicting the stale entry first).
+ *
+ * {@link StorageLocation#removeUnheldWeakEntry} is a no-op when the entry is held (in-flight query), so we
+ * detect via {@link StorageLocation#getCacheEntry} whether eviction actually happened. If not, throw a retryable
+ * {@link SegmentLoadingException}: the coordinator's load queue retries on the next sync, and by then the query
+ * should have released the hold and eviction will succeed.
+ */
+ private void evictStaleNonPartialWeakEntry(SegmentId segmentId) throws SegmentLoadingException
+ {
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId);
+ for (StorageLocation location : locations) {
+ final CacheEntry entry = location.getCacheEntry(id);
+ if (entry == null || entry instanceof PartialSegmentMetadataCacheEntry) {
+ continue;
+ }
+ location.removeUnheldWeakEntry(id);
+ final CacheEntry stillThere = location.getCacheEntry(id);
+ if (stillThere != null) {
+ throw new SegmentLoadingException(
+ "Stale non-partial cache entry[%s] at id[%s] on location[%s] blocks partial-load reservation and is "
+ + "currently held; the coordinator's load queue will retry",
+ stillThere.getClass().getSimpleName(),
+ id,
+ location.getPath()
+ );
+ }
+ log.info(
+ "Evicted stale non-partial cache entry[%s] at location[%s] to make room for partial-load rule on "
+ + "segment[%s]",
+ entry.getClass().getSimpleName(),
+ location.getPath(),
+ segmentId
+ );
+ }
+ }
+
/**
* Try to open a range reader for the given segment's {@link LoadSpec}; returns {@code null} when the backend
* doesn't support range reads (e.g. zipped storage), or when partial downloads are disabled via
@@ -790,7 +855,14 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead
FileUtils.mkdirp(partialDir);
}
catch (IOException e) {
- throw DruidException.defensive(e, "Failed to create partial cache dir for segment[%s]", dataSegment.getId());
+ // Location is unwritable, fall through to next location rather than failing the whole reservation
+ log.warn(
+ e,
+ "Failed to create partial cache dir on location[%s] for segment[%s]; trying next location",
+ location.getPath(),
+ dataSegment.getId()
+ );
+ continue;
}
final StorageLocation.ReservationHold hold = location.addWeakReservationHold(
id,
@@ -806,6 +878,7 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead
)
);
if (hold == null) {
+ atomicMoveAndDeleteCacheEntryDirectory(partialDir);
continue;
}
try {
@@ -817,19 +890,18 @@ private ReservedPartial reservePartial(DataSegment dataSegment, SegmentRangeRead
location.getPath()
);
}
- final File segmentInfoCacheFile = new File(getEffectiveInfoDir(), dataSegment.getId().toString());
- if (!segmentInfoCacheFile.exists()) {
- FileUtils.mkdirp(getEffectiveInfoDir());
- FileUtils.writeAtomically(segmentInfoCacheFile, out -> {
- jsonMapper.writeValue(out, dataSegment);
- return null;
- });
- }
+ writePartialInfoFile(dataSegment);
partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
return new ReservedPartial(partial, location, hold);
}
catch (Throwable t) {
- throw CloseableUtils.closeAndWrapInCatch(t, hold);
+ // Close the hold (removing the never-mounted weak entry), then nuke the on-disk dir.
+ try {
+ throw CloseableUtils.closeAndWrapInCatch(t, hold);
+ }
+ finally {
+ atomicMoveAndDeleteCacheEntryDirectory(partialDir);
+ }
}
}
throw DruidException.forPersona(DruidException.Persona.USER)
@@ -872,6 +944,303 @@ private record ReservedPartial(
{
}
+ /**
+ * Apply (or reconcile) a partial-load rule for a segment coming through {@link #load}. The rule state lives on the
+ * {@link PartialSegmentMetadataCacheEntry} itself; this method just drives it:
+ *
+ * - Materialize the wrapper's {@link PartialLoadSpec} and open a range reader over the backend. If the backend
+ * can't do range reads, clear any prior rule on the segment and return; the segment falls back to the
+ * weak-full-load path at query time.
+ * - Acquire (or reserve) the metadata entry via the shared on-demand path ({@link #findOrReservePartial}),
+ * take a transient eviction-protective hold across the rest of the flow.
+ * - Mount the metadata entry (idempotent).
+ * - Compute the rule's selected bundle names from the wrapper + parsed segment metadata, and call
+ * {@link PartialSegmentMetadataCacheEntry#applyRule} which installs the metadata self-hold and takes
+ * holds on any already-registered selected bundles.
+ * - Kick off async eager downloads on the loading pool for any selected bundle not yet linked. Each async task
+ * drives {@link PartialSegmentMetadataCacheEntry#ensureBundleResidentForRule}: its bundle mount will call
+ * {@link PartialSegmentMetadataCacheEntry#registerBundle}, which acquires the rule-hold on that bundle.
+ * - Release the transient hold; the metadata's self-hold keeps the entry resident under the applied rule.
+ *
+ */
+ private void loadPartial(DataSegment dataSegment) throws SegmentLoadingException
+ {
+ final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment);
+ final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment, wrapper);
+ final ReferenceCountingLock lock = lock(dataSegment);
+ synchronized (lock) {
+ try {
+ // If a stale non-partial cache entry sits at this segment id (a CompleteSegmentCacheEntry created by a prior
+ // acquireSegment while virtualStoragePartialDownloadsEnabled=false, for example), evict it before any
+ // partial-entry lookup or reservation; otherwise findExistingPartialWithHold's defensive type-check would
+ // throw, and reservePartial's addWeakReservationHold would land on the incompatible entry. If the stale
+ // entry is currently held (in-flight query), this throws a retryable SegmentLoadingException; the
+ // coordinator's load queue retries on next sync, and by then the query should have released.
+ evictStaleNonPartialWeakEntry(dataSegment.getId());
+
+ if (rangeReader == null) {
+ // Backend doesn't support range reads (e.g. zipped deep storage). The rule can't be honored as a partial
+ // load; clear any prior rule so the segment falls through to the ordinary weak-full-load path at query
+ // time.
+ final ReservedPartial existing = findExistingPartialWithHold(dataSegment.getId());
+ if (existing != null) {
+ try {
+ existing.metadata().clearRule();
+ log.warn(
+ "Backend for segment[%s] does not support range reads; released rule[fingerprint=%s], segment "
+ + "will fall back to weak full-load at query time",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(
+ existing.hold(),
+ t -> log.warn(t, "Failed to release transient hold on partial metadata for segment[%s]",
+ dataSegment.getId())
+ );
+ }
+ } else {
+ // No prior entry AND range-reader is null on this fresh load: the coordinator asked for a partial rule
+ // this historical can't honor. Log so operators can debug "why isn't my rule applied on historical X".
+ log.warn(
+ "Backend for segment[%s] does not support range reads and no prior partial entry exists; rule"
+ + "[fingerprint=%s] cannot be applied on this historical",
+ dataSegment.getId(),
+ wrapper.getFingerprint()
+ );
+ }
+ return;
+ }
+
+ final ReservedPartial reserved = findOrReservePartial(dataSegment, rangeReader);
+ try {
+ final PartialSegmentMetadataCacheEntry metadata = reserved.metadata();
+ // findOrReservePartial only invokes reservePartial (which writes the info file) on the fresh-reserve
+ // branch. On the find-existing branch the info file on disk still carries the PRIOR rule's wrapped
+ // load spec, so a rule swap here would apply in memory only. Rewrite unconditionally before mount.
+ try {
+ writePartialInfoFile(dataSegment);
+ }
+ catch (IOException e) {
+ throw new SegmentLoadingException(
+ e,
+ "Failed to write partial info file for segment[%s]",
+ dataSegment.getId()
+ );
+ }
+ try {
+ metadata.mount(reserved.location());
+ }
+ catch (IOException e) {
+ throw new SegmentLoadingException(
+ e,
+ "Failed to mount partial metadata for segment[%s]",
+ dataSegment.getId()
+ );
+ }
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ if (mapper == null) {
+ throw DruidException.defensive(
+ "Partial metadata for segment[%s] mounted without a file mapper",
+ dataSegment.getId()
+ );
+ }
+ // Register the info-file cleanup hook BEFORE anything that could throw. Any throw between mount and
+ // 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 selected = Set.copyOf(
+ wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata())
+ );
+ final String priorFingerprint = metadata.getRuleFingerprint();
+ metadata.applyRule(wrapper.getFingerprint(), selected);
+ if (priorFingerprint != null && !priorFingerprint.equals(wrapper.getFingerprint())) {
+ log.info(
+ "Reconciled partial-load rule for segment[%s]: fingerprint transitioned [%s] → [%s]",
+ dataSegment.getId(),
+ priorFingerprint,
+ wrapper.getFingerprint()
+ );
+ }
+ // Block until every eager download completes so the announcement fingerprint reflects reality: any failure
+ // clears the rule state (releasing self-hold + all bundle rule-holds) and propagates as a load failure so
+ // the coordinator's load queue can retry on its next sync. The announced fingerprint == "rule fully
+ // realized" contract stays intact.
+ awaitEagerDownloadsOrClearRule(dataSegment, metadata, selected);
+ }
+ finally {
+ CloseableUtils.closeAndSuppressExceptions(
+ reserved.hold(),
+ t -> log.warn(t, "Failed to release transient hold on partial metadata for segment[%s]",
+ dataSegment.getId())
+ );
+ }
+ }
+ finally {
+ unlock(dataSegment, lock);
+ }
+ }
+ }
+
+ /**
+ * Submit eager-download tasks to the loading pool for every rule-selected bundle not yet registered with
+ * {@code metadata}, then block until every task completes. On any failure the rule state is cleared
+ * and a {@link SegmentLoadingException} is thrown so the caller treats the load as failed and retries.
+ */
+ private void awaitEagerDownloadsOrClearRule(
+ DataSegment dataSegment,
+ PartialSegmentMetadataCacheEntry metadata,
+ Set selected
+ ) throws SegmentLoadingException
+ {
+ final List> pending = new ArrayList<>();
+ Throwable firstFailure = null;
+ try {
+ for (String bundleName : metadata.bundlesInMountOrder(selected)) {
+ // Skip only when the bundle is BOTH rule-held AND fully downloaded (every container's files present on disk).
+ if (metadata.isBundleRuleHeld(bundleName) && metadata.isBundleFullyDownloaded(bundleName)) {
+ continue;
+ }
+ pending.add(virtualStorageLoadingThreadPool.getExecutorService().submit(() -> {
+ metadata.ensureBundleResidentForRule(bundleName);
+ return null;
+ }));
+ }
+ }
+ catch (RejectedExecutionException e) {
+ // Pool shutting down or queue saturated. Any already-submitted futures may still run; track the rejection as
+ // firstFailure and fall through to the await/cancel + clearRule path so state is always cleaned up before we
+ // throw.
+ firstFailure = e;
+ }
+
+ for (Future f : pending) {
+ if (firstFailure != null) {
+ // Cancel remaining futures on any failure so we don't wait for tasks whose result we no longer intend to
+ // commit. cancel(false) — NOT cancel(true) — because a running task is mid-NIO/FileChannel read via
+ // mapper.ensureBundleDownloaded, and Thread.interrupt() on an in-flight NIO op raises
+ // ClosedByInterruptException that closes the mapper's shared underlying FD, breaking still-mounted bundles
+ // for other queries reading the same segment. Not-yet-started tasks are marked CANCELLED; running tasks
+ // continue on their own timeline, observe ruleSelectedBundleNames == {} after clearRule, and release
+ // their transient bundleAcquirer holds without acquiring rule-holds.
+ f.cancel(false);
+ // Still drain the future so an ExecutionException from a task that completed-with-failure BEFORE we
+ // canceled is logged instead of silently swallowed.
+ try {
+ f.get();
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ catch (ExecutionException e) {
+ final Throwable cause = e.getCause() != null ? e.getCause() : e;
+ log.warn(
+ cause,
+ "Eager download of a rule-selected bundle for segment[%s] failed after the primary failure",
+ dataSegment.getId()
+ );
+ }
+ catch (CancellationException e) {
+ // task never started; expected on the cancel path
+ }
+ continue;
+ }
+ try {
+ f.get();
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ firstFailure = e;
+ // Signal cancel WITHOUT interrupt — see cancel(false) comment above; interrupting mid-NIO closes the
+ // shared mapper FD.
+ f.cancel(false);
+ }
+ catch (ExecutionException e) {
+ firstFailure = e.getCause() != null ? e.getCause() : e;
+ }
+ catch (CancellationException e) {
+ // Defensive: shouldn't reach here on the first-failure branch (we only cancel after firstFailure is set),
+ // but if it does, treat it uniformly.
+ firstFailure = e;
+ }
+ }
+
+ if (firstFailure != null) {
+ // Any late-completing pool task that still succeeds after we've cleared the rule will call registerBundle →
+ // observing ruleSelectedBundleNames == {} and skipping the rule-hold acquire.
+ metadata.clearRule();
+ throw new SegmentLoadingException(
+ firstFailure,
+ "Failed eager download of rule-selected bundles for segment[%s]; cleared partial-load rule",
+ dataSegment.getId()
+ );
+ }
+ }
+
+ /**
+ * Materialize the segment's wrapped load spec to a {@link PartialLoadSpec}.
+ */
+ private PartialLoadSpec materializePartialLoadSpec(DataSegment dataSegment) throws SegmentLoadingException
+ {
+ final LoadSpec materializedLoadSpec;
+ try {
+ materializedLoadSpec = jsonMapper.convertValue(dataSegment.getLoadSpec(), LoadSpec.class);
+ }
+ catch (Exception e) {
+ throw new SegmentLoadingException(
+ e,
+ "Failed to materialize partial load spec for segment[%s]",
+ dataSegment.getId()
+ );
+ }
+ if (!(materializedLoadSpec instanceof PartialLoadSpec wrapper)) {
+ throw DruidException.defensive(
+ "Segment[%s] load spec was detected as partial but materialized to non-partial type[%s]",
+ dataSegment.getId(),
+ materializedLoadSpec.getClass().getSimpleName()
+ );
+ }
+ return wrapper;
+ }
+
+ /**
+ * Open a range reader for the segment's deep storage via the wrapper's {@link PartialLoadSpec#openRangeReader}
+ * (which delegates to the inner load spec), wrapping any {@link IOException} as a {@link SegmentLoadingException}.
+ * Returns {@code null} if the backend doesn't support range reads.
+ */
+ @Nullable
+ private SegmentRangeReader openPartialRangeReader(DataSegment dataSegment, PartialLoadSpec wrapper)
+ throws SegmentLoadingException
+ {
+ try {
+ return wrapper.openRangeReader();
+ }
+ catch (IOException e) {
+ throw new SegmentLoadingException(e, "Failed to open range reader for segment[%s]", dataSegment.getId());
+ }
+ }
+
+ /**
+ * The fingerprint of the currently-applied partial-load rule for {@code segmentId}, or {@code null} if no rule is
+ * currently applied (or the segment has no partial metadata entry (e.g. it was eager-loaded as a complete cache
+ * entry). Test-only. Delegates to {@link PartialSegmentMetadataCacheEntry#getRuleFingerprint} on the segment's
+ * cache entry.
+ */
+ @VisibleForTesting
+ @Nullable
+ String getRuleFingerprintForSegment(SegmentId segmentId)
+ {
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segmentId);
+ for (StorageLocation location : locations) {
+ final CacheEntry entry = location.getCacheEntry(id);
+ if (entry instanceof PartialSegmentMetadataCacheEntry partial) {
+ return partial.getRuleFingerprint();
+ }
+ }
+ return null;
+ }
+
@Nullable
private AcquireSegmentAction acquireExistingSegment(SegmentCacheEntryIdentifier identifier)
{
@@ -927,6 +1296,15 @@ public void load(final DataSegment dataSegment) throws SegmentLoadingException
"load() should not be called when virtualStorageIsEphemeral is true"
);
}
+ // Partial-load-rule routing: a wrapped load spec carrying a fingerprint + per-spec selection (cluster groups,
+ // projections, ...) means the coordinator wants this segment loaded with rule holds pinning the selected
+ // bundles + metadata. Install the holds via loadPartial; other segments (no wrapper, or partials disabled)
+ // take the existing weak/no-op path below.
+ if (config.isVirtualStoragePartialDownloadsEnabled()
+ && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) {
+ loadPartial(dataSegment);
+ return;
+ }
// virtual storage doesn't do anything with loading immediately, but check to see if the segment is already cached
// and if so, clear out the onUnmount action
final ReferenceCountingLock lock = lock(dataSegment);
@@ -999,6 +1377,17 @@ public void bootstrap(
dataSegment.getId()
);
}
+ // If the persisted DataSegment's loadSpec is a partial-load-rule wrapper, reapply the rule so the
+ // segment continues to be protected from eviction across the restart.
+ //
+ // Strict failure handling: reapplyRuleFromInfoFile throws on eager-download failure (and clears the
+ // rule state before throwing). The exception propagates up through SegmentManager.loadSegmentOnBootstrap,
+ // which calls cacheManager.drop and rethrows, so SegmentCacheBootstrapper marks the segment as failed
+ // and skips its announcement.
+ if (config.isVirtualStoragePartialDownloadsEnabled()
+ && PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) {
+ reapplyRuleFromInfoFile(dataSegment, partial);
+ }
} else {
throw DruidException.defensive(
"Unexpected cache entry type[%s] for segment[%s] during bootstrap",
@@ -1066,15 +1455,64 @@ public File getSegmentFiles(final DataSegment segment)
@Override
public void drop(final DataSegment segment)
{
- final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId());
- for (StorageLocation location : locations) {
- final CacheEntry entry = location.getCacheEntry(id);
- if (entry != null) {
- location.release(entry);
+ final ReferenceCountingLock lock = lock(segment);
+ synchronized (lock) {
+ try {
+
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(segment.getId());
+ for (StorageLocation location : locations) {
+ final CacheEntry entry = location.getCacheEntry(id);
+ if (entry instanceof PartialSegmentMetadataCacheEntry partial) {
+ // if the segment's cache entry is a partial metadata entry with an applied rule, clear it. Reclaim of
+ // on-disk state is left to eviction.
+ partial.clearRule();
+ }
+ if (entry != null) {
+ // static entries are removed from both cache and disk immediately (no-op for weak held entries)
+ location.release(entry);
+ }
+ }
+ }
+ finally {
+ unlock(segment, lock);
}
}
}
+ /**
+ * Reapply the persisted partial-load rule to a bootstrap-restored metadata entry. Reads the wrapper from the
+ * segment's info-file {@code loadSpec}, resolves the selected bundle names against the just-parsed on-disk
+ * metadata header, calls {@link PartialSegmentMetadataCacheEntry#applyRule}, then drives eager downloads for any
+ * selected bundle that wasn't restored from disk (via {@link #awaitEagerDownloadsOrClearRule}). On failure the
+ * exception marks the segment as failed → doesn't announce it → the coordinator's next sync re-issues load.
+ */
+ private void reapplyRuleFromInfoFile(DataSegment dataSegment, PartialSegmentMetadataCacheEntry partial)
+ throws SegmentLoadingException
+ {
+ // evict any stale non-partial cache entry at this segment id before touching the partial state.
+ // Bootstrap ordering today produces a partial entry for the passed-in `partial`, so this is expected
+ // to be a no-op except during a config switch.
+ evictStaleNonPartialWeakEntry(dataSegment.getId());
+
+ final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment);
+ final PartialSegmentFileMapperV10 mapper = partial.getFileMapper();
+ if (mapper == null) {
+ throw DruidException.defensive(
+ "Bootstrap-restored partial metadata for segment[%s] has no file mapper", dataSegment.getId()
+ );
+ }
+ // Register the info-file cleanup hook BEFORE anything that could throw. reserveFromDisk does NOT install an
+ // onUnmount hook on bootstrap-restored entries, so without this ordering a throw from getSelectedBundleNames
+ // or applyRule would leave the entry with no cleanup path. Setting the hook first ensures every teardown path
+ // deletes the info file consistently.
+ partial.setOnUnmount(() -> deleteSegmentInfoFile(dataSegment));
+ final Set selected = Set.copyOf(
+ wrapper.getSelectedBundleNames(dataSegment, mapper.getSegmentFileMetadata())
+ );
+ partial.applyRule(wrapper.getFingerprint(), selected);
+ awaitEagerDownloadsOrClearRule(dataSegment, partial, selected);
+ }
+
@Override
public void shutdownBootstrap()
{
diff --git a/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java b/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java
index 460a2fa9d2c6..49fc135a9a8e 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java
+++ b/server/src/main/java/org/apache/druid/segment/loading/StorageLocation.java
@@ -540,12 +540,11 @@ public void removeUnheldWeakEntry(CacheEntryIdentifier id)
id,
(cacheEntryIdentifier, weakCacheEntry) -> {
if (weakCacheEntry.isHeld()) {
- // a holder is responsible for cleanup when it releases; leave the entry in place
return weakCacheEntry;
}
final boolean isMounted = weakCacheEntry.cacheEntry.isMounted();
unlinkWeakEntry(weakCacheEntry);
- weakCacheEntry.unmount(); // terminate the phaser; fires cacheEntry.unmount() (idempotent)
+ weakCacheEntry.unmount();
if (isMounted) {
weakStats.getAndUpdate(s -> s.evict(weakCacheEntry.cacheEntry.getSize()));
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java
index f09ca7ba1be1..5beeb750b4e9 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/PartialSegmentMetadataCacheEntryTest.java
@@ -46,6 +46,7 @@
import java.io.IOException;
import java.util.List;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
@@ -219,6 +220,142 @@ void testOnUnmountHookRunsAfterStorageLocationCleanup() throws IOException
);
}
+ @Test
+ void testOnUnmountHookRunsOnReleaseBeforeMount()
+ {
+ // if a caller sets the onUnmount hook before mount() and then releases the reservation on a mount-failure path
+ // (or without ever calling mount), unmount() must still fire the hook — otherwise external cleanup
+ // (e.g. info-file deletion) would leak.
+ final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
+ final AtomicReference hookFired = new AtomicReference<>(false);
+ entry.setOnUnmount(() -> hookFired.set(true));
+
+ entry.unmount();
+ Assertions.assertTrue(hookFired.get(), "onUnmount hook must run even when the entry was never mounted");
+
+ // Second unmount is a no-op — the hook must not double-fire.
+ hookFired.set(false);
+ entry.unmount();
+ Assertions.assertFalse(hookFired.get(), "onUnmount hook must fire exactly once across repeated unmount() calls");
+ }
+
+ @Test
+ void testBundlesInMountOrderReturnsRootsInInputOrderWhenNoBase()
+ {
+ // With no __base bundle in the mapped segment (entry not mounted), inferBundleDependencies returns [] for every
+ // input, so the walker returns roots in their original iteration order.
+ final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
+ Assertions.assertEquals(
+ List.of("a", "b", "c"),
+ entry.bundlesInMountOrder(List.of("a", "b", "c"))
+ );
+ }
+
+ @Test
+ void testBundlesInMountOrderExpandsWithBaseWhenPresent() throws IOException
+ {
+ // A segment carrying __base: dependents' expansion places __base first, then the dependent itself.
+ final PartialSegmentMetadataCacheEntry entry = mountedEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection")
+ );
+ Assertions.assertEquals(
+ List.of(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection"),
+ entry.bundlesInMountOrder(List.of("some_projection"))
+ );
+ }
+
+ @Test
+ void testBundlesInMountOrderDedupsBaseAcrossMultipleDependents() throws IOException
+ {
+ // When multiple dependents share a common dependency, mount order emits it exactly once at the front.
+ final PartialSegmentMetadataCacheEntry entry = mountedEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b")
+ );
+ Assertions.assertEquals(
+ List.of(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b"),
+ entry.bundlesInMountOrder(List.of("projection_a", "projection_b"))
+ );
+ }
+
+ @Test
+ void testRuleStateEmptyByDefault()
+ {
+ final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
+ Assertions.assertFalse(entry.isRuleHeld());
+ Assertions.assertNull(entry.getRuleFingerprint());
+ Assertions.assertEquals(Set.of(), entry.getRuleSelectedBundleNames());
+ }
+
+ @Test
+ void testApplyRuleBeforeMountThrows()
+ {
+ final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
+ Assertions.assertThrows(
+ DruidException.class,
+ () -> entry.applyRule("fp", Set.of("bundle"))
+ );
+ }
+
+ @Test
+ void testApplyRuleSetsFingerprintAndSelection() throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a")
+ );
+ entry.applyRule("fp1", Set.of("projection_a"));
+ Assertions.assertTrue(entry.isRuleHeld());
+ Assertions.assertEquals("fp1", entry.getRuleFingerprint());
+ Assertions.assertEquals(Set.of("projection_a"), entry.getRuleSelectedBundleNames());
+ }
+
+ @Test
+ void testApplyRuleSameArgsIsIdempotent() throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a")
+ );
+ entry.applyRule("fp1", Set.of("projection_a"));
+ entry.applyRule("fp1", Set.of("projection_a"));
+ Assertions.assertEquals("fp1", entry.getRuleFingerprint());
+ Assertions.assertEquals(Set.of("projection_a"), entry.getRuleSelectedBundleNames());
+ }
+
+ @Test
+ void testApplyRuleSwapUpdatesFingerprintAndSelection() throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a", "projection_b")
+ );
+ entry.applyRule("fp1", Set.of("projection_a"));
+ entry.applyRule("fp2", Set.of("projection_b"));
+ Assertions.assertEquals("fp2", entry.getRuleFingerprint());
+ Assertions.assertEquals(Set.of("projection_b"), entry.getRuleSelectedBundleNames());
+ }
+
+ @Test
+ void testClearRuleReleasesState() throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "projection_a")
+ );
+ entry.applyRule("fp1", Set.of("projection_a"));
+ Assertions.assertTrue(entry.isRuleHeld());
+ entry.clearRule();
+ Assertions.assertFalse(entry.isRuleHeld());
+ Assertions.assertNull(entry.getRuleFingerprint());
+ Assertions.assertEquals(Set.of(), entry.getRuleSelectedBundleNames());
+ }
+
+ @Test
+ void testClearRuleOnNoRuleIsNoop() throws IOException
+ {
+ final PartialSegmentMetadataCacheEntry entry = mountedWeakEntryOver(
+ buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME)
+ );
+ entry.clearRule();
+ Assertions.assertFalse(entry.isRuleHeld());
+ }
+
@Test
void testConstructorRejectsNonPositiveEstimate()
{
@@ -385,27 +522,27 @@ void testAcquireReferenceBeforeMountReturnsEmpty()
}
@Test
- void testInferParentBundlesForBaseReturnsEmpty()
+ void testInferBundleDependenciesForBaseReturnsEmpty()
{
final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
Assertions.assertEquals(
List.of(),
- entry.inferParentBundles(Projections.BASE_TABLE_PROJECTION_NAME)
+ entry.inferBundleDependencies(Projections.BASE_TABLE_PROJECTION_NAME)
);
}
@Test
- void testInferParentBundlesForRootReturnsEmpty()
+ void testInferBundleDependenciesForRootReturnsEmpty()
{
final PartialSegmentMetadataCacheEntry entry = newEntry(ESTIMATE);
Assertions.assertEquals(
List.of(),
- entry.inferParentBundles(SegmentFileBuilder.ROOT_BUNDLE_NAME)
+ entry.inferBundleDependencies(SegmentFileBuilder.ROOT_BUNDLE_NAME)
);
}
@Test
- void testInferParentBundlesDependsOnBaseWhenBaseBundlePresent() throws IOException
+ void testInferBundleDependenciesIncludesBaseWhenBaseBundlePresent() throws IOException
{
// A segment that carries a __base bundle (the non-clustered base+projection shape): every non-base/root bundle
// depends on it. Asserted uniformly for an aggregate-projection bundle and for a cluster-group bundle name (the
@@ -414,20 +551,20 @@ void testInferParentBundlesDependsOnBaseWhenBaseBundlePresent() throws IOExcepti
buildSegmentWithBundles(Projections.BASE_TABLE_PROJECTION_NAME, "some_projection")
);
for (String dependent : List.of("some_projection", Projections.getClusterGroupBundleName(List.of(0, 1)))) {
- final List parents = entry.inferParentBundles(dependent);
- Assertions.assertEquals(1, parents.size(), "expected a __base parent for bundle[" + dependent + "]");
- Assertions.assertEquals(SEGMENT_ID, parents.getFirst().segmentId());
- Assertions.assertEquals(Projections.BASE_TABLE_PROJECTION_NAME, parents.getFirst().bundleName());
+ final List deps = entry.inferBundleDependencies(dependent);
+ Assertions.assertEquals(1, deps.size(), "expected a __base dependency for bundle[" + dependent + "]");
+ Assertions.assertEquals(SEGMENT_ID, deps.getFirst().segmentId());
+ Assertions.assertEquals(Projections.BASE_TABLE_PROJECTION_NAME, deps.getFirst().bundleName());
}
}
@Test
- void testInferParentBundlesEmptyWhenSegmentHasNoBaseBundle() throws IOException
+ void testInferBundleDependenciesEmptyWhenSegmentHasNoBaseBundle() throws IOException
{
// A clustered + aggregate-projection segment with no shared columns has no __base bundle: the base data lives in
// per-group __base$ bundles and the aggregate projection is self-contained. So neither a cluster group nor
- // the aggregate projection has a parent to depend on. (Pre-shared-columns; the old "aggregate always depends on
- // __base" rule would have wrongly tried to mount a nonexistent __base for the projection bundle.)
+ // the aggregate projection has a dependency. (Pre-shared-columns; the old "aggregate always depends on __base"
+ // rule would have wrongly tried to mount a nonexistent __base for the projection bundle.)
final PartialSegmentMetadataCacheEntry entry = mountedEntryOver(
buildSegmentWithBundles(
Projections.getClusterGroupBundleName(List.of(0)),
@@ -437,9 +574,9 @@ void testInferParentBundlesEmptyWhenSegmentHasNoBaseBundle() throws IOException
);
Assertions.assertEquals(
List.of(),
- entry.inferParentBundles(Projections.getClusterGroupBundleName(List.of(0)))
+ entry.inferBundleDependencies(Projections.getClusterGroupBundleName(List.of(0)))
);
- Assertions.assertEquals(List.of(), entry.inferParentBundles("some_projection"));
+ Assertions.assertEquals(List.of(), entry.inferBundleDependencies("some_projection"));
}
private PartialSegmentMetadataCacheEntry newEntry(long estimate)
@@ -472,8 +609,8 @@ private File buildTestSegment(int numFiles) throws IOException
/**
* Build a V10 segment whose containers are tagged with exactly the given bundle names (one column file per bundle),
- * so {@link PartialSegmentMetadataCacheEntry#inferParentBundles} can be exercised against a known bundle set without
- * a full ingestion. Returns the deep-storage directory containing the V10 file.
+ * so {@link PartialSegmentMetadataCacheEntry#inferBundleDependencies} can be exercised against a known bundle set
+ * without a full ingestion. Returns the deep-storage directory containing the V10 file.
*/
private File buildSegmentWithBundles(String... bundleNames) throws IOException
{
@@ -493,7 +630,7 @@ private File buildSegmentWithBundles(String... bundleNames) throws IOException
/**
* Reserve and mount a fresh metadata entry over the segment in {@code deepStorageDir}, into a per-call cache
- * directory. The mounted entry's file mapper is what {@code inferParentBundles} probes for the base-bundle existence.
+ * directory. The mounted entry's file mapper is what {@code inferBundleDependencies} probes for base-bundle presence.
*/
private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) throws IOException
{
@@ -515,4 +652,28 @@ private PartialSegmentMetadataCacheEntry mountedEntryOver(File deepStorageDir) t
return entry;
}
+ /**
+ * Variant of {@link #mountedEntryOver} that uses {@link StorageLocation#reserveWeak} so the mounted entry is a real
+ * weak reservation — needed by the rule-holds state machine, which calls
+ * {@link StorageLocation#addWeakReservationHoldIfExists} on itself when {@code applyRule} runs.
+ */
+ private PartialSegmentMetadataCacheEntry mountedWeakEntryOver(File deepStorageDir) throws IOException
+ {
+ final File cache = new File(tempDir, "cache_" + (fixtureSeq++));
+ FileUtils.mkdirp(cache);
+ final StorageLocation location = new StorageLocation(cache, ESTIMATE * 4, null);
+ final PartialSegmentMetadataCacheEntry entry = new PartialSegmentMetadataCacheEntry(
+ SEGMENT_ID,
+ cache,
+ IndexIO.V10_FILE_NAME,
+ List.of(),
+ new DirectoryBackedRangeReader(deepStorageDir),
+ JSON_MAPPER,
+ null,
+ ESTIMATE
+ );
+ Assertions.assertTrue(location.reserveWeak(entry));
+ entry.mount(location);
+ return entry;
+ }
}
diff --git a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
index 01862eb2b9a0..7ca9d71cc78b 100644
--- a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialAcquireTest.java
@@ -492,9 +492,9 @@ void testPartialAcquireClusteredWithProjectionMountsProjectionBundleWithoutBase(
Assertions.assertEquals(CLUSTERED_SEGMENT_ID, segment.getId());
// group-by tenant + sum(x) matches the aggregate projection. Building this cursor drives the 'proj' bundle to
- // mount through the real acquire path. inferParentBundles must return no parent for it (the clustered segment
- // has no __base bundle); the old "aggregate always depends on __base" rule would have tried to mount a
- // nonexistent __base here and failed.
+ // mount through the real acquire path. inferBundleDependencies must return no dep for it (the clustered
+ // segment has no __base bundle); the old "aggregate always depends on __base" rule would have tried to mount
+ // a nonexistent __base here and failed.
final CursorBuildSpec aggSpec = CursorBuildSpec.builder()
.setGroupingColumns(List.of("tenant"))
.setAggregators(List.of(new LongSumAggregatorFactory("sum_x", "x")))
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
new file mode 100644
index 000000000000..360fdd4cde9d
--- /dev/null
+++ b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
@@ -0,0 +1,711 @@
+/*
+ * 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.loading;
+
+import com.fasterxml.jackson.databind.InjectableValues;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.jsontype.NamedType;
+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.guice.LocalDataStorageDruidModule;
+import org.apache.druid.jackson.SegmentizerModule;
+import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.emitter.EmittingLogger;
+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.IndexBuilder;
+import org.apache.druid.segment.IndexIO;
+import org.apache.druid.segment.IndexSpec;
+import org.apache.druid.segment.SegmentLazyLoadFailCallback;
+import org.apache.druid.segment.TestHelper;
+import org.apache.druid.segment.column.ColumnConfig;
+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.incremental.IncrementalIndexSchema;
+import org.apache.druid.segment.projections.Projections;
+import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMediumFactory;
+import org.apache.druid.server.metrics.NoopServiceEmitter;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.SegmentId;
+import org.apache.druid.timeline.partition.NoneShardSpec;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Manager-level tests for the partial-load-rule {@link SegmentLocalCacheManager#load} path: a segment whose
+ * {@code loadSpec} is a {@link PartialLoadSpec} wrapper drives
+ * {@link PartialSegmentMetadataCacheEntry#applyRule} on the metadata entry, which installs a self-referential
+ * {@link StorageLocation.ReservationHold} to keep the metadata resident and holds on each selected bundle already
+ * registered. Selected bundles not yet registered get eager-downloads submitted to the loading pool; {@code load}
+ * blocks until every eager download completes so the announced fingerprint reflects "rule fully realized". On any
+ * eager-download failure the rule state is cleared and a {@link SegmentLoadingException} propagates so the
+ * coordinator's load queue retries. The cache entries themselves live in {@code weakCacheEntries} exactly as they
+ * would from an on-demand acquire; the rule-holds keep them from being SIEVE-evicted for as long as the coordinator
+ * considers the segment loaded.
+ *
+ * {@code drop} calls {@link PartialSegmentMetadataCacheEntry#clearRule} — entries become unheld weak (subject to any
+ * remaining query holds) and SIEVE reclaims them naturally.
+ */
+class SegmentLocalCacheManagerPartialRuleLoadTest
+{
+ private static final SegmentId SEGMENT_ID = SegmentId.of("test", Intervals.of("2025/2026"), "v1", 0);
+ private static final DateTime TIME = DateTimes.of("2025-01-01");
+ private static final String AGG_BUNDLE = "dim1_metric1_sum";
+ private static final String OTHER_AGG_BUNDLE = "dim1_count";
+ private static final String FINGERPRINT = "v1:rule-bundle-test";
+
+ private static final RowSignature ROW_SIGNATURE = RowSignature.builder()
+ .add("dim1", ColumnType.STRING)
+ .add("metric1", ColumnType.LONG)
+ .build();
+
+ private static final List PROJECTIONS = Arrays.asList(
+ AggregateProjectionSpec.builder(AGG_BUNDLE)
+ .groupingColumns(new StringDimensionSchema("dim1"))
+ .aggregators(
+ new LongSumAggregatorFactory("_metric1_sum", "metric1"),
+ new CountAggregatorFactory("_count")
+ )
+ .build(),
+ // A second aggregate projection so tests can verify that non-selected bundles are NOT pre-mounted while
+ // selected + parent bundles ARE.
+ AggregateProjectionSpec.builder(OTHER_AGG_BUNDLE)
+ .groupingColumns(new StringDimensionSchema("dim1"))
+ .aggregators(new CountAggregatorFactory("_count"))
+ .build()
+ );
+
+ private static final List 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))
+ );
+
+ @TempDir
+ static File SHARED_TEMP_DIR;
+
+ private static File DEEP_STORAGE_DIR;
+
+ @TempDir
+ File perTestTempDir;
+
+ private ObjectMapper jsonMapper;
+ private File cacheRoot;
+ private SegmentLocalCacheManager manager;
+
+ @BeforeAll
+ static void buildSegment()
+ {
+ final File tmp = new File(SHARED_TEMP_DIR, "build_" + ThreadLocalRandom.current().nextInt());
+ DEEP_STORAGE_DIR = IndexBuilder.create()
+ .useV10()
+ .tmpDir(tmp)
+ .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();
+ EmittingLogger.registerEmitter(new NoopServiceEmitter());
+ }
+
+ @BeforeEach
+ void setup() throws IOException
+ {
+ jsonMapper = TestHelper.makeJsonMapper();
+ jsonMapper.registerSubtypes(new NamedType(LocalLoadSpec.class, "local"));
+ jsonMapper.registerSubtypes(new NamedType(PartialProjectionLoadSpec.class, PartialProjectionLoadSpec.TYPE));
+ jsonMapper.registerModule(new SegmentizerModule());
+ jsonMapper.registerModules(new LocalDataStorageDruidModule().getJacksonModules());
+ jsonMapper.setInjectableValues(
+ new InjectableValues.Std()
+ .addValue(LocalDataSegmentPuller.class, new LocalDataSegmentPuller())
+ .addValue(IndexIO.class, TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT))
+ .addValue(ObjectMapper.class, jsonMapper)
+ .addValue(DataSegment.PruneSpecsHolder.class, DataSegment.PruneSpecsHolder.DEFAULT)
+ .addValue(ExprMacroTable.class, TestExprMacroTable.INSTANCE)
+ );
+
+ cacheRoot = new File(perTestTempDir, "cache_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE));
+ FileUtils.mkdirp(cacheRoot);
+ }
+
+ @AfterEach
+ void tearDown()
+ {
+ if (manager != null) {
+ // Drop the segment to release rule-holds and unmount mmap'd bundle files before @TempDir tries to clean up
+ // the cache directory. Without this, on some platforms the temp-dir cleanup fails on still-mapped files.
+ try {
+ manager.drop(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ }
+ catch (Throwable ignored) {
+ // best-effort — some tests never installed a rule on this segment
+ }
+ manager.shutdown();
+ }
+ }
+
+ @Test
+ void testLoadInstallsRuleHoldsOnMetadataAndSelectedBundle() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ // The metadata entry now carries the rule state.
+ final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID);
+ Assertions.assertTrue(metadata.isRuleHeld(), "rule must be applied to the metadata entry");
+ Assertions.assertEquals(FINGERPRINT, metadata.getRuleFingerprint(), "fingerprint must be stored on the entry");
+ // Metadata + selected bundle + its base parent are weak-reserved on the location; NOT static.
+ final SegmentCacheEntryIdentifier metaId = new SegmentCacheEntryIdentifier(SEGMENT_ID);
+ final PartialSegmentBundleCacheEntryIdentifier aggId =
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE);
+ final PartialSegmentBundleCacheEntryIdentifier baseId =
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME);
+ Assertions.assertTrue(location.isWeakReserved(metaId), "metadata entry should be weak-reserved");
+ Assertions.assertTrue(location.isWeakReserved(aggId), "selected bundle should be weak-reserved");
+ Assertions.assertTrue(location.isWeakReserved(baseId), "base dependency should be weak-reserved");
+ Assertions.assertFalse(location.isReserved(metaId), "metadata entry should NOT be in staticCacheEntries");
+ Assertions.assertFalse(location.isReserved(aggId), "selected bundle should NOT be in staticCacheEntries");
+ }
+
+ @Test
+ void testLoadDoesNotReserveNonSelectedBundles() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ final PartialSegmentBundleCacheEntryIdentifier otherAggId =
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE);
+ Assertions.assertFalse(
+ location.isWeakReserved(otherAggId),
+ "non-selected bundle must not be reserved by loadPartial"
+ );
+ Assertions.assertNull(
+ location.getCacheEntry(otherAggId),
+ "non-selected bundle must not be in the cache at all after load"
+ );
+ }
+
+ @Test
+ void testLoadEagerlyDownloadsSelectedBundleContainers() throws Exception
+ {
+ // ensureBundleDownloaded should be a no-op after load — the eager-download pool task already downloaded the
+ // bundle.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID);
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ Assertions.assertNotNull(mapper, "metadata mount should produce a file mapper");
+ final long downloadedBefore = mapper.getDownloadedBytes();
+ mapper.ensureBundleDownloaded(AGG_BUNDLE);
+ Assertions.assertEquals(
+ downloadedBefore,
+ mapper.getDownloadedBytes(),
+ "ensureBundleDownloaded must be a no-op — the eager task already downloaded the bundle"
+ );
+ }
+
+ @Test
+ void testLoadEagerlyDownloadsBaseDependencyOfSelectedProjection() throws Exception
+ {
+ // when the rule selects a projection bundle, __base is an implicit mount-time dependency (parent hold + sparse
+ // allocation) but its container bytes are NOT downloaded by ensureBundleDownloaded(projection).
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ final PartialSegmentMetadataCacheEntry metadata = weakReservedMetadata(location, SEGMENT_ID);
+ final PartialSegmentFileMapperV10 mapper = metadata.getFileMapper();
+ Assertions.assertNotNull(mapper, "metadata mount should produce a file mapper");
+ Assertions.assertTrue(
+ mapper.isBundleFullyDownloaded(Projections.BASE_TABLE_PROJECTION_NAME),
+ "base dependency of the selected projection must be fully downloaded eagerly, not just sparse-mounted"
+ );
+ }
+
+ @Test
+ void testLoadWithoutVirtualStorageDoesNotInstallRuleHolds() throws Exception
+ {
+ // virtualStorage=false: partial-load rules aren't a concept here. load() takes the eager full-load path (via
+ // PartialLoadSpec.loadSegment → the materialized inner LocalLoadSpec), and no rule state is created.
+ manager = makeManager(false, false);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "no rule state on the eager path"
+ );
+ }
+
+ @Test
+ void testLoadWithPartialsDisabledIsNoop() throws Exception
+ {
+ manager = makeManager(true, false);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "with partial downloads disabled, loadPartial must not run"
+ );
+ Assertions.assertFalse(
+ location.isWeakReserved(new SegmentCacheEntryIdentifier(SEGMENT_ID)),
+ "with partial downloads disabled, no reservation should be installed at load time"
+ );
+ }
+
+ @Test
+ void testSameFingerprintReIssueIsIdempotent() throws Exception
+ {
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ final PartialSegmentMetadataCacheEntry firstMeta = weakReservedMetadata(location, SEGMENT_ID);
+
+ // Second load with the same fingerprint must be idempotent — same metadata entry, unchanged rule state.
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ final PartialSegmentMetadataCacheEntry secondMeta = weakReservedMetadata(location, SEGMENT_ID);
+ Assertions.assertSame(firstMeta, secondMeta, "same-fingerprint re-issue must not create a new metadata entry");
+ Assertions.assertEquals(FINGERPRINT, secondMeta.getRuleFingerprint(), "fingerprint unchanged after re-issue");
+ }
+
+ @Test
+ void testRuleChangeSwapsHoldsWithoutDroppingSharedMetadata() throws Exception
+ {
+ // Same segment, different fingerprint: the rule's bundle selection changes but the underlying metadata entry is
+ // the same weak entry (its self-hold is never released during the swap; only bundle rule-holds diff).
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original"));
+ final PartialSegmentMetadataCacheEntry meta = weakReservedMetadata(location, SEGMENT_ID);
+ Assertions.assertEquals("v1:rule-original", meta.getRuleFingerprint());
+
+ manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated"));
+
+ // Metadata entry is the SAME instance across the swap.
+ final PartialSegmentMetadataCacheEntry metaAfter = weakReservedMetadata(location, SEGMENT_ID);
+ Assertions.assertSame(meta, metaAfter, "rule change must reuse the underlying metadata entry");
+ Assertions.assertEquals("v2:rule-updated", metaAfter.getRuleFingerprint(), "fingerprint must swap");
+ // New rule's bundle is rule-held.
+ Assertions.assertTrue(
+ location.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, OTHER_AGG_BUNDLE)),
+ "new rule's bundle should be reserved"
+ );
+ Assertions.assertFalse(metaAfter.isBundleRuleHeld(AGG_BUNDLE));
+ }
+
+ @Test
+ void testRuleSwapRewritesInfoFileOnDisk() throws Exception
+ {
+ // findOrReservePartial returns the existing entry on a rule swap, make sure we rewrite the
+ // info file.
+ manager = makeManager(true, true);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original"));
+ manager.load(partialWrapperSegment(List.of(OTHER_AGG_BUNDLE), "v2:rule-updated"));
+
+ final File infoDir = new File(cacheRoot, "info_dir");
+ final File infoFile = new File(infoDir, SEGMENT_ID.toString());
+ Assertions.assertTrue(infoFile.exists(), "info file must exist after loadPartial");
+ final DataSegment onDisk = jsonMapper.readValue(infoFile, DataSegment.class);
+ Assertions.assertEquals(
+ "v2:rule-updated",
+ onDisk.getLoadSpec().get("fingerprint"),
+ "info file on disk must carry the latest rule's fingerprint after a swap, not the prior rule's"
+ );
+ Assertions.assertEquals(
+ List.of(OTHER_AGG_BUNDLE),
+ onDisk.getLoadSpec().get("projections"),
+ "info file on disk must carry the latest rule's projection selection"
+ );
+ }
+
+ @Test
+ void testDropClearsRule() throws Exception
+ {
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+ Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID));
+
+ manager.drop(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "drop must clear the rule state on the metadata entry"
+ );
+ }
+
+ @Test
+ void testRangeReaderNullDropsPriorRule() throws Exception
+ {
+ // Prior rule installed. Second load with a wrapper whose delegate can't produce a range reader releases the rule
+ // via clearRule so subsequent queries fall through to weak-full-load semantics.
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE), "v1:rule-original"));
+ Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID));
+
+ manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v2:rule-updated"));
+
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "null range reader must clear the prior rule"
+ );
+ }
+
+ @Test
+ void testRangeReaderNullNoopWhenNoPriorRule() throws Exception
+ {
+ // No prior rule, no range reader → no rule installed, but also no failure. Segment simply doesn't get a rule;
+ // a later query would drive weak-full-load via the on-demand path.
+ manager = makeManager(true, true);
+ manager.load(partialWrapperSegmentWithNullRangeReader(List.of(AGG_BUNDLE), "v1:rule-noop"));
+ Assertions.assertNull(manager.getRuleFingerprintForSegment(SEGMENT_ID));
+ }
+
+ @Test
+ void testLoadFailureLeavesNoRuleApplied() throws Exception
+ {
+ // Wrapper referring to a non-existent projection — wrapper.getSelectedBundleNames throws before applyRule fires.
+ // In the v2 design there is no rollback-nuke of the partial dir on failure; the transient hold releases and the
+ // now-unheld weak metadata entry becomes SIEVE-eligible naturally (info-file cleanup fires via the entry's
+ // onUnmount hook when SIEVE eventually reclaims it). The important correctness invariant is: no stranded rule
+ // state, no partial rule fingerprint applied.
+ manager = makeManager(true, true);
+ final DataSegment segment = partialWrapperSegment(List.of("does_not_exist"));
+
+ final Throwable thrown = Assertions.assertThrows(
+ Throwable.class,
+ () -> manager.load(segment)
+ );
+ Assertions.assertTrue(
+ thrown.getMessage() != null && thrown.getMessage().contains("does not contain projection[does_not_exist]")
+ || thrown.getCause() != null
+ && thrown.getCause().getMessage().contains("does not contain projection[does_not_exist]"),
+ "unexpected throwable: " + thrown
+ );
+ Assertions.assertNull(
+ manager.getRuleFingerprintForSegment(SEGMENT_ID),
+ "failed load must not leave a rule applied on the metadata entry"
+ );
+ }
+
+ @Test
+ void testBootstrapReinstallsRuleHoldsFromPersistedInfoFile() throws Exception
+ {
+ // on historical restart, getCachedSegments+bootstrap must reapply the rule using the persisted DataSegment's
+ // loadSpec (the wire-form PartialLoadSpec is stored in the info file), so a rule-pinned segment is protected from
+ // eviction as soon as bootstrap completes
+ manager = makeManager(true, true);
+ final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE));
+
+ manager.load(segment);
+ Assertions.assertEquals(FINGERPRINT, manager.getRuleFingerprintForSegment(SEGMENT_ID));
+ manager.shutdown();
+ manager = null;
+
+ // Fresh manager over the same cacheRoot — simulates a restart.
+ final SegmentLocalCacheManager restarted = makeManager(true, true);
+ try {
+ final List cached = restarted.getCachedSegments();
+ Assertions.assertTrue(
+ cached.stream().anyMatch(s -> s.getId().equals(SEGMENT_ID)),
+ "restarted historical must rediscover the segment via its info file"
+ );
+
+ restarted.bootstrap(segment, SegmentLazyLoadFailCallback.NOOP);
+
+ // Fingerprint round-trips through the info file → wrapper → applyRule.
+ Assertions.assertEquals(
+ FINGERPRINT,
+ restarted.getRuleFingerprintForSegment(SEGMENT_ID),
+ "bootstrap must reapply the persisted PartialLoadSpec wrapper's fingerprint"
+ );
+ // Selected bundle + its base dependency are held after bootstrap restore (they were on disk and got
+ // restored by the metadata mount's PartialSegmentCacheBootstrap.restoreBundlesFromDisk, which register with
+ // the metadata; applyRule picks up their rule-holds from the linkedBundles state).
+ final StorageLocation loc = restarted.getLocations().get(0);
+ Assertions.assertTrue(
+ loc.isWeakReserved(new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE)),
+ "selected bundle must be reserved after bootstrap reapply"
+ );
+ Assertions.assertTrue(
+ loc.isWeakReserved(
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, Projections.BASE_TABLE_PROJECTION_NAME)
+ ),
+ "base dependency must be reserved after bootstrap reapply"
+ );
+ }
+ finally {
+ restarted.shutdown();
+ }
+ }
+
+ @Test
+ void testReserveMkdirpFailureFallsThroughToNextLocation() throws Exception
+ {
+ // Setup: writable first (so info_dir defaults there), read-only second, dummy reservation on writable so
+ // LeastBytesUsed picks the read-only location first. mkdirp on read-only fails → release + continue to writable.
+ final File writable = new File(perTestTempDir, "loc_rw_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE));
+ final File readOnly = new File(perTestTempDir, "loc_ro_" + ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE));
+ FileUtils.mkdirp(writable);
+ FileUtils.mkdirp(readOnly);
+ manager = makeManagerAtLocations(true, true, List.of(writable, readOnly));
+
+ // Bump writable's usage so LeastBytesUsed picks readOnly first.
+ final SegmentId dummy = SegmentId.of("dummy", Intervals.of("2020/2021"), "v", 0);
+ Assertions.assertTrue(
+ manager.getLocations().get(0).reserveWeak(stubCacheEntry(new SegmentCacheEntryIdentifier(dummy), 4096L))
+ );
+ Assertions.assertTrue(readOnly.setReadOnly(), "test setup must be able to make readOnly location read-only");
+ try {
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ Assertions.assertFalse(
+ new File(readOnly, SEGMENT_ID.toString()).exists(),
+ "read-only location must NOT receive the segment"
+ );
+ Assertions.assertTrue(
+ new File(writable, SEGMENT_ID.toString()).isDirectory(),
+ "load must fall through to the writable location"
+ );
+ Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID));
+ }
+ finally {
+ Assertions.assertTrue(readOnly.setWritable(true), "test teardown must restore write permission");
+ }
+ }
+
+ @Test
+ void testLoadReusesExistingWeakEntryHeldByConcurrentQuery() throws Exception
+ {
+ // A concurrent query is holding a weak metadata entry from an on-demand acquire when the coordinator issues a
+ // partial-load rule for the same segment. loadPartial must reuse that entry — addWeakReservationHoldIfExists
+ // adds another party to the same phaser rather than installing a fresh entry alongside.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+ final SegmentCacheEntryIdentifier id = new SegmentCacheEntryIdentifier(SEGMENT_ID);
+
+ final DataSegment segment = partialWrapperSegment(List.of(AGG_BUNDLE));
+ final File partialDir = new File(cacheRoot, SEGMENT_ID.toString());
+ FileUtils.mkdirp(partialDir);
+ final PartialSegmentMetadataCacheEntry preExisting = new PartialSegmentMetadataCacheEntry(
+ SEGMENT_ID,
+ partialDir,
+ IndexIO.V10_FILE_NAME,
+ List.of(),
+ new DirectoryBackedRangeReader(DEEP_STORAGE_DIR),
+ jsonMapper,
+ null,
+ 16L * 1024L * 1024L
+ );
+ final StorageLocation.ReservationHold queryHold =
+ location.addWeakReservationHold(id, () -> preExisting);
+ Assertions.assertNotNull(queryHold, "precondition: on-demand-style weak reserve must succeed");
+ try {
+ Assertions.assertTrue(location.isWeakReserved(id));
+
+ manager.load(segment);
+
+ // The rule hold reuses the pre-existing entry — same instance across the load.
+ final PartialSegmentMetadataCacheEntry ruleHeld = weakReservedMetadata(location, SEGMENT_ID);
+ Assertions.assertSame(
+ preExisting,
+ ruleHeld,
+ "loadPartial must add its hold to the pre-existing weak entry, not install a fresh one"
+ );
+ Assertions.assertNotNull(manager.getRuleFingerprintForSegment(SEGMENT_ID));
+ }
+ finally {
+ queryHold.close();
+ }
+ }
+
+ private SegmentLocalCacheManager makeManager(boolean virtualStorage, boolean partialDownloadsEnabled)
+ {
+ return makeManagerAtLocations(virtualStorage, partialDownloadsEnabled, List.of(cacheRoot));
+ }
+
+ private SegmentLocalCacheManager makeManagerAtLocations(
+ boolean virtualStorage,
+ boolean partialDownloadsEnabled,
+ List locationRoots
+ )
+ {
+ final List locConfigs = locationRoots.stream()
+ .map(root -> new StorageLocationConfig(root, 1024L * 1024L * 1024L, null))
+ .toList();
+ final SegmentLoaderConfig loaderConfig = new SegmentLoaderConfig()
+ .setLocations(locConfigs)
+ .setVirtualStorage(virtualStorage)
+ .setVirtualStoragePartialDownloadsEnabled(partialDownloadsEnabled);
+ final List storageLocations = loaderConfig.toStorageLocations();
+ return new SegmentLocalCacheManager(
+ storageLocations,
+ loaderConfig,
+ StorageLoadingThreadPool.createFromConfig(loaderConfig),
+ new LeastBytesUsedStorageLocationSelectorStrategy(storageLocations),
+ TestHelper.getTestIndexIO(jsonMapper, ColumnConfig.DEFAULT),
+ jsonMapper
+ );
+ }
+
+ private DataSegment partialWrapperSegment(List selectedProjections)
+ {
+ return partialWrapperSegment(selectedProjections, FINGERPRINT);
+ }
+
+ private DataSegment partialWrapperSegment(List selectedProjections, String fingerprint)
+ {
+ final Map delegate = Map.of(
+ "type", "local",
+ "path", DEEP_STORAGE_DIR.getAbsolutePath()
+ );
+ final Map wrapperWire =
+ PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint);
+ return DataSegment.builder(SEGMENT_ID)
+ .shardSpec(NoneShardSpec.instance())
+ .loadSpec(wrapperWire)
+ .size(0)
+ .build();
+ }
+
+ /**
+ * A wrapper whose inner LoadSpec resolves via {@code LocalLoadSpec} against a directory that holds no V10 file, so
+ * {@code openRangeReader()} returns {@code null}. Simulates the "backend doesn't support range reads" case.
+ */
+ private DataSegment partialWrapperSegmentWithNullRangeReader(List selectedProjections, String fingerprint)
+ throws IOException
+ {
+ final File noRangeReaderDir = new File(
+ perTestTempDir,
+ "no_range_reader_" + fingerprint.replace(':', '_').replace('.', '_')
+ );
+ FileUtils.mkdirp(noRangeReaderDir);
+ final Map delegate = Map.of(
+ "type", "local",
+ "path", noRangeReaderDir.getAbsolutePath()
+ );
+ final Map wrapperWire =
+ PartialProjectionLoadSpec.wireForm(delegate, selectedProjections, fingerprint);
+ return DataSegment.builder(SEGMENT_ID)
+ .shardSpec(NoneShardSpec.instance())
+ .loadSpec(wrapperWire)
+ .size(0)
+ .build();
+ }
+
+ /**
+ * Post-load lookup of the mounted partial metadata entry, verifying it's in weak — not static. loadPartial's
+ * install is synchronous by the time load() returns; the entry is either fully mounted or the call threw.
+ */
+ private static PartialSegmentMetadataCacheEntry weakReservedMetadata(StorageLocation location, SegmentId segmentId)
+ {
+ final CacheEntry entry = location.getCacheEntry(new SegmentCacheEntryIdentifier(segmentId));
+ Assertions.assertInstanceOf(PartialSegmentMetadataCacheEntry.class, entry);
+ final PartialSegmentMetadataCacheEntry partial = (PartialSegmentMetadataCacheEntry) entry;
+ Assertions.assertTrue(partial.isMounted(), "metadata entry for " + segmentId + " should be mounted after load()");
+ return partial;
+ }
+
+ private static CacheEntry stubCacheEntry(CacheEntryIdentifier id, long size)
+ {
+ return new CacheEntry()
+ {
+ @Override
+ public CacheEntryIdentifier getId()
+ {
+ return id;
+ }
+
+ @Override
+ public long getSize()
+ {
+ return size;
+ }
+
+ @Override
+ public boolean isMounted()
+ {
+ return false;
+ }
+
+ @Override
+ public void mount(StorageLocation location)
+ {
+ }
+
+ @Override
+ public void unmount()
+ {
+ }
+ };
+ }
+}