Copilot commented on code in PR #18460:
URL: https://github.com/apache/pinot/pull/18460#discussion_r3221124852
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java:
##########
@@ -415,6 +661,31 @@ private static HelixManager createHelixManager() {
return helixManager;
}
+ private static HelixManager createHelixManagerWithTiers(String tierName,
List<String> tierServers) {
+ HelixManager helixManager = mock(HelixManager.class);
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ // Build the tier instance partitions ZK record so getTierInstances() can
fetch it
+ String tierInstancePartitionsName =
+
InstancePartitionsUtils.getInstancePartitionsNameForTier(RAW_TABLE_NAME,
tierName);
+ String tierInstancePartitionsPath =
+
ZKMetadataProvider.constructPropertyStorePathForInstancePartitions(tierInstancePartitionsName);
Review Comment:
`MultiTierStrictRealtimeSegmentAssignment.getTierInstances()` builds the
tier instance-partitions name using `_tableNameWithType`, but this test stub
builds it from `RAW_TABLE_NAME`. If `getInstancePartitionsNameForTier(...)`
expects the table name *with type* (as implied by the production callsite), the
mock ZK path won’t match and `fetchInstancePartitions(...)` will return null,
meaning the tier filter won’t actually be exercised by these tests. Update the
stub to construct the tier instance-partitions name using the same “table name
with type” format used in production (e.g., via the table-name builder
utilities used elsewhere in Pinot).
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java:
##########
@@ -345,6 +390,195 @@ public void testRebalanceUpsertTableWithTiers() {
tierInstancePartitionsMap, rebalanceConfig);
}
+ @Test
+ public void testOnlyColdTierSegmentForPartitionFallsBackToComputed() {
+ // When the only segment for a partition is on cold tier,
getExistingAssignment() returns null after the
+ // tier filter skips it, and the assignment falls back to the computed
assignment from instance partitions.
+ // This covers the case where the bug previously placed a single existing
segment on cold tier (or the
+ // legitimate case where a partition's only completed segment has aged out
to cold).
+ // Only dedup tables support multi-tier assignment
(MultiTierStrictRealtimeSegmentAssignment).
+ SegmentAssignment segmentAssignment =
createSegmentAssignmentWithTiers("dedup");
+ Map<InstancePartitionsType, InstancePartitions>
onlyConsumingInstancePartitionMap =
+ Map.of(InstancePartitionsType.CONSUMING,
_instancePartitionsMap.get(InstancePartitionsType.CONSUMING));
+ Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
+
+ // Assign segments 0-3 (one per partition) to CONSUMING instances normally
+ for (int segmentId = 0; segmentId < NUM_PARTITIONS; segmentId++) {
+ String segmentName = _segments.get(segmentId);
+ List<String> instancesAssigned =
+ segmentAssignment.assignSegment(segmentName, currentAssignment,
onlyConsumingInstancePartitionMap);
+ addToAssignment(currentAssignment, segmentId, instancesAssigned);
+ }
+
+ // Replace partition 0's only segment with a cold-tier placement
+ currentAssignment.put(_segments.get(0), buildColdTierStateMap());
+
+ // Assign a new segment for partition 0 using NEW instance partitions.
With the cold-tier segment skipped
+ // by the tier filter, getExistingAssignment returns null and the
assignment falls back to the computed
+ // (new) instance partitions.
+ Map<InstancePartitionsType, InstancePartitions>
newConsumingInstancePartitionMap =
+ Map.of(InstancePartitionsType.CONSUMING,
_newConsumingInstancePartitions);
+ int newSegmentId = 4;
+ String newSegmentName = _segments.get(newSegmentId);
+ List<String> instancesAssigned =
+ segmentAssignment.assignSegment(newSegmentName, currentAssignment,
newConsumingInstancePartitionMap);
+ assertEquals(instancesAssigned.size(), NUM_REPLICAS);
+ assertEquals(instancesAssigned,
+ Arrays.asList("new_consumingInstance_0", "new_consumingInstance_3",
"new_consumingInstance_6"));
Review Comment:
This assertion hard-codes both the exact instance ids and their ordering. If
the assignment logic preserves correctness but changes ordering (or the
replica-group selection changes while still meeting invariants), this test may
become unnecessarily brittle. Consider asserting per replica-group like other
tests in this file (or asserting set-equivalence / contains-all) to validate
the intended invariant (fallback to computed assignment) without coupling to a
specific order.
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java:
##########
@@ -415,6 +661,31 @@ private static HelixManager createHelixManager() {
return helixManager;
}
+ private static HelixManager createHelixManagerWithTiers(String tierName,
List<String> tierServers) {
+ HelixManager helixManager = mock(HelixManager.class);
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ // Build the tier instance partitions ZK record so getTierInstances() can
fetch it
+ String tierInstancePartitionsName =
+
InstancePartitionsUtils.getInstancePartitionsNameForTier(RAW_TABLE_NAME,
tierName);
+ String tierInstancePartitionsPath =
+
ZKMetadataProvider.constructPropertyStorePathForInstancePartitions(tierInstancePartitionsName);
+ InstancePartitions tierInstancePartitions = new
InstancePartitions(tierInstancePartitionsName);
+ tierInstancePartitions.setInstances(0, 0, tierServers);
+ ZNRecord tierZNRecord = tierInstancePartitions.toZNRecord();
+ // Single stub that branches by path — order-independent, avoids relying
on Mockito's "last stub wins"
+ // behavior when multiple matchers could match the same call.
+ doAnswer(invocation -> {
+ String path = invocation.getArgument(0, String.class);
+ if (path.equals(tierInstancePartitionsPath)) {
+ return tierZNRecord;
+ }
+ String lastComponent = path.substring(path.lastIndexOf('/') + 1);
+ return new ZNRecord(lastComponent);
Review Comment:
For any non-tier path, the mock returns a non-null `ZNRecord`, which can
mask “path missing in ZK” behavior and potentially make unrelated
`fetchInstancePartitions(...)` calls appear to succeed. Consider returning
`null` for unknown paths (or stubbing only the specific path(s) you need) so
the test more accurately reflects HelixPropertyStore behavior and avoids false
positives.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
