You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ja...@apache.org on 2025/05/27 17:37:47 UTC

(pinot) branch master updated: [bugfix] Remove deleting segments from table status info (#15725)

This is an automated email from the ASF dual-hosted git repository.

jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new a61acb71b6 [bugfix] Remove deleting segments from table status info (#15725)
a61acb71b6 is described below

commit a61acb71b634da10e637b61fc76e0b8aaa142659
Author: Soumya Himanish Mohapatra <30...@users.noreply.github.com>
AuthorDate: Tue May 27 23:07:42 2025 +0530

    [bugfix] Remove deleting segments from table status info (#15725)
---
 .../controller/api/resources/DebugResource.java    |  2 +-
 .../api/resources/PinotSegmentRestletResource.java |  2 +-
 .../pinot/controller/api/resources/TableSize.java  |  5 ++-
 .../pinot/controller/api/resources/TableViews.java | 43 ++++++++++++++-----
 .../controller/helix/SegmentStatusChecker.java     |  2 +-
 .../helix/core/PinotHelixResourceManager.java      | 17 ++++++--
 .../helix/core/rebalance/TableRebalancer.java      |  2 +-
 .../pinot/controller/util/TableSizeReader.java     | 48 ++++++++++++++--------
 .../controller/validation/StorageQuotaChecker.java |  2 +-
 .../src/main/resources/app/requests/index.ts       |  8 ++--
 .../pinot/controller/api/TableSizeReaderTest.java  |  8 ++--
 .../resources/PinotSegmentRestletResourceTest.java |  4 +-
 .../controller/helix/SegmentStatusCheckerTest.java | 15 ++++---
 .../validation/StorageQuotaCheckerTest.java        |  2 +-
 14 files changed, 107 insertions(+), 53 deletions(-)

diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/DebugResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/DebugResource.java
index 906070a3b1..6a2235389c 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/DebugResource.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/DebugResource.java
@@ -240,7 +240,7 @@ public class DebugResource {
     TableSizeReader.TableSizeDetails tableSizeDetails;
     try {
       tableSizeDetails = _tableSizeReader
-          .getTableSizeDetails(tableNameWithType, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
+          .getTableSizeDetails(tableNameWithType, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000, true);
     } catch (Throwable t) {
       tableSizeDetails = null;
     }
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java
index 7455ac5f4e..17de0fca8f 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java
@@ -631,7 +631,7 @@ public class PinotSegmentRestletResource {
       @Nullable String instanceName) {
     if (segmentNames == null) {
       // instanceName can be null or not null, and this method below can handle both cases.
-      return _pinotHelixResourceManager.getServerToSegmentsMap(tableNameWithType, instanceName);
+      return _pinotHelixResourceManager.getServerToSegmentsMap(tableNameWithType, instanceName, true);
     }
     // Skip servers and segments not involved in the segment reloading job.
     List<String> segmnetNameList = new ArrayList<>();
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableSize.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableSize.java
index ede9d3eb73..5995450200 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableSize.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableSize.java
@@ -85,12 +85,15 @@ public class TableSize {
       @ApiParam(value = "Table name without type", required = true, example = "myTable | myTable_OFFLINE")
       @PathParam("tableName") String tableName,
       @ApiParam(value = "Provide detailed information") @DefaultValue("true") @QueryParam("verbose") boolean verbose,
+      @ApiParam(value = "Include replaced segments") @DefaultValue("true")
+      @QueryParam("includeReplacedSegments") boolean includeReplacedSegments,
       @Context HttpHeaders headers) {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     TableSizeReader.TableSizeDetails tableSizeDetails = null;
     try {
       tableSizeDetails =
-          _tableSizeReader.getTableSizeDetails(tableName, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000);
+          _tableSizeReader.getTableSizeDetails(tableName, _controllerConf.getServerAdminRequestTimeoutSeconds() * 1000,
+              includeReplacedSegments);
       if (!verbose) {
         if (tableSizeDetails._offlineSegments != null) {
           tableSizeDetails._offlineSegments._segments = new HashMap<>();
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java
index 41bda73670..9e71660dc3 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java
@@ -32,10 +32,12 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
+import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
@@ -48,6 +50,9 @@ import javax.ws.rs.core.Response;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
+import org.apache.pinot.common.lineage.SegmentLineage;
+import org.apache.pinot.common.lineage.SegmentLineageAccessHelper;
+import org.apache.pinot.common.lineage.SegmentLineageUtils;
 import org.apache.pinot.common.utils.DatabaseUtils;
 import org.apache.pinot.controller.api.exception.ControllerApplicationException;
 import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
@@ -141,28 +146,44 @@ public class TableViews {
   public String getSegmentsStatusDetails(
       @ApiParam(value = "Name of the table", required = true) @PathParam("tableName") String tableName,
       @ApiParam(value = "realtime|offline", required = false) @QueryParam("tableType") String tableTypeStr,
+      @ApiParam(value = "Include segments being replaced", required = false)
+      @QueryParam("includeReplacedSegments") @DefaultValue("true") boolean includeReplacedSegments,
       @Context HttpHeaders headers)
       throws JsonProcessingException {
     tableName = DatabaseUtils.translateTableName(tableName, headers);
     TableType tableType = validateTableType(tableTypeStr);
     TableViews.TableView externalView = getTableState(tableName, TableViews.EXTERNALVIEW, tableType);
     TableViews.TableView idealStateView = getTableState(tableName, TableViews.IDEALSTATE, tableType);
-    List<SegmentStatusInfo> segmentStatusInfoListMap = new ArrayList<>();
-    segmentStatusInfoListMap = getSegmentStatuses(externalView, idealStateView);
+
+    Map<String, Map<String, String>> externalViewStateMap = getStateMap(externalView);
+    Map<String, Map<String, String>> idealStateMap = getStateMap(idealStateView);
+    Set<String> segments = idealStateMap.keySet();
+
+    if (!includeReplacedSegments) {
+      SegmentLineage segmentLineage = SegmentLineageAccessHelper
+          .getSegmentLineage(_pinotHelixResourceManager.getPropertyStore(), tableName);
+      SegmentLineageUtils
+          .filterSegmentsBasedOnLineageInPlace(segments, segmentLineage);
+    }
+
+    List<SegmentStatusInfo> segmentStatusInfoListMap = getSegmentStatuses(externalViewStateMap, idealStateMap);
+
     return JsonUtils.objectToPrettyString(segmentStatusInfoListMap);
   }
 
-  public List<SegmentStatusInfo> getSegmentStatuses(TableViews.TableView externalView,
-      TableViews.TableView idealStateView) {
-    Map<String, Map<String, String>> idealStateMap = getStateMap(idealStateView);
-    Map<String, Map<String, String>> externalViewMap = getStateMap(externalView);
+  public List<SegmentStatusInfo> getSegmentStatuses(Map<String, Map<String, String>> externalViewMap,
+      Map<String, Map<String, String>> idealStateMap) {
     List<SegmentStatusInfo> segmentStatusInfoList = new ArrayList<>();
 
-    for (Map.Entry<String, Map<String, String>> entry : externalViewMap.entrySet()) {
+    for (Map.Entry<String, Map<String, String>> entry : idealStateMap.entrySet()) {
       String segment = entry.getKey();
-      Map<String, String> externalViewEntryValue = entry.getValue();
-      Map<String, String> idealViewEntryValue = idealStateMap.get(segment);
-      if (isErrorSegment(externalViewEntryValue)) {
+      Map<String, String> externalViewEntryValue = externalViewMap.get(segment);
+      Map<String, String> idealViewEntryValue = entry.getValue();
+
+      if (externalViewEntryValue == null) {
+        segmentStatusInfoList.add(
+            new SegmentStatusInfo(segment, CommonConstants.Helix.StateModel.DisplaySegmentStatus.UPDATING));
+      } else if (isErrorSegment(externalViewEntryValue)) {
         segmentStatusInfoList.add(
             new SegmentStatusInfo(segment, CommonConstants.Helix.StateModel.DisplaySegmentStatus.BAD));
       } else {
@@ -210,7 +231,7 @@ public class TableViews {
     return tableTypeViewResult;
   }
 
-  private Map<String, Map<String, String>> getStateMap(TableViews.TableView view) {
+  public Map<String, Map<String, String>> getStateMap(TableViews.TableView view) {
     if (view != null && view._offline != null && !view._offline.isEmpty()) {
       return view._offline;
     } else if (view != null && view._realtime != null && !view._realtime.isEmpty()) {
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java
index bc3f2fb263..a86bbaf10a 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java
@@ -202,7 +202,7 @@ public class SegmentStatusChecker extends ControllerPeriodicTask<SegmentStatusCh
 
   private void updateTableSizeMetrics(String tableNameWithType)
       throws InvalidConfigException {
-    _tableSizeReader.getTableSizeDetails(tableNameWithType, TABLE_CHECKER_TIMEOUT_MS);
+    _tableSizeReader.getTableSizeDetails(tableNameWithType, TABLE_CHECKER_TIMEOUT_MS, true);
   }
 
   /**
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
index 09bd904aee..969aab1c2f 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
@@ -3274,16 +3274,27 @@ public class PinotHelixResourceManager {
    * the ideal state because they are not supposed to be served.
    */
   public Map<String, List<String>> getServerToSegmentsMap(String tableNameWithType) {
-    return getServerToSegmentsMap(tableNameWithType, null);
+    return getServerToSegmentsMap(tableNameWithType, null, true);
   }
 
-  public Map<String, List<String>> getServerToSegmentsMap(String tableNameWithType, @Nullable String targetServer) {
+  public Map<String, List<String>> getServerToSegmentsMap(String tableNameWithType, @Nullable String targetServer,
+      boolean includeReplacedSegments) {
     Map<String, List<String>> serverToSegmentsMap = new TreeMap<>();
     IdealState idealState = _helixAdmin.getResourceIdealState(_helixClusterName, tableNameWithType);
     if (idealState == null) {
       throw new IllegalStateException("Ideal state does not exist for table: " + tableNameWithType);
     }
-    for (Map.Entry<String, Map<String, String>> entry : idealState.getRecord().getMapFields().entrySet()) {
+
+    Map<String, Map<String, String>> idealStateMap = idealState.getRecord().getMapFields();
+    Set<String> segments = idealStateMap.keySet();
+
+    if (!includeReplacedSegments) {
+      SegmentLineage segmentLineage = SegmentLineageAccessHelper
+          .getSegmentLineage(getPropertyStore(), tableNameWithType);
+      SegmentLineageUtils.filterSegmentsBasedOnLineageInPlace(segments, segmentLineage);
+    }
+
+    for (Map.Entry<String, Map<String, String>> entry : idealStateMap.entrySet()) {
       String segmentName = entry.getKey();
       for (Map.Entry<String, String> instanceStateEntry : entry.getValue().entrySet()) {
         String server = instanceStateEntry.getKey();
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java
index a89f62b41e..d96b646089 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java
@@ -658,7 +658,7 @@ public class TableRebalancer {
     tableRebalanceLogger.info("Fetching the table size");
     try {
       TableSizeReader.TableSubTypeSizeDetails sizeDetails =
-          _tableSizeReader.getTableSubtypeSize(tableNameWithType, TABLE_SIZE_READER_TIMEOUT_MS);
+          _tableSizeReader.getTableSubtypeSize(tableNameWithType, TABLE_SIZE_READER_TIMEOUT_MS, true);
       tableRebalanceLogger.info("Fetched the table size details");
       return sizeDetails;
     } catch (InvalidConfigException e) {
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableSizeReader.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableSizeReader.java
index 24148bc69d..eb134cbfbc 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableSizeReader.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/util/TableSizeReader.java
@@ -23,6 +23,7 @@ import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.BiMap;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -76,10 +77,12 @@ public class TableSizeReader {
    * Returns null if the table is not found.
    * @param tableName table name without type
    * @param timeoutMsec timeout in milliseconds for reading table sizes from server
+   * @param includeReplacedSegments include replaced segments in table size calculation
    * @return
    */
   @Nullable
-  public TableSizeDetails getTableSizeDetails(@Nonnull String tableName, @Nonnegative int timeoutMsec)
+  public TableSizeDetails getTableSizeDetails(@Nonnull String tableName, @Nonnegative int timeoutMsec,
+      boolean includeReplacedSegments)
       throws InvalidConfigException {
     Preconditions.checkNotNull(tableName, "Table name should not be null");
     Preconditions.checkArgument(timeoutMsec > 0, "Timeout value must be greater than 0");
@@ -100,7 +103,7 @@ public class TableSizeReader {
     TableSizeDetails tableSizeDetails = new TableSizeDetails(tableName);
     if (hasRealtimeTableConfig) {
       String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(tableName);
-      tableSizeDetails._realtimeSegments = getTableSubtypeSize(realtimeTableName, timeoutMsec);
+      tableSizeDetails._realtimeSegments = getTableSubtypeSize(realtimeTableName, timeoutMsec, includeReplacedSegments);
       // taking max(0,value) as values as set to -1 if all the segments are in error
       tableSizeDetails._reportedSizeInBytes += Math.max(tableSizeDetails._realtimeSegments._reportedSizeInBytes, 0L);
       tableSizeDetails._estimatedSizeInBytes += Math.max(tableSizeDetails._realtimeSegments._estimatedSizeInBytes, 0L);
@@ -126,7 +129,7 @@ public class TableSizeReader {
     }
     if (hasOfflineTableConfig) {
       String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(tableName);
-      tableSizeDetails._offlineSegments = getTableSubtypeSize(offlineTableName, timeoutMsec);
+      tableSizeDetails._offlineSegments = getTableSubtypeSize(offlineTableName, timeoutMsec, includeReplacedSegments);
       // taking max(0,value) as values as set to -1 if all the segments are in error
       tableSizeDetails._reportedSizeInBytes += Math.max(tableSizeDetails._offlineSegments._reportedSizeInBytes, 0L);
       tableSizeDetails._estimatedSizeInBytes += Math.max(tableSizeDetails._offlineSegments._estimatedSizeInBytes, 0L);
@@ -241,9 +244,11 @@ public class TableSizeReader {
     public Map<String, SegmentSizeInfo> _serverInfo = new HashMap<>();
   }
 
-  public TableSubTypeSizeDetails getTableSubtypeSize(String tableNameWithType, int timeoutMs)
+  public TableSubTypeSizeDetails getTableSubtypeSize(String tableNameWithType, int timeoutMs,
+      boolean includeReplacedSegments)
       throws InvalidConfigException {
-    Map<String, List<String>> serverToSegmentsMap = _helixResourceManager.getServerToSegmentsMap(tableNameWithType);
+    Map<String, List<String>> serverToSegmentsMap = _helixResourceManager.getServerToSegmentsMap(tableNameWithType,
+        null, includeReplacedSegments);
     ServerTableSizeReader serverTableSizeReader = new ServerTableSizeReader(_executor, _connectionManager);
     BiMap<String, String> endpoints = _helixResourceManager.getDataInstanceAdminEndpoints(serverToSegmentsMap.keySet());
     Map<String, List<SegmentSizeInfo>> serverToSegmentSizeInfoListMap =
@@ -252,23 +257,30 @@ public class TableSizeReader {
     TableSubTypeSizeDetails subTypeSizeDetails = new TableSubTypeSizeDetails();
     Map<String, SegmentSizeDetails> segmentToSizeDetailsMap = subTypeSizeDetails._segments;
 
+    Map<String, Map<String, SegmentSizeInfo>> serverToSegmentMap = new HashMap<>();
+    for (Map.Entry<String, List<SegmentSizeInfo>> entry : serverToSegmentSizeInfoListMap.entrySet()) {
+      Map<String, SegmentSizeInfo> segmentMap = new HashMap<>();
+      for (SegmentSizeInfo info : entry.getValue()) {
+        segmentMap.put(info.getSegmentName(), info);
+      }
+      serverToSegmentMap.put(entry.getKey(), segmentMap);
+    }
+
     // Convert map from (server -> List<SegmentSizeInfo>) to (segment -> SegmentSizeDetails (server -> SegmentSizeInfo))
     // If no response returned from a server, put -1 as size for all the segments on the server
-    // TODO: here we assume server contains all segments in ideal state, which might not be the case
     for (Map.Entry<String, List<String>> entry : serverToSegmentsMap.entrySet()) {
       String server = entry.getKey();
-      List<SegmentSizeInfo> segmentSizeInfoList = serverToSegmentSizeInfoListMap.get(server);
-      if (segmentSizeInfoList != null) {
-        for (SegmentSizeInfo segmentSizeInfo : segmentSizeInfoList) {
-          SegmentSizeDetails segmentSizeDetails =
-              segmentToSizeDetailsMap.computeIfAbsent(segmentSizeInfo.getSegmentName(), k -> new SegmentSizeDetails());
-          segmentSizeDetails._serverInfo.put(server, segmentSizeInfo);
-        }
-      } else {
-        List<String> segments = entry.getValue();
-        for (String segment : segments) {
-          SegmentSizeDetails segmentSizeDetails =
-              segmentToSizeDetailsMap.computeIfAbsent(segment, k -> new SegmentSizeDetails());
+      Map<String, SegmentSizeInfo> segmentToSegmentSizeInfoMap = serverToSegmentMap.getOrDefault(server,
+          Collections.emptyMap());
+
+      List<String> segments = entry.getValue();
+      for (String segment : segments) {
+        SegmentSizeDetails segmentSizeDetails =
+            segmentToSizeDetailsMap.computeIfAbsent(segment, k -> new SegmentSizeDetails());
+
+        if (segmentToSegmentSizeInfoMap.containsKey(segment)) {
+          segmentSizeDetails._serverInfo.put(server, segmentToSegmentSizeInfoMap.get(segment));
+        } else {
           segmentSizeDetails._serverInfo.put(server, new SegmentSizeInfo(segment, DEFAULT_SIZE_WHEN_MISSING_OR_ERROR));
         }
       }
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
index 37fbda8698..3c0359c960 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/validation/StorageQuotaChecker.java
@@ -110,7 +110,7 @@ public class StorageQuotaChecker {
     // read table size
     TableSizeReader.TableSubTypeSizeDetails tableSubtypeSize;
     try {
-      tableSubtypeSize = _tableSizeReader.getTableSubtypeSize(tableNameWithType, _timeoutMs);
+      tableSubtypeSize = _tableSizeReader.getTableSubtypeSize(tableNameWithType, _timeoutMs, true);
     } catch (InvalidConfigException e) {
       LOGGER.error("Failed to get table size for table {}", tableNameWithType, e);
       throw e;
diff --git a/pinot-controller/src/main/resources/app/requests/index.ts b/pinot-controller/src/main/resources/app/requests/index.ts
index dc08e09cb7..42ecfccbcb 100644
--- a/pinot-controller/src/main/resources/app/requests/index.ts
+++ b/pinot-controller/src/main/resources/app/requests/index.ts
@@ -98,8 +98,8 @@ export const putSchema = (name: string, params: string, reload?: boolean): Promi
 export const getSegmentMetadata = (tableName: string, segmentName: string): Promise<AxiosResponse<SegmentMetadata>> =>
   baseApi.get(`/segments/${tableName}/${segmentName}/metadata?columns=*`);
 
-export const getTableSize = (name: string, verbose: boolean = false): Promise<AxiosResponse<TableSize>> =>
-  baseApi.get(`/tables/${name}/size?verbose=${verbose}`);
+export const getTableSize = (name: string, verbose: boolean = false, includeReplacedSegments: boolean = false): Promise<AxiosResponse<TableSize>> =>
+  baseApi.get(`/tables/${name}/size?verbose=${verbose}&includeReplacedSegments=${includeReplacedSegments}`);
 
 export const getIdealState = (name: string): Promise<AxiosResponse<IdealState>> =>
   baseApi.get(`/tables/${name}/idealstate`);
@@ -110,8 +110,8 @@ export const getExternalView = (name: string): Promise<AxiosResponse<IdealState>
 export const getServerToSegmentsCount = (name: string, tableType: TableType, verbose: boolean = false): Promise<AxiosResponse<ServerToSegmentsCount[]>> =>
   baseApi.get(`/segments/${name}/servers?type=${tableType}&verbose=${verbose}`);
 
-export const getSegmentsStatus = (name: string): Promise<AxiosResponse<SegmentStatusInfo[]>> =>
-  baseApi.get(`/tables/${name}/segmentsStatus`);
+export const getSegmentsStatus = (name: string, includeReplacedSegments: boolean = false): Promise<AxiosResponse<SegmentStatusInfo[]>> =>
+  baseApi.get(`/tables/${name}/segmentsStatus?includeReplacedSegments=${includeReplacedSegments}`);
 
 // Fetch consuming segments information for a table
 // API: GET /tables/{tableName}/consumingSegmentsInfo
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/TableSizeReaderTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/TableSizeReaderTest.java
index 1c7731da4b..75b2b11099 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/TableSizeReaderTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/TableSizeReaderTest.java
@@ -61,6 +61,8 @@ import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -226,12 +228,12 @@ public class TableSizeReaderTest {
       throws InvalidConfigException {
     TableSizeReader reader =
         new TableSizeReader(_executor, _connectionManager, _controllerMetrics, _helix, _leadControllerManager);
-    assertNull(reader.getTableSizeDetails("mytable", 5000));
+    assertNull(reader.getTableSizeDetails("mytable", 5000, true));
   }
 
   private TableSizeReader.TableSizeDetails testRunner(final String[] servers, String table)
       throws InvalidConfigException {
-    when(_helix.getServerToSegmentsMap(anyString())).thenAnswer(new Answer<Object>() {
+    when(_helix.getServerToSegmentsMap(anyString(), any(), anyBoolean())).thenAnswer(new Answer<Object>() {
       @Override
       public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
         return subsetOfServerSegments(servers);
@@ -247,7 +249,7 @@ public class TableSizeReaderTest {
 
     TableSizeReader reader = new TableSizeReader(_executor, _connectionManager, _controllerMetrics, _helix,
         _leadControllerManager);
-    return reader.getTableSizeDetails(table, TIMEOUT_MSEC);
+    return reader.getTableSizeDetails(table, TIMEOUT_MSEC, true);
   }
 
   private Map<String, List<String>> segmentToServers(final String... servers) {
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java
index d7350a17ed..9117467bdf 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java
@@ -60,8 +60,8 @@ public class PinotSegmentRestletResourceTest {
     fullServerToSegmentsMap.put("svr01", new ArrayList<>(List.of("seg01", "seg02")));
     fullServerToSegmentsMap.put("svr02", new ArrayList<>(List.of("seg02", "seg03")));
     fullServerToSegmentsMap.put("svr03", new ArrayList<>(List.of("seg03", "seg01")));
-    when(_pinotHelixResourceManager.getServerToSegmentsMap(tableName, null)).thenReturn(fullServerToSegmentsMap);
-    when(_pinotHelixResourceManager.getServerToSegmentsMap(tableName, "svr02")).thenReturn(
+    when(_pinotHelixResourceManager.getServerToSegmentsMap(tableName, null, true)).thenReturn(fullServerToSegmentsMap);
+    when(_pinotHelixResourceManager.getServerToSegmentsMap(tableName, "svr02", true)).thenReturn(
         Map.of("svr02", new ArrayList<>(List.of("seg02", "seg03"))));
     when(_pinotHelixResourceManager.getServers(tableName, "seg01")).thenReturn(Set.of("svr01", "svr03"));
 
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/SegmentStatusCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/SegmentStatusCheckerTest.java
index 84ceb96843..f2a949051e 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/SegmentStatusCheckerTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/SegmentStatusCheckerTest.java
@@ -780,7 +780,8 @@ public class SegmentStatusCheckerTest {
     tableViewExternal._offline = tableViewExternalOffline;
     tableViewIdeal._offline = tableViewIdealOffline;
     TableViews tableviews = new TableViews();
-    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(tableViewExternal, tableViewIdeal);
+    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(
+        tableviews.getStateMap(tableViewExternal), tableviews.getStateMap(tableViewIdeal));
     assertEquals(segmentStatusInfos.get(0).getSegmentStatus(),
         CommonConstants.Helix.StateModel.DisplaySegmentStatus.GOOD);
     assertEquals(segmentStatusInfos.get(1).getSegmentStatus(),
@@ -804,7 +805,8 @@ public class SegmentStatusCheckerTest {
     tableViewExternal._offline = tableViewExternalOffline;
     tableViewIdeal._offline = tableViewIdealOffline;
     TableViews tableviews = new TableViews();
-    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(tableViewExternal, tableViewIdeal);
+    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(
+        tableviews.getStateMap(tableViewExternal), tableviews.getStateMap(tableViewIdeal));
     assertEquals(segmentStatusInfos.get(0).getSegmentStatus(),
         CommonConstants.Helix.StateModel.DisplaySegmentStatus.GOOD);
     assertEquals(segmentStatusInfos.get(1).getSegmentStatus(),
@@ -828,7 +830,8 @@ public class SegmentStatusCheckerTest {
     tableViewExternal._offline = tableViewExternalOffline;
     tableViewIdeal._offline = tableViewIdealOffline;
     TableViews tableviews = new TableViews();
-    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(tableViewExternal, tableViewIdeal);
+    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(
+        tableviews.getStateMap(tableViewExternal), tableviews.getStateMap(tableViewIdeal));
     assertEquals(segmentStatusInfos.get(0).getSegmentStatus(),
         CommonConstants.Helix.StateModel.DisplaySegmentStatus.BAD);
     assertEquals(segmentStatusInfos.get(1).getSegmentStatus(),
@@ -852,7 +855,8 @@ public class SegmentStatusCheckerTest {
     tableViewExternal._offline = tableViewExternalOffline;
     tableViewIdeal._offline = tableViewIdealOffline;
     TableViews tableviews = new TableViews();
-    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(tableViewExternal, tableViewIdeal);
+    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(
+        tableviews.getStateMap(tableViewExternal), tableviews.getStateMap(tableViewIdeal));
     assertEquals(segmentStatusInfos.get(0).getSegmentStatus(),
         CommonConstants.Helix.StateModel.DisplaySegmentStatus.UPDATING);
     assertEquals(segmentStatusInfos.get(1).getSegmentStatus(),
@@ -880,7 +884,8 @@ public class SegmentStatusCheckerTest {
     tableViewExternal._offline = tableViewExternalOffline;
     tableViewIdeal._offline = tableViewIdealOffline;
     TableViews tableviews = new TableViews();
-    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(tableViewExternal, tableViewIdeal);
+    List<SegmentStatusInfo> segmentStatusInfos = tableviews.getSegmentStatuses(
+        tableviews.getStateMap(tableViewExternal), tableviews.getStateMap(tableViewIdeal));
     assertEquals(segmentStatusInfos.get(0).getSegmentStatus(),
         CommonConstants.Helix.StateModel.DisplaySegmentStatus.GOOD);
     assertEquals(segmentStatusInfos.get(1).getSegmentStatus(),
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
index 3e42be2f50..9e6b896143 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/validation/StorageQuotaCheckerTest.java
@@ -192,6 +192,6 @@ public class StorageQuotaCheckerTest {
     tableSizeResult._estimatedSizeInBytes = tableSizeInBytes;
     tableSizeResult._segments = Collections.emptyMap();
     tableSizeResult._missingSegments = numMissingSegments;
-    when(_tableSizeReader.getTableSubtypeSize(tableName, 1000)).thenReturn(tableSizeResult);
+    when(_tableSizeReader.getTableSubtypeSize(tableName, 1000, true)).thenReturn(tableSizeResult);
   }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org