MapToMapCast.java example

Explorer
presto-master
/*
 * 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 com.facebook.presto.operator.scalar;

import com.facebook.presto.operator.aggregation.TypedSet;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.spi.StandardErrorCode;
import com.facebook.presto.spi.block.Block;
import com.facebook.presto.spi.block.BlockBuilder;
import com.facebook.presto.spi.block.BlockBuilderStatus;
import com.facebook.presto.spi.block.InterleavedBlockBuilder;
import com.facebook.presto.spi.function.OperatorDependency;
import com.facebook.presto.spi.function.ScalarOperator;
import com.facebook.presto.spi.function.SqlType;
import com.facebook.presto.spi.function.TypeParameter;
import com.facebook.presto.spi.type.StandardTypes;
import com.facebook.presto.spi.type.Type;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;

import static com.facebook.presto.spi.function.OperatorType.CAST;
import static com.facebook.presto.spi.function.OperatorType.EQUAL;
import static com.facebook.presto.spi.type.TypeUtils.readNativeValue;
import static com.facebook.presto.spi.type.TypeUtils.writeNativeValue;

@ScalarOperator(CAST)
public final class MapToMapCast
{
    private MapToMapCast() {}

    @TypeParameter("FK")
    @TypeParameter("FV")
    @TypeParameter("TK")
    @TypeParameter("TV")
    @SqlType("map(TK,TV)")
    public static Block toMap(
            @OperatorDependency(operator = EQUAL, returnType = StandardTypes.BOOLEAN, argumentTypes = {"TK", "TK"}) MethodHandle toKeyEqualsFunction,
            @OperatorDependency(operator = CAST, returnType = "TK", argumentTypes = {"FK"}) MethodHandle keyCastFunction,
            @OperatorDependency(operator = CAST, returnType = "TV", argumentTypes = {"FV"}) MethodHandle valueCastFunction,
            @TypeParameter("FK") Type fromKeyType,
            @TypeParameter("FV") Type fromValueType,
            @TypeParameter("TK") Type toKeyType,
            @TypeParameter("TV") Type toValueType,
            ConnectorSession session,
            @SqlType("map(FK,FV)") Block fromMap)
    {
        // loop over all the parameter types and bind ConnectorSession if needed
        // TODO: binding `ConnectorSession` should be done in function invocation framework
        Class<?>[] parameterTypes = keyCastFunction.type().parameterArray();
        for (int i = 0; i < parameterTypes.length; i++) {
            if (parameterTypes[i] == ConnectorSession.class) {
                keyCastFunction = MethodHandles.insertArguments(keyCastFunction, i, session);
                break;
            }
        }
        parameterTypes = valueCastFunction.type().parameterArray();
        for (int i = 0; i < parameterTypes.length; i++) {
            if (parameterTypes[i] == ConnectorSession.class) {
                valueCastFunction = MethodHandles.insertArguments(valueCastFunction, i, session);
                break;
            }
        }

        TypedSet typedSet = new TypedSet(toKeyType, fromMap.getPositionCount() / 2);
        BlockBuilder keyBlockBuilder = toKeyType.createBlockBuilder(new BlockBuilderStatus(), fromMap.getPositionCount() / 2);
        for (int i = 0; i < fromMap.getPositionCount(); i += 2) {
            Object fromKey = readNativeValue(fromKeyType, fromMap, i);
            try {
                Object toKey = keyCastFunction.invoke(fromKey);
                if (toKey == null) {
                    throw new PrestoException(StandardErrorCode.INVALID_CAST_ARGUMENT, "map key is null");
                }
                writeNativeValue(toKeyType, keyBlockBuilder, toKey);
            }
            catch (Throwable t) {
                Throwables.propagateIfInstanceOf(t, Error.class);
                Throwables.propagateIfInstanceOf(t, PrestoException.class);
                throw new PrestoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, t);
            }
        }
        Block keyBlock = keyBlockBuilder.build();
        BlockBuilder blockBuilder = new InterleavedBlockBuilder(ImmutableList.of(toKeyType, toValueType), new BlockBuilderStatus(), fromMap.getPositionCount());
        for (int i = 0; i < fromMap.getPositionCount(); i += 2) {
            if (!typedSet.contains(keyBlock, i / 2)) {
                typedSet.add(keyBlock, i / 2);
                toKeyType.appendTo(keyBlock, i / 2, blockBuilder);
                if (fromMap.isNull(i + 1)) {
                    blockBuilder.appendNull();
                    continue;
                }

                Object fromValue = readNativeValue(fromValueType, fromMap, i + 1);
                try {
                    Object toValue = valueCastFunction.invoke(fromValue);
                    writeNativeValue(toValueType, blockBuilder, toValue);
                }
                catch (Throwable t) {
                    Throwables.propagateIfInstanceOf(t, Error.class);
                    Throwables.propagateIfInstanceOf(t, PrestoException.class);
                    throw new PrestoException(StandardErrorCode.GENERIC_INTERNAL_ERROR, t);
                }
            }
            else {
                // if there are duplicated keys, fail it!
                throw new PrestoException(StandardErrorCode.INVALID_CAST_ARGUMENT, "duplicate keys");
            }
        }
        return blockBuilder.build();
    }
}