io.prestosql.plugin.kudu.KuduClientSession Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of presto-kudu Show documentation
Presto - Kudu Connector
There is a newer version: 350
/*
 * 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 io.prestosql.plugin.kudu;

import com.google.common.collect.ImmutableList;
import io.airlift.log.Logger;
import io.prestosql.plugin.kudu.properties.ColumnDesign;
import io.prestosql.plugin.kudu.properties.HashPartitionDefinition;
import io.prestosql.plugin.kudu.properties.KuduTableProperties;
import io.prestosql.plugin.kudu.properties.PartitionDesign;
import io.prestosql.plugin.kudu.properties.RangePartition;
import io.prestosql.plugin.kudu.properties.RangePartitionDefinition;
import io.prestosql.plugin.kudu.schema.SchemaEmulation;
import io.prestosql.spi.PrestoException;
import io.prestosql.spi.connector.ColumnHandle;
import io.prestosql.spi.connector.ColumnMetadata;
import io.prestosql.spi.connector.ConnectorTableMetadata;
import io.prestosql.spi.connector.SchemaNotFoundException;
import io.prestosql.spi.connector.SchemaTableName;
import io.prestosql.spi.connector.TableNotFoundException;
import io.prestosql.spi.predicate.DiscreteValues;
import io.prestosql.spi.predicate.Domain;
import io.prestosql.spi.predicate.EquatableValueSet;
import io.prestosql.spi.predicate.Marker;
import io.prestosql.spi.predicate.Range;
import io.prestosql.spi.predicate.Ranges;
import io.prestosql.spi.predicate.SortedRangeSet;
import io.prestosql.spi.predicate.TupleDomain;
import io.prestosql.spi.predicate.ValueSet;
import io.prestosql.spi.type.DecimalType;
import org.apache.kudu.ColumnSchema;
import org.apache.kudu.ColumnTypeAttributes;
import org.apache.kudu.Schema;
import org.apache.kudu.Type;
import org.apache.kudu.client.AlterTableOptions;
import org.apache.kudu.client.CreateTableOptions;
import org.apache.kudu.client.KuduClient;
import org.apache.kudu.client.KuduException;
import org.apache.kudu.client.KuduPredicate;
import org.apache.kudu.client.KuduScanToken;
import org.apache.kudu.client.KuduScanner;
import org.apache.kudu.client.KuduSession;
import org.apache.kudu.client.KuduTable;
import org.apache.kudu.client.PartialRow;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.IntStream;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static io.prestosql.spi.StandardErrorCode.QUERY_REJECTED;
import static java.util.stream.Collectors.toList;

public class KuduClientSession
{
    public static final String DEFAULT_SCHEMA = "default";
    private final Logger log = Logger.get(getClass());
    private final KuduClient client;
    private final SchemaEmulation schemaEmulation;

    public KuduClientSession(KuduClient client, SchemaEmulation schemaEmulation)
    {
        this.client = client;
        this.schemaEmulation = schemaEmulation;
    }

    public List listSchemaNames()
    {
        return schemaEmulation.listSchemaNames(client);
    }

    private List internalListTables(String prefix)
    {
        try {
            if (prefix.isEmpty()) {
                return client.getTablesList().getTablesList();
            }
            else {
                return client.getTablesList(prefix).getTablesList();
            }
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public List listTables(Optional optSchemaName)
    {
        if (optSchemaName.isPresent()) {
            return listTablesSingleSchema(optSchemaName.get());
        }

        List all = new ArrayList<>();
        for (String schemaName : listSchemaNames()) {
            List single = listTablesSingleSchema(schemaName);
            all.addAll(single);
        }
        return all;
    }

    private List listTablesSingleSchema(String schemaName)
    {
        final String prefix = schemaEmulation.getPrefixForTablesOfSchema(schemaName);

        List tables = internalListTables(prefix);
        return tables.stream()
                .map(schemaEmulation::fromRawName)
                .filter(Objects::nonNull)
                .collect(toImmutableList());
    }

    public Schema getTableSchema(KuduTableHandle tableHandle)
    {
        KuduTable table = tableHandle.getTable(this);
        return table.getSchema();
    }

    public Map getTableProperties(KuduTableHandle tableHandle)
    {
        KuduTable table = tableHandle.getTable(this);
        return KuduTableProperties.toMap(table);
    }

    public List buildKuduSplits(KuduTableHandle tableHandle)
    {
        KuduTable table = tableHandle.getTable(this);
        final int primaryKeyColumnCount = table.getSchema().getPrimaryKeyColumnCount();
        KuduScanToken.KuduScanTokenBuilder builder = client.newScanTokenBuilder(table);

        TupleDomain constraint = tableHandle.getConstraint();
        if (!addConstraintPredicates(table, builder, constraint)) {
            return ImmutableList.of();
        }

        Optional> desiredColumns = tableHandle.getDesiredColumns();

        List columnIndexes;
        if (tableHandle.isDeleteHandle()) {
            if (desiredColumns.isPresent()) {
                columnIndexes = IntStream
                        .range(0, primaryKeyColumnCount)
                        .boxed().collect(toList());
                for (ColumnHandle column : desiredColumns.get()) {
                    KuduColumnHandle k = (KuduColumnHandle) column;
                    int index = k.getOrdinalPosition();
                    if (index >= primaryKeyColumnCount) {
                        columnIndexes.add(index);
                    }
                }
                columnIndexes = ImmutableList.copyOf(columnIndexes);
            }
            else {
                columnIndexes = IntStream
                        .range(0, table.getSchema().getColumnCount())
                        .boxed().collect(toImmutableList());
            }
        }
        else {
            if (desiredColumns.isPresent()) {
                columnIndexes = desiredColumns.get().stream()
                        .map(handle -> ((KuduColumnHandle) handle).getOrdinalPosition())
                        .collect(toImmutableList());
            }
            else {
                ImmutableList.Builder columnIndexesBuilder = ImmutableList.builder();
                Schema schema = table.getSchema();
                for (int ordinal = 0; ordinal < schema.getColumnCount(); ordinal++) {
                    ColumnSchema column = schema.getColumnByIndex(ordinal);
                    // Skip hidden "row_uuid" column
                    if (!column.isKey() || !column.getName().equals(KuduColumnHandle.ROW_ID)) {
                        columnIndexesBuilder.add(ordinal);
                    }
                }
                columnIndexes = columnIndexesBuilder.build();
            }
        }

        builder.setProjectedColumnIndexes(columnIndexes);

        List tokens = builder.build();
        return tokens.stream()
                .map(token -> toKuduSplit(tableHandle, token, primaryKeyColumnCount))
                .collect(toImmutableList());
    }

    public KuduScanner createScanner(KuduSplit kuduSplit)
    {
        try {
            return KuduScanToken.deserializeIntoScanner(kuduSplit.getSerializedScanToken(), client);
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public KuduTable openTable(SchemaTableName schemaTableName)
    {
        String rawName = schemaEmulation.toRawName(schemaTableName);
        try {
            return client.openTable(rawName);
        }
        catch (KuduException e) {
            log.debug("Error on doOpenTable: " + e, e);
            if (!listSchemaNames().contains(schemaTableName.getSchemaName())) {
                throw new SchemaNotFoundException(schemaTableName.getSchemaName());
            }
            throw new TableNotFoundException(schemaTableName);
        }
    }

    public KuduSession newSession()
    {
        return client.newSession();
    }

    public void createSchema(String schemaName)
    {
        schemaEmulation.createSchema(client, schemaName);
    }

    public void dropSchema(String schemaName)
    {
        schemaEmulation.dropSchema(client, schemaName);
    }

    public void dropTable(SchemaTableName schemaTableName)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            client.deleteTable(rawName);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public void renameTable(SchemaTableName schemaTableName, SchemaTableName newSchemaTableName)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            String newRawName = schemaEmulation.toRawName(newSchemaTableName);
            AlterTableOptions alterOptions = new AlterTableOptions();
            alterOptions.renameTable(newRawName);
            client.alterTable(rawName, alterOptions);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public KuduTable createTable(ConnectorTableMetadata tableMetadata, boolean ignoreExisting)
    {
        try {
            String rawName = schemaEmulation.toRawName(tableMetadata.getTable());
            if (ignoreExisting) {
                if (client.tableExists(rawName)) {
                    return null;
                }
            }

            if (!schemaEmulation.existsSchema(client, tableMetadata.getTable().getSchemaName())) {
                throw new SchemaNotFoundException(tableMetadata.getTable().getSchemaName());
            }

            List columns = tableMetadata.getColumns();
            Map properties = tableMetadata.getProperties();

            Schema schema = buildSchema(columns, properties);
            CreateTableOptions options = buildCreateTableOptions(schema, properties);
            return client.createTable(rawName, schema, options);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public void addColumn(SchemaTableName schemaTableName, ColumnMetadata column)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            AlterTableOptions alterOptions = new AlterTableOptions();
            Type type = TypeHelper.toKuduClientType(column.getType());
            alterOptions.addNullableColumn(column.getName(), type);
            client.alterTable(rawName, alterOptions);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public void dropColumn(SchemaTableName schemaTableName, String name)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            AlterTableOptions alterOptions = new AlterTableOptions();
            alterOptions.dropColumn(name);
            client.alterTable(rawName, alterOptions);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public void renameColumn(SchemaTableName schemaTableName, String oldName, String newName)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            AlterTableOptions alterOptions = new AlterTableOptions();
            alterOptions.renameColumn(oldName, newName);
            client.alterTable(rawName, alterOptions);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    public void addRangePartition(SchemaTableName schemaTableName, RangePartition rangePartition)
    {
        changeRangePartition(schemaTableName, rangePartition, RangePartitionChange.ADD);
    }

    public void dropRangePartition(SchemaTableName schemaTableName, RangePartition rangePartition)
    {
        changeRangePartition(schemaTableName, rangePartition, RangePartitionChange.DROP);
    }

    private void changeRangePartition(SchemaTableName schemaTableName, RangePartition rangePartition,
            RangePartitionChange change)
    {
        try {
            String rawName = schemaEmulation.toRawName(schemaTableName);
            KuduTable table = client.openTable(rawName);
            Schema schema = table.getSchema();
            PartitionDesign design = KuduTableProperties.getPartitionDesign(table);
            RangePartitionDefinition definition = design.getRange();
            if (definition == null) {
                throw new PrestoException(QUERY_REJECTED, "Table " + schemaTableName + " has no range partition");
            }
            PartialRow lowerBound = KuduTableProperties.toRangeBoundToPartialRow(schema, definition, rangePartition.getLower());
            PartialRow upperBound = KuduTableProperties.toRangeBoundToPartialRow(schema, definition, rangePartition.getUpper());
            AlterTableOptions alterOptions = new AlterTableOptions();
            switch (change) {
                case ADD:
                    alterOptions.addRangePartition(lowerBound, upperBound);
                    break;
                case DROP:
                    alterOptions.dropRangePartition(lowerBound, upperBound);
                    break;
            }
            client.alterTable(rawName, alterOptions);
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }

    private Schema buildSchema(List columns, Map tableProperties)
    {
        List kuduColumns = columns.stream()
                .map(this::toColumnSchema)
                .collect(ImmutableList.toImmutableList());
        return new Schema(kuduColumns);
    }

    private ColumnSchema toColumnSchema(ColumnMetadata columnMetadata)
    {
        String name = columnMetadata.getName();
        ColumnDesign design = KuduTableProperties.getColumnDesign(columnMetadata.getProperties());
        Type ktype = TypeHelper.toKuduClientType(columnMetadata.getType());
        ColumnSchema.ColumnSchemaBuilder builder = new ColumnSchema.ColumnSchemaBuilder(name, ktype);
        builder.key(design.isPrimaryKey()).nullable(design.isNullable());
        setEncoding(name, builder, design);
        setCompression(name, builder, design);
        setTypeAttributes(columnMetadata, builder);
        return builder.build();
    }

    private void setTypeAttributes(ColumnMetadata columnMetadata, ColumnSchema.ColumnSchemaBuilder builder)
    {
        if (columnMetadata.getType() instanceof DecimalType) {
            DecimalType type = (DecimalType) columnMetadata.getType();
            ColumnTypeAttributes attributes = new ColumnTypeAttributes.ColumnTypeAttributesBuilder()
                    .precision(type.getPrecision())
                    .scale(type.getScale()).build();
            builder.typeAttributes(attributes);
        }
    }

    private void setCompression(String name, ColumnSchema.ColumnSchemaBuilder builder, ColumnDesign design)
    {
        if (design.getCompression() != null) {
            try {
                ColumnSchema.CompressionAlgorithm algorithm = KuduTableProperties.lookupCompression(design.getCompression());
                builder.compressionAlgorithm(algorithm);
            }
            catch (IllegalArgumentException e) {
                throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unknown compression algorithm " + design.getCompression() + " for column " + name);
            }
        }
    }

    private void setEncoding(String name, ColumnSchema.ColumnSchemaBuilder builder, ColumnDesign design)
    {
        if (design.getEncoding() != null) {
            try {
                ColumnSchema.Encoding encoding = KuduTableProperties.lookupEncoding(design.getEncoding());
                builder.encoding(encoding);
            }
            catch (IllegalArgumentException e) {
                throw new PrestoException(GENERIC_INTERNAL_ERROR, "Unknown encoding " + design.getEncoding() + " for column " + name);
            }
        }
    }

    private CreateTableOptions buildCreateTableOptions(Schema schema, Map properties)
    {
        CreateTableOptions options = new CreateTableOptions();

        RangePartitionDefinition rangePartitionDefinition = null;
        PartitionDesign partitionDesign = KuduTableProperties.getPartitionDesign(properties);
        if (partitionDesign.getHash() != null) {
            for (HashPartitionDefinition partition : partitionDesign.getHash()) {
                options.addHashPartitions(partition.getColumns(), partition.getBuckets());
            }
        }
        if (partitionDesign.getRange() != null) {
            rangePartitionDefinition = partitionDesign.getRange();
            options.setRangePartitionColumns(rangePartitionDefinition.getColumns());
        }

        List rangePartitions = KuduTableProperties.getRangePartitions(properties);
        if (rangePartitionDefinition != null && !rangePartitions.isEmpty()) {
            for (RangePartition rangePartition : rangePartitions) {
                PartialRow lower = KuduTableProperties.toRangeBoundToPartialRow(schema, rangePartitionDefinition, rangePartition.getLower());
                PartialRow upper = KuduTableProperties.toRangeBoundToPartialRow(schema, rangePartitionDefinition, rangePartition.getUpper());
                options.addRangePartition(lower, upper);
            }
        }

        Optional numReplicas = KuduTableProperties.getNumReplicas(properties);
        numReplicas.ifPresent(options::setNumReplicas);

        return options;
    }

    /**
     * translates TupleDomain to KuduPredicates.
     *
     * @return false if TupleDomain or one of its domains is none
     */
    private boolean addConstraintPredicates(KuduTable table, KuduScanToken.KuduScanTokenBuilder builder,
            TupleDomain constraintSummary)
    {
        if (constraintSummary.isNone()) {
            return false;
        }
        if (!constraintSummary.isAll()) {
            Schema schema = table.getSchema();
            for (TupleDomain.ColumnDomain columnDomain : constraintSummary.getColumnDomains().get()) {
                int position = ((KuduColumnHandle) columnDomain.getColumn()).getOrdinalPosition();
                ColumnSchema columnSchema = schema.getColumnByIndex(position);
                Domain domain = columnDomain.getDomain();
                if (domain.isNone()) {
                    return false;
                }
                else if (domain.isAll()) {
                    // no restriction
                }
                else if (domain.isOnlyNull()) {
                    builder.addPredicate(KuduPredicate.newIsNullPredicate(columnSchema));
                }
                else if (domain.getValues().isAll() && domain.isNullAllowed()) {
                    builder.addPredicate(KuduPredicate.newIsNotNullPredicate(columnSchema));
                }
                else if (domain.isSingleValue()) {
                    KuduPredicate predicate = createEqualsPredicate(columnSchema, domain.getSingleValue());
                    builder.addPredicate(predicate);
                }
                else {
                    ValueSet valueSet = domain.getValues();
                    if (valueSet instanceof EquatableValueSet) {
                        DiscreteValues discreteValues = valueSet.getDiscreteValues();
                        KuduPredicate predicate = createInListPredicate(columnSchema, discreteValues);
                        builder.addPredicate(predicate);
                    }
                    else if (valueSet instanceof SortedRangeSet) {
                        Ranges ranges = ((SortedRangeSet) valueSet).getRanges();
                        Range span = ranges.getSpan();
                        Marker low = span.getLow();
                        if (!low.isLowerUnbounded()) {
                            KuduPredicate.ComparisonOp op = (low.getBound() == Marker.Bound.ABOVE)
                                    ? KuduPredicate.ComparisonOp.GREATER : KuduPredicate.ComparisonOp.GREATER_EQUAL;
                            KuduPredicate predicate = createComparisonPredicate(columnSchema, op, low.getValue());
                            builder.addPredicate(predicate);
                        }
                        Marker high = span.getHigh();
                        if (!high.isUpperUnbounded()) {
                            KuduPredicate.ComparisonOp op = (low.getBound() == Marker.Bound.BELOW)
                                    ? KuduPredicate.ComparisonOp.LESS : KuduPredicate.ComparisonOp.LESS_EQUAL;
                            KuduPredicate predicate = createComparisonPredicate(columnSchema, op, high.getValue());
                            builder.addPredicate(predicate);
                        }
                    }
                    else {
                        throw new IllegalStateException("Unexpected domain: " + domain);
                    }
                }
            }
        }
        return true;
    }

    private KuduPredicate createInListPredicate(ColumnSchema columnSchema, DiscreteValues discreteValues)
    {
        io.prestosql.spi.type.Type type = TypeHelper.fromKuduColumn(columnSchema);
        List