com.facebook.presto.split.NativeSplitManager.java Source code

Java tutorial

Introduction

Here is the source code for com.facebook.presto.split.NativeSplitManager.java

Source

/*
 * 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.split;

import com.facebook.presto.metadata.Metadata;
import com.facebook.presto.metadata.NativeTableHandle;
import com.facebook.presto.metadata.Node;
import com.facebook.presto.metadata.NodeManager;
import com.facebook.presto.metadata.ShardManager;
import com.facebook.presto.metadata.TablePartition;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ConnectorSplitManager;
import com.facebook.presto.spi.HostAddress;
import com.facebook.presto.spi.Partition;
import com.facebook.presto.spi.PartitionKey;
import com.facebook.presto.spi.Split;
import com.facebook.presto.spi.TableHandle;
import com.facebook.presto.spi.TableMetadata;
import com.google.common.base.Objects;
import com.google.common.base.Stopwatch;
import com.google.common.collect.Collections2;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Multimap;
import io.airlift.log.Logger;

import javax.inject.Inject;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import static com.facebook.presto.metadata.Node.hostAndPortGetter;
import static com.google.common.base.Functions.forMap;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Iterables.transform;
import static com.google.common.collect.Maps.uniqueIndex;

public class NativeSplitManager implements ConnectorSplitManager {
    private static final Logger log = Logger.get(NativeSplitManager.class);

    private final NodeManager nodeManager;
    private final ShardManager shardManager;
    private final Metadata metadata;

    @Inject
    public NativeSplitManager(NodeManager nodeManager, ShardManager shardManager, Metadata metadata) {
        this.nodeManager = checkNotNull(nodeManager, "nodeManager is null");
        this.shardManager = checkNotNull(shardManager, "shardManager is null");
        this.metadata = checkNotNull(metadata, "metadata is null");
    }

    @Override
    public String getConnectorId() {
        return "native";
    }

    @Override
    public boolean canHandle(TableHandle handle) {
        return handle instanceof NativeTableHandle;
    }

    @Override
    public List<Partition> getPartitions(TableHandle tableHandle, Map<ColumnHandle, Object> bindings) {
        Stopwatch partitionTimer = new Stopwatch();
        partitionTimer.start();

        checkArgument(tableHandle instanceof NativeTableHandle, "Table must be a native table");

        TableMetadata tableMetadata = metadata.getTableMetadata(tableHandle);

        checkState(tableMetadata != null, "no metadata for %s found", tableHandle);

        Set<TablePartition> tablePartitions = shardManager.getPartitions(tableHandle);

        log.debug("Partition retrieval, native table %s (%d partitions): %dms", tableHandle, tablePartitions.size(),
                partitionTimer.elapsed(TimeUnit.MILLISECONDS));

        Multimap<String, ? extends PartitionKey> allPartitionKeys = shardManager.getAllPartitionKeys(tableHandle);
        Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(tableHandle);

        log.debug("Partition key retrieval, native table %s (%d keys): %dms", tableHandle, allPartitionKeys.size(),
                partitionTimer.elapsed(TimeUnit.MILLISECONDS));

        List<Partition> partitions = ImmutableList.copyOf(
                Collections2.transform(tablePartitions, new PartitionFunction(columnHandles, allPartitionKeys)));

        log.debug("Partition generation, native table %s (%d partitions): %dms", tableHandle, partitions.size(),
                partitionTimer.elapsed(TimeUnit.MILLISECONDS));

        return partitions;
    }

    @Override
    public Iterable<Split> getPartitionSplits(TableHandle tableHandle, List<Partition> partitions) {
        Stopwatch splitTimer = new Stopwatch();
        splitTimer.start();

        checkNotNull(partitions, "partitions is null");
        if (partitions.isEmpty()) {
            return ImmutableList.of();
        }

        Map<String, Node> nodesById = uniqueIndex(nodeManager.getAllNodes().getActiveNodes(),
                Node.getIdentifierFunction());

        List<Split> splits = new ArrayList<>();

        Multimap<Long, Entry<Long, String>> partitionShardNodes = shardManager
                .getCommittedPartitionShardNodes(tableHandle);

        for (Partition partition : partitions) {
            checkArgument(partition instanceof NativePartition, "Partition must be a native partition");
            NativePartition nativePartition = (NativePartition) partition;

            ImmutableMultimap.Builder<Long, String> shardNodes = ImmutableMultimap.builder();
            for (Entry<Long, String> partitionShardNode : partitionShardNodes
                    .get(nativePartition.getNativePartitionId())) {
                shardNodes.put(partitionShardNode.getKey(), partitionShardNode.getValue());
            }

            for (Map.Entry<Long, Collection<String>> entry : shardNodes.build().asMap().entrySet()) {
                List<HostAddress> addresses = getAddressesForNodes(nodesById, entry.getValue());
                checkState(addresses.size() > 0, "no host for shard %s found", entry.getKey());
                Split split = new NativeSplit(entry.getKey(), addresses);
                splits.add(split);
            }
        }

        log.debug("Split retrieval for %d partitions (%d splits): %dms", partitions.size(), splits.size(),
                splitTimer.elapsed(TimeUnit.MILLISECONDS));

        // the query engine assumes that splits are returned in a somewhat random fashion. The native split manager,
        // because it loads the data from a db table will return the splits somewhat ordered by node id so only a sub
        // set of nodes is fired up. Shuffle the splits to ensure random distribution.
        Collections.shuffle(splits);

        return ImmutableList.copyOf(splits);
    }

    private static List<HostAddress> getAddressesForNodes(Map<String, Node> nodeMap,
            Iterable<String> nodeIdentifiers) {
        return ImmutableList.copyOf(transform(transform(nodeIdentifiers, forMap(nodeMap)), hostAndPortGetter()));
    }

    public static class NativePartition implements Partition {
        private final long partitionId;
        private Map<ColumnHandle, Object> keys;

        public NativePartition(long partitionId, Map<ColumnHandle, Object> keys) {
            this.partitionId = partitionId;
            this.keys = keys;
        }

        @Override
        public String getPartitionId() {
            return Long.toString(partitionId);
        }

        public long getNativePartitionId() {
            return partitionId;
        }

        @Override
        public Map<ColumnHandle, Object> getKeys() {
            return keys;
        }

        @Override
        public int hashCode() {
            return Objects.hashCode(partitionId, keys);
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (obj == null || getClass() != obj.getClass()) {
                return false;
            }
            final NativePartition other = (NativePartition) obj;
            return this.partitionId == other.partitionId && Objects.equal(this.keys, other.keys);
        }

        @Override
        public String toString() {
            return Objects.toStringHelper(this).add("partitionId", partitionId).add("keys", keys).toString();
        }
    }
}