用 Go 写一个 Batch 函数

在 Python 下想将一个数组分成 Batch 处理是非常容易的,但是在 go 下好像就没那么容易了。

于是,自己动手,写了一个 batch 函数,支持并行操作。

func Batch(count, batchSize int, parallel int, eachFn func(idx, start, end int) error) error {
	if count <= 0 || batchSize <= 0 {
		return nil
	}

	numBatches := count / batchSize
	if 0 != count%batchSize {
		numBatches++
	}

	err := make([]error, numBatches)

	// Limit parallelism
	isParallel := parallel != 0
	if parallel < 0 || parallel > numBatches {
		parallel = numBatches
	}
	wg := sync.WaitGroup{}
	ch := make(chan struct{}, parallel)

	for idx := 0; idx < numBatches; idx++ {
		sta := idx * batchSize
		end := sta + batchSize
		if end > count {
			end = count
		}

		if isParallel {
			wg.Add(1)
			ch <- struct{}{}
			go func(index, sta, end int) {
				defer func() {
					wg.Done()
					<-ch
				}()

				err[index] = eachFn(index, sta, end)

			}(idx, sta, end)
		} else {
			err[idx] = eachFn(idx, sta, end)
		}
	}
	wg.Wait()

	for _, e := range err {
		if nil != e {
			return e
		}
	}
	return nil
}

大概用法是这样的:

someList := ...
parallel := -1
err := util.Batch(len(someList), someBatchSize, parallel,
    func(batchId, start, end int) err error {
        // do anything you want. for example:
        data := someList[start:end]
        for _, v := range data{
                // do something
        }
        return nil
    },
)

可以用 parallel 参数限制协程数量以控制并行程度,0 是不并行,-1 是不限协程数量。

还算比较好用

发表评论

您的电子邮箱地址不会被公开。