go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/collections/batch_iterator.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package collections
     9  
    10  // BatchIterator allows you to iterate over a larger
    11  // slice of items in chunsk with a given `BatchSize`.
    12  //
    13  // It is useful if you need to, for example, process
    14  // 100,000 items in 1000 item chunks, such as
    15  // for batch inserts.
    16  type BatchIterator[T any] struct {
    17  	Items     []T
    18  	BatchSize int
    19  	Cursor    int
    20  }
    21  
    22  // HasNext returns if we should process another batch.
    23  func (bi *BatchIterator[T]) HasNext() bool {
    24  	return bi.Cursor < (len(bi.Items) - 1)
    25  }
    26  
    27  // Next yields the next batch of size `BatchSize` or smaller.
    28  func (bi *BatchIterator[T]) Next() []T {
    29  	if bi.BatchSize == 0 {
    30  		return nil
    31  	}
    32  	if bi.Cursor >= len(bi.Items) {
    33  		return nil
    34  	}
    35  
    36  	if (bi.Cursor + bi.BatchSize) < len(bi.Items) {
    37  		output := bi.Items[bi.Cursor : bi.Cursor+bi.BatchSize]
    38  		bi.Cursor += len(output)
    39  		return output
    40  	}
    41  
    42  	output := bi.Items[bi.Cursor:]
    43  	bi.Cursor += len(output)
    44  	return output
    45  }