Repository: hamba/avro

Issue:
maxBufferedBytes in addition to blockLengh
Hi @nrwiersma,

First of all, thank you for your work on this project. We’ve been using this library to write terabytes of OCF files per day, and it has been running flawlessly for years - much appreciated!

We're currently encountering an issue where one of our applications hits OOM errors when writing very large individual records because the buffer can grow unbonded before writeBlocks() is triggered. (only blockLengh is taken into account)

Would it be possible to add an additional trigger in ocf.Encoder that calls writeBlocks() based on the current buffer size, not just the block lengh? This would help us better control memory usage for workloads with large record sizes.

We are using the raw ocf.Encoder.Write method.

Happy to help test or submit a PR if needed.

_edit_: Maybe blocksize is the wrong term as the block might be compressed, but maxBufferedBytes then :)

Hints from issue comments:
@nrwiersma: The issue is that the block length and size (somewhat optional) prefix the block data, so the encoder would need to know the length of the block up front. In this case it would not need to buffer. The downside is that the user must set this correctly, otherwise the encoding becomes bad and cannot be decoded. This is the only way I see to deal with your issue.@brunsgaard: Just to clarify — I’m not suggesting removing buffering entirely (I understand that would be complex, especially since the block length must be written before the data segment). What I’m proposing is simply an additional condition to trigger `writerBlock()` early, based on the buffer size, to prevent the buffer from growing to an unhealthy size.

<img width="1316" alt="Image" src="https://github.com/user-attachments/assets/33a9143e-adcc-49be-ae84-222b4b2170be" />

This is a rough proof of concept - you’re much more familiar with the Avro spec than I am, so there may be reasons this approach won’t work, but I think it’s a viable starting point.

https://github.com/hamba/avro/compare/main...brunsgaard:avro:feat/maxBufSize?expand=1


The logical change is there
```go
	if e.count >= e.blockLength || e.buf.Len() > e.maxBufSize {
		if err = e.writerBlock(); err != nil {
			return n, err
		}
```


Let me know what you think.@nrwiersma: Fair enough, I had forgotten that Write expects each call to be a full record. Should work, except that you need to check that `maxBufSize > 0` otherwise it is always true if unset.@brunsgaard: I will prepare a PR later today, do you have an opinion about what the public option should be called?@nrwiersma: Looks like `blockSize` would fit well with the existing `blockLength`.@brunsgaard: https://github.com/hamba/avro/compare/main...brunsgaard:avro:feat/maxBufSize?expand=1

Does this looks good to you?@nrwiersma: Good enough start, the rest can be worked out in the PR.

Solution patch:
diff --git a/ocf/ocf.go b/ocf/ocf.go
index ce552ec..4a1030e 100644
--- a/ocf/ocf.go
+++ b/ocf/ocf.go
@@ -217,6 +217,7 @@ func (d *Decoder) readBlock() int64 {
 
 type encoderConfig struct {
 	BlockLength     int
+	BlockSize       int
 	CodecName       CodecName
 	CodecOptions    codecOptions
 	Metadata        map[string][]byte
@@ -236,6 +237,14 @@ func WithBlockLength(length int) EncoderFunc {
 	}
 }
 
+// WithBlockSize sets the maximum uncompressed size of a buffered block before
+// it is written and flushed to the underlying io.Writer (after compression).
+func WithBlockSize(size int) EncoderFunc {
+	return func(cfg *encoderConfig) {
+		cfg.BlockSize = size
+	}
+}
+
 // WithCodec sets the compression codec on the encoder.
 func WithCodec(codec CodecName) EncoderFunc {
 	return func(cfg *encoderConfig) {
@@ -315,6 +324,7 @@ type Encoder struct {
 
 	blockLength int
 	count       int
+	blockSize   int
 }
 
 // NewEncoder returns a new encoder that writes to w using schema s.
@@ -367,6 +377,7 @@ func newEncoder(schema avro.Schema, w io.Writer, cfg encoderConfig) (*Encoder, e
 				sync:        h.Sync,
 				codec:       h.Codec,
 				blockLength: cfg.BlockLength,
+				blockSize:   cfg.BlockSize,
 			}
 			return e, nil
 		}
@@ -407,6 +418,7 @@ func newEncoder(schema avro.Schema, w io.Writer, cfg encoderConfig) (*Encoder, e
 		sync:        header.Sync,
 		codec:       codec,
 		blockLength: cfg.BlockLength,
+		blockSize:   cfg.BlockSize,
 	}
 	return e, nil
 }
@@ -429,6 +441,12 @@ func computeEncoderConfig(opts []EncoderFunc) encoderConfig {
 	return cfg
 }
 
+func (e *Encoder) shouldWriteDataBlock() bool {
+	blockLengthReached := e.blockLength != 0 && e.count >= e.blockLength
+	blockSizeReached := e.blockSize != 0 && e.buf.Len() >= e.blockSize
+	return blockLengthReached || blockSizeReached
+}
+
 // Write v to the internal buffer. This method skips the internal encoder and
 // therefore the caller is responsible for encoding the bytes. No error will be
 // thrown if the bytes does not conform to the schema given to NewEncoder, but
@@ -440,7 +458,7 @@ func (e *Encoder) Write(p []byte) (n int, err error) {
 	}
 
 	e.count++
-	if e.count >= e.blockLength {
+	if e.shouldWriteDataBlock() {
 		if err = e.writerBlock(); err != nil {
 			return n, err
 		}
@@ -456,7 +474,7 @@ func (e *Encoder) Encode(v any) error {
 	}
 
 	e.count++
-	if e.count >= e.blockLength {
+	if e.shouldWriteDataBlock() {
 		if err := e.writerBlock(); err != nil {
 			return err
 		}



Test patch:
diff --git a/ocf/ocf_test.go b/ocf/ocf_test.go
index 0831ab9..b22e6dc 100644
--- a/ocf/ocf_test.go
+++ b/ocf/ocf_test.go
@@ -987,6 +987,29 @@ func TestEncoder_EncodeWritesBlocks(t *testing.T) {
 	assert.Equal(t, 77, buf.Len())
 }
 
+func TestEncoder_EncodeWritesBlocksSize(t *testing.T) {
+	buf := &bytes.Buffer{}
+	enc, _ := ocf.NewEncoder(`"long"`, buf, ocf.WithBlockSize(2))
+	t.Cleanup(func() { _ = enc.Close() })
+
+	// Initial header should be exactly 58 bytes
+	assert.Equal(t, 58, buf.Len())
+
+	// Encode a single int64 value (1), which zig-zag encodes to 1 byte
+	err := enc.Encode(int64(1))
+	require.NoError(t, err)
+
+	// Block size threshold not reached yet; no data block should be written
+	assert.Equal(t, 58, buf.Len())
+
+	// Encode another int64 value (1), bringing buffer size to 2 bytes
+	err = enc.Encode(int64(1))
+	require.NoError(t, err)
+
+	// Block size threshold reached; data block should now be flushed
+	assert.Equal(t, 78, buf.Len())
+}
+
 func TestEncoder_EncodeHandlesWriteBlockError(t *testing.T) {
 	w := &errorBlockWriter{}
 	enc, _ := ocf.NewEncoder(`"long"`, w, ocf.WithBlockLength(1))