Go1.23のイテレータを使って実装された `slices` パッケージの `Chunk` 関数 の使い所を考えてみる
背景
2024/08/13にGo1.23がリリース[1]され、言語的にも大きな機能であるイテレータが追加されました!🎉🎉🎉
リリースから1か月経過し、実際にイテレータを使ったコードを書いてみようと考えている方も多いのではないでしょうか?
この記事では、スライスに関する汎用的な関数がまとめられているslicesパッケージに新規追加された関数Chunkからイテレータって何が良いんだっけ?どこで使えそうなんだっけ?ということを考えてみたいと思います!
使用するGoのバージョンは以下になります!
❯ go version
go version go1.23.1 darwin/arm64
また、この記事ではイテレータ自体の詳細な説明はしないので、興味がある方は以下の記事を参考にしてもらえればと思います📝
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通りの実装が可能なのでは?と筆者は考えています。
-
slices
パッケージのCollect
関数を使用する -
iter.Seq
からiter.Seq2
に変換する関数を自作する
slices
パッケージのCollect
関数を使用する
解決策1: 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
と似たような実装になっていて、イテレータのみを使用した実装と比較してパフォーマンス面で劣っている可能性が高いです。パフォーマンスに関しては後ほど検証しています。
iter.Seq
からiter.Seq2
に変換する関数を自作する
解決策2: 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
を持っているとき、その結果をスライスとして変換して使いたい」ケースに有用なのではないでしょうか。
普段コードを読んでいるだけでは気づかないような発見があって良かったです🔎
Discussion