[BugFix] Bound and expose the Iceberg partition cache memory (backport #76165) (#76181)

Co-authored-by: Youngwb <yangwenbo_mailbox@163.com>

Author: mergify[bot]

fe/fe-core/src/main/java/com/starrocks/connector/iceberg/CachingIcebergCatalog.java

@@ -144,9 +144,13 @@ public Table reload(IcebergTableName key, Table oldValue) {
                         }
                     }
                 });
-        this.partitionCache = newCacheBuilderWithMaximumSize(
-                icebergProperties.getIcebergMetaCacheTtlSec(), NEVER_CACHE,
-                enableCache ? DEFAULT_CACHE_NUM : NEVER_CACHE).build(
+        long partitionCacheSize = Math.round(Runtime.getRuntime().maxMemory() *
+                icebergProperties.getIcebergPartitionCacheMemoryUsageRatio());
+        this.partitionCache = newCacheBuilder(icebergProperties.getIcebergMetaCacheTtlSec(), NEVER_CACHE)
+                .executor(executorService)
+                .maximumWeight(partitionCacheSize)
+                .weigher((Weigher<IcebergTableName, Map<String, Partition>>) this::weighPartitionEntry)
+                .build(
                     new com.github.benmanes.caffeine.cache.CacheLoader<IcebergTableName, Map<String, Partition>>() {
                         @Override
                         public Map<String, Partition> load(IcebergTableName key) throws Exception {
@@ -610,6 +614,17 @@ private int weighTableEntry(IcebergTableName key, Table table) {
         return (int) size;
     }
 
+    // Each entry holds a whole table's partition map, which is unbounded in bytes (a table can have
+    // millions of partitions). Sample the partition values so weighing stays cheap on large maps.
+    private int weighPartitionEntry(IcebergTableName key, Map<String, Partition> partitions) {
+        long size = Estimator.estimate(key) + Estimator.estimate(partitions, 3);
+        if (size > Integer.MAX_VALUE) {
+            LOG.warn("Partition cache entry size for key {} is too large: {} bytes", key, size);
+            return Integer.MAX_VALUE;
+        }
+        return (int) size;
+    }
+
     private long estimateContentFilesSize(Set<? extends ContentFile<?>> files) {
         if (files == null || files.isEmpty()) {
             return 0L;
@@ -654,7 +669,11 @@ public long estimateSize() {
         return Estimator.estimate(tables.asMap(), 10) +
                 Estimator.estimate(databases.asMap(), 10) +
                 estimateDataFileCacheSize() +
-                estimateDeleteFileCacheSize() + Estimator.estimate(metaFileCacheMap, 10);
+                estimateDeleteFileCacheSize() +
+                Estimator.estimate(partitionCache.asMap(), 10) +
+                Estimator.estimate(metaFileCacheMap, 10) +
+                Estimator.estimate(tableLatestAccessTime, 10) +
+                Estimator.estimate(tableLatestRefreshTime, 10);
     }
 
     private long estimateDataFileCacheSize() {

fe/fe-core/src/main/java/com/starrocks/connector/iceberg/IcebergCatalogProperties.java

@@ -45,6 +45,7 @@ public class IcebergCatalogProperties {
     public static final String ICEBERG_DATA_FILE_CACHE_MEMORY_SIZE_RATIO = "iceberg_data_file_cache_memory_usage_ratio";
     public static final String ICEBERG_DELETE_FILE_CACHE_MEMORY_SIZE_RATIO = "iceberg_delete_file_cache_memory_usage_ratio";
     public static final String ICEBERG_TABLE_CACHE_MEMORY_SIZE_RATIO = "iceberg_table_cache_memory_usage_ratio";
+    public static final String ICEBERG_PARTITION_CACHE_MEMORY_SIZE_RATIO = "iceberg_partition_cache_memory_usage_ratio";
 
     // internal config
     public static final String REFRESH_ICEBERG_MANIFEST_MIN_LENGTH = "refresh_iceberg_manifest_min_length";
@@ -70,6 +71,7 @@ public class IcebergCatalogProperties {
     private double icebergDataFileCacheMemoryUsageRatio;
     private double icebergDeleteFileCacheMemoryUsageRatio;
     private double icebergTableCacheMemoryUsageRatio;
+    private double icebergPartitionCacheMemoryUsageRatio;
     private long icebergTableCacheRefreshIntervalSec;
 
     public IcebergCatalogProperties(Map<String, String> catalogProperties) {
@@ -111,6 +113,8 @@ private void initIcebergMetadataCache() {
                     properties, ICEBERG_DELETE_FILE_CACHE_MEMORY_SIZE_RATIO, 0.1);
         this.icebergTableCacheMemoryUsageRatio = PropertyUtil.propertyAsDouble(
                     properties, ICEBERG_TABLE_CACHE_MEMORY_SIZE_RATIO, 0.1);
+        this.icebergPartitionCacheMemoryUsageRatio = PropertyUtil.propertyAsDouble(
+                    properties, ICEBERG_PARTITION_CACHE_MEMORY_SIZE_RATIO, 0.1);
         this.icebergManifestCacheWithColumnStatistics = PropertyUtil.propertyAsBoolean(
                 properties, ICEBERG_MANIFEST_CACHE_WITH_COLUMN_STATISTICS, true);
         this.refreshIcebergManifestMinLength = PropertyUtil.propertyAsLong(properties, REFRESH_ICEBERG_MANIFEST_MIN_LENGTH,
@@ -193,6 +197,10 @@ public double getIcebergTableCacheMemoryUsageRatio() {
         return icebergTableCacheMemoryUsageRatio;
     }
 
+    public double getIcebergPartitionCacheMemoryUsageRatio() {
+        return icebergPartitionCacheMemoryUsageRatio;
+    }
+
     public long getRefreshIcebergManifestMinLength() {
         return refreshIcebergManifestMinLength;
     }

fe/fe-core/src/test/java/com/starrocks/connector/iceberg/CachingIcebergCatalogTest.java

@@ -214,6 +214,37 @@ public Table getTable(ConnectContext context, String dbName, String tableName) {
         }
     }
 
+    @Test
+    public void testPartitionCacheCountedInEstimateSize(@Mocked IcebergCatalog icebergCatalog) {
+        PartitionSpec spec = Mockito.mock(PartitionSpec.class);
+        Mockito.when(spec.isUnpartitioned()).thenReturn(false);
+        Table nativeTable = createBaseTableWithManifests(1, 0, spec);
+        Map<String, Partition> partitionMap = new HashMap<>();
+        for (int i = 0; i < 1000; i++) {
+            partitionMap.put("dt=part-" + i, new Partition(1234L, 1L));
+        }
+        new Expectations() {
+            {
+                icebergCatalog.getTable((ConnectContext) any, "db", "test");
+                result = nativeTable;
+                minTimes = 0;
+                icebergCatalog.getPartitions((IcebergTable) any, anyLong, null);
+                result = partitionMap;
+                minTimes = 0;
+            }
+        };
+        CachingIcebergCatalog cachingIcebergCatalog = new CachingIcebergCatalog(CATALOG_NAME, icebergCatalog,
+                DEFAULT_CATALOG_PROPERTIES, Executors.newSingleThreadExecutor());
+        IcebergTable table = IcebergTable.builder().setSrTableName("test")
+                .setCatalogDBName("db").setCatalogTableName("test").setNativeTable(nativeTable).build();
+
+        long before = cachingIcebergCatalog.estimateSize();
+        cachingIcebergCatalog.getPartitions(table, 1L, null);
+        long after = cachingIcebergCatalog.estimateSize();
+        Assertions.assertTrue(after > before,
+                "partitionCache must be counted in estimateSize; before=" + before + " after=" + after);
+    }
+
     @Test
     public void testGetDB(@Mocked IcebergCatalog icebergCatalog, @Mocked Database db) {
         new Expectations() {