Java tutorial
/* * 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 com.facebook.presto.hive; import com.facebook.presto.hive.util.AsyncRecursiveWalker; import com.facebook.presto.hive.util.BoundedExecutor; import com.facebook.presto.hive.util.FileStatusCallback; import com.facebook.presto.hive.util.SuspendingExecutor; import com.facebook.presto.spi.HostAddress; import com.facebook.presto.spi.Split; import com.google.common.base.Optional; import com.google.common.base.Splitter; import com.google.common.base.Throwables; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import io.airlift.units.DataSize; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.BlockLocation; import org.apache.hadoop.fs.FileStatus; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.hive.metastore.MetaStoreUtils; import org.apache.hadoop.hive.metastore.api.FieldSchema; import org.apache.hadoop.hive.metastore.api.Partition; import org.apache.hadoop.hive.metastore.api.Table; import org.apache.hadoop.hive.ql.io.SymlinkTextInputFormat; import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector; import org.apache.hadoop.mapred.FileInputFormat; import org.apache.hadoop.mapred.FileSplit; import org.apache.hadoop.mapred.InputFormat; import org.apache.hadoop.mapred.InputSplit; import org.apache.hadoop.mapred.JobConf; import javax.annotation.concurrent.ThreadSafe; import java.io.IOException; import java.util.Iterator; import java.util.List; import java.util.Properties; import java.util.concurrent.BlockingQueue; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import static com.facebook.presto.hive.HiveSplit.markAsLastSplit; import static com.facebook.presto.hive.HiveType.getSupportedHiveType; import static com.facebook.presto.hive.HiveUtil.convertNativeHiveType; import static com.facebook.presto.hive.HiveUtil.getInputFormat; import static com.facebook.presto.hive.HiveUtil.isSplittable; import static com.facebook.presto.hive.UnpartitionedPartition.isUnpartitioned; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; class HiveSplitIterable implements Iterable<Split> { private static final Split FINISHED_MARKER = new Split() { @Override public boolean isRemotelyAccessible() { throw new UnsupportedOperationException(); } @Override public List<HostAddress> getAddresses() { throw new UnsupportedOperationException(); } @Override public Object getInfo() { throw new UnsupportedOperationException(); } }; @SuppressWarnings("ObjectEquality") private static boolean isFinished(Split split) { return split == FINISHED_MARKER; } private final String clientId; private final Table table; private final Iterable<String> partitionNames; private final Iterable<Partition> partitions; private final Optional<Integer> bucket; private final int maxOutstandingSplits; private final int maxThreads; private final HdfsEnvironment hdfsEnvironment; private final ExecutorService executor; private final ClassLoader classLoader; private final DataSize maxSplitSize; private final int maxPartitionBatchSize; HiveSplitIterable(String clientId, Table table, Iterable<String> partitionNames, Iterable<Partition> partitions, Optional<Integer> bucket, DataSize maxSplitSize, int maxOutstandingSplits, int maxThreads, HdfsEnvironment hdfsEnvironment, ExecutorService executor, int maxPartitionBatchSize) { this.clientId = clientId; this.table = table; this.partitionNames = partitionNames; this.partitions = partitions; this.bucket = bucket; this.maxSplitSize = maxSplitSize; this.maxPartitionBatchSize = maxPartitionBatchSize; this.maxOutstandingSplits = maxOutstandingSplits; this.maxThreads = maxThreads; this.hdfsEnvironment = hdfsEnvironment; this.executor = executor; this.classLoader = Thread.currentThread().getContextClassLoader(); } @Override public Iterator<Split> iterator() { // Each iterator has its own bounded executor and can be independently suspended final SuspendingExecutor suspendingExecutor = new SuspendingExecutor( new BoundedExecutor(executor, maxThreads)); final HiveSplitQueue hiveSplitQueue = new HiveSplitQueue(maxOutstandingSplits, suspendingExecutor); executor.submit(new Callable<Void>() { @Override public Void call() throws Exception { loadPartitionSplits(hiveSplitQueue, suspendingExecutor); return null; } }); return hiveSplitQueue; } private void loadPartitionSplits(final HiveSplitQueue hiveSplitQueue, SuspendingExecutor suspendingExecutor) throws InterruptedException { final Semaphore semaphore = new Semaphore(maxPartitionBatchSize); try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { ImmutableList.Builder<ListenableFuture<Void>> futureBuilder = ImmutableList.builder(); Iterator<String> nameIterator = partitionNames.iterator(); for (Partition partition : partitions) { checkState(nameIterator.hasNext(), "different number of partitions and partition names!"); semaphore.acquire(); final String partitionName = nameIterator.next(); final Properties schema = getPartitionSchema(table, partition); final List<HivePartitionKey> partitionKeys = getPartitionKeys(table, partition); Path path = new Path(getPartitionLocation(table, partition)); final Configuration configuration = hdfsEnvironment.getConfiguration(path); final InputFormat<?, ?> inputFormat = getInputFormat(configuration, schema, false); Path partitionPath = hdfsEnvironment.getFileSystemWrapper().wrap(path); FileSystem fs = partitionPath.getFileSystem(configuration); final LastSplitMarkingQueue markerQueue = new LastSplitMarkingQueue(hiveSplitQueue); if (inputFormat instanceof SymlinkTextInputFormat) { JobConf jobConf = new JobConf(configuration); FileInputFormat.setInputPaths(jobConf, partitionPath); InputSplit[] splits = inputFormat.getSplits(jobConf, 0); for (InputSplit rawSplit : splits) { FileSplit split = ((SymlinkTextInputFormat.SymlinkTextInputSplit) rawSplit) .getTargetSplit(); // get the filesystem for the target path -- it may be a different hdfs instance FileSystem targetFilesystem = split.getPath().getFileSystem(configuration); FileStatus fileStatus = targetFilesystem.getFileStatus(split.getPath()); markerQueue.addToQueue(createHiveSplits(partitionName, fileStatus, targetFilesystem.getFileBlockLocations(fileStatus, split.getStart(), split.getLength()), split.getStart(), split.getLength(), schema, partitionKeys, false)); } markerQueue.finish(); continue; } ListenableFuture<Void> partitionFuture = new AsyncRecursiveWalker(fs, suspendingExecutor) .beginWalk(partitionPath, new FileStatusCallback() { @Override public void process(FileStatus file, BlockLocation[] blockLocations) { if (bucket.isPresent() && !fileMatchesBucket(file.getPath().getName(), bucket.get())) { return; } try { boolean splittable = isSplittable(inputFormat, file.getPath().getFileSystem(configuration), file.getPath()); markerQueue.addToQueue(createHiveSplits(partitionName, file, blockLocations, 0, file.getLen(), schema, partitionKeys, splittable)); } catch (IOException e) { hiveSplitQueue.fail(e); } } }); // release the semaphore when the partition finishes Futures.addCallback(partitionFuture, new FutureCallback<Void>() { @Override public void onSuccess(Void result) { markerQueue.finish(); semaphore.release(); } @Override public void onFailure(Throwable t) { markerQueue.finish(); semaphore.release(); } }); futureBuilder.add(partitionFuture); } // when all partitions finish, mark the queue as finished Futures.addCallback(Futures.allAsList(futureBuilder.build()), new FutureCallback<List<Void>>() { @Override public void onSuccess(List<Void> result) { hiveSplitQueue.finished(); } @Override public void onFailure(Throwable t) { hiveSplitQueue.fail(t); } }); } catch (Throwable e) { hiveSplitQueue.fail(e); Throwables.propagateIfInstanceOf(e, Error.class); } } private static boolean fileMatchesBucket(String fileName, int bucketNumber) { String currentBucket = Splitter.on('_').split(fileName).iterator().next(); try { return Integer.parseInt(currentBucket) == bucketNumber; } catch (NumberFormatException ignored) { return false; } } private List<HiveSplit> createHiveSplits(String partitionName, FileStatus file, BlockLocation[] blockLocations, long start, long length, Properties schema, List<HivePartitionKey> partitionKeys, boolean splittable) throws IOException { ImmutableList.Builder<HiveSplit> builder = ImmutableList.builder(); if (splittable) { for (BlockLocation blockLocation : blockLocations) { // get the addresses for the block List<HostAddress> addresses = toHostAddress(blockLocation.getHosts()); // divide the block into uniform chunks that are smaller than the max split size int chunks = Math.max(1, (int) (blockLocation.getLength() / maxSplitSize.toBytes())); // when block does not divide evenly into chunks, make the chunk size slightly bigger than necessary long targetChunkSize = (long) Math.ceil(blockLocation.getLength() * 1.0 / chunks); long chunkOffset = 0; while (chunkOffset < blockLocation.getLength()) { // adjust the actual chunk size to account for the overrun when chunks are slightly bigger than necessary (see above) long chunkLength = Math.min(targetChunkSize, blockLocation.getLength() - chunkOffset); builder.add(new HiveSplit(clientId, table.getDbName(), table.getTableName(), partitionName, false, file.getPath().toString(), blockLocation.getOffset() + chunkOffset, chunkLength, schema, partitionKeys, addresses)); chunkOffset += chunkLength; } checkState(chunkOffset == blockLocation.getLength(), "Error splitting blocks"); } } else { // not splittable, use the hosts from the first block builder.add(new HiveSplit(clientId, table.getDbName(), table.getTableName(), partitionName, false, file.getPath().toString(), start, length, schema, partitionKeys, toHostAddress(blockLocations[0].getHosts()))); } return builder.build(); } private static List<HostAddress> toHostAddress(String[] hosts) { ImmutableList.Builder<HostAddress> builder = ImmutableList.builder(); for (String host : hosts) { builder.add(HostAddress.fromString(host)); } return builder.build(); } /** * Buffers a single split for a given partition and when the queue * is finished, tags the final split so a reader of the stream can * know when */ @ThreadSafe private static class LastSplitMarkingQueue { private final HiveSplitQueue hiveSplitQueue; private final AtomicReference<HiveSplit> bufferedSplit = new AtomicReference<>(); private final AtomicBoolean done = new AtomicBoolean(); private LastSplitMarkingQueue(HiveSplitQueue hiveSplitQueue) { this.hiveSplitQueue = checkNotNull(hiveSplitQueue, "split is null"); } public synchronized void addToQueue(Iterable<HiveSplit> splits) { checkNotNull(splits, "splits is null"); checkState(!done.get(), "already done"); for (HiveSplit split : splits) { HiveSplit previousSplit = bufferedSplit.getAndSet(split); if (previousSplit != null) { hiveSplitQueue.addToQueue(previousSplit); } } } private synchronized void finish() { checkState(!done.getAndSet(true), "already done"); HiveSplit finalSplit = bufferedSplit.getAndSet(null); if (finalSplit != null) { hiveSplitQueue.addToQueue(markAsLastSplit(finalSplit)); } } } private static class HiveSplitQueue extends AbstractIterator<Split> { private final BlockingQueue<Split> queue = new LinkedBlockingQueue<>(); private final AtomicInteger outstandingSplitCount = new AtomicInteger(); private final AtomicReference<Throwable> throwable = new AtomicReference<>(); private final int maxOutstandingSplits; private final SuspendingExecutor suspendingExecutor; private HiveSplitQueue(int maxOutstandingSplits, SuspendingExecutor suspendingExecutor) { this.maxOutstandingSplits = maxOutstandingSplits; this.suspendingExecutor = suspendingExecutor; } private void addToQueue(Split split) { queue.add(split); if (outstandingSplitCount.incrementAndGet() == maxOutstandingSplits) { suspendingExecutor.suspend(); } } private void finished() { queue.add(FINISHED_MARKER); } private void fail(Throwable e) { throwable.set(e); queue.add(FINISHED_MARKER); } @Override protected Split computeNext() { try { Split split = queue.take(); if (isFinished(split)) { if (throwable.get() != null) { throw Throwables.propagate(throwable.get()); } return endOfData(); } if (outstandingSplitCount.getAndDecrement() == maxOutstandingSplits) { suspendingExecutor.resume(); } return split; } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw Throwables.propagate(e); } } } private static List<HivePartitionKey> getPartitionKeys(Table table, Partition partition) { if (isUnpartitioned(partition)) { return ImmutableList.of(); } ImmutableList.Builder<HivePartitionKey> partitionKeys = ImmutableList.builder(); List<FieldSchema> keys = table.getPartitionKeys(); List<String> values = partition.getValues(); checkArgument(keys.size() == values.size(), "Expected %s partition key values, but got %s", keys.size(), values.size()); for (int i = 0; i < keys.size(); i++) { String name = keys.get(i).getName(); PrimitiveObjectInspector.PrimitiveCategory primitiveCategory = convertNativeHiveType( keys.get(i).getType()); HiveType hiveType = getSupportedHiveType(primitiveCategory); String value = values.get(i); checkNotNull(value, "partition key value cannot be null for field: %s", name); partitionKeys.add(new HivePartitionKey(name, hiveType, value)); } return partitionKeys.build(); } private static Properties getPartitionSchema(Table table, Partition partition) { if (isUnpartitioned(partition)) { return MetaStoreUtils.getTableMetadata(table); } return MetaStoreUtils.getSchema(partition, table); } private static String getPartitionLocation(Table table, Partition partition) { if (isUnpartitioned(partition)) { return table.getSd().getLocation(); } return partition.getSd().getLocation(); } }