This is an automated email from the ASF dual-hosted git repository.
clintropolis pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 90563dc3b61 feat: partial load rule announcement (#19682)
90563dc3b61 is described below
commit 90563dc3b61ffbcf3ce1ea2ce79dafa63b050450
Author: Clint Wylie <[email protected]>
AuthorDate: Wed Jul 15 01:09:55 2026 -0700
feat: partial load rule announcement (#19682)
---
.../query/PartialProjectionLoadRuleQueryTest.java | 55 ++++++++++-
.../indexing/common/task/CompactionTaskTest.java | 4 +-
.../indexing/input/DruidSegmentReaderTest.java | 4 +-
.../org/apache/druid/timeline/DataSegment.java | 40 --------
.../druid/client/DataSegmentAndLoadProfile.java | 1 +
.../druid/client/HttpServerInventoryView.java | 25 ++---
.../loading/PartialSegmentMetadataCacheEntry.java | 37 +++++++
.../druid/segment/loading/SegmentCacheManager.java | 10 +-
.../segment/loading/SegmentLocalCacheManager.java | 55 ++++++++---
.../org/apache/druid/server/SegmentManager.java | 24 ++++-
.../coordination/SegmentCacheBootstrapper.java | 8 +-
.../coordination/SegmentChangeRequestLoad.java | 30 ++++--
.../coordination/SegmentLoadDropHandler.java | 10 +-
.../coordinator/loading/PartialLoadProfile.java | 43 +++------
.../server/coordinator/loading/SegmentHolder.java | 13 +--
.../druid/client/DruidServerPartialLoadTest.java | 9 +-
.../segment/loading/NoopSegmentCacheManager.java | 4 +-
...egmentLocalCacheManagerPartialRuleLoadTest.java | 106 +++++++++++++++++++++
.../coordination/SegmentChangeRequestLoadTest.java | 41 +++++++-
.../coordination/SegmentLoadDropHandlerTest.java | 13 ++-
.../loading/PartialLoadProfileTest.java | 20 +---
.../StrategicSegmentAssignerPartialTest.java | 15 ++-
.../druid/test/utils/TestSegmentCacheManager.java | 6 +-
23 files changed, 397 insertions(+), 176 deletions(-)
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
index a3efef184dc..0215e63fa63 100644
---
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/query/PartialProjectionLoadRuleQueryTest.java
@@ -38,6 +38,7 @@ 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.LongMinAggregatorFactory;
import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
import org.apache.druid.server.coordinator.rules.CannotMatchBehavior;
import org.apache.druid.server.coordinator.rules.ForeverPartialLoadRule;
@@ -84,6 +85,12 @@ import java.util.UUID;
class PartialProjectionLoadRuleQueryTest extends EmbeddedClusterTestBase
{
private static final String PROJECTION_NAME = "country_delta";
+ // A second projection ingested alongside country_delta but NOT selected by
the rule. Under the partial-load rule
+ // its container bytes stay off the historical's disk, so the historical's
realized footprint is measurably less
+ // than the full segment size — that inequality is what {@link
#testCoordinatorSeesPartialLoadFootprint} asserts
+ // against sys.servers. Uses a min-aggregator on delta so neither test
query's SUM(delta) projection-select can
+ // route through it (planner needs matching aggregator).
+ private static final String UNMATCHED_PROJECTION_NAME = "country_min_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");
@@ -248,6 +255,43 @@ class PartialProjectionLoadRuleQueryTest extends
EmbeddedClusterTestBase
);
}
+ @Test
+ void testCoordinatorSeesPartialLoadFootprint()
+ {
+ // End-to-end check that PartialLoadedDataSegment reaches the
coordinator's inventory accounting: the historical's
+ // announcement stamps realizedBytes (metadata + selected projection +
__base dep) as loadedBytes, so
+ // sys.servers.curr_size for the historical reflects the partial footprint
— strictly less than the segment's full
+ // size (which includes the unmatched projection's bytes). Without the
fix, forAnnouncement would stamp
+ // segment.getSize() and curr_size would equal the full size.
+ final long fullSize = Long.parseLong(
+ cluster.callApi().runSql(
+ "SELECT \"size\" FROM sys.segments WHERE datasource = '" +
dataSource + "'"
+ ).trim()
+ );
+ Assertions.assertTrue(fullSize > 0, "sys.segments.size must be populated
for the ingested segment");
+
+ final long currSize = Long.parseLong(
+ cluster.callApi().runSql(
+ "SELECT curr_size FROM sys.servers WHERE server_type =
'historical'"
+ ).trim()
+ );
+
+ Assertions.assertTrue(
+ currSize > 0,
+ "coordinator's historical curr_size must reflect the partial-load
footprint, got 0"
+ );
+ Assertions.assertTrue(
+ currSize < fullSize,
+ StringUtils.format(
+ "coordinator should see the partial footprint; got curr_size=%d,
full segment size=%d "
+ + "(without PartialLoadedDataSegment, forAnnouncement would stamp
segment.getSize() as loadedBytes and "
+ + "these would be equal)",
+ currSize,
+ fullSize
+ )
+ );
+ }
+
@Test
void testProjectionMissLoadsBaseBundleOnDemand()
{
@@ -344,6 +388,15 @@ class PartialProjectionLoadRuleQueryTest extends
EmbeddedClusterTestBase
.aggregators(new LongSumAggregatorFactory("sumDelta", "delta"))
.build();
+ // Second projection ingested but NOT selected by the rule. Its containers
stay off disk under the partial-load
+ // rule, giving the historical a realized footprint measurably smaller
than the full segment size. The
+ // min-aggregator prevents the planner from routing SUM(delta) queries
through this projection.
+ final AggregateProjectionSpec unmatchedProjection =
+ AggregateProjectionSpec.builder(UNMATCHED_PROJECTION_NAME)
+ .groupingColumns(new StringDimensionSchema("countryName"))
+ .aggregators(new LongMinAggregatorFactory("minDelta", "delta"))
+ .build();
+
final SegmentGranularitySpec segmentGranularitySpec = new
SegmentGranularitySpec(
Granularities.HOUR,
List.of(Intervals.of("2024-01-01/2024-01-02"))
@@ -360,7 +413,7 @@ class PartialProjectionLoadRuleQueryTest extends
EmbeddedClusterTestBase
.withTimestamp(new TimestampSpec("time", "iso", null))
.withSegmentGranularity(segmentGranularitySpec)
.withBaseTable(clusterSpec)
- .withProjections(List.of(projection))
+ .withProjections(List.of(projection, unmatchedProjection))
)
.tuningConfig(t -> t.withMaxNumConcurrentSubTasks(1))
.withId(taskId);
diff --git
a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java
b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java
index ff219260dbf..5c8a49ed482 100644
---
a/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java
+++
b/indexing-service/src/test/java/org/apache/druid/indexing/common/task/CompactionTaskTest.java
@@ -1990,9 +1990,9 @@ public class CompactionTaskTest
final SegmentCacheManager segmentCacheManager = new
NoopSegmentCacheManager()
{
@Override
- public void load(DataSegment segment)
+ public DataSegment load(DataSegment segment)
{
- // do nothing
+ return segment;
}
@Override
diff --git
a/indexing-service/src/test/java/org/apache/druid/indexing/input/DruidSegmentReaderTest.java
b/indexing-service/src/test/java/org/apache/druid/indexing/input/DruidSegmentReaderTest.java
index 54a6c2d8434..4f16fc52796 100644
---
a/indexing-service/src/test/java/org/apache/druid/indexing/input/DruidSegmentReaderTest.java
+++
b/indexing-service/src/test/java/org/apache/druid/indexing/input/DruidSegmentReaderTest.java
@@ -932,9 +932,9 @@ public class DruidSegmentReaderTest extends
InitializedNullHandlingTest
new NoopSegmentCacheManager()
{
@Override
- public void load(DataSegment segment)
+ public DataSegment load(DataSegment segment)
{
- // do nothing
+ return segment;
}
@Override
diff --git
a/processing/src/main/java/org/apache/druid/timeline/DataSegment.java
b/processing/src/main/java/org/apache/druid/timeline/DataSegment.java
index 987f2982061..0a82f8575a8 100644
--- a/processing/src/main/java/org/apache/druid/timeline/DataSegment.java
+++ b/processing/src/main/java/org/apache/druid/timeline/DataSegment.java
@@ -223,46 +223,6 @@ public class DataSegment implements
Comparable<DataSegment>, Overshadowable<Data
);
}
- /**
- * @deprecated use {@link #builder(SegmentId)} or {@link
#builder(DataSegment)} instead.
- */
- @Deprecated
- public DataSegment(
- String dataSource,
- Interval interval,
- String version,
- @Nullable Map<String, Object> loadSpec,
- @Nullable List<String> dimensions,
- @Nullable List<String> metrics,
- @Nullable List<String> projections,
- @Nullable ShardSpec shardSpec,
- @Nullable CompactionState lastCompactionState,
- Integer binaryVersion,
- long size,
- Integer totalRows,
- String indexingStateFingerprint,
- PruneSpecsHolder pruneSpecsHolder
- )
- {
- this(
- dataSource,
- interval,
- version,
- loadSpec,
- dimensions,
- metrics,
- projections,
- null,
- shardSpec,
- lastCompactionState,
- binaryVersion,
- size,
- totalRows,
- indexingStateFingerprint,
- pruneSpecsHolder
- );
- }
-
@JsonCreator
private DataSegment(
@JsonProperty("dataSource") String dataSource,
diff --git
a/server/src/main/java/org/apache/druid/client/DataSegmentAndLoadProfile.java
b/server/src/main/java/org/apache/druid/client/DataSegmentAndLoadProfile.java
index 26d35d51393..e551c6af1d0 100644
---
a/server/src/main/java/org/apache/druid/client/DataSegmentAndLoadProfile.java
+++
b/server/src/main/java/org/apache/druid/client/DataSegmentAndLoadProfile.java
@@ -53,6 +53,7 @@ public class DataSegmentAndLoadProfile extends DataSegment
delegate.getDimensions(),
delegate.getMetrics(),
delegate.getProjections(),
+ delegate.getClusterGroups(),
delegate.getShardSpec(),
delegate.getLastCompactionState(),
delegate.getBinaryVersion(),
diff --git
a/server/src/main/java/org/apache/druid/client/HttpServerInventoryView.java
b/server/src/main/java/org/apache/druid/client/HttpServerInventoryView.java
index d3063da74fc..ea81900ad4f 100644
--- a/server/src/main/java/org/apache/druid/client/HttpServerInventoryView.java
+++ b/server/src/main/java/org/apache/druid/client/HttpServerInventoryView.java
@@ -674,19 +674,8 @@ public class HttpServerInventoryView implements
ServerInventoryView, FilteredSer
/**
* Builds a {@link PartialLoadProfile} from a load announcement when the
historical populated the partial-load
* wire fields ({@code fingerprint} + {@code loadedBytes}). Returns {@code
null} for plain full-load
- * announcements.
- * <p>
- * Picks between {@link PartialLoadProfile#forLoaded forLoaded} and
- * {@link PartialLoadProfile#forFullFallback forFullFallback} based on
what the historical actually realized:
- * <ul>
- * <li>If {@code loadedBytes} is strictly less than the segment's full
size <em>and</em> the announced segment
- * carries a {@link PartialLoadSpec} wrapper on its load spec, the
historical honored the partial request —
- * preserve the wrapped load spec on the profile so per-replica
observability (e.g.
- * {@link PartialLoadProfile#isFullFallback}) reports correctly.</li>
- * <li>Otherwise the historical either fell back to full or doesn't yet
implement scheme-specific partial
- * loading; build a full-fallback profile (the wire fingerprint
still satisfies the rule, so the
- * coordinator's reconciler counts it as matching without reload
thrash).</li>
- * </ul>
+ * announcements, or defensively when a fingerprint/loadedBytes pair
arrives without a partial-load wrapper on
+ * the segment's load spec (a wire-form contract violation from the
historical).
*/
@Nullable
private static PartialLoadProfile
partialLoadProfileFor(SegmentChangeRequestLoad loadRequest)
@@ -698,10 +687,14 @@ public class HttpServerInventoryView implements
ServerInventoryView, FilteredSer
}
final DataSegment segment = loadRequest.getSegment();
final Map<String, Object> loadSpec = segment.getLoadSpec();
- if (loadedBytes < segment.getSize() &&
PartialLoadSpec.detectPartialLoadSpec(loadSpec)) {
- return PartialLoadProfile.forLoaded(loadSpec, fingerprint,
loadedBytes);
+ if (!PartialLoadSpec.detectPartialLoadSpec(loadSpec)) {
+ // fingerprint/loadedBytes on the wire without a partial-load wrapper
on the segment; shouldn't happen
+ // (SegmentChangeRequestLoad.forAnnouncement only stamps them when the
loadSpec is partial or a
+ // DataSegmentAndLoadProfile carries the profile). Fall back to null
so the inventory treats this as a
+ // plain load rather than materializing a nonsense profile.
+ return null;
}
- return PartialLoadProfile.forFullFallback(fingerprint, loadedBytes);
+ return PartialLoadProfile.forLoaded(loadSpec, fingerprint, loadedBytes);
}
private void removeSegment(final DataSegment segment, boolean fullSync)
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 9f8b53272c3..be1d871dc10 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
@@ -626,6 +626,43 @@ public class PartialSegmentMetadataCacheEntry implements
SegmentCacheEntry, Resi
}
}
+ /**
+ * The rule-declared on-disk footprint for this segment: this metadata
entry's own reservation (V10 header plus
+ * post-mount adjustment) plus the sum of every rule-held bundle's
reservation, plus every transitive dependency
+ * indirectly pinned by those rule-holds. A projection rule pins the
projection bundle in {@link #ruleBundleHolds}
+ * and indirectly pins {@code __base} too. Both contribute to what the rule
caused to be on disk, so both are
+ * included here. Deliberately excludes bundles linked for reasons unrelated
to the rule (e.g. on-demand-loaded by
+ * queries), which are evictable and would inflate the report with transient
state. Used by the historical to stamp
+ * accurate {@code loadedBytes} on partial-load announcements.
+ */
+ public long getRealizedBytes()
+ {
+ entryLock.lock();
+ try {
+ long total = currentSize;
+ // Rule-selected bundles PLUS their transitive dependency closure. Each
bundle's getSize() covers only that
+ // bundle's own containers, deps like __base are separate entries whose
sizes need to be added explicitly.
+ for (String name : bundlesInMountOrder(ruleBundleHolds.keySet())) {
+ final StorageLocation.ReservationHold<PartialSegmentBundleCacheEntry>
hold = ruleBundleHolds.get(name);
+ if (hold != null) {
+ // Rule-selected: size from the hold's entry (guaranteed non-null;
final long field, no lock nesting).
+ total += hold.getEntry().getSize();
+ } else {
+ // Transitive dependency of a rule-selected bundle: pinned via
parent-hold on the dependent, not present
+ // in ruleBundleHolds. Look up in linkedBundles.
+ final PartialSegmentBundleCacheEntry dep = linkedBundles.get(name);
+ if (dep != null) {
+ total += dep.getSize();
+ }
+ }
+ }
+ return total;
+ }
+ finally {
+ entryLock.unlock();
+ }
+ }
+
@Override
public boolean isMounted()
{
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
index c81c1d256a8..cb5570a1731 100644
---
a/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
+++
b/server/src/main/java/org/apache/druid/segment/loading/SegmentCacheManager.java
@@ -75,10 +75,14 @@ public interface SegmentCacheManager
* {@link SegmentLoadingException}.
*
* @param segment Segment to get on each download (after service bootstrap)
+ * @return the input {@code segment}, or a
+ * {@link org.apache.druid.client.DataSegmentAndLoadProfile}
wrapping it when the historical
+ * actually materialized a partial-load footprint; the announcement
path uses the wrapper's
+ * {@code realizedBytes} to stamp accurate {@code loadedBytes} on
the wire form.
* @throws SegmentLoadingException If there is an error in loading the
segment or insufficient storage space
* @see SegmentCacheManager#bootstrap(DataSegment,
SegmentLazyLoadFailCallback)
*/
- void load(DataSegment segment) throws SegmentLoadingException;
+ DataSegment load(DataSegment segment) throws SegmentLoadingException;
/**
* Similar to {@link #load(DataSegment)}, this method loads segments during
startup on data nodes. Implementations of
@@ -88,11 +92,13 @@ public interface SegmentCacheManager
* @param segment Segment to retrieve during service bootstrap
* @param loadFailed Callback to execute when segment lazy load failed. This
applies only when
* {@code lazyLoadOnStart} is enabled
+ * @return same as {@link #load(DataSegment)}: the input {@code segment} or a
+ * {@link org.apache.druid.client.DataSegmentAndLoadProfile}
wrapping it.
* @throws SegmentLoadingException - If there is an error in loading the
segment or insufficient storage space
* @see SegmentCacheManager#load(DataSegment)
* @see SegmentCacheManager#shutdownBootstrap()
*/
- void bootstrap(DataSegment segment, SegmentLazyLoadFailCallback loadFailed)
throws SegmentLoadingException;
+ DataSegment bootstrap(DataSegment segment, SegmentLazyLoadFailCallback
loadFailed) throws SegmentLoadingException;
/**
* Cleanup the segment files cache space used by the segment, releasing the
{@link StorageLocation} reservation
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 1e3e353f482..5059c932a03 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
@@ -29,6 +29,7 @@ import com.google.errorprone.annotations.concurrent.GuardedBy;
import com.google.inject.Inject;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.output.NullOutputStream;
+import org.apache.druid.client.DataSegmentAndLoadProfile;
import org.apache.druid.error.DruidException;
import org.apache.druid.guice.annotations.Json;
import org.apache.druid.java.util.common.FileUtils;
@@ -42,6 +43,7 @@ import
org.apache.druid.segment.ReferenceCountedSegmentProvider;
import org.apache.druid.segment.Segment;
import org.apache.druid.segment.SegmentLazyLoadFailCallback;
import org.apache.druid.segment.file.PartialSegmentFileMapperV10;
+import org.apache.druid.server.coordinator.loading.PartialLoadProfile;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
import org.apache.druid.utils.CloseableUtils;
@@ -967,7 +969,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
* <li>Release the transient hold; the metadata's self-hold keeps the
entry resident under the applied rule.</li>
* </ol>
*/
- private void loadPartial(DataSegment dataSegment) throws
SegmentLoadingException
+ private DataSegment loadPartial(DataSegment dataSegment) throws
SegmentLoadingException
{
final PartialLoadSpec wrapper = materializePartialLoadSpec(dataSegment);
final SegmentRangeReader rangeReader = openPartialRangeReader(dataSegment,
wrapper);
@@ -989,12 +991,18 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
final ReservedPartial existing =
findExistingPartialWithHold(dataSegment.getId());
if (existing != null) {
try {
+ // Snapshot the prior realized footprint before clearRule zeroes
out ruleBundleHolds so the log can
+ // surface the coordinator-visible curr_size transition
+ final long priorRealizedBytes =
existing.metadata().getRealizedBytes();
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",
+ "Backend for segment[%s] does not support range reads;
released rule[fingerprint=%s]. Segment "
+ + "will fall back to weak full-load at query time, and the
next announcement will report "
+ + "loadedBytes=segment.getSize()=%d (up from realized=%d).",
dataSegment.getId(),
- wrapper.getFingerprint()
+ wrapper.getFingerprint(),
+ dataSegment.getSize(),
+ priorRealizedBytes
);
}
finally {
@@ -1014,7 +1022,9 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
wrapper.getFingerprint()
);
}
- return;
+ // Rule couldn't be applied, nothing partial materialized, so hand
back the plain segment; the
+ // announcement layer will fall back to segment.getSize() for
loadedBytes.
+ return dataSegment;
}
final ReservedPartial reserved = findOrReservePartial(dataSegment,
rangeReader);
@@ -1072,6 +1082,13 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
// the coordinator's load queue can retry on its next sync. The
announced fingerprint == "rule fully
// realized" contract stays intact.
awaitEagerDownloadsOrClearRule(dataSegment, metadata, selected);
+ // Wrap the announcement with a DataSegmentAndLoadProfile carrying
the historical's realized footprint AFTER
+ // eager downloads finished. forAnnouncement reads the profile back
via profileOf() and stamps its
+ // loadedBytes + fingerprint.
+ return new DataSegmentAndLoadProfile(
+ dataSegment,
+ PartialLoadProfile.forLoaded(dataSegment.getLoadSpec(),
wrapper.getFingerprint(), metadata.getRealizedBytes())
+ );
}
finally {
CloseableUtils.closeAndSuppressExceptions(
@@ -1292,7 +1309,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
}
@Override
- public void load(final DataSegment dataSegment) throws
SegmentLoadingException
+ public DataSegment load(final DataSegment dataSegment) throws
SegmentLoadingException
{
if (config.isVirtualStorage()) {
if (config.isVirtualStorageEphemeral()) {
@@ -1306,8 +1323,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
// take the existing weak/no-op path below.
if (config.isVirtualStoragePartialDownloadsEnabled()
&& PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec()))
{
- loadPartial(dataSegment);
- return;
+ return loadPartial(dataSegment);
}
// 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
@@ -1326,7 +1342,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
unlock(dataSegment, lock);
}
}
- return;
+ return dataSegment;
}
final CompleteSegmentCacheEntry cacheEntry = new
CompleteSegmentCacheEntry(dataSegment);
final ReferenceCountingLock lock = lock(dataSegment);
@@ -1341,10 +1357,11 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
unlock(dataSegment, lock);
}
}
+ return dataSegment;
}
@Override
- public void bootstrap(
+ public DataSegment bootstrap(
final DataSegment dataSegment,
final SegmentLazyLoadFailCallback loadFailed
) throws SegmentLoadingException
@@ -1358,6 +1375,8 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
// during bootstrap, check if the segment exists in a location and mount
it; getCachedSegments already
// did the reserving for us
final SegmentCacheEntryIdentifier id = new
SegmentCacheEntryIdentifier(dataSegment.getId());
+ // Assemble the loaded-profile inside the segment lock (null = no
partial materialized)
+ PartialLoadProfile loadedProfile = null;
final ReferenceCountingLock lock = lock(dataSegment);
synchronized (lock) {
try {
@@ -1382,7 +1401,9 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
);
}
// 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.
+ // segment continues to be protected from eviction across the
restart AND build the announcement profile.
+ // Otherwise (partial downloads disabled at restart, or
persisted loadSpec no longer partial), don't
+ // advertise a rule footprint.
//
// Strict failure handling: reapplyRuleFromInfoFile throws on
eager-download failure (and clears the
// rule state before throwing). The exception propagates up
through SegmentManager.loadSegmentOnBootstrap,
@@ -1391,6 +1412,11 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
if (config.isVirtualStoragePartialDownloadsEnabled()
&&
PartialLoadSpec.detectPartialLoadSpec(dataSegment.getLoadSpec())) {
reapplyRuleFromInfoFile(dataSegment, partial);
+ loadedProfile = PartialLoadProfile.forLoaded(
+ dataSegment.getLoadSpec(),
+ (String) dataSegment.getLoadSpec().get("fingerprint"),
+ partial.getRealizedBytes()
+ );
}
} else {
throw DruidException.defensive(
@@ -1399,13 +1425,17 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
dataSegment.getId()
);
}
+ // A given segment id lives in exactly one location:
getCachedSegments assigns each cached segment to a
+ // single StorageLocation and getCacheEntry is location-local, so
hitting one location has fully handled
+ // the segment
+ break;
}
}
finally {
unlock(dataSegment, lock);
}
}
- return;
+ return loadedProfile == null ? dataSegment : new
DataSegmentAndLoadProfile(dataSegment, loadedProfile);
}
final ReferenceCountingLock lock = lock(dataSegment);
synchronized (lock) {
@@ -1419,6 +1449,7 @@ public class SegmentLocalCacheManager implements
SegmentCacheManager
unlock(dataSegment, lock);
}
}
+ return dataSegment;
}
@Nullable
diff --git a/server/src/main/java/org/apache/druid/server/SegmentManager.java
b/server/src/main/java/org/apache/druid/server/SegmentManager.java
index a3d58e10251..4bea68953ac 100644
--- a/server/src/main/java/org/apache/druid/server/SegmentManager.java
+++ b/server/src/main/java/org/apache/druid/server/SegmentManager.java
@@ -271,22 +271,30 @@ public class SegmentManager
* @param loadFailed callback to execute when segment lazy load fails. This
applies only
* when lazy loading is enabled.
*
+ * @return the input {@code dataSegment}, or a
+ * {@link org.apache.druid.client.DataSegmentAndLoadProfile}
wrapping it when the historical actually
+ * materialized a partial-load footprint. Callers pass the returned
value to the announcement layer so
+ * partial-load announcements carry accurate {@code loadedBytes}.
* @throws SegmentLoadingException if the segment cannot be loaded
* @throws IOException if the segment info cannot be cached on disk
*/
- public void loadSegmentOnBootstrap(
+ public DataSegment loadSegmentOnBootstrap(
final DataSegment dataSegment,
final SegmentLazyLoadFailCallback loadFailed
) throws SegmentLoadingException, IOException
{
+ final DataSegment loaded;
try {
- cacheManager.bootstrap(dataSegment, loadFailed);
+ loaded = cacheManager.bootstrap(dataSegment, loadFailed);
}
catch (SegmentLoadingException e) {
cacheManager.drop(dataSegment);
throw e;
}
+ // Pass the plain dataSegment (not the potentially-wrapped `loaded`) to
loadSegmentInternal: the wrapper is a
+ // load-time announcement-path artifact only
loadSegmentInternal(dataSegment);
+ return loaded;
}
@@ -298,19 +306,27 @@ public class SegmentManager
*
* @param dataSegment segment to load
*
+ * @return the input {@code dataSegment}, or a
+ * {@link org.apache.druid.client.DataSegmentAndLoadProfile}
wrapping it when the historical actually
+ * materialized a partial-load footprint. Callers pass the returned
value to the announcement layer so
+ * partial-load announcements carry accurate {@code loadedBytes}.
* @throws SegmentLoadingException if the segment cannot be loaded
* @throws IOException if the segment info cannot be cached on disk
*/
- public void loadSegment(final DataSegment dataSegment) throws
SegmentLoadingException, IOException
+ public DataSegment loadSegment(final DataSegment dataSegment) throws
SegmentLoadingException, IOException
{
+ final DataSegment loaded;
try {
- cacheManager.load(dataSegment);
+ loaded = cacheManager.load(dataSegment);
}
catch (SegmentLoadingException e) {
cacheManager.drop(dataSegment);
throw e;
}
+ // Pass the plain dataSegment (not the potentially-wrapped `loaded`) to
loadSegmentInternal: the wrapper is a
+ // load-time announcement-path artifact only
loadSegmentInternal(dataSegment);
+ return loaded;
}
private void loadSegmentInternal(
diff --git
a/server/src/main/java/org/apache/druid/server/coordination/SegmentCacheBootstrapper.java
b/server/src/main/java/org/apache/druid/server/coordination/SegmentCacheBootstrapper.java
index e9e16c9b051..c4388b9e25d 100644
---
a/server/src/main/java/org/apache/druid/server/coordination/SegmentCacheBootstrapper.java
+++
b/server/src/main/java/org/apache/druid/server/coordination/SegmentCacheBootstrapper.java
@@ -190,8 +190,9 @@ public class SegmentCacheBootstrapper
"Loading segment[%d/%d][%s]",
counter.incrementAndGet(), numSegments, segment.getId()
);
+ final DataSegment loaded;
try {
- segmentManager.loadSegmentOnBootstrap(
+ loaded = segmentManager.loadSegmentOnBootstrap(
segment,
() -> loadDropHandler.removeSegment(segment,
DataSegmentChangeCallback.NOOP, false)
);
@@ -201,7 +202,10 @@ public class SegmentCacheBootstrapper
throw new SegmentLoadingException(e, "Exception loading
segment[%s]", segment.getId());
}
try {
- backgroundSegmentAnnouncer.announceSegment(segment);
+ // loadSegmentOnBootstrap returns a PartialLoadedDataSegment
wrapper when the historical
+ // materialized a partial-load footprint on restart; enqueue
that so the announcement carries
+ // accurate loadedBytes (see PartialLoadedDataSegment +
SegmentChangeRequestLoad.forAnnouncement).
+ backgroundSegmentAnnouncer.announceSegment(loaded);
}
catch (InterruptedException e) {
Thread.currentThread().interrupt();
diff --git
a/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
b/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
index be9810b7f03..c6b5f4ac290 100644
---
a/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
+++
b/server/src/main/java/org/apache/druid/server/coordination/SegmentChangeRequestLoad.java
@@ -23,9 +23,11 @@ import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonUnwrapped;
+import org.apache.druid.client.DataSegmentAndLoadProfile;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.segment.loading.PartialLoadSpec;
+import org.apache.druid.server.coordinator.loading.PartialLoadProfile;
import org.apache.druid.timeline.DataSegment;
import javax.annotation.Nullable;
@@ -44,15 +46,20 @@ public class SegmentChangeRequestLoad implements
DataSegmentChangeRequest
private static final Logger log = new Logger(SegmentChangeRequestLoad.class);
/**
- * Builds a load announcement for a segment loaded on a historical. Any
{@link PartialLoadSpec} subtype (identified
- * by the wire-form conventions documented on that class: a {@code type}
starting with
- * {@link PartialLoadSpec#TYPE_PREFIX}, plus top-level {@code fingerprint}
and {@code delegate} fields) produces a
- * "full-fallback" announcement: the wrapper's fingerprint and {@link
DataSegment#getSize()} ride along as
- * {@code loadedBytes}, satisfying the coordinator's partial-load rule even
when the historical did a regular full
- * load via the inner delegate. Without this, the coordinator's reconciler
would treat the replica as stale and
- * re-queue the load indefinitely.
- * <p>
- * Detection is convention-based (no subtype allowlist) so future {@link
PartialLoadSpec} subtypes work
+ * Builds a load announcement for a segment loaded on a historical. Two ways
the announcement carries partial-load
+ * metadata:
+ * <ul>
+ * <li>If {@code segment} is a {@link DataSegmentAndLoadProfile}, the
historical materialized a real partial
+ * footprint and attached a {@link PartialLoadProfile}. The profile's
{@code fingerprint} and
+ * {@code loadedBytes} ride directly onto the wire form. This is the
accurate on-disk-footprint path.
+ * <li>Otherwise, if the segment's {@code loadSpec} is a {@link
PartialLoadSpec} wrapper (identified by wire-form
+ * conventions: a {@code type} starting with {@link
PartialLoadSpec#TYPE_PREFIX}, plus top-level
+ * {@code fingerprint} and {@code delegate} fields), the historical
was asked to partial-load but fell back
+ * to a full download via the inner delegate (zipped storage,
capability mismatch, etc.). {@code loadedBytes}
+ * is {@link DataSegment#getSize()} so the fingerprint still satisfies
the coordinator's partial-load rule and
+ * no reload-thrash occurs.
+ * </ul>
+ * Fallback detection is convention-based (no subtype allowlist) so future
{@link PartialLoadSpec} subtypes work
* automatically without touching this code.
* <p>
* For segments loaded without a partial-load wrapper (the common case),
this returns a bare load request with no
@@ -60,8 +67,13 @@ public class SegmentChangeRequestLoad implements
DataSegmentChangeRequest
*/
public static SegmentChangeRequestLoad forAnnouncement(DataSegment segment)
{
+ final PartialLoadProfile profile =
DataSegmentAndLoadProfile.profileOf(segment);
+ if (profile != null) {
+ return new SegmentChangeRequestLoad(segment, profile.fingerprint(),
profile.loadedBytes());
+ }
final Map<String, Object> loadSpec = segment.getLoadSpec();
if (PartialLoadSpec.detectPartialLoadSpec(loadSpec)) {
+ // Historical didn't wrap, treat as full-fallback: fingerprint from the
loadSpec, loadedBytes = full size.
return new SegmentChangeRequestLoad(segment, (String)
loadSpec.get("fingerprint"), segment.getSize());
}
if (PartialLoadSpec.hasPartialTypePrefix(loadSpec)) {
diff --git
a/server/src/main/java/org/apache/druid/server/coordination/SegmentLoadDropHandler.java
b/server/src/main/java/org/apache/druid/server/coordination/SegmentLoadDropHandler.java
index 6d8faf593cb..1e024fcc5bb 100644
---
a/server/src/main/java/org/apache/druid/server/coordination/SegmentLoadDropHandler.java
+++
b/server/src/main/java/org/apache/druid/server/coordination/SegmentLoadDropHandler.java
@@ -161,16 +161,20 @@ public class SegmentLoadDropHandler
currentDropLatch.cancelOrAwait();
}
+ final DataSegment loaded;
try {
- segmentManager.loadSegment(segment);
+ loaded = segmentManager.loadSegment(segment);
}
catch (Exception e) {
removeSegment(segment, DataSegmentChangeCallback.NOOP, false);
throw new SegmentLoadingException(e, "Exception loading segment[%s]",
segment.getId());
}
try {
- // Announce segment even if the segment file already exists.
- announcer.announceSegment(segment);
+ // Announce segment even if the segment file already exists.
loadSegment returns a
+ // PartialLoadedDataSegment wrapper when the historical actually
materialized a partial-load footprint;
+ // announcing that wrapper lets
SegmentChangeRequestLoad.forAnnouncement stamp accurate loadedBytes on
+ // the wire form, so the coordinator's inventory sees the partial
footprint rather than segment.getSize().
+ announcer.announceSegment(loaded);
}
catch (IOException e) {
throw new SegmentLoadingException(e, "Failed to announce segment[%s]",
segment.getId());
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
index 4cd1cc2e224..04e1f86f658 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/PartialLoadProfile.java
@@ -33,26 +33,22 @@ import java.util.Objects;
* matches the shape of the wire-form load-spec wrapper that gets stamped onto
outbound load requests via
* {@code DataSegment.withLoadSpec(...)}.
* <p>
- * Used in three distinct states:
+ * Used in two distinct states:
* <ul>
* <li>{@link #forRequest(Map, String) forRequest}: outbound from
coordinator to historical. Carries the wrapped
* load-spec map identifying what to load and the fingerprint
identifying that request. {@code loadedBytes} is
* null because the on-disk footprint is only known after the historical
has parsed segment metadata.</li>
- * <li>{@link #forLoaded(Map, String, long) forLoaded}: inbound on a
historical announcement after a successful
- * partial load. {@code wrappedLoadSpec} echoes the request and {@code
loadedBytes} reports the realized
- * footprint.</li>
- * <li>{@link #forFullFallback(String, long) forFullFallback}: inbound on a
historical announcement when partial
- * loading was requested but the historical fell back to a full download
(zipped-V10 segment, capability
- * mismatch on a partially-upgraded server, etc.). {@code
wrappedLoadSpec} is null to signal the fallback;
- * {@code loadedBytes} equals the full segment size. The matching
fingerprint still satisfies the rule that
- * requested the load, so no reload-thrash occurs.</li>
+ * <li>{@link #forLoaded(Map, String, long) forLoaded}: inbound on a
historical announcement after a load completed.
+ * {@code wrappedLoadSpec} echoes the request and {@code loadedBytes}
reports the realized footprint. Covers
+ * both real partial loads and historicals that fell back to a full
download. The fingerprint match is what
+ * satisfies the coordinator's rule either way; the footprint number
rides through for capacity accounting.</li>
* </ul>
* The fingerprint is derived from the resolved scheme-specific data (e.g.,
for projections, the sorted/deduped name
* list, see {@code ProjectionPartialLoadMatcher}) so that two rule
configurations that resolve to the same set on a
* segment produce the same fingerprint and don't churn replicas across
coordinator runs.
*/
public record PartialLoadProfile(
- @Nullable Map<String, Object> wrappedLoadSpec,
+ Map<String, Object> wrappedLoadSpec,
String fingerprint,
@Nullable Long loadedBytes
)
@@ -66,10 +62,9 @@ public record PartialLoadProfile(
public PartialLoadProfile
{
+ Objects.requireNonNull(wrappedLoadSpec, "wrappedLoadSpec");
Objects.requireNonNull(fingerprint, "fingerprint");
- if (wrappedLoadSpec != null) {
- wrappedLoadSpec = Map.copyOf(wrappedLoadSpec);
- }
+ wrappedLoadSpec = Map.copyOf(wrappedLoadSpec);
}
/**
@@ -86,7 +81,9 @@ public record PartialLoadProfile(
}
/**
- * Build the inbound profile for a successful partial load announcement.
+ * Build the inbound profile for a completed load announcement. Applies to
both real partial loads (whose
+ * {@code loadedBytes} reflects the rule-declared on-disk footprint) and
full-download fallbacks (whose
+ * {@code loadedBytes} equals the full segment size).
*/
public static PartialLoadProfile forLoaded(Map<String, Object>
wrappedLoadSpec, String fingerprint, long loadedBytes)
{
@@ -96,26 +93,8 @@ public record PartialLoadProfile(
return intern(new PartialLoadProfile(wrappedLoadSpec, fingerprint,
loadedBytes));
}
- /**
- * Build the inbound profile for a historical that was asked to partial-load
but fell back to a full download.
- * {@code wrappedLoadSpec} is null as a sentinel; {@code loadedBytes} should
be the full segment size so that
- * inventory accounting reflects actual on-disk footprint.
- */
- public static PartialLoadProfile forFullFallback(String fingerprint, long
fullBytes)
- {
- return intern(new PartialLoadProfile(null, fingerprint, fullBytes));
- }
-
private static PartialLoadProfile intern(PartialLoadProfile profile)
{
return INTERNER.intern(profile);
}
-
- /**
- * Whether this profile represents a historical's full-fallback (i.e., it
was asked to partial-load but couldn't).
- */
- public boolean isFullFallback()
- {
- return wrappedLoadSpec == null;
- }
}
diff --git
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentHolder.java
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentHolder.java
index 968478cd875..39f1f8a2767 100644
---
a/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentHolder.java
+++
b/server/src/main/java/org/apache/druid/server/coordinator/loading/SegmentHolder.java
@@ -119,22 +119,11 @@ public class SegmentHolder implements
Comparable<SegmentHolder>
this.profile = profile;
if (action == SegmentAction.DROP) {
this.changeRequest = new SegmentChangeRequestDrop(segment);
- } else if (profile != null && profile.wrappedLoadSpec() != null) {
+ } else if (profile != null) {
// Stamp the wrapped load-spec map onto the outbound segment so the
historical receives the partial-load
// wrapper; identity (segment field) stays original so dedup in the load
queue is unaffected.
this.changeRequest = new
SegmentChangeRequestLoad(segment.withLoadSpec(profile.wrappedLoadSpec()));
} else {
- if (profile != null) {
- // Outbound profiles (forRequest / forLoaded) must always carry a
non-null wrappedLoadSpec; only the inbound
- // forFullFallback announcement uses the null sentinel, and those
never reach this constructor. Log if it
- // happens so the bug surfaces — we still proceed with a plain load
request to avoid stalling the queue.
- log.warn(
- "Profile with null wrappedLoadSpec on outbound load for
segment[%s], action[%s]; "
- + "queuing as a regular load.",
- segment.getId(),
- action
- );
- }
this.changeRequest = new SegmentChangeRequestLoad(segment);
}
if (callback != null) {
diff --git
a/server/src/test/java/org/apache/druid/client/DruidServerPartialLoadTest.java
b/server/src/test/java/org/apache/druid/client/DruidServerPartialLoadTest.java
index e41af11672e..e4e30f1b15a 100644
---
a/server/src/test/java/org/apache/druid/client/DruidServerPartialLoadTest.java
+++
b/server/src/test/java/org/apache/druid/client/DruidServerPartialLoadTest.java
@@ -105,17 +105,20 @@ public class DruidServerPartialLoadTest
}
@Test
- void testFullFallbackProfileUsesFullSegmentSize()
+ void testFullFallbackLoadedBytesUsesFullSegmentSize()
{
// Historical was asked to partial-load but fell back to full download;
profile.loadedBytes equals
// segment.getSize, so currSize accounting reflects the full footprint.
DruidServer server = newServer();
DataSegment segment = buildSegment("ds", "v1", 1000L);
- PartialLoadProfile profile =
PartialLoadProfile.forFullFallback(FINGERPRINT, segment.getSize());
+ PartialLoadProfile profile = PartialLoadProfile.forLoaded(
+ Map.of("type", "partialProjection", "fingerprint", FINGERPRINT),
+ FINGERPRINT,
+ segment.getSize()
+ );
server.addDataSegment(segment, profile);
Assertions.assertEquals(1000L, server.getCurrSize());
Assertions.assertEquals(profile,
server.getPartialLoadProfile(segment.getId()));
- Assertions.assertTrue(profile.isFullFallback());
}
@Test
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
index 1db7a1618eb..2d9cfc39de3 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/NoopSegmentCacheManager.java
@@ -71,13 +71,13 @@ public class NoopSegmentCacheManager implements
SegmentCacheManager
}
@Override
- public void load(DataSegment segment)
+ public DataSegment load(DataSegment segment)
{
throw new UnsupportedOperationException();
}
@Override
- public void bootstrap(DataSegment segment, SegmentLazyLoadFailCallback
loadFailed)
+ public DataSegment bootstrap(DataSegment segment,
SegmentLazyLoadFailCallback loadFailed)
{
throw new UnsupportedOperationException();
}
diff --git
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
index 7edb4e7f40f..f63563474e4 100644
---
a/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
+++
b/server/src/test/java/org/apache/druid/segment/loading/SegmentLocalCacheManagerPartialRuleLoadTest.java
@@ -22,6 +22,7 @@ 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.client.DataSegmentAndLoadProfile;
import org.apache.druid.data.input.InputRow;
import org.apache.druid.data.input.ListBasedInputRow;
import org.apache.druid.data.input.impl.AggregateProjectionSpec;
@@ -33,6 +34,7 @@ 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.common.StringUtils;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.query.aggregation.CountAggregatorFactory;
@@ -51,6 +53,7 @@ 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.coordinator.loading.PartialLoadProfile;
import org.apache.druid.server.metrics.NoopServiceEmitter;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
@@ -289,6 +292,101 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
);
}
+ @Test
+ void testLoadPartialReturnsDataSegmentAndLoadProfileWithRealizedBytes()
throws Exception
+ {
+ // load() on the partial-load path returns a DataSegmentAndLoadProfile
carrying a PartialLoadProfile.forLoaded with
+ // the actual on-disk footprint. Under a projection rule, that footprint
is metadata header + rule-selected
+ // projection bundle + __base (transitive dependency pinned via the
projection's parent-hold). Anything less
+ // would under-report to the coordinator: the historical downloads __base
as part of the rule (see
+ // awaitEagerDownloadsOrClearRule's dep expansion), so its bytes belong in
the announcement's loadedBytes.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ final DataSegment loaded =
manager.load(partialWrapperSegment(List.of(AGG_BUNDLE)));
+
+ Assertions.assertInstanceOf(
+ DataSegmentAndLoadProfile.class,
+ loaded,
+ "load on the partial-rule path must return a DataSegmentAndLoadProfile
so the announcer can report actual footprint"
+ );
+ final PartialLoadProfile profile = ((DataSegmentAndLoadProfile)
loaded).profile();
+ Assertions.assertEquals(FINGERPRINT, profile.fingerprint(), "profile's
fingerprint must match the applied rule");
+
+ // Sum the parts we expect: metadata + AGG_BUNDLE (rule-selected) + __base
(transitive dep).
+ final PartialSegmentMetadataCacheEntry metadata =
weakReservedMetadata(location, SEGMENT_ID);
+ final PartialSegmentBundleCacheEntry aggEntry = location.getCacheEntry(
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, AGG_BUNDLE)
+ );
+ final PartialSegmentBundleCacheEntry baseEntry = location.getCacheEntry(
+ new PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID,
Projections.BASE_TABLE_PROJECTION_NAME)
+ );
+ Assertions.assertNotNull(aggEntry, "rule-selected projection bundle should
be present");
+ Assertions.assertNotNull(baseEntry, "transitive __base dependency should
be present");
+ final long expectedRealizedBytes = metadata.getSize() + aggEntry.getSize()
+ baseEntry.getSize();
+
+ Assertions.assertEquals(
+ Long.valueOf(expectedRealizedBytes),
+ profile.loadedBytes(),
+ "profile's loadedBytes must include metadata + rule-selected
projection + __base dependency; missing any of "
+ + "these would under-report the on-disk footprint to the coordinator"
+ );
+ Assertions.assertTrue(
+ baseEntry.getSize() > 0,
+ "sanity check: __base's size should be non-zero (so the dep-inclusion
assertion is meaningful)"
+ );
+ }
+ @Test
+ void testRealizedBytesIncludesPinnedBaseDependency() throws Exception
+ {
+ 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 long metadataBytes = metadata.getSize();
+ final long aggBytes = bundleEntry(location, AGG_BUNDLE).getSize();
+ final long baseBytes = bundleEntry(location,
Projections.BASE_TABLE_PROJECTION_NAME).getSize();
+
+ Assertions.assertTrue(baseBytes > 0, "base dependency must occupy real
on-disk bytes");
+
+ final long realized = metadata.getRealizedBytes();
+
+ Assertions.assertEquals(
+ metadataBytes + aggBytes + baseBytes,
+ realized,
+ StringUtils.format(
+ "realizedBytes must include the pinned __base dependency:
metadata=%d + selected=%d + base=%d = %d, "
+ + "but got %d (short by %d, exactly the base footprint)",
+ metadataBytes, aggBytes, baseBytes, metadataBytes + aggBytes +
baseBytes,
+ realized, (metadataBytes + aggBytes + baseBytes) - realized
+ )
+ );
+ }
+
+ @Test
+ void testRealizedBytesCountsSharedBaseDependencyOnce() throws Exception
+ {
+ // Two projections in a single rule both pin the same __base. The base
bytes must be counted exactly once.
+ manager = makeManager(true, true);
+ final StorageLocation location = manager.getLocations().get(0);
+
+ manager.load(partialWrapperSegment(List.of(AGG_BUNDLE, OTHER_AGG_BUNDLE)));
+
+ final PartialSegmentMetadataCacheEntry metadata =
weakReservedMetadata(location, SEGMENT_ID);
+ final long metadataBytes = metadata.getSize();
+ final long aggBytes = bundleEntry(location, AGG_BUNDLE).getSize();
+ final long otherBytes = bundleEntry(location, OTHER_AGG_BUNDLE).getSize();
+ final long baseBytes = bundleEntry(location,
Projections.BASE_TABLE_PROJECTION_NAME).getSize();
+
+ Assertions.assertEquals(
+ metadataBytes + aggBytes + otherBytes + baseBytes,
+ metadata.getRealizedBytes(),
+ "realizedBytes must include the shared __base exactly once across both
dependent projections"
+ );
+ }
+
@Test
void testLoadWithoutVirtualStorageDoesNotInstallRuleHolds() throws Exception
{
@@ -590,6 +688,14 @@ class SegmentLocalCacheManagerPartialRuleLoadTest
}
}
+ private static CacheEntry bundleEntry(StorageLocation location, String
bundleName)
+ {
+ final CacheEntry entry =
+ location.getCacheEntry(new
PartialSegmentBundleCacheEntryIdentifier(SEGMENT_ID, bundleName));
+ Assertions.assertNotNull(entry, "bundle[" + bundleName + "] must be
resident after load");
+ return entry;
+ }
+
private SegmentLocalCacheManager makeManager(boolean virtualStorage, boolean
partialDownloadsEnabled)
{
return makeManagerAtLocations(virtualStorage, partialDownloadsEnabled,
List.of(cacheRoot));
diff --git
a/server/src/test/java/org/apache/druid/server/coordination/SegmentChangeRequestLoadTest.java
b/server/src/test/java/org/apache/druid/server/coordination/SegmentChangeRequestLoadTest.java
index 35d5bd7f35a..74ab1f34f58 100644
---
a/server/src/test/java/org/apache/druid/server/coordination/SegmentChangeRequestLoadTest.java
+++
b/server/src/test/java/org/apache/druid/server/coordination/SegmentChangeRequestLoadTest.java
@@ -21,10 +21,12 @@ package org.apache.druid.server.coordination;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableMap;
+import org.apache.druid.client.DataSegmentAndLoadProfile;
import org.apache.druid.jackson.DefaultObjectMapper;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.jackson.JacksonUtils;
import org.apache.druid.segment.IndexIO;
+import org.apache.druid.server.coordinator.loading.PartialLoadProfile;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
import org.apache.druid.timeline.partition.NoneShardSpec;
@@ -142,9 +144,10 @@ public class SegmentChangeRequestLoadTest
@Test
public void testForAnnouncementPartialProjectionWrapperProducesFullFallback()
{
- // When the segment's loadSpec is a partialProjection wrapper, the
announcement stamps the wrapper's fingerprint
- // and the segment's full size as loadedBytes — coordinator reads this as
a full-fallback profile and counts the
- // replica as matching, avoiding reload thrash on historicals that don't
(yet) do real partial loading.
+ // Plain DataSegment carrying a partial-load wrapper (i.e. the historical
didn't attach a
+ // DataSegmentAndLoadProfile) is the full-fallback case: the wrapper's
fingerprint rides along, but loadedBytes
+ // is the segment's full size. Coordinator reads this as a fallback and
counts the replica as matching,
+ // avoiding reload thrash on historicals that don't (yet) do real partial
loading.
Map<String, Object> wrapped = Map.of(
"type", "partialProjection",
"delegate", Map.of("type", "local", "path", "/var/druid/segments/foo"),
@@ -162,6 +165,38 @@ public class SegmentChangeRequestLoadTest
Assert.assertEquals(Long.valueOf(12345L), announcement.getLoadedBytes());
}
+ @Test
+ public void
testForAnnouncementDataSegmentAndLoadProfileReportsProfileLoadedBytes()
+ {
+ // When the historical actually materialized a partial-load footprint,
load() returns a DataSegmentAndLoadProfile
+ // carrying a PartialLoadProfile.forLoaded(...) with the realized bytes.
forAnnouncement must stamp the profile's
+ // loadedBytes so HttpServerInventoryView detects it and capacity
accounting sees the actual on-disk footprint
+ // instead of the full segment size.
+ Map<String, Object> wrapped = Map.of(
+ "type", "partialProjection",
+ "delegate", Map.of("type", "local", "path", "/var/druid/segments/foo"),
+ "projections", List.of("revenue"),
+ "fingerprint", "v1:abcdef0123456789"
+ );
+ DataSegment segment = segmentBuilder("ds",
Intervals.of("2024-01-01/2024-02-01"), "v1")
+ .loadSpec(wrapped)
+ .dimensions(List.of("d"))
+ .metrics(List.of("m"))
+ .size(12345)
+ .build();
+ DataSegmentAndLoadProfile wrappedSegment = new DataSegmentAndLoadProfile(
+ segment,
+ PartialLoadProfile.forLoaded(wrapped, "v1:abcdef0123456789", 4321L)
+ );
+ SegmentChangeRequestLoad announcement =
SegmentChangeRequestLoad.forAnnouncement(wrappedSegment);
+ Assert.assertEquals("v1:abcdef0123456789", announcement.getFingerprint());
+ Assert.assertEquals(
+ "loadedBytes must come from the DataSegmentAndLoadProfile's
PartialLoadProfile, not segment.getSize()",
+ Long.valueOf(4321L),
+ announcement.getLoadedBytes()
+ );
+ }
+
@Test
public void
testForAnnouncementUnknownPartialTypeStillRecognizedByConvention()
{
diff --git
a/server/src/test/java/org/apache/druid/server/coordination/SegmentLoadDropHandlerTest.java
b/server/src/test/java/org/apache/druid/server/coordination/SegmentLoadDropHandlerTest.java
index 607e5a3fa89..4c8fc0c4079 100644
---
a/server/src/test/java/org/apache/druid/server/coordination/SegmentLoadDropHandlerTest.java
+++
b/server/src/test/java/org/apache/druid/server/coordination/SegmentLoadDropHandlerTest.java
@@ -257,10 +257,11 @@ public class SegmentLoadDropHandlerTest
public void
testProcessBatchDuplicateLoadRequestsWhenFirstRequestFailsSecondRequestShouldSucceed()
throws Exception
{
final SegmentManager segmentManager = Mockito.mock(SegmentManager.class);
- Mockito.doThrow(new RuntimeException("segment loading failure test"))
- .doNothing()
- .when(segmentManager)
- .loadSegment(ArgumentMatchers.any());
+ // loadSegment returns DataSegment, so doNothing() would be rejected by
Mockito. Throw on the first call,
+ // then return the input segment on subsequent calls (the announcement
path uses the returned segment).
+ Mockito.when(segmentManager.loadSegment(ArgumentMatchers.any()))
+ .thenThrow(new RuntimeException("segment loading failure test"))
+ .thenAnswer(invocation -> invocation.getArgument(0));
final SegmentLoadDropHandler handler =
initSegmentLoadDropHandler(segmentManager);
@@ -291,7 +292,9 @@ public class SegmentLoadDropHandlerTest
public void testProcessBatchLoadDropLoadSequenceForSameSegment() throws
Exception
{
final SegmentManager segmentManager = Mockito.mock(SegmentManager.class);
-
Mockito.doNothing().when(segmentManager).loadSegment(ArgumentMatchers.any());
+ // loadSegment returns DataSegment; return the input so the announcement
path sees a plain DataSegment.
+ Mockito.when(segmentManager.loadSegment(ArgumentMatchers.any()))
+ .thenAnswer(invocation -> invocation.getArgument(0));
Mockito.doNothing().when(segmentManager).dropSegment(ArgumentMatchers.any());
final File storageDir = temporaryFolder.newFolder();
diff --git
a/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
b/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
index f6a919f5ac4..ce0f42b7c07 100644
---
a/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
+++
b/server/src/test/java/org/apache/druid/server/coordinator/loading/PartialLoadProfileTest.java
@@ -49,7 +49,6 @@ public class PartialLoadProfileTest
Assertions.assertEquals(WRAPPED, profile.wrappedLoadSpec());
Assertions.assertEquals(FINGERPRINT, profile.fingerprint());
Assertions.assertNull(profile.loadedBytes());
- Assertions.assertFalse(profile.isFullFallback());
}
@Test
@@ -83,7 +82,6 @@ public class PartialLoadProfileTest
Assertions.assertEquals(WRAPPED, profile.wrappedLoadSpec());
Assertions.assertEquals(FINGERPRINT, profile.fingerprint());
Assertions.assertEquals(12345L, profile.loadedBytes());
- Assertions.assertFalse(profile.isFullFallback());
}
@Test
@@ -98,16 +96,6 @@ public class PartialLoadProfileTest
);
}
- @Test
- public void testForFullFallback()
- {
- PartialLoadProfile profile =
PartialLoadProfile.forFullFallback(FINGERPRINT, 99999L);
- Assertions.assertNull(profile.wrappedLoadSpec());
- Assertions.assertEquals(FINGERPRINT, profile.fingerprint());
- Assertions.assertEquals(99999L, profile.loadedBytes());
- Assertions.assertTrue(profile.isFullFallback());
- }
-
@Test
public void testFingerprintRequired()
{
@@ -131,7 +119,7 @@ public class PartialLoadProfileTest
public void testEquals()
{
EqualsVerifier.forClass(PartialLoadProfile.class)
- .withNonnullFields("fingerprint")
+ .withNonnullFields("wrappedLoadSpec", "fingerprint")
.usingGetClass()
.verify();
}
@@ -155,11 +143,5 @@ public class PartialLoadProfileTest
// Different fingerprint ⇒ different profile, no sharing.
PartialLoadProfile pd = PartialLoadProfile.forLoaded(WRAPPED,
"v1:differentfingerprint", 12345L);
Assertions.assertNotSame(pa, pd);
-
- // Full-fallback variants intern independently of forLoaded variants.
- PartialLoadProfile fb1 = PartialLoadProfile.forFullFallback(FINGERPRINT,
12345L);
- PartialLoadProfile fb2 = PartialLoadProfile.forFullFallback(FINGERPRINT,
12345L);
- Assertions.assertSame(fb1, fb2);
- Assertions.assertNotSame(pa, fb1);
}
}
diff --git
a/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
b/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
index 83fd61d047a..e1cc1fb0d39 100644
---
a/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
+++
b/server/src/test/java/org/apache/druid/server/coordinator/loading/StrategicSegmentAssignerPartialTest.java
@@ -152,13 +152,18 @@ public class StrategicSegmentAssignerPartialTest
}
@Test
- public void testFullFallbackProfileCountsAsMatching()
+ public void testFullFallbackLoadedBytesCountsAsMatching()
{
- // s1 announced with a full-fallback profile (historical was asked to
partial-load but fell back to full). The
- // fingerprint matches the rule, so the rule is satisfied; no further
action.
+ // s1 announced with a loaded profile whose loadedBytes equals the full
segment size (historical was asked to
+ // partial-load but fell back to full). The fingerprint matches the rule,
so the rule is satisfied; no further
+ // action.
final DataSegment segment = createSegment();
- final PartialLoadProfile fullFallback =
PartialLoadProfile.forFullFallback(FP_REVENUE, 1024L);
- final ServerHolder s1 = createServerWithLoaded(TIER1, segment,
fullFallback);
+ final PartialLoadProfile fullSizeLoaded = PartialLoadProfile.forLoaded(
+ profileForRevenue().wrappedLoadSpec(),
+ FP_REVENUE,
+ segment.getSize()
+ );
+ final ServerHolder s1 = createServerWithLoaded(TIER1, segment,
fullSizeLoaded);
final ServerHolder s2 = createServer(TIER1);
final DruidCluster cluster = DruidCluster.builder().addTier(TIER1, s1,
s2).build();
diff --git
a/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java
b/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java
index 06059e6cd8d..f4cbb9dcd6d 100644
---
a/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java
+++
b/server/src/test/java/org/apache/druid/test/utils/TestSegmentCacheManager.java
@@ -103,17 +103,19 @@ public class TestSegmentCacheManager extends
NoopSegmentCacheManager
}
@Override
- public void bootstrap(DataSegment segment, SegmentLazyLoadFailCallback
loadFailed)
+ public DataSegment bootstrap(DataSegment segment,
SegmentLazyLoadFailCallback loadFailed)
{
observedBootstrapSegments.add(segment);
getSegmentInternal(segment);
+ return segment;
}
@Override
- public void load(final DataSegment segment)
+ public DataSegment load(final DataSegment segment)
{
observedSegments.add(segment);
getSegmentInternal(segment);
+ return segment;
}
private ReferenceCountedSegmentProvider getSegmentInternal(final DataSegment
segment)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
