在Go中过滤范型集合：性能回顾

在一个真实的 Golang 场景中使用泛型，同时寻找与 Stream filter(Predicate predicate)和 Python list comprehension 等同的函数。我没有依赖现有的包，而是选择自己写一个过滤函数，以达到学习的目的

func filterStrings(collection []string, test func(string) bool) (f []string) {
 for _, s := range collection {
  if test(s) {
   f = append(f, s)
  }
 }

 return
}

1
2
3
4
5
6
7
8
9
10

然而，这只适用于字符串。如果我需要过滤一个整数的集合，那么我就需要另一个极其相似的函数。这对于一个范型函数来说似乎是一个完美的选择。

func filter[T any](collection []T, test func(T) bool) (f []T) {
	for _, e := range collection {
		if test(e) {
			f = append(f, e)
		}
	}

	return
}

1
2
3
4
5
6
7
8
9
10

让我们来分析类型化和范型版本之间的（少数）差异：

函数名后面是一个范型T的定义。
T被定义为任何类型。
输入 slice 中元素的类型从字符串变成了T
输入、输出的 clice 类型也从字符串变成了T

不多说了，让我们来写一些单元测试。首先，我需要一个随机集合（在我的例子中，是字符串）的生成器。

func generateStringCollection(size, strLen int) []string {
	var collection []string
	for i := 0; i < size; i++ {
		collection = append(collection, strings.Repeat(fmt.Sprintf("%c", rune('A'+(i%10))), strLen))
	}

	return collection
}
1
2
3
4
5
6
7
8

现在我可以写一个测试用例，判断 filterStrings 函数的输出与我的过滤范型器的输出相同。

func TestFilter(t *testing.T) {
	c := generateStringCollection(1000, 3)

	t.Run("the output of the typed and generic functions is the same", func(t *testing.T) {
		predicate := func(s string) bool { return s == "AAA" }

		filteredStrings := filterStrings(c, predicate)
		filteredElements := filter(c, predicate)

		if !reflect.DeepEqual(filteredStrings, filteredElements) {
			t.Errorf("the output of the two functions is not the same")
		}
	})
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14

=== RUN   TestFilter
=== RUN   TestFilter/the_output_of_the_typed_and_generic_functions_is_the_same
--- PASS: TestFilter (0.00s)
    --- PASS: TestFilter/the_output_of_the_typed_and_generic_functions_is_the_same (0.00s)
PASS
1
2
3
4
5

考虑新函数在处理大的 slice 时的性能问题。我怎样才能确保它在这种情况下也能表现良好？

答案是：基准测试。用Go编写基准测试与单元测试很相似。

const (
	CollectionSize = 1000
	ElementSize    = 3
)

func BenchmarkFilter_Typed_Copying(b *testing.B) {
	c := generateStringCollection(CollectionSize, ElementSize)

	b.Run("Equals to AAA", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			filterStrings(c, func(s string) bool { return s == "AAA" })
		}
	})
}

func BenchmarkFilter_Generics_Copying(b *testing.B) {
	c := generateStringCollection(CollectionSize, ElementSize)

	b.Run("Equals to AAA", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			filter(c, func(s string) bool { return s == "AAA" })
		}
	})
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

go test -bench=. -count=10 -benchmem
goos: darwin
goarch: arm64
pkg: github.com/timliudream/go-test/generic_test
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             718408              1641 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             718148              1640 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             732939              1655 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             723036              1639 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             699075              1639 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             707232              1643 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             616422              1652 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             730702              1649 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             691488              1700 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             717043              1646 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          428851              2754 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          428437              2762 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          430444              2800 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          429314              2757 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          430938              2754 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          429795              2754 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          426714              2755 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          418152              2755 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          431739              2761 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          412221              2755 ns/op            4080 B/op          8 allocs/op
PASS
ok      github.com/timliudream/go-test/generic_test     25.005s

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

我对这个结果并不满意。看起来我用可读性换取了性能。
此外，我对分配的数量也有点担心。你注意到我的测试名称中的_Copying后缀了吗？那是因为我的两个过滤函数都是将过滤后的项目从输入集合复制到输出集合中。为什么我必须为这样一个简单的任务占用内存？
到最后，我需要做的是过滤原始的收集。我决定先解决这个问题。

func filterInPlace[T any](collection []T, test func(T) bool) []T {
	var position, size = 0, len(collection)
	for i := 0; i < size; i++ {
		if test(collection[i]) {
			collection[position] = collection[i]
			position++
		}
	}

	return collection[:position]
}
1
2
3
4
5
6
7
8
9
10
11

我不是把过滤后的结果写到一个新的集合中，然后再返回，而是把结果写回原来的集合中，并保留一个额外的索引，以便在过滤后的项目数上 "切 "出一个片断。

我的单元测试仍然通过，在改变了下面这行之后。

filteredStrings := filterStrings(c, predicate)
//filteredElements := filter(c, predicate)
  
filteredElements := filterInPlace(c, predicate) // new memory-savvy function
1
2
3
4

再添加一个 bench 方法：

func BenchmarkFilter_Generics_InPlace(b *testing.B) {
 c := generateStringCollection(CollectionSize, 3)

 b.Run("Equals to AAA", func(b *testing.B) {
  for i := 0; i < b.N; i++ {
   filterInPlace(c, func(s string) bool { return s == "AAA" })
  }
 })
}
1
2
3
4
5
6
7
8
9

结果是出色的。

go test -bench=. -benchmem
goos: darwin
goarch: arm64
pkg: github.com/timliudream/go-test/generic_test
BenchmarkFilter_Typed_Copying/Equals_to_AAA-8             713928              1642 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Copying/Equals_to_AAA-8          426055              2787 ns/op            4080 B/op          8 allocs/op
BenchmarkFilter_Generics_Inplace/Equals_to_AAA-8          483994              2467 ns/op               0 B/op          0 allocs/op
PASS
ok      github.com/timliudream/go-test/generic_test     4.925s

1
2
3
4
5
6
7
8
9
10

不仅内存分配归零，而且性能也明显提高。