Repository: segmentio/parquet-go
Pull request #225 (2022-06-09 01:36:07)

Configure the capacity of column buffers as a number of values
The size of column buffers was configured as a number of bytes, this PR changes it to be expressed as a number of values, the intent being to improve usability of the library, and reduce memory waste.

Different column types have widely different memory footprints; considering a UUID column using 16 bytes per value, and a boolean column using 1 bit per value, there is a 128x ratio between the memory density of these columns. When the size of column buffers is expressed in bytes, we either under-allocate the buffers for columns of large types (resulting in more re-allocation and copies of the buffers), or over-allocate the buffers for columns of small types (resulting in a lot of unused memory). By expressing the capacity as a number of values, the library is able to better adjust the initial sizes of column buffers depending on the column types.

I believe that this change will also improve usability of the code: instead of having to reason in terms of how many bytes the application is expected to use, users can reason in terms of how many rows they are expecting to buffer in memory, which is a measure mapping more naturally to the application logic.

Please take a look and let me know if you have any comments or concerns about the change!


Patch:
diff --git a/buffer.go b/buffer.go
index 472e67d..1251a83 100644
--- a/buffer.go
+++ b/buffer.go
@@ -62,17 +62,17 @@ func (buf *Buffer) configure(schema *Schema) {
 		nullOrdering := nullsGoLast
 		columnIndex := int(leaf.columnIndex)
 		columnType := leaf.node.Type()
-		bufferSize := buf.config.ColumnBufferSize
+		bufferCap := buf.config.ColumnBufferCapacity
 		dictionary := (Dictionary)(nil)
 		encoding := encodingOf(leaf.node)
 
 		if isDictionaryEncoding(encoding) {
-			bufferSize /= 2
-			dictionary = columnType.NewDictionary(columnIndex, 0, make([]byte, 0, bufferSize))
+			dictBuffer := make([]byte, 0, columnType.EstimateSize(bufferCap))
+			dictionary = columnType.NewDictionary(columnIndex, 0, dictBuffer)
 			columnType = dictionary.Type()
 		}
 
-		column := columnType.NewColumnBuffer(columnIndex, bufferSize)
+		column := columnType.NewColumnBuffer(columnIndex, bufferCap)
 		switch {
 		case leaf.maxRepetitionLevel > 0:
 			column = newRepeatedColumnBuffer(column, leaf.maxRepetitionLevel, leaf.maxDefinitionLevel, nullOrdering)
diff --git a/column_buffer.go b/column_buffer.go
index db5d062..c1a540b 100644
--- a/column_buffer.go
+++ b/column_buffer.go
@@ -727,11 +727,11 @@ func (col *repeatedColumnBuffer) ReadValuesAt(values []Value, offset int64) (int
 
 type booleanColumnBuffer struct{ booleanPage }
 
-func newBooleanColumnBuffer(typ Type, columnIndex int16, bufferSize int) *booleanColumnBuffer {
+func newBooleanColumnBuffer(typ Type, columnIndex int16, numValues int32) *booleanColumnBuffer {
 	return &booleanColumnBuffer{
 		booleanPage: booleanPage{
 			typ:         typ,
-			bits:        make([]byte, 0, bufferSize),
+			bits:        make([]byte, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -916,11 +916,11 @@ func (col *booleanColumnBuffer) ReadValuesAt(values []Value, offset int64) (n in
 
 type int32ColumnBuffer struct{ int32Page }
 
-func newInt32ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *int32ColumnBuffer {
+func newInt32ColumnBuffer(typ Type, columnIndex int16, numValues int32) *int32ColumnBuffer {
 	return &int32ColumnBuffer{
 		int32Page: int32Page{
 			typ:         typ,
-			values:      make([]int32, 0, bufferSize/4),
+			values:      make([]int32, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1018,11 +1018,11 @@ func (col *int32ColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int,
 
 type int64ColumnBuffer struct{ int64Page }
 
-func newInt64ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *int64ColumnBuffer {
+func newInt64ColumnBuffer(typ Type, columnIndex int16, numValues int32) *int64ColumnBuffer {
 	return &int64ColumnBuffer{
 		int64Page: int64Page{
 			typ:         typ,
-			values:      make([]int64, 0, bufferSize/8),
+			values:      make([]int64, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1119,11 +1119,11 @@ func (col *int64ColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int,
 
 type int96ColumnBuffer struct{ int96Page }
 
-func newInt96ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *int96ColumnBuffer {
+func newInt96ColumnBuffer(typ Type, columnIndex int16, numValues int32) *int96ColumnBuffer {
 	return &int96ColumnBuffer{
 		int96Page: int96Page{
 			typ:         typ,
-			values:      make([]deprecated.Int96, 0, bufferSize/12),
+			values:      make([]deprecated.Int96, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1212,11 +1212,11 @@ func (col *int96ColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int,
 
 type floatColumnBuffer struct{ floatPage }
 
-func newFloatColumnBuffer(typ Type, columnIndex int16, bufferSize int) *floatColumnBuffer {
+func newFloatColumnBuffer(typ Type, columnIndex int16, numValues int32) *floatColumnBuffer {
 	return &floatColumnBuffer{
 		floatPage: floatPage{
 			typ:         typ,
-			values:      make([]float32, 0, bufferSize/4),
+			values:      make([]float32, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1313,11 +1313,11 @@ func (col *floatColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int,
 
 type doubleColumnBuffer struct{ doublePage }
 
-func newDoubleColumnBuffer(typ Type, columnIndex int16, bufferSize int) *doubleColumnBuffer {
+func newDoubleColumnBuffer(typ Type, columnIndex int16, numValues int32) *doubleColumnBuffer {
 	return &doubleColumnBuffer{
 		doublePage: doublePage{
 			typ:         typ,
-			values:      make([]float64, 0, bufferSize/8),
+			values:      make([]float64, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1417,14 +1417,14 @@ type byteArrayColumnBuffer struct {
 	offsets []uint32
 }
 
-func newByteArrayColumnBuffer(typ Type, columnIndex int16, bufferSize int) *byteArrayColumnBuffer {
+func newByteArrayColumnBuffer(typ Type, columnIndex int16, numValues int32) *byteArrayColumnBuffer {
 	return &byteArrayColumnBuffer{
 		byteArrayPage: byteArrayPage{
 			typ:         typ,
-			values:      make([]byte, 0, bufferSize/2),
+			values:      make([]byte, 0, typ.EstimateSize(int(numValues))),
 			columnIndex: ^columnIndex,
 		},
-		offsets: make([]uint32, 0, bufferSize/8),
+		offsets: make([]uint32, 0, numValues),
 	}
 }
 
@@ -1566,13 +1566,13 @@ type fixedLenByteArrayColumnBuffer struct {
 	tmp []byte
 }
 
-func newFixedLenByteArrayColumnBuffer(typ Type, columnIndex int16, bufferSize int) *fixedLenByteArrayColumnBuffer {
+func newFixedLenByteArrayColumnBuffer(typ Type, columnIndex int16, numValues int32) *fixedLenByteArrayColumnBuffer {
 	size := typ.Length()
 	return &fixedLenByteArrayColumnBuffer{
 		fixedLenByteArrayPage: fixedLenByteArrayPage{
 			typ:         typ,
 			size:        size,
-			data:        make([]byte, 0, bufferSize),
+			data:        make([]byte, 0, typ.EstimateSize(int(numValues))),
 			columnIndex: ^columnIndex,
 		},
 		tmp: make([]byte, size),
@@ -1691,11 +1691,11 @@ func (col *fixedLenByteArrayColumnBuffer) ReadValuesAt(values []Value, offset in
 
 type uint32ColumnBuffer struct{ uint32Page }
 
-func newUint32ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *uint32ColumnBuffer {
+func newUint32ColumnBuffer(typ Type, columnIndex int16, numValues int32) *uint32ColumnBuffer {
 	return &uint32ColumnBuffer{
 		uint32Page: uint32Page{
 			typ:         typ,
-			values:      make([]uint32, 0, bufferSize/4),
+			values:      make([]uint32, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1792,11 +1792,11 @@ func (col *uint32ColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int
 
 type uint64ColumnBuffer struct{ uint64Page }
 
-func newUint64ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *uint64ColumnBuffer {
+func newUint64ColumnBuffer(typ Type, columnIndex int16, numValues int32) *uint64ColumnBuffer {
 	return &uint64ColumnBuffer{
 		uint64Page: uint64Page{
 			typ:         typ,
-			values:      make([]uint64, 0, bufferSize/8),
+			values:      make([]uint64, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
@@ -1893,11 +1893,11 @@ func (col *uint64ColumnBuffer) ReadValuesAt(values []Value, offset int64) (n int
 
 type be128ColumnBuffer struct{ be128Page }
 
-func newBE128ColumnBuffer(typ Type, columnIndex int16, bufferSize int) *be128ColumnBuffer {
+func newBE128ColumnBuffer(typ Type, columnIndex int16, numValues int32) *be128ColumnBuffer {
 	return &be128ColumnBuffer{
 		be128Page: be128Page{
 			typ:         typ,
-			values:      make([][16]byte, 0, bufferSize/16),
+			values:      make([][16]byte, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
diff --git a/config.go b/config.go
index 079ec9a..6cdddb0 100644
--- a/config.go
+++ b/config.go
@@ -10,8 +10,8 @@ import (
 const (
 	DefaultCreatedBy            = "github.com/segmentio/parquet-go"
 	DefaultColumnIndexSizeLimit = 16
-	DefaultColumnBufferSize     = 1 * 1024 * 1024
-	DefaultPageBufferSize       = 1 * 1024 * 1024
+	DefaultColumnBufferCapacity = 16 * 1024
+	DefaultPageBufferSize       = 256 * 1024
 	DefaultWriteBufferSize      = 32 * 1024
 	DefaultDataPageVersion      = 2
 	DefaultDataPageStatistics   = false
@@ -223,20 +223,20 @@ func (c *WriterConfig) Validate() error {
 // directly as argument to the NewBuffer function when needed, for example:
 //
 //	buffer := parquet.NewBuffer(&parquet.RowGroupConfig{
-//		ColumnBufferSize: 8 * 1024 * 1024,
+//		ColumnBufferCapacity: 10_000,
 //	})
 //
 type RowGroupConfig struct {
-	ColumnBufferSize int
-	SortingColumns   []SortingColumn
-	Schema           *Schema
+	ColumnBufferCapacity int
+	SortingColumns       []SortingColumn
+	Schema               *Schema
 }
 
 // DefaultRowGroupConfig returns a new RowGroupConfig value initialized with the
 // default row group configuration.
 func DefaultRowGroupConfig() *RowGroupConfig {
 	return &RowGroupConfig{
-		ColumnBufferSize: DefaultColumnBufferSize,
+		ColumnBufferCapacity: DefaultColumnBufferCapacity,
 	}
 }
 
@@ -255,7 +255,7 @@ func NewRowGroupConfig(options ...RowGroupOption) (*RowGroupConfig, error) {
 func (c *RowGroupConfig) Validate() error {
 	const baseName = "parquet.(*RowGroupConfig)."
 	return errorInvalidConfiguration(
-		validatePositiveInt(baseName+"ColumnBufferSize", c.ColumnBufferSize),
+		validatePositiveInt(baseName+"ColumnBufferCapacity", c.ColumnBufferCapacity),
 	)
 }
 
@@ -267,9 +267,9 @@ func (c *RowGroupConfig) Apply(options ...RowGroupOption) {
 
 func (c *RowGroupConfig) ConfigureRowGroup(config *RowGroupConfig) {
 	*config = RowGroupConfig{
-		ColumnBufferSize: coalesceInt(c.ColumnBufferSize, config.ColumnBufferSize),
-		SortingColumns:   coalesceSortingColumns(c.SortingColumns, config.SortingColumns),
-		Schema:           coalesceSchema(c.Schema, config.Schema),
+		ColumnBufferCapacity: coalesceInt(c.ColumnBufferCapacity, config.ColumnBufferCapacity),
+		SortingColumns:       coalesceSortingColumns(c.SortingColumns, config.SortingColumns),
+		Schema:               coalesceSchema(c.Schema, config.Schema),
 	}
 }
 
@@ -325,7 +325,7 @@ func SkipBloomFilters(skip bool) FileOption {
 // read and write pages rather than controlling the space used by the encoded
 // representation on disk.
 //
-// Defaults to 1 MiB.
+// Defaults to 256KiB.
 func PageBufferSize(size int) WriterOption {
 	return writerOption(func(config *WriterConfig) { config.PageBufferSize = size })
 }
@@ -335,7 +335,7 @@ func PageBufferSize(size int) WriterOption {
 // Setting the writer buffer size to zero deactivates buffering, all writes are
 // immediately sent to the output io.Writer.
 //
-// Defaults to 32 KiB.
+// Defaults to 32KiB.
 func WriteBufferSize(size int) WriterOption {
 	return writerOption(func(config *WriterConfig) { config.WriteBufferSize = size })
 }
@@ -425,12 +425,12 @@ func Compression(codec compress.Codec) WriterOption {
 	return writerOption(func(config *WriterConfig) { config.Compression = codec })
 }
 
-// ColumnBufferSize creates a configuration option which defines the size of
+// ColumnBufferCapacity creates a configuration option which defines the size of
 // row group column buffers.
 //
-// Defaults to 1 MiB.
-func ColumnBufferSize(size int) RowGroupOption {
-	return rowGroupOption(func(config *RowGroupConfig) { config.ColumnBufferSize = size })
+// Defaults to 16384.
+func ColumnBufferCapacity(size int) RowGroupOption {
+	return rowGroupOption(func(config *RowGroupConfig) { config.ColumnBufferCapacity = size })
 }
 
 // SortingColumns creates a configuration option which defines the sorting order
diff --git a/dictionary.go b/dictionary.go
index 3404f8d..ad2e656 100644
--- a/dictionary.go
+++ b/dictionary.go
@@ -1079,8 +1079,8 @@ func newIndexedType(typ Type, dict Dictionary) *indexedType {
 	return &indexedType{Type: typ, dict: dict}
 }
 
-func (t *indexedType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newIndexedColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *indexedType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newIndexedColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *indexedType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -1215,11 +1215,11 @@ func (r *indexedPageValues) ReadValues(values []Value) (n int, err error) {
 // builds a page of indexes into a parent dictionary when values are written.
 type indexedColumnBuffer struct{ indexedPage }
 
-func newIndexedColumnBuffer(typ *indexedType, columnIndex int16, bufferSize int) *indexedColumnBuffer {
+func newIndexedColumnBuffer(typ *indexedType, columnIndex int16, numValues int32) *indexedColumnBuffer {
 	return &indexedColumnBuffer{
 		indexedPage: indexedPage{
 			typ:         typ,
-			values:      make([]int32, 0, bufferSize/4),
+			values:      make([]int32, 0, numValues),
 			columnIndex: ^columnIndex,
 		},
 	}
diff --git a/type.go b/type.go
index 65f29af..790e7f7 100644
--- a/type.go
+++ b/type.go
@@ -62,6 +62,12 @@ type Type interface {
 	// For other types, the value is zero.
 	Length() int
 
+	// Returns an estimation of the number of bytes required to hold the given
+	// number of values of this type in memory.
+	//
+	// The method returns zero for group types.
+	EstimateSize(numValues int) int64
+
 	// Compares two values and returns a negative integer if a < b, positive if
 	// a > b, or zero if a == b.
 	//
@@ -121,18 +127,16 @@ type Type interface {
 	// accumulating values in memory for (relative to the parent schema),
 	// and the size of their memory buffer.
 	//
-	// The buffer size is given in bytes, because we want to control memory
-	// consumption of the application, which is simpler to achieve with buffer
-	// size expressed in bytes rather than number of elements.
-	//
-	// Note that the buffer size is not a hard limit, it defines the initial
-	// capacity of the column buffer, but may grow as needed. Programs can use
-	// the Size method of the column buffer (or the parent row group, when
-	// relevant) to determine how many bytes are being used, and perform a flush
-	// of the buffers to a storage layer.
+	// The application may give an estimate of the number of values it expects
+	// to write to the buffer as second argument. This estimate helps set the
+	// initialize buffer capacity but is not a hard limit, the underlying memory
+	// buffer will grown as needed to allow more values to be written. Programs
+	// may use the Size method of the column buffer (or the parent row group,
+	// when relevant) to determine how many bytes are being used, and perform a
+	// flush of the buffers to a storage layer.
 	//
 	// The method panics if it is called on a group type.
-	NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer
+	NewColumnBuffer(columnIndex, numValues int) ColumnBuffer
 
 	// Creates a dictionary holding values of this type.
 	//
@@ -232,6 +236,7 @@ type booleanType struct{}
 func (t booleanType) String() string                           { return "BOOLEAN" }
 func (t booleanType) Kind() Kind                               { return Boolean }
 func (t booleanType) Length() int                              { return 1 }
+func (t booleanType) EstimateSize(n int) int64                 { return (int64(n) + 7) / 8 }
 func (t booleanType) Compare(a, b Value) int                   { return compareBool(a.Boolean(), b.Boolean()) }
 func (t booleanType) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t booleanType) LogicalType() *format.LogicalType         { return nil }
@@ -242,8 +247,8 @@ func (t booleanType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newBooleanColumnIndexer()
 }
 
-func (t booleanType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newBooleanColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t booleanType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newBooleanColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t booleanType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -267,6 +272,7 @@ type int32Type struct{}
 func (t int32Type) String() string                           { return "INT32" }
 func (t int32Type) Kind() Kind                               { return Int32 }
 func (t int32Type) Length() int                              { return 32 }
+func (t int32Type) EstimateSize(n int) int64                 { return 4 * int64(n) }
 func (t int32Type) Compare(a, b Value) int                   { return compareInt32(a.Int32(), b.Int32()) }
 func (t int32Type) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t int32Type) LogicalType() *format.LogicalType         { return nil }
@@ -277,8 +283,8 @@ func (t int32Type) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newInt32ColumnIndexer()
 }
 
-func (t int32Type) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t int32Type) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t int32Type) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -302,6 +308,7 @@ type int64Type struct{}
 func (t int64Type) String() string                           { return "INT64" }
 func (t int64Type) Kind() Kind                               { return Int64 }
 func (t int64Type) Length() int                              { return 64 }
+func (t int64Type) EstimateSize(n int) int64                 { return 8 * int64(n) }
 func (t int64Type) Compare(a, b Value) int                   { return compareInt64(a.Int64(), b.Int64()) }
 func (t int64Type) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t int64Type) LogicalType() *format.LogicalType         { return nil }
@@ -312,8 +319,8 @@ func (t int64Type) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newInt64ColumnIndexer()
 }
 
-func (t int64Type) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t int64Type) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t int64Type) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -338,6 +345,7 @@ func (t int96Type) String() string { return "INT96" }
 
 func (t int96Type) Kind() Kind                               { return Int96 }
 func (t int96Type) Length() int                              { return 96 }
+func (t int96Type) EstimateSize(n int) int64                 { return 12 * int64(n) }
 func (t int96Type) Compare(a, b Value) int                   { return compareInt96(a.Int96(), b.Int96()) }
 func (t int96Type) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t int96Type) LogicalType() *format.LogicalType         { return nil }
@@ -348,8 +356,8 @@ func (t int96Type) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newInt96ColumnIndexer()
 }
 
-func (t int96Type) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newInt96ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t int96Type) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newInt96ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t int96Type) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -373,6 +381,7 @@ type floatType struct{}
 func (t floatType) String() string                           { return "FLOAT" }
 func (t floatType) Kind() Kind                               { return Float }
 func (t floatType) Length() int                              { return 32 }
+func (t floatType) EstimateSize(n int) int64                 { return 4 * int64(n) }
 func (t floatType) Compare(a, b Value) int                   { return compareFloat32(a.Float(), b.Float()) }
 func (t floatType) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t floatType) LogicalType() *format.LogicalType         { return nil }
@@ -383,8 +392,8 @@ func (t floatType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newFloatColumnIndexer()
 }
 
-func (t floatType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newFloatColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t floatType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newFloatColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t floatType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -408,6 +417,7 @@ type doubleType struct{}
 func (t doubleType) String() string                           { return "DOUBLE" }
 func (t doubleType) Kind() Kind                               { return Double }
 func (t doubleType) Length() int                              { return 64 }
+func (t doubleType) EstimateSize(n int) int64                 { return 8 * int64(n) }
 func (t doubleType) Compare(a, b Value) int                   { return compareFloat64(a.Double(), b.Double()) }
 func (t doubleType) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t doubleType) LogicalType() *format.LogicalType         { return nil }
@@ -418,8 +428,8 @@ func (t doubleType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newDoubleColumnIndexer()
 }
 
-func (t doubleType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newDoubleColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t doubleType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newDoubleColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t doubleType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -443,6 +453,7 @@ type byteArrayType struct{}
 func (t byteArrayType) String() string                           { return "BYTE_ARRAY" }
 func (t byteArrayType) Kind() Kind                               { return ByteArray }
 func (t byteArrayType) Length() int                              { return 0 }
+func (t byteArrayType) EstimateSize(n int) int64                 { return 10 * int64(n) }
 func (t byteArrayType) Compare(a, b Value) int                   { return bytes.Compare(a.ByteArray(), b.ByteArray()) }
 func (t byteArrayType) ColumnOrder() *format.ColumnOrder         { return &typeDefinedColumnOrder }
 func (t byteArrayType) LogicalType() *format.LogicalType         { return nil }
@@ -453,8 +464,8 @@ func (t byteArrayType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newByteArrayColumnIndexer(sizeLimit)
 }
 
-func (t byteArrayType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t byteArrayType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t byteArrayType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -483,6 +494,8 @@ func (t fixedLenByteArrayType) Kind() Kind { return FixedLenByteArray }
 
 func (t fixedLenByteArrayType) Length() int { return t.length }
 
+func (t fixedLenByteArrayType) EstimateSize(n int) int64 { return int64(t.length) * int64(n) }
+
 func (t fixedLenByteArrayType) Compare(a, b Value) int {
 	return bytes.Compare(a.ByteArray(), b.ByteArray())
 }
@@ -499,8 +512,8 @@ func (t fixedLenByteArrayType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newFixedLenByteArrayColumnIndexer(t.length, sizeLimit)
 }
 
-func (t fixedLenByteArrayType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newFixedLenByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t fixedLenByteArrayType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newFixedLenByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t fixedLenByteArrayType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -538,6 +551,8 @@ func (t be128Type) Kind() Kind { return FixedLenByteArray }
 
 func (t be128Type) Length() int { return 16 }
 
+func (t be128Type) EstimateSize(n int) int64 { return 16 * int64(n) }
+
 func (t be128Type) Compare(a, b Value) int {
 	return compareBE128((*[16]byte)(a.ByteArray()), (*[16]byte)(b.ByteArray()))
 }
@@ -554,8 +569,8 @@ func (t be128Type) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newBE128ColumnIndexer()
 }
 
-func (t be128Type) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newBE128ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t be128Type) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newBE128ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t be128Type) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -644,6 +659,8 @@ func (t *intType) Kind() Kind {
 
 func (t *intType) Length() int { return int(t.BitWidth) }
 
+func (t *intType) EstimateSize(n int) int64 { return int64(t.BitWidth/8) * int64(n) }
+
 func (t *intType) Compare(a, b Value) int {
 	if t.BitWidth == 64 {
 		i1 := a.Int64()
@@ -706,18 +723,18 @@ func (t *intType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	}
 }
 
-func (t *intType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
+func (t *intType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
 	if t.IsSigned {
 		if t.BitWidth == 64 {
-			return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+			return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 		} else {
-			return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+			return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 		}
 	} else {
 		if t.BitWidth == 64 {
-			return newUint64ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+			return newUint64ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 		} else {
-			return newUint32ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+			return newUint32ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 		}
 	}
 }
@@ -817,6 +834,8 @@ func (t *stringType) Kind() Kind { return ByteArray }
 
 func (t *stringType) Length() int { return 0 }
 
+func (t *stringType) EstimateSize(n int) int64 { return 10 * int64(n) }
+
 func (t *stringType) Compare(a, b Value) int {
 	return bytes.Compare(a.ByteArray(), b.ByteArray())
 }
@@ -845,8 +864,8 @@ func (t *stringType) NewDictionary(columnIndex, numValues int, data []byte) Dict
 	return newByteArrayDictionary(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *stringType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *stringType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *stringType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -874,6 +893,8 @@ func (t *uuidType) Kind() Kind { return FixedLenByteArray }
 
 func (t *uuidType) Length() int { return 16 }
 
+func (t *uuidType) EstimateSize(n int) int64 { return 16 * int64(n) }
+
 func (t *uuidType) Compare(a, b Value) int {
 	return compareBE128((*[16]byte)(a.ByteArray()), (*[16]byte)(b.ByteArray()))
 }
@@ -896,8 +917,8 @@ func (t *uuidType) NewDictionary(columnIndex, numValues int, data []byte) Dictio
 	return newBE128Dictionary(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *uuidType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newBE128ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *uuidType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newBE128ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *uuidType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -925,6 +946,8 @@ func (t *enumType) Kind() Kind { return ByteArray }
 
 func (t *enumType) Length() int { return 0 }
 
+func (t *enumType) EstimateSize(n int) int64 { return 10 * int64(n) }
+
 func (t *enumType) Compare(a, b Value) int {
 	return bytes.Compare(a.ByteArray(), b.ByteArray())
 }
@@ -953,8 +976,8 @@ func (t *enumType) NewDictionary(columnIndex, numValues int, data []byte) Dictio
 	return newByteArrayDictionary(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *enumType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *enumType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *enumType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -982,6 +1005,8 @@ func (t *jsonType) Kind() Kind { return ByteArray }
 
 func (t *jsonType) Length() int { return 0 }
 
+func (t *jsonType) EstimateSize(n int) int64 { return 10 * int64(n) }
+
 func (t *jsonType) Compare(a, b Value) int {
 	return bytes.Compare(a.ByteArray(), b.ByteArray())
 }
@@ -1010,8 +1035,8 @@ func (t *jsonType) NewDictionary(columnIndex, numValues int, data []byte) Dictio
 	return newByteArrayDictionary(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *jsonType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *jsonType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *jsonType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -1039,6 +1064,8 @@ func (t *bsonType) Kind() Kind { return ByteArray }
 
 func (t *bsonType) Length() int { return 0 }
 
+func (t *bsonType) EstimateSize(n int) int64 { return 10 * int64(n) }
+
 func (t *bsonType) Compare(a, b Value) int {
 	return bytes.Compare(a.ByteArray(), b.ByteArray())
 }
@@ -1067,8 +1094,8 @@ func (t *bsonType) NewDictionary(columnIndex, numValues int, data []byte) Dictio
 	return newByteArrayDictionary(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *bsonType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *bsonType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newByteArrayColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *bsonType) NewPage(columnIndex, numValues int, data []byte) Page {
@@ -1096,6 +1123,8 @@ func (t *dateType) Kind() Kind { return Int32 }
 
 func (t *dateType) Length() int { return 32 }
 
+func (t *dateType) EstimateSize(n int) int64 { return 4 * int64(n) }
+
 func (t *dateType) Compare(a, b Value) int { return compareInt32(a.Int32(), b.Int32()) }
 
 func (t *dateType) ColumnOrder() *format.ColumnOrder {
@@ -1116,8 +1145,8 @@ func (t *dateType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newInt32ColumnIndexer()
 }
 
-func (t *dateType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *dateType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *dateType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -1209,6 +1238,14 @@ func (t *timeType) Length() int {
 	}
 }
 
+func (t *timeType) EstimateSize(n int) int64 {
+	if t.useInt32() {
+		return 4 * int64(n)
+	} else {
+		return 8 * int64(n)
+	}
+}
+
 func (t *timeType) Compare(a, b Value) int {
 	if t.useInt32() {
 		return compareInt32(a.Int32(), b.Int32())
@@ -1252,11 +1289,11 @@ func (t *timeType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	}
 }
 
-func (t *timeType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
+func (t *timeType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
 	if t.useInt32() {
-		return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+		return newInt32ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 	} else {
-		return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+		return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 	}
 }
 
@@ -1307,6 +1344,8 @@ func (t *timestampType) Kind() Kind { return Int64 }
 
 func (t *timestampType) Length() int { return 64 }
 
+func (t *timestampType) EstimateSize(n int) int64 { return 8 * int64(n) }
+
 func (t *timestampType) Compare(a, b Value) int { return compareInt64(a.Int64(), b.Int64()) }
 
 func (t *timestampType) ColumnOrder() *format.ColumnOrder { return &typeDefinedColumnOrder }
@@ -1332,8 +1371,8 @@ func (t *timestampType) NewColumnIndexer(sizeLimit int) ColumnIndexer {
 	return newInt64ColumnIndexer()
 }
 
-func (t *timestampType) NewColumnBuffer(columnIndex, bufferSize int) ColumnBuffer {
-	return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), bufferSize)
+func (t *timestampType) NewColumnBuffer(columnIndex, numValues int) ColumnBuffer {
+	return newInt64ColumnBuffer(t, makeColumnIndex(columnIndex), makeNumValues(numValues))
 }
 
 func (t *timestampType) NewDictionary(columnIndex, numValues int, data []byte) Dictionary {
@@ -1371,6 +1410,8 @@ func (t *listType) Kind() Kind { panic("cannot call Kind on parquet LIST type")
 
 func (t *listType) Length() int { return 0 }
 
+func (t *listType) EstimateSize(int) int64 { return 0 }
+
 func (t *listType) Compare(Value, Value) int { panic("cannot compare values on parquet LIST type") }
 
 func (t *listType) ColumnOrder() *format.ColumnOrder { return nil }
@@ -1433,6 +1474,8 @@ func (t *mapType) Kind() Kind { panic("cannot call Kind on parquet MAP type") }
 
 func (t *mapType) Length() int { return 0 }
 
+func (t *mapType) EstimateSize(int) int64 { return 0 }
+
 func (t *mapType) Compare(Value, Value) int { panic("cannot compare values on parquet MAP type") }
 
 func (t *mapType) ColumnOrder() *format.ColumnOrder { return nil }
@@ -1479,6 +1522,8 @@ func (t *nullType) Kind() Kind { return -1 }
 
 func (t *nullType) Length() int { return 0 }
 
+func (t *nullType) EstimateSize(int) int64 { return 0 }
+
 func (t *nullType) Compare(Value, Value) int { panic("cannot compare values on parquet NULL type") }
 
 func (t *nullType) ColumnOrder() *format.ColumnOrder { return nil }
@@ -1553,6 +1598,8 @@ func (groupType) Decode(_, _ []byte, _ encoding.Encoding) ([]byte, error) {
 
 func (groupType) Length() int { return 0 }
 
+func (groupType) EstimateSize(int) int64 { return 0 }
+
 func (groupType) ColumnOrder() *format.ColumnOrder { return nil }
 
 func (groupType) PhysicalType() *format.Type { return nil }
diff --git a/writer.go b/writer.go
index e82dd6a..aecac13 100644
--- a/writer.go
+++ b/writer.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"hash/crc32"
 	"io"
+	"math"
 	"math/bits"
 	"sort"
 
@@ -927,7 +928,10 @@ func (c *writerColumn) resizeBloomFilter(numValues int64) {
 }
 
 func (c *writerColumn) newColumnBuffer() ColumnBuffer {
-	column := c.columnType.NewColumnBuffer(int(c.bufferIndex), int(c.bufferSize))
+	columnBufferCapacity := sort.Search(math.MaxInt32, func(i int) bool {
+		return c.columnType.EstimateSize(i) >= int64(c.bufferSize)
+	})
+	column := c.columnType.NewColumnBuffer(int(c.bufferIndex), columnBufferCapacity)
 	switch {
 	case c.maxRepetitionLevel > 0:
 		column = newRepeatedColumnBuffer(column, c.maxRepetitionLevel, c.maxDefinitionLevel, nullsGoLast)


Test patch:
diff --git a/buffer_test.go b/buffer_test.go
index bb62989..924b278 100644
--- a/buffer_test.go
+++ b/buffer_test.go
@@ -206,7 +206,7 @@ func TestBuffer(t *testing.T) {
 
 									options := []parquet.RowGroupOption{
 										schema,
-										parquet.ColumnBufferSize(1024),
+										parquet.ColumnBufferCapacity(100),
 									}
 									if ordering.sorting != nil {
 										options = append(options, parquet.SortingColumns(ordering.sorting))