-
Notifications
You must be signed in to change notification settings - Fork 27
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add batch.BySize
function
#735
Changes from 3 commits
cbd1eb0
424f24c
9719319
9dcb0eb
7b06869
2dd4abd
ce3c838
e5019fa
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,48 @@ | ||
package batch | ||
|
||
import "fmt" | ||
|
||
// BySize takes a series of elements, encodes them, groups them into batches of bytes that sum to at most [maxSize], and | ||
// then passes each of those batches to the [yield] function. | ||
func BySize[T any](in []T, maxSize int, encode func(T) ([]byte, error), yield func([][]byte) error) error { | ||
var buffer [][]byte | ||
var currentSize int | ||
|
||
for i, item := range in { | ||
bytes, err := encode(item) | ||
if err != nil { | ||
return fmt.Errorf("failed to encode item %d: %w", i, err) | ||
} | ||
|
||
if len(bytes) > maxSize { | ||
return fmt.Errorf("item %d is larger (%d bytes) than max size (%d bytes)", i, len(bytes), maxSize) | ||
} | ||
|
||
currentSize += len(bytes) | ||
|
||
if currentSize < maxSize { | ||
buffer = append(buffer, bytes) | ||
} else if currentSize == maxSize { | ||
buffer = append(buffer, bytes) | ||
if err := yield(buffer); err != nil { | ||
return err | ||
} | ||
buffer = [][]byte{} | ||
currentSize = 0 | ||
} else { | ||
if err := yield(buffer); err != nil { | ||
return err | ||
} | ||
buffer = [][]byte{bytes} | ||
currentSize = len(bytes) | ||
} | ||
} | ||
|
||
if len(buffer) > 0 { | ||
if err := yield(buffer); err != nil { | ||
return err | ||
} | ||
} | ||
|
||
return nil | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,90 @@ | ||
package batch | ||
|
||
import ( | ||
"fmt" | ||
"testing" | ||
|
||
"github.com/stretchr/testify/assert" | ||
) | ||
|
||
func TestBySize(t *testing.T) { | ||
goodEncoder := func(value string) ([]byte, error) { | ||
return []byte(value), nil | ||
} | ||
|
||
panicEncoder := func(value string) ([]byte, error) { | ||
panic("should not be called") | ||
} | ||
|
||
badEncoder := func(value string) ([]byte, error) { | ||
return nil, fmt.Errorf("failed to encode %q", value) | ||
} | ||
|
||
testBySize := func(in []string, maxSize int, encoder func(value string) ([]byte, error)) ([][][]byte, error) { | ||
out := [][][]byte{} | ||
err := BySize(in, maxSize, encoder, func(batch [][]byte) error { out = append(out, batch); return nil }) | ||
return out, err | ||
} | ||
|
||
{ | ||
// Empty slice: | ||
batches, err := testBySize([]string{}, 0, panicEncoder) | ||
assert.NoError(t, err) | ||
assert.Empty(t, batches) | ||
} | ||
{ | ||
// Non-empty slice + bad encoder: | ||
_, err := testBySize([]string{"foo", "bar"}, 10, badEncoder) | ||
assert.ErrorContains(t, err, `failed to encode item 0: failed to encode "foo"`) | ||
} | ||
{ | ||
// Non-empty slice + two items that are < maxSize + yield returns error. | ||
err := BySize([]string{"foo", "bar"}, 10, goodEncoder, func(batch [][]byte) error { return fmt.Errorf("yield failed") }) | ||
assert.ErrorContains(t, err, "yield failed") | ||
} | ||
{ | ||
// Non-empty slice + two items that are = maxSize + yield returns error. | ||
err := BySize([]string{"foo", "bar"}, 6, goodEncoder, func(batch [][]byte) error { return fmt.Errorf("yield failed") }) | ||
assert.ErrorContains(t, err, "yield failed") | ||
} | ||
{ | ||
// Non-empty slice + two items that are > maxSize + yield returns error. | ||
err := BySize([]string{"foo", "bar-baz"}, 8, goodEncoder, func(batch [][]byte) error { return fmt.Errorf("yield failed") }) | ||
assert.ErrorContains(t, err, "yield failed") | ||
} | ||
{ | ||
// Non-empty slice + item is larger than max size: | ||
_, err := testBySize([]string{"foo", "i-am-23-characters-long", "bar"}, 20, goodEncoder) | ||
assert.ErrorContains(t, err, "item 1 is larger (23 bytes) than max size (20 bytes)") | ||
} | ||
{ | ||
// Non-empty slice + item equal to max size: | ||
batches, err := testBySize([]string{"foo", "i-am-23-characters-long", "bar"}, 23, goodEncoder) | ||
assert.NoError(t, err) | ||
assert.Len(t, batches, 3) | ||
assert.Equal(t, [][]byte{[]byte("foo")}, batches[0]) | ||
assert.Equal(t, [][]byte{[]byte("i-am-23-characters-long")}, batches[1]) | ||
assert.Equal(t, [][]byte{[]byte("bar")}, batches[2]) | ||
} | ||
{ | ||
// Non-empty slice + one item: | ||
batches, err := testBySize([]string{"foo"}, 100, goodEncoder) | ||
assert.NoError(t, err) | ||
assert.Len(t, batches, 1) | ||
assert.Equal(t, [][]byte{[]byte("foo")}, batches[0]) | ||
} | ||
{ | ||
// Non-empty slice + all items exactly fit into one batch: | ||
batches, err := testBySize([]string{"foo", "bar", "baz", "qux"}, 12, goodEncoder) | ||
assert.NoError(t, err) | ||
assert.Len(t, batches, 1) | ||
assert.Equal(t, [][]byte{[]byte("foo"), []byte("bar"), []byte("baz"), []byte("qux")}, batches[0]) | ||
} | ||
{ | ||
// Non-empty slice + all items exactly fit into just under one batch: | ||
batches, err := testBySize([]string{"foo", "bar", "baz", "qux"}, 13, goodEncoder) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Worth adding a check on how many times There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In this case |
||
assert.NoError(t, err) | ||
assert.Len(t, batches, 1) | ||
assert.Equal(t, [][]byte{[]byte("foo"), []byte("bar"), []byte("baz"), []byte("qux")}, batches[0]) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maxSizeBytes?