flink-connector-kudu/src/main/java/org/apache/flink/connectors/kudu/table/dynamic/KuduDynamicTableSource.java

/*
 * 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.connectors.kudu.table.dynamic;

import org.apache.commons.collections.CollectionUtils;
import org.apache.flink.connectors.kudu.connector.KuduFilterInfo;
import org.apache.flink.connectors.kudu.connector.KuduTableInfo;
import org.apache.flink.connectors.kudu.connector.convertor.RowResultRowDataConvertor;
import org.apache.flink.connectors.kudu.connector.reader.KuduReaderConfig;
import org.apache.flink.connectors.kudu.format.KuduRowDataInputFormat;
import org.apache.flink.connectors.kudu.table.function.lookup.KuduLookupOptions;
import org.apache.flink.connectors.kudu.table.function.lookup.KuduRowDataLookupFunction;
import org.apache.flink.connectors.kudu.table.utils.KuduTableUtils;
import org.apache.flink.table.api.TableSchema;
import org.apache.flink.table.connector.ChangelogMode;
import org.apache.flink.table.connector.source.*;
import org.apache.flink.table.connector.source.abilities.SupportsFilterPushDown;
import org.apache.flink.table.connector.source.abilities.SupportsLimitPushDown;
import org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown;
import org.apache.flink.table.expressions.ResolvedExpression;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.FieldsDataType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.utils.DataTypeUtils;
import org.apache.flink.util.Preconditions;
import org.apache.kudu.shaded.com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

import static org.apache.flink.table.utils.TableSchemaUtils.containsPhysicalColumnsOnly;

/**
 * A {@link DynamicTableSource} for Kudu.
 */
public class KuduDynamicTableSource implements ScanTableSource, SupportsProjectionPushDown,
        SupportsLimitPushDown, LookupTableSource, SupportsFilterPushDown {

    private static final Logger LOG = LoggerFactory.getLogger(KuduDynamicTableSource.class);
    private final KuduTableInfo tableInfo;
    private final KuduLookupOptions kuduLookupOptions;
    private final KuduRowDataInputFormat kuduRowDataInputFormat;
    private final transient List<KuduFilterInfo> predicates = Lists.newArrayList();
    private KuduReaderConfig.Builder configBuilder;
    private TableSchema physicalSchema;
    private String[] projectedFields;
    private transient List<ResolvedExpression> filters;

    public KuduDynamicTableSource(KuduReaderConfig.Builder configBuilder, KuduTableInfo tableInfo,
                                  TableSchema physicalSchema, String[] projectedFields,
                                  KuduLookupOptions kuduLookupOptions) {
        this.configBuilder = configBuilder;
        this.tableInfo = tableInfo;
        this.physicalSchema = physicalSchema;
        this.projectedFields = projectedFields;
        this.kuduRowDataInputFormat = new KuduRowDataInputFormat(configBuilder.build(),
                new RowResultRowDataConvertor(), tableInfo,
                predicates,
                projectedFields == null ? null : Lists.newArrayList(projectedFields));
        this.kuduLookupOptions = kuduLookupOptions;
    }

    @Override
    public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext context) {
        int keysLen = context.getKeys().length;
        String[] keyNames = new String[keysLen];
        for (int i = 0; i < keyNames.length; ++i) {
            int[] innerKeyArr = context.getKeys()[i];
            Preconditions.checkArgument(innerKeyArr.length == 1, "Kudu only support non-nested look up keys");
            keyNames[i] = this.physicalSchema.getFieldNames()[innerKeyArr[0]];
        }
        KuduRowDataLookupFunction rowDataLookupFunction = KuduRowDataLookupFunction.Builder.options()
                .keyNames(keyNames)
                .kuduReaderConfig(configBuilder.build())
                .projectedFields(projectedFields)
                .tableInfo(tableInfo)
                .kuduLookupOptions(kuduLookupOptions)
                .build();
        return TableFunctionProvider.of(rowDataLookupFunction);
    }

    @Override
    public ChangelogMode getChangelogMode() {
        return ChangelogMode.insertOnly();
    }

    @Override
    public ScanRuntimeProvider getScanRuntimeProvider(ScanContext runtimeProviderContext) {
        if (CollectionUtils.isNotEmpty(this.filters)) {
            for (ResolvedExpression filter : this.filters) {
                Optional<KuduFilterInfo> kuduFilterInfo = KuduTableUtils.toKuduFilterInfo(filter);
                if (kuduFilterInfo != null && kuduFilterInfo.isPresent()) {
                    this.predicates.add(kuduFilterInfo.get());
                }

            }
        }
        KuduRowDataInputFormat inputFormat = new KuduRowDataInputFormat(configBuilder.build(),
                new RowResultRowDataConvertor(), tableInfo,
                this.predicates,
                projectedFields == null ? null : Lists.newArrayList(projectedFields));
        return InputFormatProvider.of(inputFormat);
    }

    @Override
    public DynamicTableSource copy() {
        return new KuduDynamicTableSource(this.configBuilder, this.tableInfo, this.physicalSchema,
                this.projectedFields, this.kuduLookupOptions);
    }

    @Override
    public String asSummaryString() {
        return "kudu";
    }

    @Override
    public boolean supportsNestedProjection() {
        //  planner doesn't support nested projection push down yet.
        return false;
    }

    @Override
    public void applyProjection(int[][] projectedFields, DataType producedDataType) {
        // parser projectFields
        this.physicalSchema = projectSchema(this.physicalSchema, projectedFields);
        this.projectedFields = physicalSchema.getFieldNames();
    }

    private TableSchema projectSchema(TableSchema tableSchema, int[][] projectedFields) {
        Preconditions.checkArgument(
                containsPhysicalColumnsOnly(tableSchema),
                "Projection is only supported for physical columns.");
        TableSchema.Builder builder = TableSchema.builder();

        FieldsDataType fields =
                (FieldsDataType)
                        DataTypeUtils.projectRow(tableSchema.toRowDataType(), projectedFields);
        RowType topFields = (RowType) fields.getLogicalType();
        for (int i = 0; i < topFields.getFieldCount(); i++) {
            builder.field(topFields.getFieldNames().get(i), fields.getChildren().get(i));
        }
        return builder.build();
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        KuduDynamicTableSource that = (KuduDynamicTableSource) o;
        return Objects.equals(configBuilder, that.configBuilder) && Objects.equals(tableInfo, that.tableInfo) && Objects.equals(physicalSchema, that.physicalSchema) && Arrays.equals(projectedFields, that.projectedFields) && Objects.equals(kuduLookupOptions, that.kuduLookupOptions) && Objects.equals(kuduRowDataInputFormat, that.kuduRowDataInputFormat) && Objects.equals(filters, that.filters) && Objects.equals(predicates, that.predicates);
    }

    @Override
    public int hashCode() {
        int result = Objects.hash(configBuilder, tableInfo, physicalSchema,
                kuduLookupOptions, kuduRowDataInputFormat, filters, predicates);
        result = 31 * result + Arrays.hashCode(projectedFields);
        return result;
    }

    @Override
    public void applyLimit(long limit) {
        this.configBuilder = this.configBuilder.setRowLimit((int) limit);
    }

    @Override
    public Result applyFilters(List<ResolvedExpression> filters) {
        this.filters = filters;
        return Result.of(Collections.emptyList(), filters);
    }
}