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

org.apache.flink.table.planner.plan.nodes.exec.common.CommonExecLegacySink Maven / Gradle / Ivy

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.flink.table.planner.plan.nodes.exec.common;

import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.dag.Transformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSink;
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
import org.apache.flink.table.api.TableException;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.planner.codegen.CodeGenUtils;
import org.apache.flink.table.planner.codegen.CodeGeneratorContext;
import org.apache.flink.table.planner.codegen.SinkCodeGenerator;
import org.apache.flink.table.planner.delegation.PlannerBase;
import org.apache.flink.table.planner.plan.nodes.exec.ExecEdge;
import org.apache.flink.table.planner.plan.nodes.exec.ExecNode;
import org.apache.flink.table.planner.plan.nodes.exec.ExecNodeBase;
import org.apache.flink.table.planner.plan.nodes.exec.InputProperty;
import org.apache.flink.table.planner.plan.nodes.exec.MultipleTransformationTranslator;
import org.apache.flink.table.planner.sinks.DataStreamTableSink;
import org.apache.flink.table.planner.sinks.TableSinkUtils;
import org.apache.flink.table.runtime.operators.CodeGenOperatorFactory;
import org.apache.flink.table.runtime.typeutils.TypeCheckUtils;
import org.apache.flink.table.sinks.AppendStreamTableSink;
import org.apache.flink.table.sinks.RetractStreamTableSink;
import org.apache.flink.table.sinks.StreamTableSink;
import org.apache.flink.table.sinks.TableSink;
import org.apache.flink.table.sinks.UpsertStreamTableSink;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;

import javax.annotation.Nullable;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * Base {@link ExecNode} to to write data into an external sink defined by a {@link TableSink}.
 *
 * @param  The return type of the {@link TableSink}.
 */
public abstract class CommonExecLegacySink extends ExecNodeBase
        implements MultipleTransformationTranslator {
    protected final TableSink tableSink;
    protected final @Nullable String[] upsertKeys;
    protected final boolean needRetraction;
    protected final boolean isStreaming;

    public CommonExecLegacySink(
            TableSink tableSink,
            @Nullable String[] upsertKeys,
            boolean needRetraction,
            boolean isStreaming,
            InputProperty inputProperty,
            LogicalType outputType,
            String description) {
        super(Collections.singletonList(inputProperty), outputType, description);
        this.tableSink = tableSink;
        this.upsertKeys = upsertKeys;
        this.needRetraction = needRetraction;
        this.isStreaming = isStreaming;
    }

    @SuppressWarnings("unchecked")
    @Override
    protected Transformation translateToPlanInternal(PlannerBase planner) {
        if (tableSink instanceof StreamTableSink) {
            final Transformation transform;
            if (tableSink instanceof RetractStreamTableSink) {
                transform = translateToTransformation(planner, true);
            } else if (tableSink instanceof UpsertStreamTableSink) {
                UpsertStreamTableSink upsertSink = (UpsertStreamTableSink) tableSink;
                final boolean isAppendOnlyTable = !needRetraction;
                upsertSink.setIsAppendOnly(isAppendOnlyTable);
                if (upsertKeys != null) {
                    upsertSink.setKeyFields(upsertKeys);
                } else {
                    if (isAppendOnlyTable) {
                        upsertSink.setKeyFields(null);
                    } else {
                        throw new TableException(
                                "UpsertStreamTableSink requires that Table has a full primary keys if it is updated.");
                    }
                }

                transform = translateToTransformation(planner, true);
            } else if (tableSink instanceof AppendStreamTableSink) {
                // verify table is an insert-only (append-only) table
                if (needRetraction) {
                    throw new TableException(
                            "AppendStreamTableSink requires that Table has only insert changes.");
                }
                transform = translateToTransformation(planner, false);
            } else {
                if (isStreaming) {
                    throw new TableException(
                            "Stream Tables can only be emitted by AppendStreamTableSink, "
                                    + "RetractStreamTableSink, or UpsertStreamTableSink.");
                } else {
                    transform = translateToTransformation(planner, false);
                }
            }

            final DataStream dataStream = new DataStream(planner.getExecEnv(), transform);
            final DataStreamSink dsSink =
                    ((StreamTableSink) tableSink).consumeDataStream(dataStream);
            if (dsSink == null) {
                throw new TableException(
                        String.format(
                                "The StreamTableSink#consumeDataStream(DataStream) must be implemented "
                                        + "and return the sink transformation DataStreamSink. "
                                        + "However, %s doesn't implement this method.",
                                tableSink.getClass().getCanonicalName()));
            }
            return dsSink.getTransformation();
        } else if (tableSink instanceof DataStreamTableSink) {
            // In case of table to DataStream through
            // StreamTableEnvironment#toAppendStream/toRetractStream,
            // we insert a DataStreamTableSink that wraps the given DataStream as a LogicalSink. It
            // is no real table sink, so we just need translate its input to Transformation.
            return (Transformation)
                    translateToTransformation(
                            planner, ((DataStreamTableSink) tableSink).withChangeFlag());
        } else {
            throw new TableException(
                    String.format(
                            "Only Support StreamTableSink! However %s is not a StreamTableSink.",
                            tableSink.getClass().getCanonicalName()));
        }
    }

    /** Check whether the given row type is legal and do some conversion if needed. */
    protected abstract RowType checkAndConvertInputTypeIfNeeded(RowType inputRowType);

    /**
     * Translates {@link TableSink} into a {@link Transformation}.
     *
     * @param withChangeFlag Set to true to emit records with change flags.
     * @return The {@link Transformation} that corresponds to the translated {@link TableSink}.
     */
    @SuppressWarnings("unchecked")
    private Transformation translateToTransformation(
            PlannerBase planner, boolean withChangeFlag) {
        // if no change flags are requested, verify table is an insert-only (append-only) table.
        if (!withChangeFlag && needRetraction) {
            throw new TableException(
                    "Table is not an append-only table. "
                            + "Use the toRetractStream() in order to handle add and retract messages.");
        }

        final ExecEdge inputEdge = getInputEdges().get(0);
        final Transformation inputTransform =
                (Transformation) inputEdge.translateToPlan(planner);
        final RowType inputRowType = (RowType) inputEdge.getOutputType();
        final RowType convertedInputRowType = checkAndConvertInputTypeIfNeeded(inputRowType);
        final DataType resultDataType = tableSink.getConsumedDataType();
        if (CodeGenUtils.isInternalClass(resultDataType)) {
            return (Transformation) inputTransform;
        } else {
            final int rowtimeIndex = getRowtimeIndex(inputRowType);

            final DataType physicalOutputType =
                    TableSinkUtils.inferSinkPhysicalDataType(
                            resultDataType, convertedInputRowType, withChangeFlag);

            final TypeInformation outputTypeInfo =
                    SinkCodeGenerator.deriveSinkOutputTypeInfo(
                            tableSink, physicalOutputType, withChangeFlag);

            final CodeGenOperatorFactory converterOperator =
                    SinkCodeGenerator.generateRowConverterOperator(
                            new CodeGeneratorContext(planner.getTableConfig()),
                            convertedInputRowType,
                            tableSink,
                            physicalOutputType,
                            withChangeFlag,
                            "SinkConversion",
                            rowtimeIndex);
            return new OneInputTransformation<>(
                    inputTransform,
                    "SinkConversionTo" + resultDataType.getConversionClass().getSimpleName(),
                    converterOperator,
                    outputTypeInfo,
                    inputTransform.getParallelism());
        }
    }

    private int getRowtimeIndex(RowType inputRowType) {
        int rowtimeIndex = -1;
        final List rowtimeFieldIndices = new ArrayList<>();
        for (int i = 0; i < inputRowType.getFieldCount(); ++i) {
            if (TypeCheckUtils.isRowTime(inputRowType.getTypeAt(i))) {
                rowtimeFieldIndices.add(i);
            }
        }
        if (rowtimeFieldIndices.size() == 1) {
            rowtimeIndex = rowtimeFieldIndices.get(0);
        }
        return rowtimeIndex;
    }
}