All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.trino.plugin.hive.coercions.CoercionUtils 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.plugin.hive.coercions;

import com.google.common.collect.ImmutableList;
import io.trino.metastore.HiveType;
import io.trino.metastore.type.Category;
import io.trino.metastore.type.ListTypeInfo;
import io.trino.metastore.type.MapTypeInfo;
import io.trino.metastore.type.StructTypeInfo;
import io.trino.plugin.hive.HiveStorageFormat;
import io.trino.plugin.hive.HiveTimestampPrecision;
import io.trino.plugin.hive.coercions.BooleanCoercer.BooleanToVarcharCoercer;
import io.trino.plugin.hive.coercions.DateCoercer.DateToVarcharCoercer;
import io.trino.plugin.hive.coercions.DateCoercer.VarcharToDateCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToDateCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.LongTimestampToVarcharCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToLongTimestampCoercer;
import io.trino.plugin.hive.coercions.TimestampCoercer.VarcharToShortTimestampCoercer;
import io.trino.spi.TrinoException;
import io.trino.spi.block.ArrayBlock;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.block.ColumnarArray;
import io.trino.spi.block.ColumnarMap;
import io.trino.spi.block.DictionaryBlock;
import io.trino.spi.block.LazyBlock;
import io.trino.spi.block.RowBlock;
import io.trino.spi.block.RunLengthEncodedBlock;
import io.trino.spi.type.ArrayType;
import io.trino.spi.type.BigintType;
import io.trino.spi.type.BooleanType;
import io.trino.spi.type.CharType;
import io.trino.spi.type.DateType;
import io.trino.spi.type.DecimalType;
import io.trino.spi.type.IntegerType;
import io.trino.spi.type.MapType;
import io.trino.spi.type.RowType;
import io.trino.spi.type.RowType.Field;
import io.trino.spi.type.SmallintType;
import io.trino.spi.type.TimestampType;
import io.trino.spi.type.TinyintType;
import io.trino.spi.type.Type;
import io.trino.spi.type.TypeManager;
import io.trino.spi.type.VarbinaryType;
import io.trino.spi.type.VarcharType;

import java.util.List;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkArgument;
import static io.trino.metastore.HiveType.HIVE_BOOLEAN;
import static io.trino.metastore.HiveType.HIVE_BYTE;
import static io.trino.metastore.HiveType.HIVE_DOUBLE;
import static io.trino.metastore.HiveType.HIVE_FLOAT;
import static io.trino.metastore.HiveType.HIVE_INT;
import static io.trino.metastore.HiveType.HIVE_LONG;
import static io.trino.metastore.HiveType.HIVE_SHORT;
import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT;
import static io.trino.plugin.hive.HiveStorageFormat.ORC;
import static io.trino.plugin.hive.coercions.BooleanCoercer.createVarcharToBooleanCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToDecimalCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToDoubleCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToInteger;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToRealCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDecimalToVarcharCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createDoubleToDecimalCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createIntegerNumberToDecimalCoercer;
import static io.trino.plugin.hive.coercions.DecimalCoercers.createRealToDecimalCoercer;
import static io.trino.plugin.hive.coercions.DoubleToVarcharCoercers.createDoubleToVarcharCoercer;
import static io.trino.plugin.hive.coercions.FloatToVarcharCoercers.createFloatToVarcharCoercer;
import static io.trino.plugin.hive.coercions.VarbinaryToVarcharCoercers.createVarbinaryToVarcharCoercer;
import static io.trino.plugin.hive.coercions.VarcharToIntegralNumericCoercers.createVarcharToIntegerNumberCoercer;
import static io.trino.plugin.hive.util.HiveTypeTranslator.toHiveType;
import static io.trino.plugin.hive.util.HiveTypeUtil.getType;
import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
import static io.trino.spi.block.ColumnarArray.toColumnarArray;
import static io.trino.spi.block.ColumnarMap.toColumnarMap;
import static io.trino.spi.type.DoubleType.DOUBLE;
import static io.trino.spi.type.RealType.REAL;
import static io.trino.spi.type.TimestampType.TIMESTAMP_NANOS;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

public final class CoercionUtils
{
    private CoercionUtils() {}

    public static Type createTypeFromCoercer(TypeManager typeManager, HiveType fromHiveType, HiveType toHiveType, CoercionContext coercionContext)
    {
        return createCoercer(typeManager, fromHiveType, toHiveType, coercionContext)
                .map(TypeCoercer::getFromType)
                .orElseGet(() -> getType(fromHiveType, typeManager, coercionContext.timestampPrecision()));
    }

    public static Optional> createCoercer(TypeManager typeManager, HiveType fromHiveType, HiveType toHiveType, CoercionContext coercionContext)
    {
        if (fromHiveType.equals(toHiveType)) {
            return Optional.empty();
        }

        Type fromType = getType(fromHiveType, typeManager, coercionContext.timestampPrecision());
        Type toType = getType(toHiveType, typeManager, coercionContext.timestampPrecision());
        boolean isOrcFile = coercionContext.storageFormat() == ORC;

        if (toType instanceof VarcharType toVarcharType && (fromHiveType.equals(HIVE_BYTE) || fromHiveType.equals(HIVE_SHORT) || fromHiveType.equals(HIVE_INT) || fromHiveType.equals(HIVE_LONG))) {
            return Optional.of(new IntegerNumberToVarcharCoercer<>(fromType, toVarcharType));
        }
        if (fromType instanceof VarcharType fromVarcharType && (toHiveType.equals(HIVE_BYTE) || toHiveType.equals(HIVE_SHORT) || toHiveType.equals(HIVE_INT) || toHiveType.equals(HIVE_LONG))) {
            return Optional.of(createVarcharToIntegerNumberCoercer(fromVarcharType, toType, isOrcFile));
        }
        if (fromType instanceof VarcharType fromVarcharType && toHiveType.equals(HIVE_BOOLEAN)) {
            return Optional.of(createVarcharToBooleanCoercer(fromVarcharType, isOrcFile));
        }
        if (fromType instanceof VarcharType varcharType && toHiveType.equals(HIVE_DOUBLE)) {
            return Optional.of(new VarcharToDoubleCoercer(varcharType, isOrcFile));
        }
        if (fromType instanceof VarcharType varcharType && toHiveType.equals(HIVE_FLOAT)) {
            return Optional.of(new VarcharToFloatCoercer(varcharType, isOrcFile));
        }
        if (fromType instanceof VarcharType varcharType && toType instanceof TimestampType timestampType) {
            if (timestampType.isShort()) {
                return Optional.of(new VarcharToShortTimestampCoercer(varcharType, timestampType));
            }
            return Optional.of(new VarcharToLongTimestampCoercer(varcharType, timestampType));
        }
        if (fromType instanceof VarcharType fromVarcharType && toType instanceof VarcharType toVarcharType) {
            if (narrowerThan(toVarcharType, fromVarcharType)) {
                return Optional.of(new VarcharCoercer(fromVarcharType, toVarcharType));
            }
            return Optional.empty();
        }
        if (fromType instanceof VarcharType fromVarcharType && toType instanceof DateType toDateType) {
            return Optional.of(new VarcharToDateCoercer(fromVarcharType, toDateType));
        }
        if (fromType instanceof BooleanType && toType instanceof VarcharType toVarcharType) {
            return Optional.of(new BooleanToVarcharCoercer(toVarcharType));
        }
        if (fromType instanceof VarcharType fromVarcharType && toType instanceof CharType toCharType) {
            return Optional.of(new VarcharToCharCoercer(fromVarcharType, toCharType));
        }
        if (fromType instanceof CharType fromCharType && toType instanceof CharType toCharType) {
            if (narrowerThan(toCharType, fromCharType)) {
                return Optional.of(new CharCoercer(fromCharType, toCharType));
            }
            return Optional.empty();
        }
        if (fromType instanceof CharType fromCharType && toType instanceof VarcharType toVarcharType) {
            if (!toVarcharType.isUnbounded() && toVarcharType.getBoundedLength() < fromCharType.getLength()) {
                return Optional.of(new CharToVarcharCoercer(fromCharType, toVarcharType));
            }
            return Optional.empty();
        }
        if (fromHiveType.equals(HIVE_BYTE)) {
            if (toHiveType.equals(HIVE_SHORT) || toHiveType.equals(HIVE_INT) || toHiveType.equals(HIVE_LONG)) {
                return Optional.of(new IntegerNumberUpscaleCoercer<>(fromType, toType));
            }
            if (toHiveType.equals(HIVE_DOUBLE)) {
                return Optional.of(new IntegerNumberToDoubleCoercer<>(fromType));
            }
            if (toType instanceof DecimalType toDecimalType) {
                return Optional.of(createIntegerNumberToDecimalCoercer(fromType, toDecimalType));
            }
        }
        if (fromHiveType.equals(HIVE_SHORT)) {
            if (toHiveType.equals(HIVE_INT) || toHiveType.equals(HIVE_LONG)) {
                return Optional.of(new IntegerNumberUpscaleCoercer<>(fromType, toType));
            }
            if (toHiveType.equals(HIVE_DOUBLE)) {
                return Optional.of(new IntegerNumberToDoubleCoercer<>(fromType));
            }
            if (toType instanceof DecimalType toDecimalType) {
                return Optional.of(createIntegerNumberToDecimalCoercer(fromType, toDecimalType));
            }
        }
        if (fromHiveType.equals(HIVE_INT)) {
            if (toHiveType.equals(HIVE_LONG)) {
                return Optional.of(new IntegerToBigintCoercer());
            }
            if (toHiveType.equals(HIVE_DOUBLE)) {
                return Optional.of(new IntegerNumberToDoubleCoercer<>(fromType));
            }
            if (toType instanceof DecimalType toDecimalType) {
                return Optional.of(createIntegerNumberToDecimalCoercer(fromType, toDecimalType));
            }
        }
        if (fromHiveType.equals(HIVE_LONG)) {
            if (toHiveType.equals(HIVE_DOUBLE)) {
                return Optional.of(new IntegerNumberToDoubleCoercer<>(fromType));
            }
            if (toType instanceof DecimalType toDecimalType) {
                return Optional.of(createIntegerNumberToDecimalCoercer(fromType, toDecimalType));
            }
        }
        if (fromHiveType.equals(HIVE_FLOAT) && toHiveType.equals(HIVE_DOUBLE)) {
            return Optional.of(new FloatToDoubleCoercer());
        }
        if (fromHiveType.equals(HIVE_DOUBLE) && toHiveType.equals(HIVE_FLOAT)) {
            return Optional.of(new DoubleToFloatCoercer());
        }
        if (fromType instanceof DecimalType fromDecimalType && toType instanceof DecimalType toDecimalType) {
            return Optional.of(createDecimalToDecimalCoercer(fromDecimalType, toDecimalType));
        }
        if (fromType instanceof DecimalType fromDecimalType && toType == DOUBLE) {
            return Optional.of(createDecimalToDoubleCoercer(fromDecimalType));
        }
        if (fromType instanceof DecimalType fromDecimalType && toType == REAL) {
            return Optional.of(createDecimalToRealCoercer(fromDecimalType));
        }
        if (fromType instanceof DecimalType fromDecimalType && toType instanceof VarcharType toVarcharType) {
            return Optional.of(createDecimalToVarcharCoercer(fromDecimalType, toVarcharType));
        }
        if (fromType instanceof DecimalType fromDecimalType &&
                (toType instanceof TinyintType ||
                        toType instanceof SmallintType ||
                        toType instanceof IntegerType ||
                        toType instanceof BigintType)) {
            return Optional.of(createDecimalToInteger(fromDecimalType, toType));
        }
        if (fromType == DOUBLE && toType instanceof DecimalType toDecimalType) {
            return Optional.of(createDoubleToDecimalCoercer(toDecimalType));
        }
        if (fromType == REAL && toType instanceof DecimalType toDecimalType) {
            return Optional.of(createRealToDecimalCoercer(toDecimalType));
        }
        if (fromType instanceof TimestampType) {
            if (toType instanceof VarcharType varcharType) {
                return Optional.of(new LongTimestampToVarcharCoercer(TIMESTAMP_NANOS, varcharType));
            }
            if (toType instanceof DateType toDateType) {
                return Optional.of(new LongTimestampToDateCoercer(TIMESTAMP_NANOS, toDateType));
            }
            return Optional.empty();
        }
        if (fromType instanceof DateType && toType instanceof VarcharType toVarcharType) {
            return Optional.of(new DateToVarcharCoercer(toVarcharType));
        }
        if (fromType == REAL && toType instanceof VarcharType toVarcharType) {
            return Optional.of(createFloatToVarcharCoercer(toVarcharType, isOrcFile));
        }
        if (fromType == DOUBLE && toType instanceof VarcharType toVarcharType) {
            return Optional.of(createDoubleToVarcharCoercer(toVarcharType, isOrcFile));
        }
        if ((fromType instanceof ArrayType) && (toType instanceof ArrayType)) {
            return createCoercerForList(
                    typeManager,
                    (ListTypeInfo) fromHiveType.getTypeInfo(),
                    (ListTypeInfo) toHiveType.getTypeInfo(),
                    coercionContext);
        }
        if ((fromType instanceof MapType) && (toType instanceof MapType)) {
            return createCoercerForMap(
                    typeManager,
                    (MapTypeInfo) fromHiveType.getTypeInfo(),
                    (MapTypeInfo) toHiveType.getTypeInfo(),
                    coercionContext);
        }
        if ((fromType instanceof RowType) && (toType instanceof RowType)) {
            HiveType fromHiveTypeStruct = (fromHiveType.getCategory() == Category.UNION) ? toHiveType(fromType) : fromHiveType;
            HiveType toHiveTypeStruct = (toHiveType.getCategory() == Category.UNION) ? toHiveType(toType) : toHiveType;

            return createCoercerForStruct(
                    typeManager,
                    (StructTypeInfo) fromHiveTypeStruct.getTypeInfo(),
                    (StructTypeInfo) toHiveTypeStruct.getTypeInfo(),
                    coercionContext);
        }

        if (fromType instanceof VarbinaryType && toType instanceof VarcharType varcharType) {
            return Optional.of(createVarbinaryToVarcharCoercer(varcharType, coercionContext.storageFormat()));
        }

        throw new TrinoException(NOT_SUPPORTED, format("Unsupported coercion from %s to %s", fromHiveType, toHiveType));
    }

    public static boolean narrowerThan(VarcharType first, VarcharType second)
    {
        requireNonNull(first, "first is null");
        requireNonNull(second, "second is null");
        if (first.isUnbounded() || second.isUnbounded()) {
            return !first.isUnbounded();
        }
        return first.getBoundedLength() < second.getBoundedLength();
    }

    public static boolean narrowerThan(CharType first, CharType second)
    {
        requireNonNull(first, "first is null");
        requireNonNull(second, "second is null");
        return first.getLength() < second.getLength();
    }

    private static Optional> createCoercerForList(
            TypeManager typeManager,
            ListTypeInfo fromListTypeInfo,
            ListTypeInfo toListTypeInfo,
            CoercionContext coercionContext)
    {
        HiveType fromElementHiveType = HiveType.valueOf(fromListTypeInfo.getListElementTypeInfo().getTypeName());
        HiveType toElementHiveType = HiveType.valueOf(toListTypeInfo.getListElementTypeInfo().getTypeName());

        return createCoercer(typeManager, fromElementHiveType, toElementHiveType, coercionContext)
                .map(elementCoercer -> new ListCoercer(new ArrayType(elementCoercer.getFromType()), new ArrayType(elementCoercer.getToType()), elementCoercer));
    }

    private static Optional> createCoercerForMap(
            TypeManager typeManager,
            MapTypeInfo fromMapTypeInfo,
            MapTypeInfo toMapTypeInfo,
            CoercionContext coercionContext)
    {
        HiveType fromKeyHiveType = HiveType.valueOf(fromMapTypeInfo.getMapKeyTypeInfo().getTypeName());
        HiveType fromValueHiveType = HiveType.valueOf(fromMapTypeInfo.getMapValueTypeInfo().getTypeName());
        HiveType toKeyHiveType = HiveType.valueOf(toMapTypeInfo.getMapKeyTypeInfo().getTypeName());
        HiveType toValueHiveType = HiveType.valueOf(toMapTypeInfo.getMapValueTypeInfo().getTypeName());
        Optional> keyCoercer = createCoercer(typeManager, fromKeyHiveType, toKeyHiveType, coercionContext);
        Optional> valueCoercer = createCoercer(typeManager, fromValueHiveType, toValueHiveType, coercionContext);
        MapType fromType = new MapType(
                keyCoercer.map(TypeCoercer::getFromType).orElseGet(() -> getType(fromKeyHiveType, typeManager, coercionContext.timestampPrecision())),
                valueCoercer.map(TypeCoercer::getFromType).orElseGet(() -> getType(fromValueHiveType, typeManager, coercionContext.timestampPrecision())),
                typeManager.getTypeOperators());

        MapType toType = new MapType(
                keyCoercer.map(TypeCoercer::getToType).orElseGet(() -> getType(toKeyHiveType, typeManager, coercionContext.timestampPrecision())),
                valueCoercer.map(TypeCoercer::getToType).orElseGet(() -> getType(toValueHiveType, typeManager, coercionContext.timestampPrecision())),
                typeManager.getTypeOperators());

        return Optional.of(new MapCoercer(fromType, toType, keyCoercer, valueCoercer));
    }

    private static Optional> createCoercerForStruct(
            TypeManager typeManager,
            StructTypeInfo fromStructTypeInfo,
            StructTypeInfo toStructTypeInfo,
            CoercionContext coercionContext)
    {
        ImmutableList.Builder>> coercers = ImmutableList.builder();
        ImmutableList.Builder fromField = ImmutableList.builder();
        ImmutableList.Builder toField = ImmutableList.builder();

        List fromStructFieldName = fromStructTypeInfo.getAllStructFieldNames();
        List toStructFieldNames = toStructTypeInfo.getAllStructFieldNames();

        for (int i = 0; i < toStructFieldNames.size(); i++) {
            HiveType toStructFieldType = HiveType.valueOf(toStructTypeInfo.getAllStructFieldTypeInfos().get(i).getTypeName());
            if (i >= fromStructFieldName.size()) {
                toField.add(new Field(
                        Optional.of(toStructFieldNames.get(i)),
                        getType(toStructFieldType, typeManager, coercionContext.timestampPrecision())));
                coercers.add(Optional.empty());
            }
            else {
                HiveType fromStructFieldType = HiveType.valueOf(fromStructTypeInfo.getAllStructFieldTypeInfos().get(i).getTypeName());

                Optional> coercer = createCoercer(typeManager, fromStructFieldType, toStructFieldType, coercionContext);

                fromField.add(new Field(
                        Optional.of(fromStructFieldName.get(i)),
                        coercer.map(TypeCoercer::getFromType).orElseGet(() -> getType(fromStructFieldType, typeManager, coercionContext.timestampPrecision()))));
                toField.add(new Field(
                        Optional.of(toStructFieldNames.get(i)),
                        coercer.map(TypeCoercer::getToType).orElseGet(() -> getType(toStructFieldType, typeManager, coercionContext.timestampPrecision()))));

                coercers.add(coercer);
            }
        }

        return Optional.of(new StructCoercer(RowType.from(fromField.build()), RowType.from(toField.build()), coercers.build()));
    }

    public static class ListCoercer
            extends TypeCoercer
    {
        private final TypeCoercer elementCoercer;

        public ListCoercer(ArrayType fromType, ArrayType toType, TypeCoercer elementCoercer)
        {
            super(fromType, toType);
            this.elementCoercer = requireNonNull(elementCoercer, "elementCoercer is null");
        }

        @Override
        public Block apply(Block block)
        {
            ColumnarArray arrayBlock = toColumnarArray(block);
            Block elementsBlock = elementCoercer.apply(arrayBlock.getElementsBlock());
            boolean[] valueIsNull = new boolean[arrayBlock.getPositionCount()];
            int[] offsets = new int[arrayBlock.getPositionCount() + 1];
            for (int i = 0; i < arrayBlock.getPositionCount(); i++) {
                valueIsNull[i] = arrayBlock.isNull(i);
                offsets[i + 1] = offsets[i] + arrayBlock.getLength(i);
            }
            return ArrayBlock.fromElementBlock(arrayBlock.getPositionCount(), Optional.of(valueIsNull), offsets, elementsBlock);
        }

        @Override
        protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position)
        {
            throw new UnsupportedOperationException("Not supported");
        }
    }

    public static class MapCoercer
            extends TypeCoercer
    {
        private final Optional> keyCoercer;
        private final Optional> valueCoercer;

        public MapCoercer(
                MapType fromType,
                MapType toType,
                Optional> keyCoercer,
                Optional> valueCoercer)
        {
            super(fromType, toType);
            this.keyCoercer = requireNonNull(keyCoercer, "keyCoercer is null");
            this.valueCoercer = requireNonNull(valueCoercer, "valueCoercer is null");
        }

        @Override
        public Block apply(Block block)
        {
            ColumnarMap mapBlock = toColumnarMap(block);
            Block keysBlock = keyCoercer.isEmpty() ? mapBlock.getKeysBlock() : keyCoercer.get().apply(mapBlock.getKeysBlock());
            Block valuesBlock = valueCoercer.isEmpty() ? mapBlock.getValuesBlock() : valueCoercer.get().apply(mapBlock.getValuesBlock());
            boolean[] valueIsNull = new boolean[mapBlock.getPositionCount()];
            int[] offsets = new int[mapBlock.getPositionCount() + 1];
            for (int i = 0; i < mapBlock.getPositionCount(); i++) {
                valueIsNull[i] = mapBlock.isNull(i);
                offsets[i + 1] = offsets[i] + mapBlock.getEntryCount(i);
            }
            return toType.createBlockFromKeyValue(Optional.of(valueIsNull), offsets, keysBlock, valuesBlock);
        }

        @Override
        protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position)
        {
            throw new UnsupportedOperationException("Not supported");
        }
    }

    public static class StructCoercer
            extends TypeCoercer
    {
        private final List>> coercers;

        public StructCoercer(RowType fromType, RowType toType, List>> coercers)
        {
            super(fromType, toType);
            checkArgument(toType.getTypeParameters().size() == coercers.size());
            checkArgument(fromType.getTypeParameters().size() <= coercers.size());
            this.coercers = ImmutableList.copyOf(requireNonNull(coercers, "coercers is null"));
        }

        @Override
        public Block apply(Block block)
        {
            if (block instanceof LazyBlock lazyBlock) {
                // only load the top level block so non-coerced fields are not loaded
                block = lazyBlock.getBlock();
            }

            if (block instanceof RunLengthEncodedBlock runLengthEncodedBlock) {
                RowBlock rowBlock = (RowBlock) runLengthEncodedBlock.getValue();
                RowBlock newRowBlock = RowBlock.fromNotNullSuppressedFieldBlocks(
                        1,
                        rowBlock.isNull(0) ? Optional.of(new boolean[] {true}) : Optional.empty(),
                        coerceFields(rowBlock.getFieldBlocks()));
                return RunLengthEncodedBlock.create(newRowBlock, runLengthEncodedBlock.getPositionCount());
            }
            if (block instanceof DictionaryBlock dictionaryBlock) {
                RowBlock rowBlock = (RowBlock) dictionaryBlock.getDictionary();
                // create a dictionary block for each field, by rewraping the nested fields in a new dictionary
                List fieldBlocks = rowBlock.getFieldBlocks().stream()
                        .map(dictionaryBlock::createProjection)
                        .toList();
                // coerce the wrapped fields, so only the used dictionary values are coerced
                Block[] newFields = coerceFields(fieldBlocks);
                return RowBlock.fromNotNullSuppressedFieldBlocks(
                        dictionaryBlock.getPositionCount(),
                        getNulls(dictionaryBlock),
                        newFields);
            }
            RowBlock rowBlock = (RowBlock) block;
            return RowBlock.fromNotNullSuppressedFieldBlocks(
                    rowBlock.getPositionCount(),
                    getNulls(rowBlock),
                    coerceFields(rowBlock.getFieldBlocks()));
        }

        private static Optional getNulls(Block rowBlock)
        {
            if (!rowBlock.mayHaveNull()) {
                return Optional.empty();
            }

            boolean[] valueIsNull = new boolean[rowBlock.getPositionCount()];
            for (int i = 0; i < rowBlock.getPositionCount(); i++) {
                valueIsNull[i] = rowBlock.isNull(i);
            }
            return Optional.of(valueIsNull);
        }

        private Block[] coerceFields(List fields)
        {
            Block[] newFields = new Block[coercers.size()];
            for (int i = 0; i < coercers.size(); i++) {
                Optional> coercer = coercers.get(i);
                if (coercer.isPresent()) {
                    newFields[i] = coercer.get().apply(fields.get(i));
                }
                else if (i < fields.size()) {
                    newFields[i] = fields.get(i);
                }
                else {
                    newFields[i] = RunLengthEncodedBlock.create(toType.getTypeParameters().get(i), null, fields.get(0).getPositionCount());
                }
            }
            return newFields;
        }

        @Override
        protected void applyCoercedValue(BlockBuilder blockBuilder, Block block, int position)
        {
            throw new UnsupportedOperationException("Not supported");
        }
    }

    public record CoercionContext(HiveTimestampPrecision timestampPrecision, HiveStorageFormat storageFormat)
    {
        public CoercionContext
        {
            requireNonNull(storageFormat, "storageFormat is null");
            requireNonNull(timestampPrecision, "timestampPrecision is null");
        }
    }

    public static HiveStorageFormat extractHiveStorageFormat(String deserializedClass)
    {
        for (HiveStorageFormat format : HiveStorageFormat.values()) {
            if (format.getSerde().equals(deserializedClass)) {
                return format;
            }
        }
        throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, format("SerDe %s is not supported", deserializedClass));
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy