Starrocks 主键查询和分区裁剪/bucket裁剪

背景

在之前的Starrocks ShortCircuit短路径的调度,我们只是提到经过了SQL优化模块，就直接进入到了发送BE阶段，其实在这里还有一个很重要的部分，就是分区裁剪和tablet裁剪,经过这两个步骤，才会发送到对应的BE节点上,也从侧面解释了StarRocks 各类索引以及存储位置详解中的主键为啥会在tablet层级（经过了tablet裁剪后，就可以直接通过主键索引获取到对应的数据），这里也顺便了解一下 分区裁剪和tablet裁剪 的实现。

分区裁剪PartitionPruneRule

• 作用：在“分区（Partition）”维度剔除不相关的粗粒度数据集合。
• 原理： 该 Rule 提取查询中的过滤条件（例如 WHERE pk = 123 或者 date = ‘2023-01-01’），并将这些条件（Predicate）与表的分区信息（Range 或 List）进行求交集运算。
• 在主键查询中的表现：

1. 如果你的表是非分区表（即整表只有一个默认分区），或者查询条件中未带上分区列，这一步会原封不动，保留全部分区或者这单一分区。
2. 如果查询条件中包含了分区列（例如通常按时间范围分区，查询包含时间条件），PartitionPruneRule会计算出特定的值落在哪个分区，剔除（Prune）掉所有不符合条件的 Partition。

• 收益：避免扫描几百到上千个无关的分区元数据和底层文件。
  这里的关键代码是在OptOlapPartitionPruner.prunePartitions:

 public static LogicalOlapScanOperator prunePartitions(LogicalOlapScanOperator logicalOlapScanOperator) {
        List<Long> selectedPartitionIds = null;
        OlapTable table = (OlapTable) logicalOlapScanOperator.getTable();

        PartitionInfo partitionInfo = table.getPartitionInfo();

        if (partitionInfo.isRangePartition()) {
            selectedPartitionIds = rangePartitionPrune(table, (RangePartitionInfo) partitionInfo, logicalOlapScanOperator);
        } else if (partitionInfo.getType() == PartitionType.LIST) {
            selectedPartitionIds = listPartitionPrune(table, (ListPartitionInfo) partitionInfo, logicalOlapScanOperator);
        }

这里会根据分区类型不同，选择不同的处理路径：

• Range 分区裁剪 (rangePartitionPrune)：内部利用查询中解析出的 where col1 >= ‘X’ and col1 <= ‘Y’，构建出一个 Range 过滤器，遍历所有 Partition 的区间 [low, high)，找交集，交集不为空则保留这个分区。
• List 分区裁剪 (listPartitionPrune)：如果是 city in (‘beijing’, ‘shanghai’) 这种分区，它会直接找到这几个离散的分区 ID
  这里的代码和其他引擎的大同小异，读者自行查看即可

tablet裁剪DistributionPruneRule

• 作用：在裁剪后剩余的分区中，进一步在“分桶 / 分片（Tablet）”维度精确定位数据。这个规则必须在 PartitionPruneRule 执行完成之后进行。

• 原理： StarRocks 的分布式存储架构中，每个分区会被切分成多个分桶（Tablet），通常是通过主键 Hash 的方式来切分（即主键表的 hash distribution 配置）。
  DistributionPruneRule根据查询的分布列（Distribution columns，主键表里即对应主键字段）条件进行哈希计算。如果查询谓词给出的是主键的精确匹配等值条件（如 id = 1）：规则类底层的 HashDistributionPruner 会对查询值 “1” 运行与导入时一样的数据 Hash 函数，直接计算出该数据必然落在该分区内的**某 1 个特定的 Tablet（Bucket）**上。

• 在主键查询中的表现： 经过这一步，原本需要在选定分区下扫描的所有 Tablet (通常会有几到几百个)，会被硬核裁剪到只剩 1 个指定的 Tablet ID。

• 收益：彻底消除无需读取的分桶 I/O 开销和分布式执行资源调度的开销。

这里只说核心部分：
以 Hash 分桶为例,关键的代码在于HashDistributionPruner.prune

public Collection<Long> prune(int columnId, HashDistributionKey hashKey, int complex) {
        if (columnId == distributionColumns.size()) {
            // compute Hash Key
            long hashValue = hashKey.getHashValue();
            return Lists.newArrayList(tabletIdsInOrder.get((int) ((hashValue & 0xffffffff) % tabletIdsInOrder.size())));
        }
        Column keyColumn = distributionColumns.get(columnId);
...
  public long getHashValue() {
        CRC32 hashValue = new CRC32();
        int i = 0;
        for (LiteralExpr expr : keys) {
            ByteBuffer buffer = expr.getHashValue(types.get(i));
            hashValue.update(buffer.array(), 0, buffer.limit());
            i++;
        }
        return hashValue.getValue();
    }

bucketId是通过如下方法获取的bucketId = (hashValue & 0xffffffff) % bucketNum ,用这个 hashValue 对 bucket 数取模，得到 bucket 序号，然后用该序号去 tabletIdsInOrder 里取对应的 tabletId
而这里的分桶算法和当初写入数据时计算的bucketId是一致的，核心代码在OlapTablePartitionParam::find_tablets():

Status OlapTablePartitionParam::find_tablets(Chunk* chunk, std::vector<OlapTablePartition*>* partitions,
                                             std::vector<uint32_t>* hashes, std::vector<uint8_t>* selection,
                                             std::vector<int>* invalid_row_indexs, int64_t txn_id,
                                             std::vector<std::vector<std::string>>* partition_not_exist_row_values) {
    size_t num_rows = chunk->num_rows();
    partitions->resize(num_rows);

    _compute_hashes(chunk, hashes);
...

void OlapTablePartitionParam::_compute_hashes(const Chunk* chunk, std::vector<uint32_t>* hashes) {
    size_t num_rows = chunk->num_rows();
    hashes->assign(num_rows, 0);

    if (is_hash_distribution()) {
        for (size_t i = 0; i < _distributed_slot_descs.size(); ++i) {
            _distributed_columns[i] = chunk->get_column_by_slot_id(_distributed_slot_descs[i]->id()).get();
            _distributed_columns[i]->crc32_hash(&(*hashes)[0], 0, num_rows);
        }
    } else if (is_random_distribution()) {
        uint32_t r = _rand.Next();
        for (auto i = 0; i < num_rows; ++i) {
            (*hashes)[i] = r++;
        }
    }
}

对于这里的Hash算法，采用的是CRC32,和tablet裁剪的时候是一致的：

• 单列分桶：hash = CRC32(col)
• 多列分桶：hash = CRC32( CRC32(0, col1), col2, … )（按列顺序累积）