Flink SQL自定義TableSource讀取Kudu數據

本文基於Flink1.9，之前文章基於Flink1.6。在Flink的官方文檔中提供了很多connector用於連接外部系統數據源。如果提供的connector不能滿足需要，還可以通過自定義方式定義讀取外部數據源的邏輯。本文的背景就是使用SQL查詢批量數據，但是批量數據存儲在kudu中，由於沒有提供connector所以需要自定義讀取數據邏輯。

官方文檔中給出了自定義批量數據讀取的實現方式：

然而實際API卻顯示該類已經不推薦使用了：

於是使用InputFormatTableSource抽象類來實現此功能，我們繼承該類

package com.xxx.bigdata;

import org.apache.flink.api.common.io.InputFormat;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.TableSchema;
import org.apache.flink.table.sources.InputFormatTableSource;
import org.apache.flink.table.types.DataType;

public class KuduTableSource extends InputFormatTableSource {


    @Override
    public InputFormat getInputFormat() {
        return new KuduInputFormat();
    }

    @Override
    public TableSchema getTableSchema() {
        String[] names = {"info"};
        DataType[] dataTypes = {DataTypes.STRING()};
        TableSchema tableSchema = TableSchema.builder().fields(names, dataTypes).build();
        return tableSchema;
    }

    @Override
    public DataType getProducedDataType() {
        return getTableSchema().getFieldDataTypes()[0];
    }
}

上面的代碼中，getTableSchema是定義Table的schema描述，getProducedDataType定義數據輸出類型，而父類InputFormatTableSource裏面的一個核心抽象方法是：

public abstract InputFormat<T, ?> getInputFormat();

該方法的功能就是讀數的邏輯，我們需要重寫此方法。

package com.xxx.bigdata;

import org.apache.flink.api.common.io.GenericInputFormat;
import org.apache.flink.core.io.GenericInputSplit;
import org.apache.kudu.client.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.StringJoiner;

public class KuduInputFormat extends GenericInputFormat<String> {
    private static Logger LOG = LoggerFactory.getLogger(KuduInputFormat.class);

    private KuduClient client;
    private KuduTable table;
    private KuduScanner scanner;
    private static final String IMPALA_TABLE_PREFIX = "impala::";
    private static final String KUDU_MASTER_ADDRESS = "192.168.1.1:7051";
    private String tableName = "class";
    List<String> columnsList = new ArrayList<>();
    private RowResultIterator iterator;

    @Override
    public void open(GenericInputSplit split) throws IOException {
        super.open(split);
        client = new KuduClient.KuduClientBuilder(KUDU_MASTER_ADDRESS).build();
        table = client.openTable(IMPALA_TABLE_PREFIX + tableName);
        columnsList.add("id");
        columnsList.add("teacher_id");
        scanner = client.newScannerBuilder(table).setProjectedColumnNames(columnsList).build();
        iterator = scanner.nextRows();
    }

    @Override
    public boolean reachedEnd() {
        return iterator.hasNext() == false;
    }

    @Override
    public String nextRecord(String reuse) {
        StringJoiner joiner = new StringJoiner(":");
        RowResult rowResult = iterator.next();
        String id = String.valueOf(rowResult.getLong("id"));
        String teacher_id = String.valueOf(rowResult.getLong("teacher_id"));
        joiner.add(id).add(teacher_id);
        return joiner.toString();
    }

    @Override
    public GenericInputSplit[] createInputSplits(int numSplits) {
        GenericInputSplit[] splits = new GenericInputSplit[1];
        splits[0] = new GenericInputSplit(0, 1);
        return splits;
    }

    @Override
    public void close() throws IOException {
        scanner.close();
        client.close();
    }
}

1.open用於打開一些連接，如數據庫，文件系統等，reachedEnd判斷數據是否已經讀完，nextRecord用於讀取下條數據並返回處理結果數據，createInputSplits方法是創建數據輸入分片，這裏不重寫，默認調用父類的同名方法，產生8個分片，即8個並行度。close裏進行一些資源的釋放。

2.這裏繼承的是GenericInputFormat類，其實還有其他的類可以使用：

所有這些類都有一個共有的頂層接口InputFormat：

不同的InputFormat都是基於不同的功能實現不同的方法，有興趣的可以打開看一下。那麼這些方法之間的調用鏈是怎樣的呢？如下所示：

每次讀取下一條記錄之前都會判斷記錄是否讀完，如果讀完直接執行close方法，否則讀取下一條記錄。這裏重點說一下createInputSplits創建輸入分片，其實在Kudu(類似HBase)的scan讀取中，我們要做到並行讀取，在open方法構建scanner時，應該每一個scanner對應一段範圍內的rowkey，這樣才能並行讀取，不重複讀取。因此，基於此，需要重寫createInputSplits和getInputSplitAssigner，創建scanner的時候也需要做相應的改動，這裏不做詳細說明，程序入口如下：

package com.xxx.bigdata;

import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.java.BatchTableEnvironment;

public class BatchEngine {
    public static void main(String[] args) throws Exception {
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        BatchTableEnvironment batchEnv = BatchTableEnvironment.create(env);
        batchEnv.registerTableSource("class", new KuduTableSource());
        Table result = batchEnv.sqlQuery("select * from class_room");
        batchEnv.toDataSet(result, String.class).print();
        batchEnv.execute("kudu job");
    }
}

本地測試：

由於程序裏沒有定義sink，我們只是print了一下，因此會報上述錯誤，無關緊要。