Go1.23のイテレータを使って実装された `slices` パッケージの `Chunk` 関数 の使い所を考えてみる

2024/10/07に公開

背景

2024/08/13にGo1.23がリリース[1]され、言語的にも大きな機能であるイテレータが追加されました!🎉🎉🎉

リリースから1か月経過し、実際にイテレータを使ったコードを書いてみようと考えている方も多いのではないでしょうか?

この記事では、スライスに関する汎用的な関数がまとめられているslicesパッケージに新規追加された関数Chunkからイテレータって何が良いんだっけ?どこで使えそうなんだっけ?ということを考えてみたいと思います!

使用するGoのバージョンは以下になります!

❯ go version
go version go1.23.1 darwin/arm64

また、この記事ではイテレータ自体の詳細な説明はしないので、興味がある方は以下の記事を参考にしてもらえればと思います📝

https://zenn.dev/kkkxxx/articles/d9505540581b5d

Chunk関数の説明

まずはChunk関数のシグネチャや使い方を見てみます。

シグネチャは以下のようになっており、引数であるスライスのsと整数nを渡すと、sを要素数nごとに分割したスライスがiter.Seq[Slice]で提供されるようになっています

func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice]

実際に使用してみるとこんな感じでしょうか。

package main

import (
	"fmt"
	"slices"
)

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	for chunk := range slices.Chunk(nums, 3) {
		fmt.Println(chunk)
	}
}

動かしてみると確かに要素数3ごとのスライスに分割されています👏

❯ go run .
[1 2 3]
[4 5 6]
[7]

実際何が嬉しいの?

そもそもイテレータは何が良いのでしょうか?

Chunk関数の実装以前にはジェネリクスを使って、スライスのスライスを返す関数を自作しているケースが多いのではと思います。

たとえば、筆者は以下のようなコードを何回か見たことがあります。

func ChunkWithGenerics[S ~[]T, T any](slice S, size int) [][]T {
	if len(slice) == 0 {
		return [][]T{}
	}

	var chunks [][]T
	for i := 0; i <= len(slice); i += size {
		end := i + size
		if end > len(slice) {
			end = len(slice)
		}
		chunks = append(chunks, slice[i:end])
	}
	return chunks
}

このようなコードとイテレータを使ったChunk関数ではどこに差分が出るのでしょうか?

差分を考えるためにChunk関数の実装を見てみます。

func Chunk[Slice ~[]E, E any](s Slice, n int) iter.Seq[Slice] {
	if n < 1 {
		panic("cannot be less than 1")
	}

	return func(yield func(Slice) bool) {
		for i := 0; i < len(s); i += n {
			end := min(n, len(s[i:]))
			if !yield(s[i : i+end : i+end]) {
				return
			}
		}
	}
}

パッとみる感じだとそこまで大きな差分はないように見えます。
しかし、ジェネリクスを使った実装はvar chunk [][]E を定義し、そのスライスにappendしていますが、イテレータは元のスライスの要素の一部をそのまま関数の返り値にしています。

つまり、差分としては「var chunk [][]Eを定義している」「chunk変数にappendの処理が行われている」という2点がありそうです。


実際にベンチマークを取ってみたいと思います!

ベンチマークのテストコードは以下です。

package main

import (
	"slices"
	"testing"
)

func BenchmarkChunk_BySlicePkg(b *testing.B) {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	size := 3

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = slices.Chunk(nums, size)
	}
}

func BenchmarkChunk_ByGenerics(b *testing.B) {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	size := 3

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = ChunkWithGenerics(nums, size)
	}
}

計測してみると、Chunk関数の方は実行速度が早く、メモリのアロケーションも行われていないことがわかると思います!👏

イテレータにすることでパフォーマンス向上の恩恵がありそうです!

❯ go test -bench . -benchmem
goos: darwin
goarch: arm64
pkg: github.com/k3forx/iter_research
cpu: Apple M3 Pro
BenchmarkChunk_BySlicePkg-11            1000000000               0.2565 ns/op          0 B/op          0 allocs/op
BenchmarkChunk_ByGenerics-11            15918478                75.72 ns/op          168 B/op          3 allocs/op
PASS
ok      github.com/k3forx/iter_research 1.827s

実装のケースを考えてみる

上記のChunk関数を実際に使うケースを考えてみます。

ケース1: 単純にchunkしたスライスのみ必要な場合

単純にchunkしたスライスのみが必要な場合は、Chunk関数を使えば問題ないと思います。

先ほどの例と同じような感じですね!

ケース2: chunkしたスライスとそのindexが必要な場合

では、次にchunkしたスライスとそのindexが必要なケースはどうでしょうか?

先ほどの例でもう少し具体的に説明すると、[1, 2, 3]の時は0[4, 5, 6]の時は1[7]の時は2を取得するという場合です。

package main

import (
	"fmt"
	"slices"
)

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	for chunk := range slices.Chunk(nums, 3) {
		fmt.Println(chunk) // この時にchunkのindex(0,1,2)も取得したい
	}
}

個人的にはこのようなケースの場合だとChunk関数がうまく使えないと考えています。なぜならChunk関数のシグネチャがiter.Seq[Slice]だからです。

むしろChunk関数が実装される以前は、[][]Tのようなスライスのスライスを返すジェネリクス関数を自作しているケースが多いかと思われるので、そういった関数をそのまま使うほうが良さそうでもあります。

頑張ってなんとかしてみる

ここで終わっては面白くないので、もう少しなんとかしてみようと思います!

今のGoの仕様だと以下の2通りの実装が可能なのでは?と筆者は考えています。

  1. slicesパッケージのCollect関数を使用する
  2. iter.Seqからiter.Seq2に変換する関数を自作する

解決策1: slicesパッケージのCollect関数を使用する

slicesパッケージの関数を眺めてみるとCollect関数というのがあります。引数のiter.Seq[E]を渡すと、それを実行してスライス[]Eを返してくれる関数のようです。

func Collect[E any](seq iter.Seq[E]) []E

この関数を使えば、chunkを提供するiter.Seq[Slice]をスライスのスライスとして返してくれるようになるので、chunkとindexを両方取得できそうな予感がします!

実際に試してみます。コードは以下のような感じになると思います。

package main

import (
	"fmt"
	"iter"
	"slices"
)

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	for index, chunk := range slices.Collect(slices.Chunk(nums, 3)) {
		fmt.Println(index, chunk)
	}
}

実際に動かしてみると、確かに意図した挙動になっていそうです🎉🎉🎉

❯ go run .
0 [1 2 3]
1 [4 5 6]
2 [7]

一方で、このような形での実装はイテレータの恩恵をあまり受けられていないと感じます。なぜなら、一度イテレータを実行して、その結果を別のスライスとして取得しているからです。

具体的にCollect関数の実装を深掘りしてみます。実装は以下のようになっています。

func Collect[E any](seq iter.Seq[E]) []E {
	return AppendSeq([]E(nil), seq)
}

func AppendSeq[Slice ~[]E, E any](s Slice, seq iter.Seq[E]) Slice {
	for v := range seq {
		s = append(s, v)
	}
	return s
}

Collect関数は内部的にAppendSeq関数を呼んでおり、そのAppendSeq関数の中でスライス[]E(nil)にイテレータの結果をappendしている実装になっています。

つまり、冒頭で出てきた関数ChunkWithGenericsと似たような実装になっていて、イテレータのみを使用した実装と比較してパフォーマンス面で劣っている可能性が高いです。パフォーマンスに関しては後ほど検証しています。

解決策2: iter.Seqからiter.Seq2に変換する関数を自作する

iterパッケージにはiter.Seq2[K, V any]という型も定義されており、その定義は以下のようになっています。

Seq2[K, V any] func(yield func(K, V) bool)

このKにindexを、Vにchunkされたスライスを渡すことによって、スライスのスライスを作ることなくイテレータのみでやりたいことが実現できそうです!

具体的に iter.Seq[Slice]iter.Seq2[int, Slice]に変換するような関数ConvertToIterSeq2を自作 してみます。

実装は以下のような形になると思います!実装自体はかなりシンプルになりそうです。

func ConvertToIterSeq2[S ~[]E, E any](seq iter.Seq[S]) iter.Seq2[int, S] {
	return func(yield func(int, S) bool) {
		var index int
		for v := range seq {
			if !yield(index, v) {
				return
			}
			index++
		}
	}
}

イテレータを連結させているような形になりますね。このような実装をすれば、スライスのスライスを定義することなく、イテレータのみでやりたいことを完結できて良さそうです👍


コードの全体像は以下のようになります。

package main

import (
	"fmt"
	"iter"
	"slices"
)

func main() {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	for index, chunk := range ConvertToIterSeq2(slices.Chunk(nums, 3)) {
		fmt.Println(index, chunk)
	}
}

func ConvertToIterSeq2[S ~[]E, E any](seq iter.Seq[S]) iter.Seq2[int, S] {
	return func(yield func(int, S) bool) {
		var index int
		for v := range seq {
			if !yield(index, v) {
				return
			}
			index++
		}
	}
}

実際に動かしてみると、意図した挙動になっていました🎉🎉🎉

❯ go run .
0 [1 2 3]
1 [4 5 6]
2 [7]

実際にイテレータに適用するイテレータのような便利関数は61898のproposalで議論がされています。まだproposal-acceptedの状況にはなっていないものの、近い将来に実装される可能性があります。

パフォーマンスを比較してみる

最後にslices.Collect関数と自作したConvertToIterSeq2関数でどちらのパフォーマンスが良いのか計測してみたいと思います。

ベンチマークのテストコードは以下です。

package main

import (
	"slices"
	"testing"
)

func BenchmarkChunk_ByIterSeq2(b *testing.B) {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	size := 3

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = ConvertToIterSeq2(slices.Chunk(nums, size))
	}
}

func BenchmarkChunk_ByCollect(b *testing.B) {
	nums := []int{1, 2, 3, 4, 5, 6, 7}
	size := 3

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = slices.Collect(slices.Chunk(nums, size))
	}
}

ベンチマークを取ってみると、やはりイテレータのみを使った実装であるConvertToIterSeq2関数の処理が早く、メモリのアロケーションも少ないですね!!

❯ go test -bench . -benchmem
goos: darwin
goarch: arm64
pkg: github.com/k3forx/iter_research
cpu: Apple M3 Pro
BenchmarkChunk_ByIterSeq2-11            1000000000               0.2560 ns/op          0 B/op          0 allocs/op
BenchmarkChunk_ByCollect-11             10219161               117.5 ns/op           232 B/op          6 allocs/op
PASS
ok      github.com/k3forx/iter_research 1.877s

まとめ

これまでの議論をまとめてみると以下のようになります!

  • スライスをchunkに分割し、そのchunkのみを使用したい場合はslicesパッケージのChunk関数が使える
    • ジェネリクスを使ったスライスのスライスを返す関数を使用するよりパフォーマンスが良い
  • 一方で、分割したchunkとそのindex自体が必要な場合、標準ライブラリで実装できるがイテレータの恩恵がなくなる
    • イテレータのみを使用した自作の関数を使ったほうが良い
  • イテレータ自体に適応されるイテレータのプロポーザルが61898でされている

イテレータが実装されて色々な関数が便利に書けるようになった反面、まだ痒いところに手が届かないケースもあるように感じました!Goのチームがイテレータ関連の便利関数の実装を議論中なので、近い将来にそういった関数が実装されるのではと予想されます!

おまけの感想

余談ですが、筆者はCollect関数をどういう場面で使うのだろうとずっと疑問に思っていました。

今回はiter.Seqをスライスに変換する目的で使用しました。Collect関数は「外部パッケージなどに自分たちが手を加えづらい関数があり、その返り値がiter.Seqを持っているとき、その結果をスライスとして変換して使いたい」ケースに有用なのではないでしょうか。

普段コードを読んでいるだけでは気づかないような発見があって良かったです🔎

脚注
  1. https://go.dev/blog/go1.23 ↩︎

Canary Tech Blog

Discussion