io.trino.plugin.accumulo.AccumuloSplitManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trino-accumulo Show documentation
Trino - Accumulo connector
There is a newer version: 463
/*
 * 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.trino.plugin.accumulo;

import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import io.trino.plugin.accumulo.model.AccumuloColumnConstraint;
import io.trino.plugin.accumulo.model.AccumuloColumnHandle;
import io.trino.plugin.accumulo.model.AccumuloSplit;
import io.trino.plugin.accumulo.model.AccumuloTableHandle;
import io.trino.plugin.accumulo.model.SerializedRange;
import io.trino.plugin.accumulo.model.TabletSplitMetadata;
import io.trino.spi.connector.ColumnHandle;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.connector.ConnectorSplit;
import io.trino.spi.connector.ConnectorSplitManager;
import io.trino.spi.connector.ConnectorSplitSource;
import io.trino.spi.connector.ConnectorTableHandle;
import io.trino.spi.connector.ConnectorTransactionHandle;
import io.trino.spi.connector.Constraint;
import io.trino.spi.connector.DynamicFilter;
import io.trino.spi.connector.FixedSplitSource;
import io.trino.spi.predicate.Domain;
import io.trino.spi.predicate.TupleDomain;

import java.util.List;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;

public class AccumuloSplitManager
        implements ConnectorSplitManager
{
    private final AccumuloMetadataManager metadataManager;

    @Inject
    public AccumuloSplitManager(AccumuloMetadataManager metadataManager)
    {
        this.metadataManager = requireNonNull(metadataManager, "metadataManager is null");
    }

    @Override
    public ConnectorSplitSource getSplits(
            ConnectorTransactionHandle transactionHandle,
            ConnectorSession session,
            ConnectorTableHandle tableHandle,
            DynamicFilter dynamicFilter,
            Constraint constraint)
    {
        AccumuloTableHandle handle = (AccumuloTableHandle) tableHandle;

        String schemaName = handle.getSchema();
        String tableName = handle.getTable();
        String rowIdName = handle.getRowId();

        // Get non-row ID column constraints
        List constraints = getColumnConstraints(rowIdName, handle.getConstraint());

        // Get the row domain column range
        Optional rDom = getRangeDomain(rowIdName, handle.getConstraint());

        // Call out to our client to retrieve all tablet split metadata using the row ID domain and the secondary index
        List tabletSplits = metadataManager.getTabletSplits(session, schemaName, tableName, rDom, constraints, handle.getSerializerInstance());

        // Pack the tablet split metadata into a connector split
        ImmutableList.Builder cSplits = ImmutableList.builder();
        for (TabletSplitMetadata splitMetadata : tabletSplits) {
            AccumuloSplit split = new AccumuloSplit(
                    splitMetadata.ranges().stream().map(SerializedRange::serialize).collect(Collectors.toList()),
                    splitMetadata.hostPort());
            cSplits.add(split);
        }

        return new FixedSplitSource(cSplits.build());
    }

    private static Optional getRangeDomain(String rowIdName, TupleDomain constraint)
    {
        if (constraint.getDomains().isPresent()) {
            for (Entry columnDomain : constraint.getDomains().get().entrySet()) {
                AccumuloColumnHandle col = (AccumuloColumnHandle) columnDomain.getKey();
                if (col.name().equals(rowIdName)) {
                    return Optional.of(columnDomain.getValue());
                }
            }
        }

        return Optional.empty();
    }

    /**
     * Gets a list of {@link AccumuloColumnConstraint} based on the given constraint ID, excluding the row ID column
     *
     * @param rowIdName Trino column name mapping to the Accumulo row ID
     * @param constraint Set of query constraints
     * @return List of all column constraints
     */
    private static List getColumnConstraints(String rowIdName, TupleDomain constraint)
    {
        ImmutableList.Builder constraintBuilder = ImmutableList.builder();
        constraint.getDomains().orElseThrow().forEach((handle, domain) -> {
            AccumuloColumnHandle columnHandle = (AccumuloColumnHandle) handle;

            if (!columnHandle.name().equals(rowIdName)) {
                // Family and qualifier will exist for non-row ID columns
                constraintBuilder.add(new AccumuloColumnConstraint(
                        columnHandle.name(),
                        columnHandle.family().get(),
                        columnHandle.qualifier().get(),
                        Optional.of(domain),
                        columnHandle.indexed()));
            }
        });

        return constraintBuilder.build();
    }
}