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 3e526ad48c86..d69225bdc038 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 @@ -19,17 +19,20 @@ package org.apache.pinot.controller.validation; import com.google.common.base.Preconditions; +import java.util.Set; +import javax.inject.Inject; import org.apache.pinot.common.exception.InvalidConfigException; import org.apache.pinot.common.metrics.ControllerGauge; import org.apache.pinot.common.metrics.ControllerMetrics; +import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.util.TableSizeReader; import org.apache.pinot.spi.config.table.QuotaConfig; import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TenantConfig; import org.apache.pinot.spi.utils.DataSizeUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; - /** * Class to check if a new segment is within the configured storage quota for the table * @@ -37,6 +40,9 @@ public class StorageQuotaChecker { private static final Logger LOGGER = LoggerFactory.getLogger(StorageQuotaChecker.class); + @Inject + PinotHelixResourceManager _pinotHelixResourceManager; + private final TableSizeReader _tableSizeReader; private final TableConfig _tableConfig; private final ControllerMetrics _controllerMetrics; @@ -80,7 +86,18 @@ public QuotaCheckerResponse isSegmentStorageWithinQuota(String segmentName, long // 3. update predicted segment sizes // 4. is the updated size within quota QuotaConfig quotaConfig = _tableConfig.getQuotaConfig(); - int numReplicas = _tableConfig.getValidationConfig().getReplicationNumber(); + int numReplicas; + + if (_tableConfig.isDimTable()) { + // If the table is a dimension table then fetch the tenant config and get the number of server belonging + // to the tenant + TenantConfig tenantConfig = _tableConfig.getTenantConfig(); + Set serverInstances = _pinotHelixResourceManager.getAllInstancesForServerTenant(tenantConfig.getServer()); + numReplicas = serverInstances.size(); + } else { + numReplicas = _tableConfig.getValidationConfig().getReplicationNumber(); + } + final String tableNameWithType = _tableConfig.getTableName(); if (quotaConfig == null || quotaConfig.getStorage() == null) {