presto连接器-grouped_execution功能要实现的接口

原理

分组执行是presto在2017年就支持的功能，它的原理是：

根据相同字段（orderid）分桶（bucketing）且分桶数量相同的两个表（orders，orders_item），

在通过orderid进行join的时候，由于两个表相同的orderid都分到相同id的桶里，所以是可以独立进行join以及聚合计算的（参考MapReduer的partition过程）。

这样，每当一个桶的数据计算完成后，可以立即把这个桶所占用的内存释放掉，因此，通过控制并行处理桶的数量可以限制内存的占用。

计算理论占用的内存：优化后的内存占用=原内存占用/表的桶数量*并行处理桶的数量

完整的示例参考：

https://archongum.cn/2019/08/21/presto-query-memory-optimze/

要实现的接口

我是在基于legacy机制实现的kudu connector的基础上实现的分组执行。下面是要实现的接口：

ConnectorPartitioningHandle

这个接口可以存放我们在处理分桶逻辑时需要用到的参数。比如数据库名、表名、分桶个数、分桶字段明细信息等等。

equals
该方法用来判断join操作的左右表是否可以执行分桶操作，所以一定不要根据表名是否相等来判断true和false，否则会导致join操作永远也无法做grouped_execution。可以使用分桶个数、分桶字段类型等属性来判断。

ConnectorNodePartitioningProvider

listPartitionHandles
列出当前表所有的bucket，入参ConnectorPartitioningHandle就是我们封装的KuduPartitioningHandle，可以把接口实现需要用到的参数放到这个类中。返回的bucket只要有一个bucket_number就行了。这样在ConnectorSplitSource中会依次处理这些bucket_number。
getBucketNodeMap
仿照hive的实现即可，他的作用是构建bucket和node的映射关系，供调度使用。相关逻辑可以参考FixedSourcePartitionedScheduler的构造函数。
getSplitBucketFunction
获取bucket的bucket_number。
getBucketFunction
封装BucketFunction对象。这个对象的作用参考下文。

BucketFunction

这个类只在join操作时会使用，作用还不能确定，理论上说是获取一个page是属于哪一个bucket的。他的作用体现在唯一的接口getBucket中。

getBucket
入参page是上层封装的某个operator要处理的page中的所有分桶字段的值。每一列的值以page中一个block的形式存在。入参position是行的index。这一部分的实现可以参考HiveBucketFunction和IcebergBucketFunction。

ConnectorMetadata#getTableLayouts

这个接口即使不实现grouped_execution功能我们也要实现。通常情况只会封装一个ConnectorTableLayoutHandle对象，然后调用下图api返回一个ConnectorTableLayout对象。

    public ConnectorTableLayout(ConnectorTableLayoutHandle handle)
    {
        this(handle,
                Optional.empty(),
                TupleDomain.all(),
                Optional.empty(),
                Optional.empty(),
                Optional.empty(),
                emptyList());
    }

但是如果要实现grouped_execution就需要通过这个构造函数来将其他参数实例化：

    public ConnectorTableLayout(
            ConnectorTableLayoutHandle handle,
            Optional<List<ColumnHandle>> columns,
            TupleDomain<ColumnHandle> predicate,
            Optional<ConnectorTablePartitioning> tablePartitioning,
            Optional<Set<ColumnHandle>> streamPartitioningColumns,
            Optional<DiscretePredicates> discretePredicates,
            List<LocalProperty<ColumnHandle>> localProperties)

因为在PlanFragmenter#visitTableScan方法中会判断tablePartitioning是否存在，如果不存在则返回GroupedExecutionProperties.notCapable()，这将导致presto认定这个sql不符合grouped_execution的要求。

注：非legacy的实现会自动调用getTableProperties方法，就不需要再从getTableLayouts方法中做初始化了（非legacy也不会有getTableLayouts方法）。

PlanFragmenter#visitTableScan方法中判断tablePartitioning是否存在的逻辑如下图：

@Override
public GroupedExecutionProperties visitTableScan(TableScanNode node, Void context)
{
    Optional<TablePartitioning> tablePartitioning = metadata.getTableProperties(session, node.getTable()).getTablePartitioning();
    if (!tablePartitioning.isPresent()) {
        return GroupedExecutionProperties.notCapable();
    }
    List<ConnectorPartitionHandle> partitionHandles = nodePartitioningManager.listPartitionHandles(session, tablePartitioning.get().getPartitioningHandle());
    if (ImmutableList.of(NOT_PARTITIONED).equals(partitionHandles)) {
        return new GroupedExecutionProperties(false, false, ImmutableList.of());
    }
    else {
        return new GroupedExecutionProperties(true, false, ImmutableList.of(node.getId()));
    }
}

所以我们必须保证返回的tableProperties方法包含完整的tablePartitioning信息。

封装ConnectorTableLayout用到的其他几个参数需要通过ConnectorMetadata#getTableProperties方法一并返回。该方法的说明如下：

ConnectorMetadata#getTableProperties

该接口用来返回表的扩展信息，例如ConnectorPartitioningHandle、切分bucket的字段（streamingPartitionColumns）、分桶个数等等。

同时在这里也可以使用自定义的session参数（或者表名）判断是否要走分组执行模式。如果不走只需要将tablePartitioning信息返回为optional.empty()即可。

注意：方法返回的ConnectorTableProperties对象的localProperties属性可以设置分桶内数据的排序方式。这里如果设置的与实际不匹配，会导致group by操作中，同一个组中的数据被拆分到多个组中。Presto的逻辑我没细看，跟同事讨论认为有可能是因为因为排序被预设了，所以当处理到其他key的数据时就认为当前正在统计的分桶结束了，会自动开始统计一个新的组。

ConnectorMetadata#getTableLayout

返回一个带有ConnectorTablePartitioning等信息的ConnectorTableLayout即可。可以调用getTableProperties方法获取相关信息。

ConnectorSplitSource

getNextBatch
KuduSplitSource中有一个集合用来存放所有Splits（或者有一个方法能获取所有splits），这些Splits都是与Group（即入参partitionHandle）相关联的。

外界会遍历所有的Group，依次获取每个Group下所有的split，

如果调用一次getNextBatch方法无法获取完当前Group的所有split，可以在返回的ConnectorSplitBatch中把noMoreSplits设置为false，这样外界还会再基于同样的Group重新调用这个接口:

GetNextBatch#fetchSplits -> ConnectorAwareSplitSource#getNextBatch

直到这个Group所有的split都被获取到之后，我们把返回的ConnectorSplitBatch的noMoreSplits设置为true，外界就会再继续获取其他Group的splits了。而kudu是一个bucket切分一个split，所以每次返回一个split，并把noMoreSplits设置为true。如下：

@Override
public CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize) {
    KuduPartitionHandle kuduPartitionHandle = (KuduPartitionHandle) partitionHandle;

    KuduSplit kuduSplit = splitMap.remove(kuduPartitionHandle.getBucket());

    // 一个bucket只有一个split，所以noMoreSplits一定是true
    return completedFuture(new ConnectorSplitBatch(kuduSplit == null ? ImmutableList.of() : ImmutableList.of(kuduSplit), true));
}

外界调用getNextBatch获取所有Group的Split的逻辑，可以从SourcePartitionedScheduler#schedule开始阅读。

isFinished
是否表中所有的split都被获取完了。

    /**
     * Returns whether any more {@link ConnectorSplit} may be produced.
     * <p>
     * This method should only be called when there has been no invocation of getNextBatch,
     * or result Future of previous getNextBatch is done.
     * Calling this method at other time is not useful because the contract of such an invocation
     * will be inherently racy.
     */
    boolean isFinished();

ConnectorSplitManager#getSplits

该方法中，需要根据入参SplitSchedulingStrategy的值判断应该走自定义的ConnectorSplitSource还是FixedSplitSource

switch (splitSchedulingStrategy) {
    case UNGROUPED_SCHEDULING:
        return new FixedSplitSource(splits);
    case GROUPED_SCHEDULING:
        return new KuduSplitSource(splits);
    default:
        throw new IllegalArgumentException("Unknown splitSchedulingStrategy: " + splitSchedulingStrategy);
}

Connector

在该实现类的构造函数中增加ConnectorNodePartitioningProvider参数，并赋值给本地变量。

getNodePartitioningProvider

@Override
public ConnectorNodePartitioningProvider getNodePartitioningProvider()
{
    return nodePartitioningProvider;
}

ConnectorHandleResolver

getPartitioningHandleClass

@Override
public Class<? extends ConnectorPartitioningHandle> getPartitioningHandleClass()
{
    return KuduPartitioningHandle.class;
}

AbstractModule

configure
在该方法中设置ConnectorNodePartitioningProvider

bind(ConnectorNodePartitioningProvider.class).to(KuduNodePartitioningProvider.class).in(Scopes.SINGLETON);

注：presto的dynamic filtering和grouped_execution不能同时采用，所以需要通过以下设置将dynamic filtering关闭
set session enable_dynamic_filtering=false;

Presto连接器-grouped_execution功能要实现的接口