• Home
• io.trino
• trino-spi
• 468
• source code
• SortedRangeSet.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

io.trino.spi.predicate.SortedRangeSet Maven / Gradle / Ivy

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

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.errorprone.annotations.DoNotCall;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.block.DictionaryBlock;
import io.trino.spi.block.LazyBlock;
import io.trino.spi.block.RunLengthEncodedBlock;
import io.trino.spi.block.ValueBlock;
import io.trino.spi.connector.ConnectorSession;
import io.trino.spi.type.Type;

import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.StringJoiner;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.IntStream;
import java.util.stream.Stream;

import static io.airlift.slice.SizeOf.instanceSize;
import static io.airlift.slice.SizeOf.sizeOf;
import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION;
import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION_NOT_NULL;
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.DEFAULT_ON_NULL;
import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
import static io.trino.spi.function.InvocationConvention.simpleConvention;
import static io.trino.spi.predicate.SortedRangeSet.DiscreteSetMarker.DISCRETE;
import static io.trino.spi.predicate.SortedRangeSet.DiscreteSetMarker.NON_DISCRETE;
import static io.trino.spi.predicate.SortedRangeSet.DiscreteSetMarker.UNKNOWN;
import static io.trino.spi.predicate.Utils.TUPLE_DOMAIN_TYPE_OPERATORS;
import static io.trino.spi.predicate.Utils.handleThrowable;
import static io.trino.spi.predicate.Utils.nativeValueToBlock;
import static io.trino.spi.type.TypeUtils.isFloatingPointNaN;
import static io.trino.spi.type.TypeUtils.readNativeValue;
import static io.trino.spi.type.TypeUtils.writeNativeValue;
import static java.lang.Math.max;
import static java.lang.Math.min;
import static java.lang.String.format;
import static java.util.Arrays.asList;
import static java.util.Collections.unmodifiableList;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.joining;

/**
 * A set containing zero or more Ranges of the same type over a continuous space of possible values.
 * Ranges are coalesced into the most compact representation of non-overlapping Ranges. This structure
 * allows iteration across these compacted Ranges in increasing order, as well as other common
 * set-related operation.
 */
public final class SortedRangeSet
        implements ValueSet
{
    private static final int INSTANCE_SIZE = instanceSize(SortedRangeSet.class);

    private final Type type;
    private final MethodHandle equalOperator;
    private final MethodHandle hashCodeOperator;
    private final MethodHandle comparisonOperator;
    private final MethodHandle rangeComparisonOperator;

    private final boolean[] inclusive;
    private final Block sortedRanges;
    private volatile DiscreteSetMarker discreteSetMarker;

    private int lazyHash;

    public enum DiscreteSetMarker
    {
        DISCRETE,
        // empty set is also considered non discrete
        NON_DISCRETE,
        UNKNOWN
    }

    private SortedRangeSet(Type type, boolean[] inclusive, Block sortedRanges, DiscreteSetMarker discreteSetMarker)
    {
        requireNonNull(type, "type is null");
        if (!type.isOrderable()) {
            throw new IllegalArgumentException("Type is not orderable: " + type);
        }
        this.type = type;
        this.equalOperator = TUPLE_DOMAIN_TYPE_OPERATORS.getEqualOperator(type, simpleConvention(DEFAULT_ON_NULL, BLOCK_POSITION_NOT_NULL, BLOCK_POSITION_NOT_NULL));
        this.hashCodeOperator = TUPLE_DOMAIN_TYPE_OPERATORS.getHashCodeOperator(type, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION_NOT_NULL));
        // choice of placing unordered values first or last does not matter for this code
        this.comparisonOperator = TUPLE_DOMAIN_TYPE_OPERATORS.getComparisonUnorderedLastOperator(type, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION));
        // Calculating the comparison operator once instead of per range to avoid hitting TypeOperators cache
        this.rangeComparisonOperator = Range.getComparisonOperator(type);

        requireNonNull(inclusive, "inclusive is null");
        requireNonNull(sortedRanges, "sortedRanges is null");
        if (inclusive.length % 2 != 0) {
            throw new IllegalArgumentException("Malformed inclusive markers");
        }
        if (inclusive.length != sortedRanges.getPositionCount()) {
            throw new IllegalArgumentException(format("Size mismatch between inclusive markers and sortedRanges block: %s, %s", inclusive.length, sortedRanges.getPositionCount()));
        }
        for (int position = 0; position < sortedRanges.getPositionCount(); position++) {
            if (sortedRanges.isNull(position)) {
                if (inclusive[position]) {
                    throw new IllegalArgumentException("Invalid inclusive marker for null value at position " + position);
                }
                if (position != 0 && position != sortedRanges.getPositionCount() - 1) {
                    throw new IllegalArgumentException(format("Invalid null value at position %s of %s", position, sortedRanges.getPositionCount()));
                }
            }
        }
        this.inclusive = inclusive;
        this.sortedRanges = sortedRanges;
        this.discreteSetMarker = requireNonNull(discreteSetMarker, "discreteSetMarker is null");
    }

    static SortedRangeSet none(Type type)
    {
        return new SortedRangeSet(
                type,
                new boolean[0],
                // TODO This can perhaps use an empty block singleton
                type.createBlockBuilder(null, 0).build(),
                // empty => no discrete set
                NON_DISCRETE);
    }

    static SortedRangeSet all(Type type)
    {
        return new SortedRangeSet(
                type,
                new boolean[] {false, false},
                // TODO This can perhaps use a "block with two nulls" singleton
                type.createBlockBuilder(null, 2)
                        .appendNull()
                        .appendNull()
                        .build(),
                NON_DISCRETE);
    }

    @JsonCreator
    @DoNotCall // For JSON deserialization only
    @Deprecated // Discourage usages in SPI consumers
    public static SortedRangeSet fromJson(
            @JsonProperty("type") Type type,
            @JsonProperty("inclusive") boolean[] inclusive,
            @JsonProperty("sortedRanges") Block sortedRanges,
            @JsonProperty("discreteSetMarker") DiscreteSetMarker discreteSetMarker)
    {
        if (sortedRanges instanceof BlockBuilder) {
            throw new IllegalArgumentException("sortedRanges must be a block: " + sortedRanges);
        }
        return new SortedRangeSet(type, inclusive.clone(), sortedRanges, discreteSetMarker);
    }

    /**
     * Provided discrete values that are unioned together to form the SortedRangeSet
     */
    static SortedRangeSet of(Type type, Object first, Object... rest)
    {
        if (rest.length == 0) {
            return of(type, first);
        }

        BlockBuilder blockBuilder = type.createBlockBuilder(null, 1 + rest.length);
        checkNotNaN(type, first);
        writeNativeValue(type, blockBuilder, first);
        for (Object value : rest) {
            checkNotNaN(type, value);
            writeNativeValue(type, blockBuilder, value);
        }
        Block block = blockBuilder.build();

        return fromUnorderedValuesBlock(type, block);
    }

    static SortedRangeSet of(Type type, Collection values)
    {
        if (values.isEmpty()) {
            return none(type);
        }

        BlockBuilder blockBuilder = type.createBlockBuilder(null, values.size());
        for (Object value : values) {
            checkNotNaN(type, value);
            writeNativeValue(type, blockBuilder, value);
        }
        Block block = blockBuilder.build();

        return fromUnorderedValuesBlock(type, block);
    }

    private static SortedRangeSet fromUnorderedValuesBlock(Type type, Block block)
    {
        // choice of placing unordered values first or last does not matter for this code
        MethodHandle comparisonOperator = TUPLE_DOMAIN_TYPE_OPERATORS.getComparisonUnorderedLastOperator(type, simpleConvention(FAIL_ON_NULL, BLOCK_POSITION, BLOCK_POSITION));

        List indexes = new ArrayList<>(block.getPositionCount());
        for (int position = 0; position < block.getPositionCount(); position++) {
            indexes.add(position);
        }
        indexes.sort((left, right) -> compareValues(comparisonOperator, block, left, block, right));

        int[] dictionary = new int[block.getPositionCount() * 2];
        dictionary[0] = indexes.get(0);
        dictionary[1] = indexes.get(0);
        int dictionaryIndex = 2;

        for (int i = 1; i < indexes.size(); i++) {
            int compare = compareValues(comparisonOperator, block, indexes.get(i - 1), block, indexes.get(i));
            if (compare > 0) {
                throw new IllegalStateException("Values not sorted");
            }
            if (compare == 0) {
                // equal, skip
                continue;
            }
            dictionary[dictionaryIndex] = indexes.get(i);
            dictionaryIndex++;
            dictionary[dictionaryIndex] = indexes.get(i);
            dictionaryIndex++;
        }

        boolean[] inclusive = new boolean[dictionaryIndex];
        Arrays.fill(inclusive, true);

        return new SortedRangeSet(
                type,
                inclusive,
                DictionaryBlock.create(dictionaryIndex, block, dictionary),
                DISCRETE);
    }

    /**
     * Provided Ranges are unioned together to form the SortedRangeSet
     */
    static SortedRangeSet of(Range first, Range... rest)
    {
        if (rest.length == 0 && first.isSingleValue()) {
            return of(first.getType(), first.getSingleValue());
        }

        List rangeList = new ArrayList<>(rest.length + 1);
        rangeList.add(first);
        rangeList.addAll(asList(rest));
        return copyOf(first.getType(), rangeList);
    }

    static SortedRangeSet of(List rangeList)
    {
        if (rangeList.isEmpty()) {
            throw new IllegalArgumentException("cannot use empty rangeList");
        }
        return copyOf(rangeList.get(0).getType(), rangeList);
    }

    private static SortedRangeSet of(Type type, Object value)
    {
        checkNotNaN(type, value);
        Block block = nativeValueToBlock(type, value);
        return new SortedRangeSet(
                type,
                new boolean[] {true, true},
                RunLengthEncodedBlock.create(block, 2),
                DISCRETE);
    }

    static SortedRangeSet copyOf(Type type, Collection ranges)
    {
        return buildFromUnsortedRanges(type, ranges);
    }

    /**
     * Provided Ranges are unioned together to form the SortedRangeSet
     */
    public static SortedRangeSet copyOf(Type type, List ranges)
    {
        return copyOf(type, (Collection) ranges);
    }

    @Override
    @JsonProperty
    public Type getType()
    {
        return type;
    }

    @JsonProperty
    public boolean[] getInclusive()
    {
        return inclusive;
    }

    @JsonProperty
    public Block getSortedRanges()
    {
        return sortedRanges;
    }

    public List getOrderedRanges()
    {
        List ranges = new ArrayList<>(getRangeCount());
        for (int rangeIndex = 0; rangeIndex < getRangeCount(); rangeIndex++) {
            ranges.add(getRange(rangeIndex));
        }
        return unmodifiableList(ranges);
    }

    public int getRangeCount()
    {
        return inclusive.length / 2;
    }

    @Override
    public boolean isNone()
    {
        return getRangeCount() == 0;
    }

    @Override
    public boolean isAll()
    {
        if (getRangeCount() != 1) {
            return false;
        }
        return isRangeLowUnbounded(0) && isRangeHighUnbounded(0);
    }

    @Override
    public boolean isSingleValue()
    {
        return getRangeCount() == 1 && getRangeView(0).isSingleValue();
    }

    @Override
    public Object getSingleValue()
    {
        if (getRangeCount() == 1) {
            Optional