Support embedded alluxio cache in hive

docs/src/main/sphinx/connector/filesystem-cache.md

@@ -12,6 +12,7 @@ source [Alluxio](https://github.com/Alluxio/alluxio) libraries with catalogs
 using the following connectors:
 
 * [](/connector/delta-lake)
+* [](/connector/hive)
 
 (fs-cache-distributed)=
 ## Distributed caching

lib/trino-filesystem-cache-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioCoordinatorNoOpFileSystemCache.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.filesystem.alluxio;
+
+import alluxio.client.file.cache.CacheManager;
+import alluxio.conf.AlluxioProperties;
+import alluxio.conf.InstancedConfiguration;
+import com.google.inject.Inject;
+import io.trino.filesystem.Location;
+import io.trino.filesystem.TrinoInput;
+import io.trino.filesystem.TrinoInputFile;
+import io.trino.filesystem.TrinoInputStream;
+import io.trino.filesystem.cache.TrinoFileSystemCache;
+
+import java.io.IOException;
+
+/**
+ * Used to skip caching data on coordinator while still registering alluxio metrics so
+ * that JMX queries for metrics can succeed on the coordinator
+ */
+public class AlluxioCoordinatorNoOpFileSystemCache
+        implements TrinoFileSystemCache
+{
+    @Inject
+    public AlluxioCoordinatorNoOpFileSystemCache()
+    {
+        try {
+            CacheManager cacheManager = CacheManager.Factory.create(new InstancedConfiguration(new AlluxioProperties()));
+            cacheManager.close();
+        }
+        catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public TrinoInput cacheInput(TrinoInputFile delegate, String key)
+            throws IOException
+    {
+        return delegate.newInput();
+    }
+
+    @Override
+    public TrinoInputStream cacheStream(TrinoInputFile delegate, String key)
+            throws IOException
+    {
+        return delegate.newStream();
+    }
+
+    @Override
+    public void expire(Location source)
+            throws IOException
+    {
+    }
+}

lib/trino-filesystem-cache-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystemCacheModule.java

@@ -32,6 +32,13 @@
 public class AlluxioFileSystemCacheModule
         extends AbstractConfigurationAwareModule
 {
+    private final boolean isCoordinator;
+
+    public AlluxioFileSystemCacheModule(boolean isCoordinator)
+    {
+        this.isCoordinator = isCoordinator;
+    }
+
     @Override
     protected void setup(Binder binder)
     {
@@ -40,8 +47,13 @@ protected void setup(Binder binder)
         binder.bind(AlluxioCacheStats.class).in(SINGLETON);
         newExporter(binder).export(AlluxioCacheStats.class).as(generator -> generator.generatedNameOf(AlluxioCacheStats.class));
 
-        binder.bind(TrinoFileSystemCache.class).to(AlluxioFileSystemCache.class).in(SINGLETON);
-        newOptionalBinder(binder, CachingHostAddressProvider.class).setBinding().to(ConsistentHashingHostAddressProvider.class).in(SINGLETON);
+        if (isCoordinator) {
+            binder.bind(TrinoFileSystemCache.class).to(AlluxioCoordinatorNoOpFileSystemCache.class).in(SINGLETON);
+            newOptionalBinder(binder, CachingHostAddressProvider.class).setBinding().to(ConsistentHashingHostAddressProvider.class).in(SINGLETON);
+        }
+        else {
+            binder.bind(TrinoFileSystemCache.class).to(AlluxioFileSystemCache.class).in(SINGLETON);
+        }
 
         Properties metricProps = new Properties();
         metricProps.put("sink.jmx.class", "alluxio.metrics.sink.JmxSink");

lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java

@@ -29,7 +29,7 @@
 import io.trino.filesystem.cache.CacheKeyProvider;
 import io.trino.filesystem.cache.CachingHostAddressProvider;
 import io.trino.filesystem.cache.DefaultCacheKeyProvider;
-import io.trino.filesystem.cache.NoneCachingHostAddressProvider;
+import io.trino.filesystem.cache.DefaultCachingHostAddressProvider;
 import io.trino.filesystem.cache.TrinoFileSystemCache;
 import io.trino.filesystem.gcs.GcsFileSystemFactory;
 import io.trino.filesystem.gcs.GcsFileSystemModule;
@@ -101,7 +101,7 @@ protected void setup(Binder binder)
             factories.addBinding("gs").to(GcsFileSystemFactory.class);
         }
 
-        newOptionalBinder(binder, CachingHostAddressProvider.class).setDefault().to(NoneCachingHostAddressProvider.class).in(Scopes.SINGLETON);
+        newOptionalBinder(binder, CachingHostAddressProvider.class).setDefault().to(DefaultCachingHostAddressProvider.class).in(Scopes.SINGLETON);
         newOptionalBinder(binder, CacheKeyProvider.class).setDefault().to(DefaultCacheKeyProvider.class).in(Scopes.SINGLETON);
         newMapBinder(binder, FileSystemConfig.CacheType.class, TrinoFileSystemCache.class);
 
@@ -110,7 +110,7 @@ protected void setup(Binder binder)
         install(conditionalModule(
                 FileSystemConfig.class,
                 cache -> cache.getCacheType() == FileSystemConfig.CacheType.ALLUXIO,
-                new AlluxioFileSystemCacheModule()));
+                new AlluxioFileSystemCacheModule(nodeManager.getCurrentNode().isCoordinator())));
     }
 
     @Provides

lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/CachingHostAddressProvider.java

@@ -22,5 +22,5 @@ public interface CachingHostAddressProvider
     /**
      * Returns a lists of hosts which are preferred to cache the split with the given path.
      */
-    List<HostAddress> getHosts(String splitPath);
+    List<HostAddress> getHosts(String splitPath, List<HostAddress> defaultAddresses);
 }

lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/ConsistentHashingHostAddressProvider.java

@@ -61,7 +61,7 @@ public ConsistentHashingHostAddressProvider(NodeManager nodeManager, ConsistentH
     }
 
     @Override
-    public List<HostAddress> getHosts(String splitPath)
+    public List<HostAddress> getHosts(String splitPath, List<HostAddress> defaultAddresses)
     {
         return consistentHashRing.locate(splitPath, replicationFactor)
                 .stream()

lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/NoneCachingHostAddressProvider.java → lib/trino-filesystem/src/main/java/io/trino/filesystem/cache/DefaultCachingHostAddressProvider.java

@@ -13,17 +13,16 @@
  */
 package io.trino.filesystem.cache;
 
-import com.google.common.collect.ImmutableList;
 import io.trino.spi.HostAddress;
 
 import java.util.List;
 
-public class NoneCachingHostAddressProvider
+public class DefaultCachingHostAddressProvider
         implements CachingHostAddressProvider
 {
     @Override
-    public List<HostAddress> getHosts(String splitPath)
+    public List<HostAddress> getHosts(String splitPath, List<HostAddress> defaultAddresses)
     {
-        return ImmutableList.of();
+        return defaultAddresses;
     }
 }

lib/trino-filesystem/src/test/java/io/trino/filesystem/cache/TestConsistentHashingCacheHostAddressProvider.java

@@ -13,6 +13,7 @@
  */
 package io.trino.filesystem.cache;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Sets;
 import io.trino.client.NodeVersion;
 import io.trino.metadata.InternalNode;
@@ -85,7 +86,7 @@ private static void assertFairDistribution(CachingHostAddressProvider cachingHos
         int n = 1000;
         Map<String, Integer> counts = new HashMap<>();
         for (int i = 0; i < n; i++) {
-            counts.merge(cachingHostAddressProvider.getHosts(String.valueOf(i)).get(0).getHostText(), 1, Math::addExact);
+            counts.merge(cachingHostAddressProvider.getHosts(String.valueOf(i), ImmutableList.of()).get(0).getHostText(), 1, Math::addExact);
         }
         assertThat(nodeNames.stream().map(m -> m.getHostAndPort().getHostText()).collect(Collectors.toSet())).isEqualTo(counts.keySet());
         counts.values().forEach(c -> assertThat(abs(c - n / nodeNames.size()) < 0.1 * n).isTrue());
@@ -105,7 +106,7 @@ private Map<String, Set<Integer>> getDistribution(ConsistentHashingHostAddressPr
         int n = 1000;
         Map<String, Set<Integer>> distribution = new HashMap<>();
         for (int i = 0; i < n; i++) {
-            String host = provider.getHosts(String.valueOf(i)).get(0).getHostText();
+            String host = provider.getHosts(String.valueOf(i), ImmutableList.of()).get(0).getHostText();
             distribution.computeIfAbsent(host, (k) -> new HashSet<>()).add(i);
         }
         return distribution;

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeSplitManager.java

@@ -333,7 +333,7 @@ private List<DeltaLakeSplit> splitsForFile(
                     addFileEntry.getStats().flatMap(DeltaLakeFileStatistics::getNumRecords),
                     addFileEntry.getModificationTime(),
                     addFileEntry.getDeletionVector(),
-                    cachingHostAddressProvider.getHosts(splitPath),
+                    cachingHostAddressProvider.getHosts(splitPath, ImmutableList.of()),
                     SplitWeight.standard(),
                     statisticsPredicate,
                     partitionKeys));
@@ -359,7 +359,7 @@ private List<DeltaLakeSplit> splitsForFile(
                     Optional.empty(),
                     addFileEntry.getModificationTime(),
                     addFileEntry.getDeletionVector(),
-                    cachingHostAddressProvider.getHosts(splitPath),
+                    cachingHostAddressProvider.getHosts(splitPath, ImmutableList.of()),
                     SplitWeight.fromProportion(clamp((double) splitSize / maxSplitSize, minimumAssignedSplitWeight, 1.0)),
                     statisticsPredicate,
                     partitionKeys));

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeAlluxioCacheFileOperations.java

@@ -71,9 +71,10 @@ protected DistributedQueryRunner createQueryRunner()
                 .buildOrThrow();
 
         DistributedQueryRunner queryRunner = DeltaLakeQueryRunner.builder(session)
+                .setCoordinatorProperties(ImmutableMap.of("node-scheduler.include-coordinator", "false"))
                 .setDeltaProperties(deltaLakeProperties)
                 .setCatalogName(DELTA_CATALOG)
-                .setNodeCount(1)
+                .setNodeCount(2)
                 .build();
 
         queryRunner.execute("CREATE SCHEMA " + session.getSchema().orElseThrow());

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeMetadata.java

@@ -25,7 +25,7 @@
 import io.opentelemetry.api.trace.Tracer;
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.filesystem.cache.CachingHostAddressProvider;
-import io.trino.filesystem.cache.NoneCachingHostAddressProvider;
+import io.trino.filesystem.cache.DefaultCachingHostAddressProvider;
 import io.trino.filesystem.hdfs.HdfsFileSystemFactory;
 import io.trino.hdfs.HdfsEnvironment;
 import io.trino.hdfs.TrinoHdfsFileSystemStats;
@@ -207,7 +207,7 @@ public void setUp()
                     binder.bind(HdfsEnvironment.class).toInstance(HDFS_ENVIRONMENT);
                     binder.bind(TrinoHdfsFileSystemStats.class).toInstance(HDFS_FILE_SYSTEM_STATS);
                     binder.bind(TrinoFileSystemFactory.class).to(HdfsFileSystemFactory.class).in(Scopes.SINGLETON);
-                    binder.bind(CachingHostAddressProvider.class).to(NoneCachingHostAddressProvider.class).in(Scopes.SINGLETON);
+                    binder.bind(CachingHostAddressProvider.class).to(DefaultCachingHostAddressProvider.class).in(Scopes.SINGLETON);
                 },
                 new AbstractModule()
                 {

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSplitManager.java

@@ -20,7 +20,7 @@
 import io.airlift.json.JsonCodecFactory;
 import io.airlift.units.DataSize;
 import io.trino.filesystem.Location;
-import io.trino.filesystem.cache.NoneCachingHostAddressProvider;
+import io.trino.filesystem.cache.DefaultCachingHostAddressProvider;
 import io.trino.filesystem.hdfs.HdfsFileSystemFactory;
 import io.trino.filesystem.memory.MemoryFileSystemFactory;
 import io.trino.plugin.deltalake.statistics.CachingExtendedStatisticsAccess;
@@ -241,7 +241,7 @@ public Stream<AddFileEntry> getActiveFiles(
                 deltaLakeConfig,
                 HDFS_FILE_SYSTEM_FACTORY,
                 deltaLakeTransactionManager,
-                new NoneCachingHostAddressProvider());
+                new DefaultCachingHostAddressProvider());
     }
 
     private AddFileEntry addFileEntryOfSize(long fileSize)

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java

@@ -24,6 +24,7 @@
 import io.airlift.stats.CounterStat;
 import io.airlift.units.DataSize;
 import io.trino.filesystem.TrinoFileSystemFactory;
+import io.trino.filesystem.cache.CachingHostAddressProvider;
 import io.trino.plugin.hive.metastore.Column;
 import io.trino.plugin.hive.metastore.Partition;
 import io.trino.plugin.hive.metastore.SemiTransactionalHiveMetastore;
@@ -114,6 +115,7 @@ public class HiveSplitManager
     private final boolean recursiveDfsWalkerEnabled;
     private final CounterStat highMemorySplitSourceCounter;
     private final TypeManager typeManager;
+    private final CachingHostAddressProvider cachingHostAddressProvider;
     private final int maxPartitionsPerScan;
 
     @Inject
@@ -124,7 +126,8 @@ public HiveSplitManager(
             TrinoFileSystemFactory fileSystemFactory,
             ExecutorService executorService,
             VersionEmbedder versionEmbedder,
-            TypeManager typeManager)
+            TypeManager typeManager,
+            CachingHostAddressProvider cachingHostAddressProvider)
     {
         this(
                 transactionManager,
@@ -141,6 +144,7 @@ public HiveSplitManager(
                 hiveConfig.getMaxSplitsPerSecond(),
                 hiveConfig.getRecursiveDirWalkerEnabled(),
                 typeManager,
+                cachingHostAddressProvider,
                 hiveConfig.getMaxPartitionsPerScan());
     }
 
@@ -159,6 +163,7 @@ public HiveSplitManager(
             @Nullable Integer maxSplitsPerSecond,
             boolean recursiveDfsWalkerEnabled,
             TypeManager typeManager,
+            CachingHostAddressProvider cachingHostAddressProvider,
             int maxPartitionsPerScan)
     {
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
@@ -176,6 +181,7 @@ public HiveSplitManager(
         this.maxSplitsPerSecond = firstNonNull(maxSplitsPerSecond, Integer.MAX_VALUE);
         this.recursiveDfsWalkerEnabled = recursiveDfsWalkerEnabled;
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
+        this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null");
         this.maxPartitionsPerScan = maxPartitionsPerScan;
     }
 
@@ -275,6 +281,7 @@ public ConnectorSplitSource getSplits(
                 hiveSplitLoader,
                 executor,
                 highMemorySplitSourceCounter,
+                cachingHostAddressProvider,
                 hiveTable.isRecordScannedFiles());
         hiveSplitLoader.start(splitSource);
 

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitSource.java

@@ -19,6 +19,7 @@
 import io.airlift.log.Logger;
 import io.airlift.stats.CounterStat;
 import io.airlift.units.DataSize;
+import io.trino.filesystem.cache.CachingHostAddressProvider;
 import io.trino.plugin.hive.InternalHiveSplit.InternalHiveBlock;
 import io.trino.plugin.hive.util.AsyncQueue;
 import io.trino.plugin.hive.util.AsyncQueue.BorrowResult;
@@ -84,6 +85,7 @@ class HiveSplitSource
     private final CounterStat highMemorySplitSourceCounter;
     private final AtomicBoolean loggedHighMemoryWarning = new AtomicBoolean();
     private final HiveSplitWeightProvider splitWeightProvider;
+    private final CachingHostAddressProvider cachingHostAddressProvider;
 
     private final boolean recordScannedFiles;
     private final ImmutableList.Builder<Object> scannedFilePaths = ImmutableList.builder();
@@ -98,6 +100,7 @@ private HiveSplitSource(
             HiveSplitLoader splitLoader,
             AtomicReference<State> stateReference,
             CounterStat highMemorySplitSourceCounter,
+            CachingHostAddressProvider cachingHostAddressProvider,
             boolean recordScannedFiles)
     {
         requireNonNull(session, "session is null");
@@ -114,6 +117,7 @@ private HiveSplitSource(
         this.maxInitialSplitSize = getMaxInitialSplitSize(session);
         this.remainingInitialSplits = new AtomicInteger(maxInitialSplits);
         this.splitWeightProvider = isSizeBasedSplitWeightsEnabled(session) ? new SizeBasedSplitWeightProvider(getMinimumAssignedSplitWeight(session), maxSplitSize) : HiveSplitWeightProvider.uniformStandardWeightProvider();
+        this.cachingHostAddressProvider = requireNonNull(cachingHostAddressProvider, "cachingHostAddressProvider is null");
         this.recordScannedFiles = recordScannedFiles;
     }
 
@@ -128,6 +132,7 @@ public static HiveSplitSource allAtOnce(
             HiveSplitLoader splitLoader,
             Executor executor,
             CounterStat highMemorySplitSourceCounter,
+            CachingHostAddressProvider cachingHostAddressProvider,
             boolean recordScannedFiles)
     {
         AtomicReference<State> stateReference = new AtomicReference<>(State.initial());
@@ -168,6 +173,7 @@ public boolean isFinished()
                 splitLoader,
                 stateReference,
                 highMemorySplitSourceCounter,
+                cachingHostAddressProvider,
                 recordScannedFiles);
     }
 
@@ -305,7 +311,7 @@ else if (maxSplitBytes * 2 >= remainingBlockBytes) {
                         internalSplit.getFileModifiedTime(),
                         internalSplit.getSchema(),
                         internalSplit.getPartitionKeys(),
-                        block.getAddresses(),
+                        cachingHostAddressProvider.getHosts(internalSplit.getPath(), block.getAddresses()),
                         internalSplit.getReadBucketNumber(),
                         internalSplit.getTableBucketNumber(),
                         internalSplit.isForceLocalScheduling(),

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java

@@ -29,6 +29,7 @@
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.filesystem.TrinoInputFile;
 import io.trino.filesystem.TrinoOutputFile;
+import io.trino.filesystem.cache.DefaultCachingHostAddressProvider;
 import io.trino.filesystem.memory.MemoryFileSystemFactory;
 import io.trino.plugin.hive.HiveColumnHandle.ColumnType;
 import io.trino.plugin.hive.fs.CachingDirectoryLister;
@@ -1230,6 +1231,7 @@ private HiveSplitSource hiveSplitSource(HiveSplitLoader hiveSplitLoader)
                 hiveSplitLoader,
                 executor,
                 new CounterStat(),
+                new DefaultCachingHostAddressProvider(),
                 false);
     }