Fix listing of files in AlluxioFileSystem

lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileIterator.java

@@ -24,19 +24,18 @@
 import java.util.List;
 import java.util.Optional;
 
-import static io.trino.filesystem.alluxio.AlluxioUtils.convertToLocation;
 import static java.util.Objects.requireNonNull;
 
 public class AlluxioFileIterator
         implements FileIterator
 {
     private final Iterator<URIStatus> files;
-    private final String mountRoot;
+    private final String basePath;
 
-    public AlluxioFileIterator(List<URIStatus> files, String mountRoot)
+    public AlluxioFileIterator(List<URIStatus> files, String basePath)
     {
         this.files = requireNonNull(files.iterator(), "files is null");
-        this.mountRoot = requireNonNull(mountRoot, "mountRoot is null");
+        this.basePath = requireNonNull(basePath, "basePath is null");
     }
 
     @Override
@@ -54,7 +53,8 @@ public FileEntry next()
             return null;
         }
         URIStatus fileStatus = files.next();
-        Location location = convertToLocation(fileStatus.getPath(), mountRoot);
+        String filePath = fileStatus.getPath();
+        Location location = Location.of(basePath + filePath);
         return new FileEntry(
                 location,
                 fileStatus.getLength(),

lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystem.java

@@ -38,6 +38,7 @@
 import java.util.stream.Collectors;
 
 import static io.trino.filesystem.alluxio.AlluxioUtils.convertToAlluxioURI;
+import static io.trino.filesystem.alluxio.AlluxioUtils.getAlluxioBase;
 import static java.util.Objects.requireNonNull;
 import static java.util.UUID.randomUUID;
 
@@ -189,20 +190,20 @@ public FileIterator listFiles(Location location)
         try {
             URIStatus status = alluxioClient.getStatus(convertToAlluxioURI(location, mountRoot));
             if (status == null) {
-                new AlluxioFileIterator(Collections.emptyList(), mountRoot);
+                new AlluxioFileIterator(Collections.emptyList(), getAlluxioBase(location.toString()));
             }
             if (!status.isFolder()) {
                 throw new IOException("Location is not a directory: %s".formatted(location));
             }
         }
         catch (NotFoundRuntimeException | AlluxioException e) {
-            return new AlluxioFileIterator(Collections.emptyList(), mountRoot);
+            return new AlluxioFileIterator(Collections.emptyList(), getAlluxioBase(location.toString()));
         }
 
         try {
             List<URIStatus> filesStatus = alluxioClient.listStatus(convertToAlluxioURI(location, mountRoot),
                     ListStatusPOptions.newBuilder().setRecursive(true).build());
-            return new AlluxioFileIterator(filesStatus.stream().filter(status -> !status.isFolder() & status.isCompleted()).toList(), mountRoot);
+            return new AlluxioFileIterator(filesStatus.stream().filter(status -> !status.isFolder() & status.isCompleted()).toList(), getAlluxioBase(location.toString()));
         }
         catch (AlluxioException e) {
             throw new IOException("Error listFiles %s".formatted(location), e);

lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioUtils.java

@@ -50,6 +50,16 @@ public static Location convertToLocation(String path, String mountRoot)
         return Location.of(schema + mountRootWithSlash + path);
     }
 
+    public static String getAlluxioBase(String path)
+    {
+        requireNonNull(path, "path is null");
+        if (!path.startsWith("alluxio://")) {
+            throw new IllegalArgumentException("path is not an alluxio://");
+        }
+        int index = path.indexOf('/', "alluxio://".length());
+        return path.substring(0, index);
+    }
+
     public static String simplifyPath(String path)
     {
         // Use a deque to store the path components

lib/trino-filesystem-alluxio/src/test/java/io/trino/filesystem/alluxio/TestAlluxioFileIterator.java

@@ -0,0 +1,62 @@
+/*
+ * 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.URIStatus;
+import alluxio.wire.FileInfo;
+import io.trino.filesystem.FileEntry;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+final class TestAlluxioFileIterator
+{
+    private final String fileName = "000000_0";
+    private final String filePath = "/s3a/tables/sales/000000_0";
+    private final String ufsFilePath = "s3a://test-bucket/tables/sales/000000_0";
+
+    @Test
+    void testNext()
+            throws IOException
+    {
+        String alluxioBasePath = "alluxio://master:19998";
+        FileInfo fileInfo = new FileInfo();
+        fileInfo.setName(fileName);
+        fileInfo.setPath(filePath);
+        fileInfo.setUfsPath(ufsFilePath);
+        URIStatus fileStatus = new URIStatus(fileInfo);
+        AlluxioFileIterator iterator = new AlluxioFileIterator(
+                List.of(fileStatus),
+                alluxioBasePath);
+        FileEntry fileEntry = iterator.next();
+        assertThat(fileEntry.location().toString())
+                .isEqualTo(alluxioBasePath + filePath);
+
+        alluxioBasePath = "alluxio:/";
+        fileInfo = new FileInfo();
+        fileInfo.setName(fileName);
+        fileInfo.setPath(filePath);
+        fileInfo.setUfsPath(ufsFilePath);
+        fileStatus = new URIStatus(fileInfo);
+        iterator = new AlluxioFileIterator(
+                List.of(fileStatus),
+                alluxioBasePath);
+        fileEntry = iterator.next();
+        assertThat(fileEntry.location().toString())
+                .isEqualTo(alluxioBasePath + filePath);
+    }
+}