package com.facebook.presto.elasticsearch;

import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ConnectorPartition;
import com.facebook.presto.spi.ConnectorPartitionResult;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.ConnectorSplit;
import com.facebook.presto.spi.ConnectorSplitSource;
import com.facebook.presto.spi.ConnectorTableHandle;
import com.facebook.presto.spi.FixedSplitSource;
import com.facebook.presto.spi.connector.ConnectorSplitManager;
import com.facebook.presto.spi.predicate.TupleDomain;
import com.google.common.collect.ImmutableList;

import javax.inject.Inject;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static com.facebook.presto.elasticsearch.Types.checkType;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;

public class ElasticsearchSplitManager
        implements ConnectorSplitManager
{
    private final String connectorId;
    private final ElasticsearchClient elasticsearchClient;

    @Inject
    public ElasticsearchSplitManager(ElasticsearchConnectorId connectorId, ElasticsearchClient elasticsearchClient)
    {
        this.connectorId = requireNonNull(connectorId, "connectorId is null").toString();
        this.elasticsearchClient = requireNonNull(elasticsearchClient, "client is null");
    }

    @Override
    public ConnectorPartitionResult getPartitions(ConnectorSession session, ConnectorTableHandle table, TupleDomain<ColumnHandle> tupleDomain)
    {
        return getPartitions(table, tupleDomain);
    }

    private ConnectorPartitionResult getPartitions(ConnectorTableHandle tableHandle, TupleDomain<ColumnHandle> tupleDomain)
    {
        ElasticsearchTableHandle elasticsearchTableHandle = checkType(tableHandle, ElasticsearchTableHandle.class, "tableHandle");

        // elasticsearch connector has only one partition
        List<ConnectorPartition> partitions = ImmutableList.<ConnectorPartition>of(new ElasticsearchPartition(elasticsearchTableHandle.getSchemaName(), elasticsearchTableHandle.getTableName()));
        // elasticsearch connector does not do any additional processing/filtering with the TupleDomain, so just return the whole TupleDomain
        return new ConnectorPartitionResult(partitions, tupleDomain);
    }

    @Override
    public ConnectorSplitSource getPartitionSplits(ConnectorSession session, ConnectorTableHandle table, List<ConnectorPartition> partitions)
    {
        return getPartitionSplits(table, partitions);
    }

    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List<ConnectorPartition> partitions)
    {
        requireNonNull(partitions, "partitions is null");
        checkArgument(partitions.size() == 1, "Expected one partition but got %s", partitions.size());
        ConnectorPartition partition = partitions.get(0);

        ElasticsearchPartition elasticsearchPartition = checkType(partition, ElasticsearchPartition.class, "partition");

        ElasticsearchTableHandle elasticsearchTableHandle = (ElasticsearchTableHandle) tableHandle;
        ElasticsearchTable table = elasticsearchClient.getTable(elasticsearchTableHandle.getSchemaName(), elasticsearchTableHandle.getTableName());
        // this can happen if table is removed during a query
        checkState(table != null, "Table %s.%s no longer exists", elasticsearchTableHandle.getSchemaName(), elasticsearchTableHandle.getTableName());

        List<ConnectorSplit> splits = new ArrayList<>();
        for (ElasticsearchTableSource uri : table.getSources()) {
            int clmsCount = table.getColumns().size();
            splits.add(new ElasticsearchSplit(connectorId, elasticsearchPartition.getSchemaName(), elasticsearchPartition.getTableName(), uri));
        }
        Collections.shuffle(splits);

        return new FixedSplitSource(connectorId, splits);
    }
}

@@ -41,161 +41,8 @@
 h;%0A%0A
-import com.facebook.presto.spi.ColumnHandle;%0Aimport com.facebook.presto.spi.ConnectorPartition;%0Aimport com.facebook.presto.spi.ConnectorPartitionResult;%0A
 impo
@@ -227,24 +227,30 @@
 nnectorTable
+Layout
 Handle;%0Aimpo
@@ -393,77 +393,44 @@
 spi.
-predicate.TupleDomain;%0Aimport com.google.common.collect.ImmutableList
+connector.ConnectorTransactionHandle
 ;%0A%0Ai
@@ -608,74 +608,8 @@
 pe;%0A
-import static com.google.common.base.Preconditions.checkArgument;%0A
 impo
@@ -1213,38 +1213,24 @@
 l%22);%0A    %7D%0A%0A
-    @Override%0A
     public C
@@ -1241,993 +1241,76 @@
 ctor
-PartitionResult getPartitions(ConnectorSession session, ConnectorTableHandle table, TupleDomain%3CColumnHandle%3E tupleDomain)%0A    %7B%0A        return getPartitions(table, tupleDomain);%0A    %7D%0A%0A    private ConnectorPartitionResult getPartitions(ConnectorTableHandle tableHandle, TupleDomain%3CColumnHandle%3E tupleDomain)%0A    %7B%0A        ElasticsearchTableHandle elasticsearchTableHandle = checkType(tableHandle, ElasticsearchTableHandle.class, %22tableHandle%22);%0A%0A        // elasticsearch connector has only one partition%0A        List%3CConnectorPartition%3E partitions = ImmutableList.%3CConnectorPartition%3Eof(new ElasticsearchPartition(elasticsearchTableHandle.getSchemaName(), elasticsearchTableHandle.getTableName()));%0A        // elasticsearch connector does not do any additional processing/filtering with the TupleDomain, so just return the whole TupleDomain%0A        return new ConnectorPartitionResult(partitions, tupleDomain);%0A    %7D%0A%0A    @Override%0A    public ConnectorSplitSource getPartitionSplits(
+SplitSource getSplits(ConnectorTransactionHandle transactionHandle, 
 Conn
@@ -1349,573 +1349,137 @@
 able
+Layout
 Handle 
-table, List%3CConnectorPartition%3E partitions)%0A    %7B%0A        return getPartitionSplits(table, partitions);%0A    %7D%0A%0A    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List%3CConnectorPartition%3E partitions)%0A    %7B%0A        requireNonNull(partitions, %22partitions is null%22);%0A        checkArgument(partitions.size() == 1, %22Expected one partition but got %25s%22, partitions.size());%0A        ConnectorPartition partition = partitions.get(0);%0A%0A        ElasticsearchPartition elasticsearchPartition = checkType(partition, ElasticsearchPartition
+layout)%0A    %7B%0A        ElasticsearchTableLayoutHandle layoutHandle = checkType(layout, ElasticsearchTableLayoutHandle
 .cla
@@ -1496,17 +1496,16 @@
 tion%22);%0A
-%0A
         
@@ -1533,73 +1533,45 @@
 dle 
-elasticsearchTableHandle = (ElasticsearchTableHandle) tableHandle
+tableHandle = layoutHandle.getTable()
 ;%0A  
@@ -1632,30 +1632,17 @@
 etTable(
-elasticsearchT
+t
 ableHand
@@ -1657,38 +1657,25 @@
 hemaName(), 
-elasticsearchT
+t
 ableHandle.g
@@ -1818,30 +1818,17 @@
 xists%22, 
-elasticsearchT
+t
 ableHand
@@ -1851,22 +1851,9 @@
 (), 
-elasticsearchT
+t
 able
@@ -2118,70 +2118,48 @@
 Id, 
-elasticsearchPartition.getSchemaName(), elasticsearchPartition
+tableHandle.getSchemaName(), tableHandle
 .get

Convert elasticsearch connector to TableLayout API