Move Hive directory recursion into HiveFileIterator

presto-hive/src/main/java/com/facebook/presto/hive/BackgroundHiveSplitLoader.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.hive.metastore.Partition;
 import com.facebook.presto.hive.metastore.Table;
 import com.facebook.presto.hive.util.HiveFileIterator;
+import com.facebook.presto.hive.util.HiveFileIterator.NestedDirectoryNotAllowedException;
 import com.facebook.presto.hive.util.InternalHiveSplitFactory;
 import com.facebook.presto.hive.util.ResumableTask;
 import com.facebook.presto.hive.util.ResumableTasks;
@@ -26,6 +27,8 @@
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.predicate.TupleDomain;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterators;
+import com.google.common.collect.Streams;
 import com.google.common.io.CharStreams;
 import com.google.common.util.concurrent.ListenableFuture;
 import org.apache.hadoop.conf.Configuration;
@@ -49,6 +52,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Deque;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
 import java.util.Properties;
@@ -58,7 +62,6 @@
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.IntPredicate;
 
-import static com.facebook.presto.hadoop.HadoopFileStatus.isDirectory;
 import static com.facebook.presto.hive.HiveBucketing.HiveBucket;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_BAD_DATA;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_INVALID_BUCKET_FILES;
@@ -69,6 +72,9 @@
 import static com.facebook.presto.hive.HiveUtil.getInputFormat;
 import static com.facebook.presto.hive.metastore.MetastoreUtil.getHiveSchema;
 import static com.facebook.presto.hive.util.ConfigurationUtils.toJobConf;
+import static com.facebook.presto.hive.util.HiveFileIterator.NestedDirectoryPolicy.FAIL;
+import static com.facebook.presto.hive.util.HiveFileIterator.NestedDirectoryPolicy.IGNORED;
+import static com.facebook.presto.hive.util.HiveFileIterator.NestedDirectoryPolicy.RECURSE;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
@@ -96,7 +102,7 @@ public class BackgroundHiveSplitLoader
     private final Executor executor;
     private final ConnectorSession session;
     private final ConcurrentLazyQueue<HivePartitionMetadata> partitions;
-    private final Deque<HiveFileIterator> fileIterators = new ConcurrentLinkedDeque<>();
+    private final Deque<Iterator<InternalHiveSplit>> fileIterators = new ConcurrentLinkedDeque<>();
 
     // Purpose of this lock:
     // * When write lock is acquired, except the holder, no one can do any of the following:
@@ -210,32 +216,20 @@ private void invokeNoMoreSplitsIfNecessary()
     private ListenableFuture<?> loadSplits()
             throws IOException
     {
-        HiveFileIterator files = fileIterators.poll();
-        if (files == null) {
+        Iterator<InternalHiveSplit> splits = fileIterators.poll();
+        if (splits == null) {
             HivePartitionMetadata partition = partitions.poll();
             if (partition == null) {
                 return COMPLETED_FUTURE;
             }
             return loadPartition(partition);
         }
 
-        while (files.hasNext() && !stopped) {
-            LocatedFileStatus file = files.next();
-            if (isDirectory(file)) {
-                if (recursiveDirWalkerEnabled) {
-                    fileIterators.add(files.withPath(file.getPath()));
-                }
-            }
-            else {
-                Optional<InternalHiveSplit> internalHiveSplit = files.getSplitFactory().createInternalHiveSplit(file);
-                if (!internalHiveSplit.isPresent()) {
-                    continue;
-                }
-                ListenableFuture<?> future = hiveSplitSource.addToQueue(internalHiveSplit.get());
-                if (!future.isDone()) {
-                    fileIterators.addFirst(files);
-                    return future;
-                }
+        while (splits.hasNext() && !stopped) {
+            ListenableFuture<?> future = hiveSplitSource.addToQueue(splits.next());
+            if (!future.isDone()) {
+                fileIterators.addFirst(splits);
+                return future;
             }
         }
 
@@ -301,7 +295,7 @@ private ListenableFuture<?> loadPartition(HivePartitionMetadata partition)
             return hiveSplitSource.addToQueue(getBucketedSplits(path, fs, splitFactory, bucketSplitInfo.get()));
         }
 
-        fileIterators.addLast(new HiveFileIterator(path, fs, directoryLister, namenodeStats, partitionName, splitFactory));
+        fileIterators.addLast(createInternalHiveSplitIterator(path, fs, splitFactory));
         return COMPLETED_FUTURE;
     }
 
@@ -329,37 +323,44 @@ private static boolean shouldUseFileSplitsFromInputFormat(InputFormat<?, ?> inpu
                 .anyMatch(name -> name.equals("UseFileSplitsFromInputFormat"));
     }
 
+    private Iterator<InternalHiveSplit> createInternalHiveSplitIterator(Path path, FileSystem fileSystem, InternalHiveSplitFactory splitFactory)
+    {
+        return Streams.stream(new HiveFileIterator(path, fileSystem, directoryLister, namenodeStats, recursiveDirWalkerEnabled ? RECURSE : IGNORED))
+                .map(splitFactory::createInternalHiveSplit)
+                .filter(Optional::isPresent)
+                .map(Optional::get)
+                .iterator();
+    }
+
     private List<InternalHiveSplit> getBucketedSplits(Path path, FileSystem fileSystem, InternalHiveSplitFactory splitFactory, BucketSplitInfo bucketSplitInfo)
     {
         int bucketCount = bucketSplitInfo.getBucketCount();
 
         // list all files in the partition
-        ArrayList<LocatedFileStatus> list = new ArrayList<>(bucketCount);
-        HiveFileIterator hiveFileIterator = new HiveFileIterator(path, fileSystem, directoryLister, namenodeStats, splitFactory.getPartitionName(), splitFactory);
-        while (hiveFileIterator.hasNext()) {
-            LocatedFileStatus next = hiveFileIterator.next();
-            if (isDirectory(next)) {
-                // Fail here to be on the safe side. This seems to be the same as what Hive does
-                throw new PrestoException(HIVE_INVALID_BUCKET_FILES, format("%s Found sub-directory in bucket directory for partition: %s", CORRUPT_BUCKETING, splitFactory.getPartitionName()));
-            }
-            list.add(next);
+        ArrayList<LocatedFileStatus> files = new ArrayList<>(bucketCount);
+        try {
+            Iterators.addAll(files, new HiveFileIterator(path, fileSystem, directoryLister, namenodeStats, FAIL));
+        }
+        catch (NestedDirectoryNotAllowedException e) {
+            // Fail here to be on the safe side. This seems to be the same as what Hive does
+            throw new PrestoException(HIVE_INVALID_BUCKET_FILES, format("%s Found sub-directory in bucket directory for partition: %s", CORRUPT_BUCKETING, splitFactory.getPartitionName()));
         }
 
         // verify we found one file per bucket
-        if (list.size() != bucketCount) {
-            throw new PrestoException(HIVE_INVALID_BUCKET_FILES, format("%s The number of files in the directory (%s) does not match the declared bucket count (%s) for partition: %s", CORRUPT_BUCKETING, list.size(),
+        if (files.size() != bucketCount) {
+            throw new PrestoException(HIVE_INVALID_BUCKET_FILES, format("%s The number of files in the directory (%s) does not match the declared bucket count (%s) for partition: %s", CORRUPT_BUCKETING, files.size(),
                     bucketCount,
                     splitFactory.getPartitionName()));
         }
 
         // Sort FileStatus objects (instead of, e.g., fileStatus.getPath().toString). This matches org.apache.hadoop.hive.ql.metadata.Table.getSortedPaths
-        list.sort(null);
+        files.sort(null);
 
         // convert files internal splits
         List<InternalHiveSplit> splitList = new ArrayList<>();
         for (int bucketNumber = 0; bucketNumber < bucketCount; bucketNumber++) {
             if (bucketSplitInfo.isBucketEnabled(bucketNumber)) {
-                LocatedFileStatus file = list.get(bucketNumber);
+                LocatedFileStatus file = files.get(bucketNumber);
                 splitFactory.createInternalHiveSplit(file, bucketNumber)
                         .ifPresent(splitList::add);
             }

presto-hive/src/main/java/com/facebook/presto/hive/util/HiveFileIterator.java

@@ -25,103 +25,154 @@
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Collections;
+import java.util.Deque;
+import java.util.Iterator;
 
+import static com.facebook.presto.hadoop.HadoopFileStatus.isDirectory;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_FILE_NOT_FOUND;
 import static java.util.Objects.requireNonNull;
 
 public class HiveFileIterator
         extends AbstractIterator<LocatedFileStatus>
 {
+    public enum NestedDirectoryPolicy
+    {
+        IGNORED,
+        RECURSE,
+        FAIL
+    }
+
+    private final Deque<Path> paths = new ArrayDeque<>();
     private final FileSystem fileSystem;
     private final DirectoryLister directoryLister;
     private final NamenodeStats namenodeStats;
-    private final Path path;
-    private final String partitionName;
-    private final InternalHiveSplitFactory splitFactory;
+    private final NestedDirectoryPolicy nestedDirectoryPolicy;
 
-    private RemoteIterator<LocatedFileStatus> remoteIterator;
+    private Iterator<LocatedFileStatus> remoteIterator = Collections.emptyIterator();
 
     public HiveFileIterator(
             Path path,
             FileSystem fileSystem,
             DirectoryLister directoryLister,
             NamenodeStats namenodeStats,
-            String partitionName,
-            InternalHiveSplitFactory splitFactory)
+            NestedDirectoryPolicy nestedDirectoryPolicy)
     {
-        this.partitionName = requireNonNull(partitionName, "partitionName is null");
-        this.path = requireNonNull(path, "path is null");
+        paths.addLast(requireNonNull(path, "path is null"));
         this.fileSystem = requireNonNull(fileSystem, "fileSystem is null");
         this.directoryLister = requireNonNull(directoryLister, "directoryLister is null");
         this.namenodeStats = requireNonNull(namenodeStats, "namenodeStats is null");
-        this.splitFactory = requireNonNull(splitFactory, "splitFactory is null");
+        this.nestedDirectoryPolicy = requireNonNull(nestedDirectoryPolicy, "nestedDirectoryPolicy is null");
     }
 
     @Override
     protected LocatedFileStatus computeNext()
     {
-        try {
-            if (remoteIterator == null) {
-                remoteIterator = getLocatedFileStatusRemoteIterator(path);
-            }
-
+        while (true) {
             while (remoteIterator.hasNext()) {
                 LocatedFileStatus status = getLocatedFileStatus(remoteIterator);
 
-                // ignore hidden files. Hive ignores files starting with _ and . as well.
+                // Ignore hidden files and directories. Hive ignores files starting with _ and . as well.
                 String fileName = status.getPath().getName();
                 if (fileName.startsWith("_") || fileName.startsWith(".")) {
                     continue;
                 }
+
+                if (isDirectory(status)) {
+                    switch (nestedDirectoryPolicy) {
+                        case IGNORED:
+                            continue;
+                        case RECURSE:
+                            paths.add(status.getPath());
+                            continue;
+                        case FAIL:
+                            throw new NestedDirectoryNotAllowedException();
+                    }
+                }
+
                 return status;
             }
-            return endOfData();
-        }
-        catch (FileNotFoundException e) {
-            throw new PrestoException(HIVE_FILE_NOT_FOUND, "Partition location does not exist: " + path);
-        }
-        catch (IOException e) {
-            throw new PrestoException(HIVE_FILESYSTEM_ERROR, "Failed to list directory: " + path, e);
+
+            if (paths.isEmpty()) {
+                return endOfData();
+            }
+            remoteIterator = getLocatedFileStatusRemoteIterator(paths.removeFirst());
         }
     }
 
-    private RemoteIterator<LocatedFileStatus> getLocatedFileStatusRemoteIterator(Path path)
-            throws IOException
+    private Iterator<LocatedFileStatus> getLocatedFileStatusRemoteIterator(Path path)
     {
         try (TimeStat.BlockTimer ignored = namenodeStats.getListLocatedStatus().time()) {
-            return directoryLister.list(fileSystem, path);
-        }
-        catch (IOException | RuntimeException e) {
-            namenodeStats.getListLocatedStatus().recordException(e);
-            throw e;
+            return new FileStatusIterator(path, fileSystem, directoryLister, namenodeStats);
         }
     }
 
-    private LocatedFileStatus getLocatedFileStatus(RemoteIterator<LocatedFileStatus> iterator)
-            throws IOException
+    private LocatedFileStatus getLocatedFileStatus(Iterator<LocatedFileStatus> iterator)
     {
         try (TimeStat.BlockTimer ignored = namenodeStats.getRemoteIteratorNext().time()) {
             return iterator.next();
         }
-        catch (IOException | RuntimeException e) {
-            namenodeStats.getRemoteIteratorNext().recordException(e);
-            throw e;
-        }
     }
 
-    public String getPartitionName()
+    private static class FileStatusIterator
+            implements Iterator<LocatedFileStatus>
     {
-        return partitionName;
-    }
+        private final Path path;
+        private final NamenodeStats namenodeStats;
+        private final RemoteIterator<LocatedFileStatus> fileStatusIterator;
 
-    public InternalHiveSplitFactory getSplitFactory()
-    {
-        return splitFactory;
+        private FileStatusIterator(Path path, FileSystem fileSystem, DirectoryLister directoryLister, NamenodeStats namenodeStats)
+        {
+            this.path = path;
+            this.namenodeStats = namenodeStats;
+            try {
+                this.fileStatusIterator = directoryLister.list(fileSystem, path);
+            }
+            catch (IOException e) {
+                throw processException(e);
+            }
+        }
+
+        @Override
+        public boolean hasNext()
+        {
+            try {
+                return fileStatusIterator.hasNext();
+            }
+            catch (IOException e) {
+                throw processException(e);
+            }
+        }
+
+        @Override
+        public LocatedFileStatus next()
+        {
+            try {
+                return fileStatusIterator.next();
+            }
+            catch (IOException e) {
+                throw processException(e);
+            }
+        }
+
+        private PrestoException processException(IOException exception)
+        {
+            namenodeStats.getRemoteIteratorNext().recordException(exception);
+            if (exception instanceof FileNotFoundException) {
+                throw new PrestoException(HIVE_FILE_NOT_FOUND, "Partition location does not exist: " + path);
+            }
+            return new PrestoException(HIVE_FILESYSTEM_ERROR, "Failed to list directory: " + path, exception);
+        }
     }
 
-    public HiveFileIterator withPath(Path path)
+    public static class NestedDirectoryNotAllowedException
+            extends RuntimeException
     {
-        return new HiveFileIterator(path, fileSystem, directoryLister, namenodeStats, partitionName, splitFactory);
+        public NestedDirectoryNotAllowedException()
+        {
+            super("Nested sub-directories are not allowed");
+        }
     }
 }