Drill storage plugin實現原理分析

Drill Storage Plugin介紹

Drill是一個交互式SQL查詢引擎，官方默認支持的數據源有hive、hbase、kafka、kudu、mongo、opentsdb、jdbc等，其中jdbc storage plugin可以覆蓋所有支持jdbc協議的數據源，如：mysql、oracle等關係型數據庫。所有數據源的接入都是通過drill的storage plugin實現的，理論上Drill通過storage plugin機制可以支持對任何數據源進行異構查詢。

Drill作爲一個SQL查詢引擎，它跟傳統數據庫有着很多相似之處，主要包括SQL Parser、SQL Validator、Query Optimizer、Data Flow Operators等幾部分。如下圖所示，SQL Parser階段會把SQL語句解析爲SQL查詢語法樹，這個階段Storage Plugin沒有介入。 從SQL Valiator階段Storage Plugin開始介入，在這個階段會通過Storage Plugin獲取Schema信息對SQL進行校驗，如判斷表、字段是否存在等。Query Optimzer階段會把SqlNode轉換爲PhysicalPlan，在這個過程中會通過Storage Plugin獲取Planner Rule對SQL進行優化。Data Flow Operators階段是對目標數據源進行數據讀取，這部分操作是通過Storage Plugin的RecordReader實現的。

圖 1

Drill Storage Plugin加載機制

DrillBit爲drill的主類，drill啓動時會自動加載所有有效的Storage Plugin，加載時序圖如圖2所示。

圖 2

Plugin的註冊主要是在類StoragePluginRegistryImpl中完成，插件註冊主要分爲以下幾步。
第一步是加載classpath下所有drill-module.conf文件，這個文件配置了需要掃描的包路徑，在這個包路徑下接口StoragePlugin所有實現類都會被加載；第二步是校驗，首先校驗的是接口StoragePlugin的實現類的構造方法是否符合標準要求，構造方法參數必須爲3個，且三個參數的類型必須分別爲StoragePluginConfig，DrillbitContext、String。其次是校驗plugin的配置是否有效，加載plugin配置，如果是首次啓動，會讀取classpath下bootstrap-storage-plugins.json文件,每個plugin都會對應一個這樣的json文件。這個json文件最終會反序列爲StoragePluginConfig實現類對象。非首次啓動bootstrap-storage-plugins.json文件不會被加載。drill會以本地/tmp/drill/sys.storage_plugins目錄下配置文件爲準，集羣模式配置信息保存在zookeeper /drill/sys.storage_plugins目錄下。第三步是通過發射的方式進行插件實例化並註冊。整個Plugin的註冊流程如圖 3所示

圖 3

Drill查詢流程分析

在正式介紹Drill Storage Plugin開發實戰之前我們先了解下Drill的查詢流程，這樣有助於對Storage Plugin進行深入理解，而不是簡單的依葫蘆畫瓢。Drill是分佈式的，並且節點之間是對等的，所有drill節點都可以對外提供服務，當節點接收到sql查詢請求之後，在UserWorker中會拉起一個Foreman線程來單獨處理這個請求，Foreman會完成sql到物理執行計劃的轉換，並根據物理執行計劃切分成可並行執行的Fragment，Foreman根據一定的算法把Fragment分發到本機或者其他drill節點進行執行，執行完之後會在接收初始請求的Drill節點中進行結果合併，然後返回給客戶端。如圖4 所示。

圖 4

一條SQL到物理執行計劃，會經過SqlNode（sql節點解析樹）、RelNode（關係表達式）、DrillRel（drill關係表達式）、Prel（物理關係表達式）、PhysicalPlan（物理執行計劃）幾個步驟的轉換。如圖 5所示

圖 5

SqlNode、RelNode、DrillRel、Prel都是樹形結構，以一條簡單查詢druid數據源的SQL爲例，SQL->SqlNode->RelNode這兩個階段只會存在節點之間的轉換，不會考慮目標數據源之間的差異進行SQL優化和改寫之類的動作。RelNode->DrilRel節點會進行邏輯執行計劃的優化，示例中對filter進行了下推操作。DrillRel->Prel節點會進行物理執行計劃的優化。   各階段文本化之後如下所示。
原始SQL

select * from hbase.staff t1 where row_key='10000'

RelNode(關係表達式節點樹)，有3個節點分別爲LogicalProject、LogicalFilter、EnumerableTableScan

LogicalProject(row_key=[$0], f1=[$1])
  LogicalFilter(condition=[=($0, '10000')])
    EnumerableTableScan(table=[[hbase, staff]])

DrillRel（Drill關係表達式節點樹),轉換爲drill中關係表達式節點

DrillScreenRel
  DrillFilterRel(condition=[=($0, '10000')])
    DrillScanRel(table=[[hbase, staff]], groupscan=[HBaseGroupScan [HBaseScanSpec=HBaseScanSpec [
             tableName=staff, startRow=null, stopRow=null, filter=null], columns=[`row_key`, `f1`, `**`]]])

Prel（物理關係表達式節點樹），這一步應用了物理優化規則，把filter下推到scan裏面了

ScreenPrel
  ProjectPrel(row_key=[$0], f1=[$1])
    ScanPrel(groupscan=[HBaseGroupScan [HBaseScanSpec=HBaseScanSpec [
             tableName=staff, startRow=10000, stopRow=10000\x00, filter=null], columns=[`row_key`, `f1`, `**`]]])

Drill Storage Plugin開發實戰

經過前面的介紹，大家對Drill Storage Plugin作用與原理應該已經有一個比較全面的瞭解。接下來以hbase爲例詳細介紹drill storage plugin開發流程。 Hbase是一個分佈式列存數據庫，默認是不支持SQL查詢的。爲了實現在Drill中用SQL對Hbase進行異構查詢，需要實現一個Hbase的storage plugin。 下面以Hbase storage plugin爲例介紹storage plugin的開發流程。
1、在目錄contrib新建mvn模塊，如： stroage-hbase
2、在新建的模塊resource目錄新建兩個文件drill-module.conf和bootstrap-storage-plugins.json。drill-module.conf定義plugin所在的包路徑，在plugin加載的時候會用到。bootstrap-storage-plugins.json文件是一些必要連接Hbase的配置。Drill首次啓動時會用這個文件作爲Plugin的初始配置。

{
  "type": "hbase",
  "config": {
    "hbase.zookeeper.quorum": "172.168.1.100",
    "hbase.zookeeper.property.clientPort": "2181"
  },
  "size.calculator.enabled": false,
  "enabled": true
}

3、修改UserBitShared.proto文件，在CoreOpertorType對象裏面新增一行HBASE_SUB_SCAN = 33，33這個數字需根據自身實際情況進行遞增。修改proto文件之後需要重新編譯，具體參考protocol模塊下的readme.txt
4、修改distribution模塊下的bin.xml, 新增org.apache.drill.contrib:drill-hbase-storage
5、代碼實現，部分代碼剖析如下

HbaseStoragePlugin: 相當於plugin的總入口，對scheme進行註冊，加載插件配置，指定優化規則等

public class HBaseStoragePlugin extends AbstractStoragePlugin {
  private static final HBaseConnectionManager hbaseConnectionManager = HBaseConnectionManager.INSTANCE;

  private final HBaseStoragePluginConfig storeConfig;
  private final HBaseSchemaFactory schemaFactory;
  private final HBaseConnectionKey connectionKey;

  private final String name;

  //構造方法，參數必須是3個，且類型需要匹配
  public HBaseStoragePlugin(HBaseStoragePluginConfig storeConfig, DrillbitContext context, String name)
      throws IOException {
    super(context, name);
    this.schemaFactory = new HBaseSchemaFactory(this, name);
    this.storeConfig = storeConfig;
    this.name = name;
    this.connectionKey = new HBaseConnectionKey();
  }

  //註冊schema 
  @Override
  public void registerSchemas(SchemaConfig schemaConfig, SchemaPlus parent) throws IOException {
    schemaFactory.registerSchemas(schemaConfig, parent);
  }
  //添加物理優化規則
  @Override
  public Set<StoragePluginOptimizerRule> getPhysicalOptimizerRules(OptimizerRulesContext optimizerRulesContext) {
    return ImmutableSet.of(HBasePushFilterIntoScan.FILTER_ON_SCAN, HBasePushFilterIntoScan.FILTER_ON_PROJECT);
  }
}

HbaseStoragePluginConfig: Plugin配置，參數與bootstrap-storage-plugins.json對應

HBaseSchemaFactory: Schema工廠，Schema相當於一個表元數據，包括表名、字段、以及字段類型等信息

public class HBaseSchemaFactory extends AbstractSchemaFactory {
  //註冊schema，schema是有層級，查詢時每層之間用.分隔
  @Override
  public void registerSchemas(SchemaConfig schemaConfig, SchemaPlus parent) throws IOException {
    HBaseSchema schema = new HBaseSchema(getName());
    SchemaPlus hPlus = parent.add(getName(), schema);
    schema.setHolder(hPlus);
  }

  class HBaseSchema extends AbstractSchema {

    HBaseSchema(String name) {
      super(Collections.emptyList(), name);
    }
    // hbase schema只有一層
    @Override
    public AbstractSchema getSubSchema(String name) {
      return null;
    }

    @Override
    public Table getTable(String name) {
      HBaseScanSpec scanSpec = new HBaseScanSpec(name);
      try {
        return new DrillHBaseTable(getName(), plugin, scanSpec);
      } catch (Exception e) {
        // Calcite firstly looks for a table in the default schema, if the table was not found,
        // it looks in the root schema.
        // If the table does not exist, a query will fail at validation stage,
        // so the error should not be thrown here.
        logger.warn("Failure while loading table '{}' for database '{}'.", name, getName(), e.getCause());
        return null;
      }
    }

    //調用hbase提供api，獲取表信息
    @Override
    public Set<String> getTableNames() {
      try(Admin admin = plugin.getConnection().getAdmin()) {
        HTableDescriptor[] tables = admin.listTables();
        Set<String> tableNames = Sets.newHashSet();
        for (HTableDescriptor table : tables) {
          tableNames.add(new String(table.getTableName().getNameAsString()));
        }
        return tableNames;
      } catch (Exception e) {
        logger.warn("Failure while loading table names for database '{}'.", getName(), e.getCause());
        return Collections.emptySet();
      }
    }

    @Override
    public String getTypeName() {
      return HBaseStoragePluginConfig.NAME;
    }
  }
}

public abstract class AbstractHBaseDrillTable extends DrillTable {
  protected HTableDescriptor tableDesc;

  public AbstractHBaseDrillTable(String storageEngineName, StoragePlugin plugin, Object selection) {
    super(storageEngineName, plugin, selection);
  }

  //字段類型轉換，把hbase中的字段類型映射爲SQL類型
  @Override
  public RelDataType getRowType(RelDataTypeFactory typeFactory) {
    ArrayList<RelDataType> typeList = new ArrayList<>();
    ArrayList<String> fieldNameList = new ArrayList<>();

    fieldNameList.add(ROW_KEY);
    typeList.add(typeFactory.createSqlType(SqlTypeName.ANY));

    Set<byte[]> families = tableDesc.getFamiliesKeys();
    for (byte[] family : families) {
      fieldNameList.add(Bytes.toString(family));
      //family映射爲map結構
      typeList.add(typeFactory.createMapType(typeFactory.createSqlType(SqlTypeName.VARCHAR), typeFactory.createSqlType(SqlTypeName.ANY)));
    }
    return typeFactory.createStructType(typeList, fieldNameList);
  }
}

HbaseSubScan: 關係表達式的葉子節點，目標數據源能夠識別的查詢語言會在這裏面定義
HbaseGroupScan: SubScan的一個超集
HbaseScanBatchCreator:根據節點泛型HbaseSubScan反射獲取，獲取HbaseSubScan參數並創建HbaseRecordReader對象
HbaseRecordReader:實現對目標數據源的進行記錄讀取,setup方法是在讀取記錄之前進行一些初始化工作， next方法中會調用hbase的api獲取數據並放入OutputMutator對象中。

Rule: drill的優化規則，可用在邏輯計劃、物理計劃等優化階段

實現一個Storage Plugin主要難點是在如何實現優化規則，where條件、聚合函數、分組、排序等是否可以下推都是由優化規則決定。下面以一個where條件下推爲例介紹如何實現一個Rule。如圖7所示，Filter經過下推轉換爲一顆 等價的查詢樹

圖 6

Drill中優化規則很多，所有規則都是StoragePluginOptimizerRule類的子類，在進行邏輯計劃和物理計劃優化時並不是所有規則都會應用，只有匹配上的規則纔會應用。匹配策略分兩級，一級匹配比較粗略，只要查詢節點樹最小子樹與規則類的構造放中操作類型class匹配就算匹配。如圖7左邊圈中部分和圖8圈中部分所示。二級匹配是在matches方法，返回true纔會執行onMatch方法進行關係表達式等價轉換，這個方法默認是返回true，需要根據實際情況決定是否重寫。在這個列子中我們進一步判斷GroupScan是否是HbaseGroupScan實例，也就是說只有查詢Hbase數據源的查詢纔會匹配這個規則。這裏要說明一點的是，你在其中一個Storage Plugin中寫的規則，對其他Storage Plugin來說都是可以使用的。

public abstract class HBasePushFilterIntoScan extends StoragePluginOptimizerRule {

  private HBasePushFilterIntoScan(RelOptRuleOperand operand, String description) {
    super(operand, description);
  }

  //FilterPrel.class、ScanPrel.class與圖6圈中的部分匹配
  public static final StoragePluginOptimizerRule FILTER_ON_SCAN = 
      new HBasePushFilterIntoScan(RelOptHelper.some(
            FilterPrel.class, RelOptHelper.any(ScanPrel.class)), "HBasePushFilterIntoScan:Filter_On_Scan") {

    @Override
    public void onMatch(RelOptRuleCall call) {
      final ScanPrel scan = (ScanPrel) call.rel(1);
      final FilterPrel filter = (FilterPrel) call.rel(0);
      final RexNode condition = filter.getCondition();

      HBaseGroupScan groupScan = (HBaseGroupScan)scan.getGroupScan();
      if (groupScan.isFilterPushedDown()) {
        /*
         * The rule can get triggered again due to the transformed "scan => filter" sequence
         * created by the earlier execution of this rule when we could not do a complete
         * conversion of Optiq Filter's condition to HBase Filter. In such cases, we rely upon
         * this flag to not do a re-processing of the rule on the already transformed call.
         */
        return;
      }

      doPushFilterToScan(call, filter, null, scan, groupScan, condition);
    }

    //二級匹配
    @Override
    public boolean matches(RelOptRuleCall call) {
      final ScanPrel scan = (ScanPrel) call.rel(1);
      //hbase數據源纔會匹配
      if (scan.getGroupScan() instanceof HBaseGroupScan) {
        return super.matches(call);
      }
      return false;
    }
  };

  protected void doPushFilterToScan(final RelOptRuleCall call, final FilterPrel filter, final ProjectPrel project, final ScanPrel scan, final HBaseGroupScan groupScan, final RexNode condition) {

    final LogicalExpression conditionExp = DrillOptiq.toDrill(new DrillParseContext(PrelUtil.getPlannerSettings(call.getPlanner())), scan, condition);
    final HBaseFilterBuilder hbaseFilterBuilder = new HBaseFilterBuilder(groupScan, conditionExp);
    final HBaseScanSpec newScanSpec = hbaseFilterBuilder.parseTree();
    if (newScanSpec == null) {
      return; //no filter pushdown ==> No transformation.
    }

    final HBaseGroupScan newGroupsScan = new HBaseGroupScan(groupScan.getUserName(), groupScan.getStoragePlugin(),
        newScanSpec, groupScan.getColumns());
    newGroupsScan.setFilterPushedDown(true);

    //filter下推至scan中
    final ScanPrel newScanPrel = ScanPrel.create(scan, filter.getTraitSet(), newGroupsScan, scan.getRowType());

    // Depending on whether is a project in the middle, assign either scan or copy of project to childRel.
    final RelNode childRel = project == null ? newScanPrel : project.copy(project.getTraitSet(), ImmutableList.of(newScanPrel));

    if (hbaseFilterBuilder.isAllExpressionsConverted()) {
        /*
         * Since we could convert the entire filter condition expression into an HBase filter,
         * we can eliminate the filter operator altogether.
         */
      call.transformTo(childRel);
    } else {
      call.transformTo(filter.copy(filter.getTraitSet(), ImmutableList.of(childRel)));
    }
  }

}
 

效果演示

演示數據準備，表staff包含一個列簇f1, 數據詳細信息如下

row_key	f1:name	f1:sex	f1:age
10000	張三	男	18
10001	李四	男	28
10002	王五	男	38

演示SQL 1

select * from hbase.staff t1 where row_key='10000'

結果1

從上圖結果可看出，同一個列簇f1是在同一個字段顯示的，並且是一個json格式，列值都是經過編碼的，這些結果的輸出形式都是在HbaseRecordReader類中指定的，在類HbaseRecordReader中指定了row_key的輸出類型爲VarBinary, 列簇的輸出類型爲Map，Map中value爲VarBinary類型。如果想要個列單獨顯示，SQL可以按以下方式書寫。

演示SQL 2

select cast(row_key as varchar) row_key, 
cast(t1.f1.name as varchar) name, 
cast(t1.f1.sex as varchar) sex,
cast(t1.f1.age as varchar) age
from hbase.staff t1 limit 10

結果2

演示SQL 3

按列條件查詢

select cast(row_key as varchar) row_key, 
cast(t1.f1.name as varchar) name, 
cast(t1.f1.sex as varchar) sex,
cast(t1.f1.age as varchar) age
from hbase.staff t1 where t1.f1.name='張三'

結果3

小結

本文對Drill SQL查詢流程、Storage Plugin加載機制、以及Storage Plugin實現原理進行了分析。 希望對讀者自己實現一個Storage Plugin有所幫助