feat: add dataset caching in worker node (#24)

jackye1995 · web-flow · commit c60f84e43624 · 2025-12-17T11:25:27.000-08:00
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -18,8 +18,8 @@ jobs:
         distribution: 'temurin'
         cache: maven
 
-    - name: Build and test with Maven
-      run: ./mvnw clean install -Dair.check.skip-enforcer=true
+    - name: Build and test
+      run: make install
 
     - name: Upload test results
       uses: actions/upload-artifact@v4
diff --git a/Makefile b/Makefile
@@ -0,0 +1,59 @@
+.PHONY: build test clean install compile package help run lint format check
+
+# Default target
+help:
+	@echo "Available targets:"
+	@echo "  build     - Build the project (compile + package)"
+	@echo "  test      - Run tests"
+	@echo "  install   - Build and install to local Maven repository"
+	@echo "  compile   - Compile source code"
+	@echo "  package   - Package the plugin"
+	@echo "  clean     - Clean build artifacts"
+	@echo "  verify    - Run full verification (compile, test, package)"
+	@echo "  run       - Run the development query runner server"
+	@echo "  lint      - Run all code style checks (checkstyle, modernizer, sortpom)"
+	@echo "  format    - Format pom.xml files"
+	@echo "  check     - Run lint checks without tests"
+
+# Build the project
+build: compile package
+
+# Run tests
+test:
+	./mvnw test -Dair.check.skip-enforcer=true
+
+# Install to local repository
+install:
+	./mvnw clean install -Dair.check.skip-enforcer=true
+
+# Compile source code
+compile:
+	./mvnw compile -Dair.check.skip-enforcer=true
+
+# Package the plugin
+package:
+	./mvnw package -Dair.check.skip-enforcer=true -DskipTests
+
+# Clean build artifacts
+clean:
+	./mvnw clean
+
+# Full verification
+verify:
+	./mvnw verify -Dair.check.skip-enforcer=true
+
+# Run the development query runner server
+run:
+	./mvnw exec:java -pl plugin/trino-lance -Dexec.mainClass="io.trino.plugin.lance.LanceQueryRunner" -Dair.check.skip-enforcer=true
+
+# Run all code style checks (Trino standard)
+lint:
+	./mvnw checkstyle:check sortpom:verify modernizer:modernizer -Dair.check.skip-enforcer=true
+
+# Format pom.xml files
+format:
+	./mvnw sortpom:sort
+
+# Run all checks without tests
+check:
+	./mvnw compile checkstyle:check sortpom:verify modernizer:modernizer -Dair.check.skip-enforcer=true -DskipTests
diff --git a/README.md b/README.md
@@ -72,7 +72,23 @@ GROUP BY category;
 Build the connector:
 
 ```bash
-./mvnw clean install
+make install
+```
+
+Other available commands:
+
+```bash
+make help      # Show all available commands
+make build     # Compile and package
+make test      # Run tests
+make compile   # Compile only
+make package   # Package without tests
+make clean     # Clean build artifacts
+make verify    # Full verification
+make run       # Run development server
+make lint      # Run code style checks (checkstyle, modernizer, sortpom)
+make format    # Format pom.xml files
+make check     # Run all checks without tests
 ```
 
 ## Installation
@@ -93,16 +109,15 @@ We periodically upgrade the Trino version to stay up to date with the latest Tri
 ### Running Tests
 
 ```bash
-./mvnw test -pl plugin/trino-lance
+make test
 ```
 
 ### Running the Query Runner (Development Server)
 
 You can run a local Trino server for development:
 
 ```bash
-cd plugin/trino-lance
-mvn exec:java -Dexec.mainClass="io.trino.plugin.lance.LanceQueryRunner"
+make run
 ```
 
 This starts a Trino server on port 8080 with the Lance connector configured.
diff --git a/plugin/trino-lance/pom.xml b/plugin/trino-lance/pom.xml
@@ -53,6 +53,11 @@
             <artifactId>units</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-cache</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-plugin-toolkit</artifactId>
diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/LanceFragmentPageSource.java
@@ -17,7 +17,6 @@
 import io.trino.plugin.lance.internal.LanceReader;
 import io.trino.plugin.lance.internal.ScannerFactory;
 import org.apache.arrow.memory.BufferAllocator;
-import org.lance.Dataset;
 import org.lance.Fragment;
 import org.lance.ipc.LanceScanner;
 import org.lance.ipc.ScanOptions;
@@ -46,7 +45,6 @@ public static class FragmentScannerFactory
             implements ScannerFactory
     {
         private final int fragmentId;
-        private Dataset lanceDataset;
         private Fragment lanceFragment;
         private LanceScanner lanceScanner;
 
@@ -58,12 +56,11 @@ public FragmentScannerFactory(int fragmentId)
         @Override
         public LanceScanner open(String tablePath, BufferAllocator allocator, List<String> columns)
         {
-            this.lanceDataset = Dataset.open(tablePath, allocator);
-            // Find fragment by ID, not by list index (fragment IDs may not match list positions)
-            this.lanceFragment = lanceDataset.getFragments().stream()
-                    .filter(f -> f.getId() == this.fragmentId)
-                    .findFirst()
-                    .orElseThrow(() -> new RuntimeException("Fragment not found: " + this.fragmentId));
+            // Use cached fragment lookup instead of opening dataset and filtering
+            this.lanceFragment = LanceReader.getFragment(tablePath, this.fragmentId);
+            if (this.lanceFragment == null) {
+                throw new RuntimeException("Fragment not found: " + this.fragmentId);
+            }
             ScanOptions.Builder optionsBuilder = new ScanOptions.Builder();
             // Only set columns if non-empty; empty list means read all columns
             if (!columns.isEmpty()) {
@@ -76,6 +73,7 @@ public LanceScanner open(String tablePath, BufferAllocator allocator, List<Strin
         @Override
         public void close()
         {
+            // Only close the scanner; the dataset is managed by LanceReader's cache
             try {
                 if (lanceScanner != null) {
                     lanceScanner.close();
@@ -84,9 +82,6 @@ public void close()
             catch (Exception e) {
                 log.warn("error while closing lance scanner, Exception: %s", e.getMessage());
             }
-            if (lanceDataset != null) {
-                lanceDataset.close();
-            }
         }
     }
 }
diff --git a/plugin/trino-lance/src/main/java/io/trino/plugin/lance/internal/LanceReader.java b/plugin/trino-lance/src/main/java/io/trino/plugin/lance/internal/LanceReader.java
@@ -13,7 +13,12 @@
  */
 package io.trino.plugin.lance.internal;
 
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
 import com.google.inject.Inject;
+import io.airlift.log.Logger;
+import io.trino.cache.NonEvictableLoadingCache;
+import io.trino.cache.SafeCaches;
 import io.trino.plugin.lance.LanceColumnHandle;
 import io.trino.plugin.lance.LanceConfig;
 import io.trino.plugin.lance.LanceNamespaceProperties;
@@ -38,20 +43,62 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 public class LanceReader
         implements Closeable
 {
+    private static final Logger log = Logger.get(LanceReader.class);
+
     // TODO: support multiple schemas
     public static final String SCHEMA = "default";
     private static final String TABLE_PATH_SUFFIX = ".lance";
     private static final BufferAllocator allocator = new RootAllocator(
             RootAllocator.configBuilder().from(RootAllocator.defaultConfig()).maxAllocation(4 * 1024 * 1024).build());
 
+    // Cache for dataset metadata (fragments) - shared across all LanceReader instances per worker JVM
+    // Maximum 100 entries, expires 1 hour after last access (same as lance-spark)
+    private static final NonEvictableLoadingCache<CacheKey, Map<Integer, Fragment>> FRAGMENT_CACHE =
+            SafeCaches.buildNonEvictableCache(
+                    CacheBuilder.newBuilder()
+                            .maximumSize(100)
+                            .expireAfterAccess(1, TimeUnit.HOURS),
+                    new CacheLoader<>()
+                    {
+                        @Override
+                        public Map<Integer, Fragment> load(CacheKey key)
+                        {
+                            log.debug("Loading fragments for table: %s", key.getTablePath());
+                            Dataset dataset = Dataset.open(key.getTablePath(), allocator);
+                            return dataset.getFragments().stream()
+                                    .collect(Collectors.toMap(Fragment::getId, f -> f));
+                        }
+                    });
+
+    // Cache for schema metadata - shared across all LanceReader instances per worker JVM
+    private static final NonEvictableLoadingCache<CacheKey, Schema> SCHEMA_CACHE =
+            SafeCaches.buildNonEvictableCache(
+                    CacheBuilder.newBuilder()
+                            .maximumSize(100)
+                            .expireAfterAccess(1, TimeUnit.HOURS),
+                    new CacheLoader<>()
+                    {
+                        @Override
+                        public Schema load(CacheKey key)
+                        {
+                            log.debug("Loading schema for table: %s", key.getTablePath());
+                            try (Dataset dataset = Dataset.open(key.getTablePath(), allocator)) {
+                                return dataset.getSchema();
+                            }
+                        }
+                    });
+
     private final String root;
     private final LanceNamespace namespace;
 
@@ -138,17 +185,46 @@ public List<Fragment> getFragments(LanceTableHandle tableHandle)
         return getFragments(getTablePath(tableHandle.getTableName()));
     }
 
+    /**
+     * Get a specific fragment by ID from the cache.
+     * This is useful for workers that need to access a specific fragment for data reading.
+     *
+     * @param tablePath the path to the lance table
+     * @param fragmentId the fragment ID to retrieve
+     * @return the Fragment object, or null if not found
+     */
+    public static Fragment getFragment(String tablePath, int fragmentId)
+    {
+        try {
+            CacheKey key = new CacheKey(tablePath);
+            Map<Integer, Fragment> fragments = FRAGMENT_CACHE.get(key);
+            return fragments.get(fragmentId);
+        }
+        catch (ExecutionException e) {
+            throw new RuntimeException("Failed to get fragment from cache for table: " + tablePath, e);
+        }
+    }
+
     private static List<Fragment> getFragments(String tablePath)
     {
-        try (Dataset dataset = Dataset.open(tablePath, allocator)) {
-            return dataset.getFragments();
+        try {
+            CacheKey key = new CacheKey(tablePath);
+            Map<Integer, Fragment> fragmentMap = FRAGMENT_CACHE.get(key);
+            return List.copyOf(fragmentMap.values());
+        }
+        catch (ExecutionException e) {
+            throw new RuntimeException("Failed to get fragments from cache for table: " + tablePath, e);
         }
     }
 
     private static Schema getSchema(String tablePath)
     {
-        try (Dataset dataset = Dataset.open(tablePath, allocator)) {
-            return dataset.getSchema();
+        try {
+            CacheKey key = new CacheKey(tablePath);
+            return SCHEMA_CACHE.get(key);
+        }
+        catch (ExecutionException e) {
+            throw new RuntimeException("Failed to get schema from cache for table: " + tablePath, e);
         }
     }
 
@@ -176,4 +252,39 @@ public void close()
             }
         }
     }
+
+    /**
+     * Cache key for dataset metadata caching.
+     * Uses table path as the unique identifier.
+     */
+    private static class CacheKey
+    {
+        private final String tablePath;
+
+        CacheKey(String tablePath)
+        {
+            this.tablePath = tablePath;
+        }
+
+        public String getTablePath()
+        {
+            return tablePath;
+        }
+
+        @Override
+        public boolean equals(Object o)
+        {
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            CacheKey cacheKey = (CacheKey) o;
+            return Objects.equals(tablePath, cacheKey.tablePath);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(tablePath);
+        }
+    }
 }
