From 26d7b0e565d2144218edb085208886d64af27a5d Mon Sep 17 00:00:00 2001 From: tanyajun <2498618501@qq.com> Date: Mon, 4 Nov 2024 09:14:43 +0800 Subject: [PATCH] Fix listing of files in AlluxioFileSystem Co-authored by: JiamingMai --- .../alluxio/AlluxioFileIterator.java | 10 +-- .../filesystem/alluxio/AlluxioFileSystem.java | 7 ++- .../filesystem/alluxio/AlluxioUtils.java | 10 +++ .../alluxio/TestAlluxioFileIterator.java | 62 +++++++++++++++++++ 4 files changed, 81 insertions(+), 8 deletions(-) create mode 100644 lib/trino-filesystem-alluxio/src/test/java/io/trino/filesystem/alluxio/TestAlluxioFileIterator.java 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 e67a51f16e90..d20bc4cc2aff 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 files; - private final String mountRoot; + private final String basePath; - public AlluxioFileIterator(List files, String mountRoot) + public AlluxioFileIterator(List 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(), 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 b6a949ee44d6..94654da3f1b2 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 @@ -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 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); diff --git a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioUtils.java b/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioUtils.java index 49ece36a8b93..ac95c970dd36 100644 --- a/lib/trino-filesystem-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioUtils.java +++ b/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 diff --git a/lib/trino-filesystem-alluxio/src/test/java/io/trino/filesystem/alluxio/TestAlluxioFileIterator.java b/lib/trino-filesystem-alluxio/src/test/java/io/trino/filesystem/alluxio/TestAlluxioFileIterator.java new file mode 100644 index 000000000000..6ae971673887 --- /dev/null +++ b/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); + } +}