Examples with HostAddressFactory io.trino.plugin.cassandra.util.HostAddressFactory used on opensource projects

Example 1 with HostAddressFactory

use of io.trino.plugin.cassandra.util.HostAddressFactory in project trino by trinodb.

the class CassandraSplitManager method getSplitsForPartitions.

private List<ConnectorSplit> getSplitsForPartitions(CassandraTableHandle cassTableHandle, List<CassandraPartition> partitions, String clusteringPredicates) {
    String schema = cassTableHandle.getSchemaName();
    HostAddressFactory hostAddressFactory = new HostAddressFactory();
    ImmutableList.Builder<ConnectorSplit> builder = ImmutableList.builder();
    // For single partition key column table, we can merge multiple partitions into a single split
    // by using IN CLAUSE in a single select query if the partitions have the same host list.
    // For multiple partition key columns table, we can't merge them into a single select query, so
    // keep them in a separate split.
    boolean singlePartitionKeyColumn = true;
    String partitionKeyColumnName = null;
    if (!partitions.isEmpty()) {
        singlePartitionKeyColumn = partitions.get(0).getTupleDomain().getDomains().get().size() == 1;
        if (singlePartitionKeyColumn) {
            String partitionId = partitions.get(0).getPartitionId();
            partitionKeyColumnName = partitionId.substring(0, partitionId.lastIndexOf('=') - 1);
        }
    }
    Map<Set<String>, Set<String>> hostsToPartitionKeys = new HashMap<>();
    Map<Set<String>, List<HostAddress>> hostMap = new HashMap<>();
    for (CassandraPartition cassandraPartition : partitions) {
        Set<Host> hosts = cassandraSession.getReplicas(schema, cassandraPartition.getKeyAsByteBuffer());
        List<HostAddress> addresses = hostAddressFactory.toHostAddressList(hosts);
        if (singlePartitionKeyColumn) {
            // host ip addresses
            ImmutableSet.Builder<String> sb = ImmutableSet.builder();
            for (HostAddress address : addresses) {
                sb.add(address.getHostText());
            }
            Set<String> hostAddresses = sb.build();
            // partition key values
            Set<String> values = hostsToPartitionKeys.get(hostAddresses);
            if (values == null) {
                values = new HashSet<>();
            }
            String partitionId = cassandraPartition.getPartitionId();
            values.add(partitionId.substring(partitionId.lastIndexOf('=') + 2));
            hostsToPartitionKeys.put(hostAddresses, values);
            hostMap.put(hostAddresses, addresses);
        } else {
            builder.add(createSplitForClusteringPredicates(cassandraPartition.getPartitionId(), addresses, clusteringPredicates));
        }
    }
    if (singlePartitionKeyColumn) {
        for (Map.Entry<Set<String>, Set<String>> entry : hostsToPartitionKeys.entrySet()) {
            StringBuilder sb = new StringBuilder(partitionSizeForBatchSelect);
            int size = 0;
            for (String value : entry.getValue()) {
                if (size > 0) {
                    sb.append(",");
                }
                sb.append(value);
                size++;
                if (size > partitionSizeForBatchSelect) {
                    String partitionId = format("%s in (%s)", partitionKeyColumnName, sb);
                    builder.add(createSplitForClusteringPredicates(partitionId, hostMap.get(entry.getKey()), clusteringPredicates));
                    size = 0;
                    sb.setLength(0);
                    sb.trimToSize();
                }
            }
            if (size > 0) {
                String partitionId = format("%s in (%s)", partitionKeyColumnName, sb);
                builder.add(createSplitForClusteringPredicates(partitionId, hostMap.get(entry.getKey()), clusteringPredicates));
            }
        }
    }
    return builder.build();
}

Example 2 with HostAddressFactory

use of io.trino.plugin.cassandra.util.HostAddressFactory in project trino by trinodb.

the class CassandraSplitManager method getSplitsByTokenRange.

private List<ConnectorSplit> getSplitsByTokenRange(CassandraTable table, String partitionId, Optional<Long> sessionSplitsPerNode) {
    String schema = table.getTableHandle().getSchemaName();
    String tableName = table.getTableHandle().getTableName();
    String tokenExpression = table.getTokenExpression();
    ImmutableList.Builder<ConnectorSplit> builder = ImmutableList.builder();
    List<CassandraTokenSplitManager.TokenSplit> tokenSplits = tokenSplitMgr.getSplits(schema, tableName, sessionSplitsPerNode);
    for (CassandraTokenSplitManager.TokenSplit tokenSplit : tokenSplits) {
        String condition = buildTokenCondition(tokenExpression, tokenSplit.getStartToken(), tokenSplit.getEndToken());
        List<HostAddress> addresses = new HostAddressFactory().hostAddressNamesToHostAddressList(tokenSplit.getHosts());
        CassandraSplit split = new CassandraSplit(partitionId, condition, addresses);
        builder.add(split);
    }
    return builder.build();
}