From 93e7556460b99fe693c993e9c56d75350c684e9e Mon Sep 17 00:00:00 2001
From: JingsongLi <jingsonglee0@gmail.com>
Date: Mon, 27 Jul 2026 23:47:56 +0800
Subject: [PATCH] [core] Support off-heap lookup cache

---
 .../apache/paimon/io/cache/CacheManager.java  | 35 +++++++++++++++++--
 .../paimon/io/cache/CacheManagerTest.java     | 17 +++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/paimon-common/src/main/java/org/apache/paimon/io/cache/CacheManager.java b/paimon-common/src/main/java/org/apache/paimon/io/cache/CacheManager.java
index 073c3db6c32e..e8cd314725d5 100644
--- a/paimon-common/src/main/java/org/apache/paimon/io/cache/CacheManager.java
+++ b/paimon-common/src/main/java/org/apache/paimon/io/cache/CacheManager.java
@@ -31,7 +31,7 @@
 import static org.apache.paimon.utils.Preconditions.checkNotNull;
 
 /** Cache manager to cache bytes to paged {@link MemorySegment}s. */
-public class CacheManager {
+public class CacheManager implements AutoCloseable {
 
     private static final Logger LOG = LoggerFactory.getLogger(CacheManager.class);
 
@@ -43,8 +43,13 @@ public class CacheManager {
 
     private final Cache dataCache;
     private final Cache indexCache;
+    private final boolean offHeap;
 
     public CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio) {
+        this(maxMemorySize, highPriorityPoolRatio, false);
+    }
+
+    private CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio, boolean offHeap) {
         Preconditions.checkArgument(
                 highPriorityPoolRatio >= 0 && highPriorityPoolRatio < 1,
                 "The high priority pool ratio should in the range [0, 1).");
@@ -58,12 +63,19 @@ public CacheManager(MemorySize maxMemorySize, double highPriorityPoolRatio) {
         } else {
             this.indexCache = CacheBuilder.newBuilder().maximumWeight(indexCacheSize).build();
         }
+        this.offHeap = offHeap;
         LOG.info(
-                "Initialize cache manager with data cache of {} and index cache of {}.",
+                "Initialize {} cache manager with data cache of {} and index cache of {}.",
+                offHeap ? "off-heap" : "heap",
                 dataCacheSize,
                 indexCacheSize);
     }
 
+    public static CacheManager createOffHeap(
+            MemorySize maxMemorySize, double highPriorityPoolRatio) {
+        return new CacheManager(maxMemorySize, highPriorityPoolRatio, true);
+    }
+
     @VisibleForTesting
     public Cache dataCache() {
         return dataCache;
@@ -82,7 +94,7 @@ public MemorySegment getPage(CacheKey key, CacheReader reader, CacheCallback cal
                         k -> {
                             try {
                                 return new Cache.CacheValue(
-                                        MemorySegment.wrap(reader.read(key)), callback);
+                                        toMemorySegment(reader.read(key)), callback);
                             } catch (IOException e) {
                                 throw new RuntimeException(e);
                             }
@@ -98,6 +110,23 @@ public void invalidPage(CacheKey key) {
         }
     }
 
+    private MemorySegment toMemorySegment(byte[] bytes) {
+        if (!offHeap) {
+            return MemorySegment.wrap(bytes);
+        }
+        MemorySegment segment = MemorySegment.allocateOffHeapMemory(bytes.length);
+        segment.put(0, bytes);
+        return segment;
+    }
+
+    @Override
+    public void close() {
+        dataCache.invalidateAll();
+        if (indexCache != dataCache) {
+            indexCache.invalidateAll();
+        }
+    }
+
     /** The container for the segment. */
     public static class SegmentContainer {
 
diff --git a/paimon-common/src/test/java/org/apache/paimon/io/cache/CacheManagerTest.java b/paimon-common/src/test/java/org/apache/paimon/io/cache/CacheManagerTest.java
index 6045956a6d72..e8ca35271e35 100644
--- a/paimon-common/src/test/java/org/apache/paimon/io/cache/CacheManagerTest.java
+++ b/paimon-common/src/test/java/org/apache/paimon/io/cache/CacheManagerTest.java
@@ -66,4 +66,21 @@ void testCaffeineCache() throws Exception {
             }
         }
     }
+
+    @Test
+    void testOffHeapCache() throws Exception {
+        File file = new File(tempDir.toFile(), "test.off-heap");
+        assertThat(file.createNewFile()).isTrue();
+        CacheKey key = CacheKey.forPageIndex(new RandomAccessFile(file, "r"), 0, 0);
+
+        try (CacheManager cacheManager = CacheManager.createOffHeap(MemorySize.ofBytes(10), 0)) {
+            MemorySegment segment =
+                    cacheManager.getPage(key, ignored -> new byte[] {1, 2, 3}, ignored -> {});
+
+            assertThat(segment.isOffHeap()).isTrue();
+            byte[] bytes = new byte[3];
+            segment.get(0, bytes);
+            assertThat(bytes).containsExactly(1, 2, 3);
+        }
+    }
 }