From 30ac4c8fc6d3b511c87d89b9e7d274d554bd1930 Mon Sep 17 00:00:00 2001
From: tanyajun <tanyajun538@gmail.com>
Date: Thu, 17 Oct 2024 15:41:05 +0800
Subject: [PATCH] fix the data pull problem when use alluxio as fs.

---
 .../trino/filesystem/alluxio/AlluxioFileIterator.java | 11 ++++++-----
 .../trino/filesystem/alluxio/AlluxioFileSystem.java   |  6 +++---
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileIterator.java b/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileIterator.java
index e67a51f16e90ed..105e7158d311c2 100644
--- a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileIterator.java
+++ b/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 Location pathLocation;
 
-    public AlluxioFileIterator(List<URIStatus> files, String mountRoot)
+    public AlluxioFileIterator(List<URIStatus> files, Location pathLocation)
     {
         this.files = requireNonNull(files.iterator(), "files is null");
-        this.mountRoot = requireNonNull(mountRoot, "mountRoot is null");
+        this.pathLocation = requireNonNull(pathLocation, "pathLocation is null");
     }
 
     @Override
@@ -54,7 +53,9 @@ public FileEntry next()
             return null;
         }
         URIStatus fileStatus = files.next();
-        Location location = convertToLocation(fileStatus.getPath(), mountRoot);
+        String filePath = fileStatus.getPath();
+        String fileName = filePath.substring(filePath.lastIndexOf('/') + 1);
+        Location location = pathLocation.appendSuffix("/" + fileName);
         return new FileEntry(
                 location,
                 fileStatus.getLength(),
diff --git a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystem.java b/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystem.java
index b6a949ee44d629..73c3dfe0b56e6a 100644
--- a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystem.java
+++ b/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioFileSystem.java
@@ -189,20 +189,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(), location);
             }
             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(), location);
         }
 
         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(), location);
         }
         catch (AlluxioException e) {
             throw new IOException("Error listFiles %s".formatted(location), e);