09. 范围查询

我们已经在键值存储（key-value store，KV store）之上实现了表结构，并且能够通过主键检索记录。在本章中，我们将增加按排序顺序检索一系列记录的功能。

9.1 B树迭代器

第一步是在B树中添加范围查询功能。BIter 类型允许我们以迭代方式遍历B树。

// B树迭代器
type BIter struct {
    tree *BTree
    path []BNode // 从根节点到叶节点的路径
    pos  []uint16 // 节点中的索引
}

// 获取当前键值对
func (iter *BIter) Deref() ([]byte, []byte)
// Deref() 的前置条件
func (iter *BIter) Valid() bool
// 向后和向前移动
func (iter *BIter) Prev()
func (iter *BIter) Next()

BIter 是从根节点到叶节点中键值对的一条路径。移动迭代器只需将位置或节点移动到兄弟节点即可。

func iterPrev(iter *BIter, level int) {
    if iter.pos[level] > 0 {
       iter.pos[level]-- // 在当前节点内移动
    } else if level > 0 {
       iterPrev(iter, level - 1) // 移动到兄弟节点
    } else {
       return // 伪键
    }

    if level + 1 < len(iter.pos) {
       // 更新子节点
       node := iter.path[level]
       kid := iter.tree.get(node.getPtr(iter.pos[level]))
       iter.path[level + 1] = kid
       iter.pos[level + 1] = kid.nkeys() - 1
    }
}

func (iter *BIter) Prev() {
    iterPrev(iter, len(iter.path) - 1)
}

BTree.SeekLE 是在范围查询中查找初始位置的函数。它只是一个普通的B树查找，并记录路径。

// 查找小于或等于输入键的最接近位置
func (tree *BTree) SeekLE(key []byte) *BIter {
    iter := &BIter{tree: tree}
    for ptr := tree.root; ptr != 0; {
       node := tree.get(ptr)
       idx := nodeLookupLE(node, key)
       iter.path = append(iter.path, node)
       iter.pos = append(iter.pos, idx)
       if node.btype() == BNODE_NODE {
          ptr = node.getPtr(idx)
       } else {
          ptr = 0
       }
    }
    return iter
}

nodeLookupLE 函数仅适用于范围查询中的 “小于或等于” 运算符，对于其他3个运算符（小于、大于、大于或等于），结果可能会偏差一个。我们将通过 BTree.Seek 函数来修复这个问题。

const (
    CMP_GE = +3 // >=
    CMP_GT = +2 // >
    CMP_LT = -2 // <
    CMP_LE = -3 // <=
)

// 根据比较关系查找与键最接近的位置
func (tree *BTree) Seek(key []byte, cmp int) *BIter {
    iter := tree.SeekLE(key)
    if cmp != CMP_LE && iter.Valid() {
       cur, _ := iter.Deref()
       if!cmpOK(cur, cmp, key) {
          // 偏差一个
          if cmp > 0 {
             iter.Next()
          } else {
             iter.Prev()
          }
       }
    }
    return iter
}

// 键比较参考
func cmpOK(key []byte, cmp int, ref []byte) bool {
    r := bytes.Compare(key, ref)
    switch cmp {
    case CMP_GE:
       return r >= 0
    case CMP_GT:
       return r > 0
    case CMP_LT:
       return r < 0
    case CMP_LE:
       return r <= 0
    default:
       panic("what?")
    }
}

9.2 数据序列化

为了支持范围查询，序列化的主键必须在键值存储中能够正确比较。一种方法是反序列化主键并逐列进行比较。我们将采用另一种方法，使序列化的键字节反映其字典顺序，也就是说，无需先反序列化，就可以通过 bytes.Compare 或 memcmp 正确比较键。我们将这种技术称为 “保序编码（order-preserving encoding）”，它可以在不控制底层键值存储的键比较函数的情况下使用。

对于整数，很容易看出无符号大端序（big-endian）整数是保序的 —— 大端序格式中最高有效位排在前面。以空字符结尾的字符串也是保序的。

对于有符号整数，问题在于负数的最高有效位（符号位）被设置。在以大端序编码之前，我们需要翻转符号位，使负数的值更低。

// 保序编码
func encodeValues(out []byte, vals []Value) []byte {
    for _, v := range vals {
       switch v.Type {
       case TYPE_INT64:
          var buf [8]byte
          u := uint64(v.I64) + (1 << 63)
          binary.BigEndian.PutUint64(buf[:], u)
          out = append(out, buf[:]...)
       case TYPE_BYTES:
          out = append(out, escapeString(v.Str)...)
          out = append(out, 0) // 以空字符结尾
       default:
          panic("what?")
       }
    }
    return out
}

func decodeValues(in []byte, out []Value) {
    // 省略...
}

以空字符结尾的字符串的问题在于它们不能包含空字节。我们将通过 “转义” 空字节来解决这个问题。\x00 被替换为 \x01\x01，转义字节 \x01 本身被替换为 \x01\x02，这样仍然保留了排序顺序。

// 字符串被编码为以空字符结尾的字符串，
// 转义空字节，使字符串不包含空字节。
func escapeString(in []byte) []byte {
    zeros := bytes.Count(in, []byte{0})
    ones := bytes.Count(in, []byte{1})
    if zeros + ones == 0 {
       return in
    }

    out := make([]byte, len(in) + zeros + ones)
    pos := 0
    for _, ch := range in {
       if ch <= 1 {
          out[pos + 0] = 0x01
          out[pos + 1] = ch + 1
          pos += 2
       } else {
          out[pos] = ch
          pos += 1
       }
    }
    return out
}

9.3 范围查询

最后，我们将添加 Scanner 类型，它允许我们按排序顺序遍历一系列记录。

// 范围查询的迭代器
type Scanner struct {
    // 范围，从Key1到Key2
    Cmp1 int // CMP_??
    Cmp2 int
    Key1 Record
    Key2 Record
    // 内部
    tdef   *TableDef
    iter   *BIter // 底层B树迭代器
    keyEnd []byte // 编码后的Key2
}

// 是否在范围内？
func (sc *Scanner) Valid() bool
// 移动底层B树迭代器
func (sc *Scanner) Next()
// 获取当前行
func (sc *Scanner) Deref(rec *Record)

func (db *DB) Scan(table string, req *Scanner) error {
    tdef := getTableDef(db, table)
    if tdef == nil {
       return fmt.Errorf("table not found: %s", table)
    }
    return dbScan(db, tdef, req)
}

初始化迭代器：

func dbScan(db *DB, tdef *TableDef, req *Scanner) error {
    // 合理性检查
    switch {
    case req.Cmp1 > 0 && req.Cmp2 < 0:
    case req.Cmp2 > 0 && req.Cmp1 < 0:
    default:
       return fmt.Errorf("bad range")
    }

    values1, err := checkRecord(tdef, req.Key1, tdef.PKeys)
    if err != nil {
       return err
    }
    values2, err := checkRecord(tdef, req.Key2, tdef.PKeys)
    if err != nil {
       return err
    }

    req.tdef = tdef

    // 定位到起始键
    keyStart := encodeKey(nil, tdef.Prefix, values1[:tdef.PKeys])
    req.keyEnd = encodeKey(nil, tdef.Prefix, values2[:tdef.PKeys])
    req.iter = db.kv.tree.Seek(keyStart, req.Cmp1)
    return nil
}

移动迭代器：

// 是否在范围内？
func (sc *Scanner) Valid() bool {
    if!sc.iter.Valid() {
       return false
    }
    key, _ := sc.iter.Deref()
    return cmpOK(key, sc.Cmp2, sc.keyEnd)
}

// 移动底层B树迭代器
func (sc *Scanner) Next() {
    assert(sc.Valid())
    if sc.Cmp1 > 0 {
       sc.iter.Next()
    } else {
       sc.iter.Prev()
    }
}

点查询（Point queries）只是范围查询的特殊情况，那么为什么不把它们去掉呢？

// 通过主键获取单行记录
func dbGet(db *DB, tdef *TableDef, rec *Record) (bool, error) {
    // 只是扫描操作的快捷方式
    sc := Scanner{
       Cmp1: CMP_GE,
       Cmp2: CMP_LE,
       Key1: *rec,
       Key2: *rec,
    }

    if err := dbScan(db, tdef, &sc); err != nil {
       return false, err
    }

    if sc.Valid() {
       sc.Deref(rec)
       return true, nil
    } else {
       return false, nil
    }
}

我们只允许对完整主键进行范围查询，但对主键前缀进行范围查询也是合理的。我们将在下一章连同二级索引（secondary indexes）一起解决这个问题。