Go1.23 から導入される Iterator を使って Chunk 処理を書く

Overview

Go1.23 から導入される Iterator を使って同様に Chunk 処理を実装してみます。

Sample

 func chunk[T any](src []T, chunkSize int) iter.Seq[[]T] {
  return func(yeild func([]T) bool) {
      for i := 0; i < len(src); i += chunkSize {
          end := i + chunkSize
          if end > len(src) {
              end = len(src)
          }
          if !yeild(src[i:end]) {
              break
          }
      }
  }
 }

利用方法

func main() {
    arr := make([]int64, 10000)
    for i := range arr {
        arr[i] = int64(i + 1)
    }
    for chunkedIDs := range chunk(arr, 1000) {
        echo(len(chunkedIDs))
    }
}

func echo[T any](v T) {
    fmt.Println(v)
}

Benchmark

以下のような従来の一般的な方法で実装された Chunk 処理とベンチマークの比較をしてみる。

func chunkBy[T any](items []T, chunkSize int) [][]T {
    chunk := make([][]T, 0, (len(items)/chunkSize)+1)
    for chunkSize < len(items) {
        chunk, items = append(chunk, items[0:chunkSize:chunkSize]), items[chunkSize:]
    }
    return append(chunk, items)
}

Benchmark のコードは以下

func BenchmarkChunk(b *testing.B) {
    // chunk と chunkBy のベンチを取る
    arr := make([]int64, 10000)
    for i := range arr {
        arr[i] = int64(i + 1)
    }

    b.Run("chunk", func(b *testing.B) {
        b.ResetTimer()
        for i := 0; i < b.N; i++ {
            for range chunk(arr, 1000) {
            }
        }
    })

    b.Run("chunkBy", func(b *testing.B) {
        b.ResetTimer()
        for i := 0; i < b.N; i++ {
            for range chunkBy(arr, 1000) {
            }
        }
    })
}

ベンチマーク結果は以下

go test -bench=. -benchmem
goos: darwin
goarch: amd64
pkg: emahiro.dev/adhoc/go-sandbox
cpu: VirtualApple @ 2.50GHz
BenchmarkChunk/chunk-8          37948653                32.04 ns/op            0 B/op          0 allocs/op
BenchmarkChunk/chunkBy-8        16093440                69.57 ns/op          288 B/op          1 allocs/op
PASS
ok      emahiro.dev/adhoc/go-sandbox    3.678s

※ このベンチマークは Go1.23rc2 を使っています。

chunk する処理の中で chunk した結果を入れる箱を用意しない分 Allocation が発生せず iterator を使うほうがパフォーマンスが良いことがわかりました。

See Also

zenn.dev