聊一聊 go 的字典 map

492 阅读13分钟

什么是 map

map 是 go 中内置的一种数据类型,用来存储无序的键值对「key-value」集合。在这个集合中,key 只能出现一次。map 又叫哈希表、字典。

在 map 中,对于键值对的操作都是基于 key 来完成。常见的操作包括写入、更新、读取、删除等。

  • 写入:增加一个键值对
  • 更新:更新 key 对应的 value
  • 读取:获取 key 对应的 value,同时也可以遍历 map 中的键值对
  • 删除:删除指定 key 的键值对

map 的使用

声明 & 初始化

如果 map 中存放的是标量类型,会将数据存储在栈上,需要注意栈空间溢出

var m map[int]int             // 声明一个 map 为 nil,可读不可写
var m1 = make(map[int]int)    // 使用 make 初始化一个 map, 非 nil
var m2 = make(map[int]int, 8) // 初始化一个 map 并指定容量
var m3 = map[int]int{0: 0}    // 初始化一个 map 并赋值
// _ = map[int][1024 * 1024 * 1024]byte{0: {}}    ./main.go:5:6: stack frame too large (>1GB): 1024 MB locals + 0 MB args
fmt.Println(m == nil, m1 == nil, m2 == nil, m3 == nil) // true false false false

增删改查

func main() {
    var m map[int]int
    var m1 = make(map[int]int)

    // 插入一条记录
    // m[0] = 0; panic: assignment to entry in nil map, 对未初始化的 map 插入数据会 panic, 读取没事
    // 插入和删除操作相同,如果 key 存在则插入,key 不存在则删除
    m1[1] = 0
    m1[1] = 1

    // 删除 key, key 存在直接删除, key 不存在直接跳过
    delete(m1, 0)

    // 遍历键值对
    for key, value := range m1 {
       println(key, " : ", value)
    }

    // 只遍历 key
    for key := range m1 {
       println("key : ", key)
    }

    // 获取指定 key
    v := m[1]           // 如果 key 不存在,返回零值
    fmt.Println(v)      // 0
    v2, ok := m[1]      // 标识 key 是否存在
    fmt.Println(v2, ok) // 0, false
}

实现一个简单的 map

map 本质上就是对键值对处理的一种数据类型,且键只能存在一次。其特点就是可以通过 key 快速进行查找对应的 value。常见的实现方式就是利用数组存储键值对,通过对 key 取哈希来计算对应的数组下标。下面我们来尝试实现一个简单的 map。

结构体设计

使用一个切片存储 key-value 数据

type simpleMap[K comparable, V any] struct {
    data []*Pair[K, V]
    size int // 支持存储的容量
    hash func(k K) int
}

type Pair[K comparable, V any] struct {
    key   K
    value V
}

增删改查

func newSimpleMap[K comparable, V any](h func(k K) int, size int) *simpleMap[K, V] {
    return &simpleMap[K, V]{
       data: make([]*Pair[K, V], size),
       size: size,
       hash: h,
    }
}

func (m *simpleMap[K, V]) set(k K, v V) {
    m.data[m.getIndex(k)] = &Pair[K, V]{key: k, value: v}
}

func (m *simpleMap[K, V]) get(k K) (v V, exist bool) {
    index := m.getIndex(k)
    if m.data[index] != nil {
       return m.data[index].value, true
    }
    return v, false
}

func (m *simpleMap[K, V]) getIndex(k K) int {
    return m.hash(k) % m.size
}

使用例子

func main() {
    m := newSimpleMap[int, int](
       func(k int) int {
          return k
       },
       5,
    )
    m.set(1, 1)
    m.set(2, 2)
    k, ok := m.get(1)
    fmt.Println("get k: ", k, " get ok: ", ok)
    k, ok = m.get(0)
    fmt.Println("get k: ", k, " get ok: ", ok)
}

哈希冲突

在例子中,我们通过对 hash(key) 取余来计算存储的下标。如果不同 key 的 hash 结果相同,就会导致数据丢失的问题。常见的解决方法包括:

  • 链地址法:如果计算的数组下标结果相同,则在 Pair 结构体中维护一个链表。遍历链表并与需要存储的 key 做比较,key 相同则更新数据;key 不同则链表中追加一个节点。
  • 开放地址法:如果计算的结果相同,且 Pair 中的 key 与需要 set 的 key 不同。则向数组的下一个下标遍历,如果 key 不同则继续遍历;key 相同则更新数据;下标为空则存储数据。

这里使用链地址法来解决冲突,更新存储结构。

type Pair[K comparable, V any] struct {
    key      K
    value    V
    overflow *Pair[K, V]  // 拉链法解决冲突
}

func (m *simpleMap[K, V]) set(k K, v V) {
    pair := m.data[m.getIndex(k)]
    for pair != nil {
       if pair.key == k { // 更新数据
          pair.value = v
          return
       }
       pair = pair.overflow
    }
    m.data[m.getIndex(k)] = &Pair[K, V]{
       key:      k,
       value:    v,
       overflow: m.data[m.getIndex(k)],
    }
}

func (m *simpleMap[K, V]) get(k K) (v V, exist bool) {
    pair := m.data[m.getIndex(k)]
    for pair != nil {
       if pair.key == k {
          return pair.value, true
       }
       pair = pair.overflow
    }
    return v, false
}

func (m *simpleMap[K, V]) del(k K) {
    pre, node := m.data[m.getIndex(k)], m.data[m.getIndex(k)]
    for node != nil {
       if node.key == k {
          if node == pre {
             m.data[m.getIndex(k)] = node.overflow
             return
          }
          pre.overflow = node.overflow
          return
       }
       pre, node = node, node.overflow
    }
}

扩缩容

在上面的例子中, map 初始化时会确定数组的大小。当存储的数据不断增加时,会导致拉链越来越长,从而将时间复杂度退化至 o(n)。可以考虑对 map 里面的数据进行扩容操作,从而减少性能劣化。

负载因子: 哈希表中已经存储的数据和哈希表长度的比值。它是判断哈希表是否需要扩容或缩容的重要指标。

当插入新的元素时,如果负载因子超过一定的阈值(比如0.75),那么哈希表通常需要进行扩容操作,也就是增加哈希表的长度,并将所有的元素重新进行哈希,存入新的位置,以保证哈希表的查找、插入和删除操作的时间复杂度。

同样的,当删除元素时,如果负载因子低于一定的阈值(比如0.25),那么哈希表可能需要进行缩容操作,减少哈希表的长度,减小空间占用。

下面我们修改代码,支持进行扩容:

const maxLoadFactor = 0.75

type simpleMap[K comparable, V any] struct {
    data  []*Pair[K, V]
    size  int
    count int
    hash  func(k K) int
}

type Pair[K comparable, V any] struct {
    key      K
    value    V
    overflow *Pair[K, V] // 拉链法解决冲突
}

func newSimpleMap[K comparable, V any](h func(k K) int, size int) *simpleMap[K, V] {
    return &simpleMap[K, V]{
       data:  make([]*Pair[K, V], size),
       size:  size,
       count: 0,
       hash:  h,
    }
}

func (m *simpleMap[K, V]) set(k K, v V) {
    pair := m.data[m.getIndex(k)]
    for pair != nil {
       if pair.key == k { // 更新数据
          pair.value = v
          return
       }
       pair = pair.overflow
    }
    m.data[m.getIndex(k)] = &Pair[K, V]{
       key:      k,
       value:    v,
       overflow: m.data[m.getIndex(k)],
    }
    m.count++

    if m.needReHash() {
       m.reHash()
    }
}

func (m *simpleMap[K, V]) get(k K) (v V, exist bool) {
    pair := m.data[m.getIndex(k)]
    for pair != nil {
       if pair.key == k {
          return pair.value, true
       }
       pair = pair.overflow
    }
    return v, false
}

func (m *simpleMap[K, V]) del(k K) {
    pre, node := m.data[m.getIndex(k)], m.data[m.getIndex(k)]
    for node != nil {
       if node.key == k {
          m.count--
          if node == pre {
             m.data[m.getIndex(k)] = node.overflow
             return
          }
          pre.overflow = node.overflow
          return
       }
       pre, node = node, node.overflow
    }
}

func (m *simpleMap[K, V]) getIndex(k K) int {
    return m.hash(k) % m.size
}

func (m *simpleMap[K, V]) needReHash() bool {
    loadFactor := float64(m.count) / float64(m.size)
    return loadFactor > maxLoadFactor
}

func (m *simpleMap[K, V]) reHash() {
    m.size = m.size * 2
    newData := make([]*Pair[K, V], m.size)
    for _, pair := range m.data {
       for pair != nil {
          newData[m.getIndex(pair.key)] = pair
          pair = pair.overflow
       }
    }
    m.data = newData
}

其他问题优化

到目前为止,我们的哈希表已经可以存储大量数据,且查询复杂度较低。然而,在这个哈希表中,会存在性能尖刺。即发生 reHash 时,所有的 key 需要重新计算下标,进行移动。在 go 的官方实现中,在 map 和 key-value 之间增加一层 bucket 的概念,即 key-value 存储在 bucket 中,一个 bucket 在不发生溢出的情况下存储 8 个 key-value。当发生 reHash 时,只需要搬迁 bucket 即可,无需按 key 搬迁。其次,在 reHash 过程中,采用的是渐进式扩容,即一次至多迁移两个 bucket,避免尖刺。

image.png 当然,在 go 的实现中,也增加了其他优化。例如:索引计算的优化,bucket key 查询的优化,溢出桶的预分配等。具体可以看下一节的介绍。与我们当前的思路区别不大,核心是增加了 bucket 和渐进式扩容。

Go 中 map 的设计与实现

本文使用的 go 版本 go1.22.1 darwin/arm64

结构设计

map 的整体设计如下所示。由两个核心结构体组成 hmap 和 buckets。hmap 存储 map 的基础信息,例如 key 的数量,桶数量 2**B,指向桶的指针等。buckets 也就是我们所说的桶,每个桶是一个 bamp 类型的数据,用于存储一批 key-value。引入桶的好处在于每次 rehash 只需要迁移桶即可,不用逐个迁移每个 key。

每个桶「bmap」由三部分组成:

  • tophash: 存储桶中已经存储的 hash(key) 的低八位,便于快速查找 key 是否存在桶中
  • keys: 实际存储的 key
  • values: 实际存储 value
  • overflow: 指向溢出桶的指针,即使用拉链法来解决哈希冲突。当桶中的数据存满时,会创建一个链表存储 key-value。

 // map 结构体的定义
type hmap struct {
    count     int // map 中 key 的数量,可使用 len() 获取
    flags     uint8 // 标记位,如读、写标记
    B         uint8  // map 桶的数量 2^B
    noverflow uint16 // 溢出桶的数量
    hash0     uint32 // 计算哈希值的哈希种子

    buckets    unsafe.Pointer // 指向 bucket 的指针
    oldbuckets unsafe.Pointer // 指向旧 bucket 的指针,扩容的时候会用到
    nevacuate  uintptr        // 表示已经迁移了多少桶

    extra *mapextra // 额外信息
}

 // bucket,桶,实际存储 key-value 数据的地方.
type bmap struct {
    // 存储桶中 key 的 hash 高八位,方便快速查找 key 是否存在桶中。 If tophash[0] < minTopHash,
    // tophash[0] 为桶已经迁移的状态。
    tophash [bucketCnt]uint8
      
    keys [bucketCnt]keytype // 8 个键   
    values [bucketCnt]valuetype // 8 个值   
    overflow *bmap // 指向溢出桶的指针
}

初始化

当使用 make(map[k]v, hint) 初始化 map 时,调用的 makemap 函数,生成 *hmap

makemap

func makemap(t *maptype, hint int, h *hmap) *hmap {
    // 初始化 Hmap
    if h == nil {
       h = new(hmap)
    }
    h.hash0 = uint32(rand())
    
    // 通过传入的 hint 计算 B, 即计算 map 的容量 2**B
    B := uint8(0)
    for overLoadFactor(hint, B) { // 通过负载因子和 hit 计算最终的 B
       B++
    }
    h.B = B

     // allocate initial hash table
     // if B == 0, the buckets field is allocated lazily later (in mapassign)
     // If hint is large zeroing this memory could take a while.
    if h.B != 0 {
       var nextOverflow *bmap
       h.buckets, nextOverflow = makeBucketArray(t, h.B, nil)
       if nextOverflow != nil {
          h.extra = new(mapextra)
          h.extra.nextOverflow = nextOverflow
       }
    }

    return h
}

overLoadFactor

判断存储数量为 count 时,是否超过 map 的负载因子。如果超过负载因子,需要扩容。

  • count 数量小于 bucketCnt「一个 bucket 存储的数据总量,即一个 bucket 就可以存储全部数据」 时,返回 false。
  • count 数量 > 桶数量 * 负载因子「6.5」时,意味着平均每个桶需要存储多于 6.5 个 key,此时需要扩容。增加 bucket,减少哈希冲突。
func overLoadFactor(count int, B uint8) bool {
    return count > bucketCnt && uintptr(count) > loadFactorNum*(bucketShift(B)/loadFactorDen)
}

makeBucketArray

另一个核心的函数是创建 buckets 的函数,其核心流程如下「这里只讨论 b> 4 的情况」:

  • 通过 b 计算需要分配的桶的数量,如果 b>4,则会预先分配 2**(b-4) 个溢出桶。此时,buckets 由两部分组成:
  • 普通桶:数量为 2**b 个,用于索引并存储 key。
  • 溢出桶:数量为 2**(b-4) 个,当发生哈希冲突时会使用预先分配的溢出桶。
  • 在 hmap 中,将 mapextra 的 nextOverflow 指向第一个溢出桶的位置。
  • 最后一个溢出桶的 overflow 指针指向 buckets 的入口,作为一个标记作用。这样就可以通过判断溢出桶的 overflow 是不是 nil,来确定是不是最后一个溢出桶。在使用溢出桶的时候,如果 overflow 不是 nil,则意味着需要创建新的溢出桶。

// 初始化 map 的 buckets, 数量是 1<<b。dirtyalloc 是相同的 t 和 b 之间申请的 buckets,如果其不为空则清空数据,并复用空间。
func makeBucketArray(t *maptype, b uint8, dirtyalloc unsafe.Pointer) (buckets unsafe.Pointer, nextOverflow *bmap) {
    base := bucketShift(b)
    nbuckets := base
    // 如果分配的 buckets 数量大量 16「1<<4」,则需要分配溢出桶。b 过大时,意味着 map 需要存储大量的 key。因此,设计时会提前申请溢出桶,当发生哈希冲突时,可直接使用分配的溢出桶。
    if b >= 4 {
       // 溢出桶的数量为 2**(b-4)
nbuckets += bucketShift(b - 4)
       // 溢出桶所需要的内存
       sz := t.Bucket.Size_ * nbuckets
       // 计算实际内存
       up := roundupsize(sz, t.Bucket.PtrBytes == 0)
       if up != sz {
          nbuckets = up / t.Bucket.Size_
       }
    }

    if dirtyalloc == nil {
        // 创建新的 buckets 数组
       buckets = newarray(t.Bucket, int(nbuckets))
    } else {
       // 晴空 dirtyalloc 的内存
buckets = dirtyalloc
       size := t.Bucket.Size_ * nbuckets
       if t.Bucket.PtrBytes != 0 {
          memclrHasPointers(buckets, size)
       } else {
          memclrNoHeapPointers(buckets, size)
       }
    }

    // 当 b >=4 时,即 buckets 的数量为 2**4,溢出桶的数量为 2**(b-4)
    if base != nbuckets {
       // 处理预先分配的溢出桶,将最后一个溢出桶的 overflow 指向第一个普通桶
       nextOverflow = (*bmap)(add(buckets, base*uintptr(t.BucketSize)))
       last := (*bmap)(add(buckets, (nbuckets-1)*uintptr(t.BucketSize)))
       last.setoverflow(t, (*bmap)(buckets))
    }
    return buckets, nextOverflow
}

读取

map 的读取有三种:

  • mapaccess1: 只返回的对应的值,对应 v := map[k]
  • mapaccess2: 返回对应的值,已经 key 是否存在,对应 v,ok := map[k]
  • mapaccessK:用于遍历 map,返回 key 和 value。对应 for k,v := range map{} 这里只介绍一下 mapacesss2,整体逻辑都是相同的。整理流程如下:
    1. 通过 key 计算 hash
    2. 根据 hash 值索引到对应的 bucket。使用位运算「hash & (2^B-1)」,优化计算速度。其实就是取哈希值的低 B 位来确定最终的 bucket
    3. 遍历 bucket 的每一个槽「即 tophash 数组」,比较哈希值的高 8 位是否相同,从而快速判断 key 是否存在 bucket 中。如果 tophash 相同,则查找 key,如果 key 相同,则找到数据;key 不相同则遍历下一个槽。
    4. 如果 bucket 中不包含对应的 key,则去 bucket 的溢出桶中去查找。直到所有溢出桶都查找完毕。

写入 & 更新

map 的写入是通过 mapassign 实现的。整体流程如下:

  1. 通过 key 计算 hash,并标记为 writing
  2. 通过 hash 低 B 位找到对应的 bucket
  3. 如果 bucket 正在扩容,则迁移要操作的 bucket 和按下标迁移「扩容逻辑下面会讲」
  4. 如果 bucket 中存在要写入的 key 则更新,不存在的话,依据条件插入数据,或者插入到溢出桶中
  5. 再次判断是否需要扩容,如果需要扩容则返回 2. 执行一个 bucket 迁移逻辑
  6. 更新标记位和 count

image.png

渐进式迁移

在写入过程中,有一个重要的逻辑就是扩容。当 map 没有在扩容、超过负载因子、太多溢出桶,三个条件同时满足时触发扩容。扩容分为两类:

  • 2 倍扩容:map 中的 key 数量过多,避免溢出桶太多而导致查询性能退化
  • 等量扩容:map 中删除的 key 过多,存在大量空数据,优化存储空间 扩容意味着旧的数据需要存储到新的 buckets 中,即 rehash。在 go 中采用渐进式扩容的方式,避免一次性的数据迁移而导致性能降低。会有两条线同时迁移,其流程如下:
    1. 从第一个 bucket 进行迁移
    2. 插入、删除、修改的时候,key 所在的 bucket 会被迁移

迁移示例

如下图所示,第一次插入 KEY,触发 map 扩容。第一条线从 nevacuate = 0 开始迁移,迁移一个 bucket,更新 nevacuate = 1;第二条线计算 KEY 所在的 bucket 为 2,迁移对应的 bucket。
第二次更新 KEY,map 处于迁移中。第一条线开始迁移 bueckt 1;第二条线的 bucket 2 已经迁移,需要操作。
这样可以保证,最多 B-1 次操作可以完成 map 的迁移。

image.png

删除

删除逻辑比较简单,调用的是 mapdelete 函数。这里需要注意的是,删除只是清空对应 key 和 value ,并不会释放内存。只能通过 GC 清理空间。

总结

本文介绍了 go 中 map 实现和常见用法。在 map 的实现中,采用的拉链法解决哈希冲突。同时,为了避免 reHash 操作时的性能下降,go 采用了 bucket 和渐进式迁移的思路,降低了迁移操作的开销。除此之外,也进行了其他的优化。例如:预分配空间、通过位运算计算下标、使用标记位快速终止查询等。