Repository: segmentio/parquet-go
Pull request #434 (2022-12-07 05:08:45)

better estimation of the size of decode output buffers
This PR adds a new method to `parquet.Type` named `EstimateDecodeSize` which is similar in functionality to `EstimateSize` but provides more precise results when working with pages of variable size values (e.g. BYTE_ARRAY).

The use of this new method allows us to better size the output buffer where pages are decode, which greatly reduces memory pressure by avoiding reallocation of the buffer since the one that was selected to decode BYTE_ARRAY pages was almost always too small.


Patch:
diff --git a/column.go b/column.go
index 1e25258..2cccadd 100644
--- a/column.go
+++ b/column.go
@@ -647,14 +647,12 @@ func (c *Column) decodeDataPage(header DataPageHeader, numValues int, repetition
 	var pageValues []byte
 	var pageOffsets []uint32
 
-	vbuf = buffers.get(int(pageType.EstimateSize(numValues)))
+	vbuf = buffers.get(pageType.EstimateDecodeSize(numValues, page.data, pageEncoding))
 	defer vbuf.unref()
 	pageValues = vbuf.data
 
-	pageKind := pageType.Kind()
-
 	// Page offsets not needed when dictionary-encoded
-	if pageKind == ByteArray && !isDictionaryEncoding(pageEncoding) {
+	if pageType.Kind() == ByteArray && !isDictionaryEncoding(pageEncoding) {
 		obuf = buffers.get(4 * (numValues + 1))
 		defer obuf.unref()
 		pageOffsets = unsafecast.BytesToUint32(obuf.data)
diff --git a/dictionary.go b/dictionary.go
index 69c2f06..eae3592 100644
--- a/dictionary.go
+++ b/dictionary.go
@@ -1225,10 +1225,6 @@ func (t *indexedType) NewPage(columnIndex, numValues int, data encoding.Values)
 	return newIndexedPage(t, makeColumnIndex(columnIndex), makeNumValues(numValues), data)
 }
 
-func (t *indexedType) NewValues(values []byte, _ []uint32) encoding.Values {
-	return encoding.Int32ValuesFromBytes(values)
-}
-
 // indexedPage is an implementation of the Page interface which stores
 // indexes instead of plain value. The indexes reference the values in a
 // dictionary that the page was created for.
@@ -1311,6 +1307,10 @@ func (page *indexedPage) Slice(i, j int64) Page {
 // its dictionary instead of plain values.
 type indexedPageType struct{ *indexedType }
 
+func (t indexedPageType) NewValues(values []byte, _ []uint32) encoding.Values {
+	return encoding.Int32ValuesFromBytes(values)
+}
+
 func (t indexedPageType) Encode(dst []byte, src encoding.Values, enc encoding.Encoding) ([]byte, error) {
 	return encoding.EncodeInt32(dst, src, enc)
 }
@@ -1319,6 +1319,10 @@ func (t indexedPageType) Decode(dst encoding.Values, src []byte, enc encoding.En
 	return encoding.DecodeInt32(dst, src, enc)
 }
 
+func (t indexedPageType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return Int32Type.EstimateDecodeSize(numValues, src, enc)
+}
+
 type indexedPageValues struct {
 	page   *indexedPage
 	offset int
diff --git a/encoding/delta/byte_array.go b/encoding/delta/byte_array.go
index d8cee61..f260a48 100644
--- a/encoding/delta/byte_array.go
+++ b/encoding/delta/byte_array.go
@@ -177,6 +177,15 @@ func (e *ByteArrayEncoding) DecodeFixedLenByteArray(dst []byte, src []byte, size
 	return decodeFixedLenByteArray(dst[:0], src, size, prefix.values, suffix.values)
 }
 
+func (e *ByteArrayEncoding) EstimateDecodeByteArraySize(src []byte) int {
+	length := getInt32Buffer()
+	defer putInt32Buffer(length)
+	src, _ = length.decode(src)
+	sum := int(length.sum())
+	length.decode(src)
+	return sum + int(length.sum())
+}
+
 func (e *ByteArrayEncoding) wrap(err error) error {
 	if err != nil {
 		err = encoding.Error(e, err)
diff --git a/encoding/delta/delta.go b/encoding/delta/delta.go
index f05f087..e68cefe 100644
--- a/encoding/delta/delta.go
+++ b/encoding/delta/delta.go
@@ -25,6 +25,13 @@ func (buf *int32Buffer) decode(src []byte) ([]byte, error) {
 	return remain, err
 }
 
+func (buf *int32Buffer) sum() (sum int32) {
+	for _, v := range buf.values {
+		sum += v
+	}
+	return sum
+}
+
 var (
 	int32BufferPool sync.Pool // *int32Buffer
 )
diff --git a/encoding/delta/length_byte_array.go b/encoding/delta/length_byte_array.go
index 5cac384..22af240 100644
--- a/encoding/delta/length_byte_array.go
+++ b/encoding/delta/length_byte_array.go
@@ -62,6 +62,13 @@ func (e *LengthByteArrayEncoding) DecodeByteArray(dst []byte, src []byte, offset
 	return append(dst, src[:lastOffset]...), offsets, nil
 }
 
+func (e *LengthByteArrayEncoding) EstimateDecodeByteArraySize(src []byte) int {
+	length := getInt32Buffer()
+	defer putInt32Buffer(length)
+	length.decode(src)
+	return int(length.sum())
+}
+
 func (e *LengthByteArrayEncoding) wrap(err error) error {
 	if err != nil {
 		err = encoding.Error(e, err)
diff --git a/encoding/encoding.go b/encoding/encoding.go
index 03bb252..c199dea 100644
--- a/encoding/encoding.go
+++ b/encoding/encoding.go
@@ -54,4 +54,15 @@ type Encoding interface {
 	DecodeDouble(dst []float64, src []byte) ([]float64, error)
 	DecodeByteArray(dst []byte, src []byte, offsets []uint32) ([]byte, []uint32, error)
 	DecodeFixedLenByteArray(dst []byte, src []byte, size int) ([]byte, error)
+
+	// Computes an estimation of the output size of decoding the encoded page
+	// of values passed as argument.
+	//
+	// Note that this is an estimate, it is useful to preallocate the output
+	// buffer that will be passed to the decode method, but the actual output
+	// size may be different.
+	//
+	// The estimate never errors since it is not intended to be used as an
+	// input validation method.
+	EstimateDecodeByteArraySize(src []byte) int
 }
diff --git a/encoding/notsupported.go b/encoding/notsupported.go
index 07712e0..530be12 100644
--- a/encoding/notsupported.go
+++ b/encoding/notsupported.go
@@ -196,6 +196,10 @@ func (NotSupported) DecodeFixedLenByteArray(dst []byte, src []byte, size int) ([
 	return dst, errNotSupported("FIXED_LEN_BYTE_ARRAY")
 }
 
+func (NotSupported) EstimateDecodeByteArraySize(src []byte) int {
+	return 0
+}
+
 func errNotSupported(typ string) error {
 	return fmt.Errorf("%w for type %s", ErrNotSupported, typ)
 }
diff --git a/encoding/plain/plain.go b/encoding/plain/plain.go
index 7f1da9c..e382cce 100644
--- a/encoding/plain/plain.go
+++ b/encoding/plain/plain.go
@@ -147,6 +147,10 @@ func (e *Encoding) DecodeFixedLenByteArray(dst []byte, src []byte, size int) ([]
 	return append(dst[:0], src...), nil
 }
 
+func (e *Encoding) EstimateDecodeByteArraySize(src []byte) int {
+	return len(src)
+}
+
 func Boolean(v bool) []byte { return AppendBoolean(nil, 0, v) }
 
 func Int32(v int32) []byte { return AppendInt32(nil, v) }
diff --git a/type.go b/type.go
index 82a06f9..52dda73 100644
--- a/type.go
+++ b/type.go
@@ -69,7 +69,7 @@ type Type interface {
 	// number of values of this type in memory.
 	//
 	// The method returns zero for group types.
-	EstimateSize(numValues int) int64
+	EstimateSize(numValues int) int
 
 	// Compares two values and returns a negative integer if a < b, positive if
 	// a > b, or zero if a == b.
@@ -188,6 +188,14 @@ type Type interface {
 	// encoding methods.
 	Decode(dst encoding.Values, src []byte, enc encoding.Encoding) (encoding.Values, error)
 
+	// Returns an estimation of the output size after decoding the values passed
+	// as first argument with the given encoding.
+	//
+	// For most types, this is similar to calling EstimateSize with the known
+	// number of encoded values. For variable size types, using this method may
+	// provide a more precise result since it can inspect the input buffer.
+	EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int
+
 	// Assigns a Parquet value to a Go value. Returns an error if assignment is
 	// not possible. The source Value must be an expected logical type for the
 	// receiver. This can be accomplished using ConvertValue.
@@ -256,7 +264,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) EstimateSize(n int) int                   { return (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 }
@@ -291,6 +299,10 @@ func (t booleanType) Decode(dst encoding.Values, src []byte, enc encoding.Encodi
 	return encoding.DecodeBoolean(dst, src, enc)
 }
 
+func (t booleanType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t booleanType) AssignValue(dst reflect.Value, src Value) error {
 	v := src.boolean()
 	switch dst.Kind() {
@@ -312,7 +324,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) EstimateSize(n int) int                   { return 4 * 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 }
@@ -347,6 +359,10 @@ func (t int32Type) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeInt32(dst, src, enc)
 }
 
+func (t int32Type) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t int32Type) AssignValue(dst reflect.Value, src Value) error {
 	v := src.int32()
 	switch dst.Kind() {
@@ -370,7 +386,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) EstimateSize(n int) int                   { return 8 * 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 }
@@ -405,6 +421,10 @@ func (t int64Type) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeInt64(dst, src, enc)
 }
 
+func (t int64Type) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t int64Type) AssignValue(dst reflect.Value, src Value) error {
 	v := src.int64()
 	switch dst.Kind() {
@@ -428,7 +448,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) EstimateSize(n int) int                   { return 12 * 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 }
@@ -463,6 +483,10 @@ func (t int96Type) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeInt96(dst, src, enc)
 }
 
+func (t int96Type) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t int96Type) AssignValue(dst reflect.Value, src Value) error {
 	v := src.Int96()
 	dst.Set(reflect.ValueOf(v))
@@ -478,7 +502,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) EstimateSize(n int) int                   { return 4 * 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 }
@@ -513,6 +537,10 @@ func (t floatType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeFloat(dst, src, enc)
 }
 
+func (t floatType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t floatType) AssignValue(dst reflect.Value, src Value) error {
 	v := src.float()
 	switch dst.Kind() {
@@ -533,7 +561,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) EstimateSize(n int) int                   { return 8 * 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 }
@@ -568,6 +596,10 @@ func (t doubleType) Decode(dst encoding.Values, src []byte, enc encoding.Encodin
 	return encoding.DecodeDouble(dst, src, enc)
 }
 
+func (t doubleType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t doubleType) AssignValue(dst reflect.Value, src Value) error {
 	v := src.double()
 	switch dst.Kind() {
@@ -589,7 +621,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) EstimateSize(n int) int                   { return 10 * 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 }
@@ -624,6 +656,10 @@ func (t byteArrayType) Decode(dst encoding.Values, src []byte, enc encoding.Enco
 	return encoding.DecodeByteArray(dst, src, enc)
 }
 
+func (t byteArrayType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return enc.EstimateDecodeByteArraySize(src)
+}
+
 func (t byteArrayType) AssignValue(dst reflect.Value, src Value) error {
 	v := src.byteArray()
 	switch dst.Kind() {
@@ -653,7 +689,7 @@ 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) EstimateSize(n int) int { return t.length * n }
 
 func (t fixedLenByteArrayType) Compare(a, b Value) int {
 	return bytes.Compare(a.byteArray(), b.byteArray())
@@ -695,6 +731,10 @@ func (t fixedLenByteArrayType) Decode(dst encoding.Values, src []byte, enc encod
 	return encoding.DecodeFixedLenByteArray(dst, src, enc)
 }
 
+func (t fixedLenByteArrayType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t fixedLenByteArrayType) AssignValue(dst reflect.Value, src Value) error {
 	v := src.byteArray()
 	switch dst.Kind() {
@@ -743,7 +783,7 @@ 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) EstimateSize(n int) int { return 16 * n }
 
 func (t be128Type) Compare(a, b Value) int { return compareBE128(a.be128(), b.be128()) }
 
@@ -783,6 +823,10 @@ func (t be128Type) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeFixedLenByteArray(dst, src, enc)
 }
 
+func (t be128Type) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return t.EstimateSize(numValues)
+}
+
 func (t be128Type) AssignValue(dst reflect.Value, src Value) error {
 	return fixedLenByteArrayType{length: 16}.AssignValue(dst, src)
 }
@@ -861,7 +905,7 @@ 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) EstimateSize(n int) int { return (int(t.BitWidth) / 8) * n }
 
 func (t *intType) Compare(a, b Value) int {
 	if t.BitWidth == 64 {
@@ -997,6 +1041,14 @@ func (t *intType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding)
 	}
 }
 
+func (t *intType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	if t.BitWidth == 64 {
+		return Int64Type.EstimateDecodeSize(numValues, src, enc)
+	} else {
+		return Int32Type.EstimateDecodeSize(numValues, src, enc)
+	}
+}
+
 func (t *intType) AssignValue(dst reflect.Value, src Value) error {
 	if t.BitWidth == 64 {
 		return Int64Type.AssignValue(dst, src)
@@ -1060,7 +1112,7 @@ 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) EstimateSize(n int) int { return 10 * n }
 
 func (t *stringType) Compare(a, b Value) int {
 	return bytes.Compare(a.byteArray(), b.byteArray())
@@ -1110,6 +1162,10 @@ func (t *stringType) Decode(dst encoding.Values, src []byte, enc encoding.Encodi
 	return encoding.DecodeByteArray(dst, src, enc)
 }
 
+func (t *stringType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return ByteArrayType.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *stringType) AssignValue(dst reflect.Value, src Value) error {
 	return ByteArrayType.AssignValue(dst, src)
 }
@@ -1131,7 +1187,7 @@ 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) EstimateSize(n int) int { return 16 * n }
 
 func (t *uuidType) Compare(a, b Value) int { return compareBE128(a.be128(), b.be128()) }
 
@@ -1173,12 +1229,16 @@ func (t *uuidType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeFixedLenByteArray(dst, src, enc)
 }
 
+func (t *uuidType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return be128Type{}.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *uuidType) AssignValue(dst reflect.Value, src Value) error {
-	return fixedLenByteArrayType{length: 16}.AssignValue(dst, src)
+	return be128Type{}.AssignValue(dst, src)
 }
 
 func (t *uuidType) ConvertValue(val Value, typ Type) (Value, error) {
-	return fixedLenByteArrayType{length: 16}.ConvertValue(val, typ)
+	return be128Type{}.ConvertValue(val, typ)
 }
 
 // Enum constructs a leaf node with a logical type representing enumerations.
@@ -1194,7 +1254,7 @@ 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) EstimateSize(n int) int { return 10 * n }
 
 func (t *enumType) Compare(a, b Value) int {
 	return bytes.Compare(a.byteArray(), b.byteArray())
@@ -1244,6 +1304,10 @@ func (t *enumType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeByteArray(dst, src, enc)
 }
 
+func (t *enumType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return ByteArrayType.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *enumType) AssignValue(dst reflect.Value, src Value) error {
 	return ByteArrayType.AssignValue(dst, src)
 }
@@ -1265,7 +1329,7 @@ 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) EstimateSize(n int) int { return 10 * n }
 
 func (t *jsonType) Compare(a, b Value) int {
 	return bytes.Compare(a.byteArray(), b.byteArray())
@@ -1315,6 +1379,10 @@ func (t *jsonType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeByteArray(dst, src, enc)
 }
 
+func (t *jsonType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return ByteArrayType.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *jsonType) AssignValue(dst reflect.Value, src Value) error {
 	return ByteArrayType.AssignValue(dst, src)
 }
@@ -1336,7 +1404,7 @@ 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) EstimateSize(n int) int { return 10 * n }
 
 func (t *bsonType) Compare(a, b Value) int {
 	return bytes.Compare(a.byteArray(), b.byteArray())
@@ -1386,6 +1454,10 @@ func (t *bsonType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeByteArray(dst, src, enc)
 }
 
+func (t *bsonType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return ByteArrayType.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *bsonType) AssignValue(dst reflect.Value, src Value) error {
 	return ByteArrayType.AssignValue(dst, src)
 }
@@ -1407,7 +1479,7 @@ 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) EstimateSize(n int) int { return 4 * n }
 
 func (t *dateType) Compare(a, b Value) int { return compareInt32(a.int32(), b.int32()) }
 
@@ -1453,6 +1525,10 @@ func (t *dateType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	return encoding.DecodeInt32(dst, src, enc)
 }
 
+func (t *dateType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return Int32Type.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *dateType) AssignValue(dst reflect.Value, src Value) error {
 	return Int32Type.AssignValue(dst, src)
 }
@@ -1534,11 +1610,11 @@ func (t *timeType) Length() int {
 	}
 }
 
-func (t *timeType) EstimateSize(n int) int64 {
+func (t *timeType) EstimateSize(n int) int {
 	if t.useInt32() {
-		return 4 * int64(n)
+		return 4 * n
 	} else {
-		return 8 * int64(n)
+		return 8 * n
 	}
 }
 
@@ -1633,6 +1709,14 @@ func (t *timeType) Decode(dst encoding.Values, src []byte, enc encoding.Encoding
 	}
 }
 
+func (t *timeType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	if t.useInt32() {
+		return Int32Type.EstimateDecodeSize(numValues, src, enc)
+	} else {
+		return Int64Type.EstimateDecodeSize(numValues, src, enc)
+	}
+}
+
 func (t *timeType) AssignValue(dst reflect.Value, src Value) error {
 	if t.useInt32() {
 		return Int32Type.AssignValue(dst, src)
@@ -1664,7 +1748,7 @@ 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) EstimateSize(n int) int { return 8 * n }
 
 func (t *timestampType) Compare(a, b Value) int { return compareInt64(a.int64(), b.int64()) }
 
@@ -1715,6 +1799,10 @@ func (t *timestampType) Decode(dst encoding.Values, src []byte, enc encoding.Enc
 	return encoding.DecodeInt64(dst, src, enc)
 }
 
+func (t *timestampType) EstimateDecodeSize(numValues int, src []byte, enc encoding.Encoding) int {
+	return Int64Type.EstimateDecodeSize(numValues, src, enc)
+}
+
 func (t *timestampType) AssignValue(dst reflect.Value, src Value) error {
 	switch dst.Type() {
 	case reflect.TypeOf(time.Time{}):
@@ -1788,7 +1876,7 @@ 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) EstimateSize(int) int { return 0 }
 
 func (t *listType) Compare(Value, Value) int { panic("cannot compare values on parquet LIST type") }
 
@@ -1832,6 +1920,10 @@ func (t *listType) Decode(_ encoding.Values, _ []byte, _ encoding.Encoding) (enc
 	panic("cannot decode parquet LIST type")
 }
 
+func (t *listType) EstimateDecodeSize(_ int, _ []byte, _ encoding.Encoding) int {
+	panic("cannot estimate decode size of parquet LIST type")
+}
+
 func (t *listType) AssignValue(reflect.Value, Value) error {
 	panic("cannot assign value to a parquet LIST type")
 }
@@ -1864,7 +1956,7 @@ 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) EstimateSize(int) int { return 0 }
 
 func (t *mapType) Compare(Value, Value) int { panic("cannot compare values on parquet MAP type") }
 
@@ -1908,6 +2000,10 @@ func (t *mapType) Decode(_ encoding.Values, _ []byte, _ encoding.Encoding) (enco
 	panic("cannot decode parquet MAP type")
 }
 
+func (t *mapType) EstimateDecodeSize(_ int, _ []byte, _ encoding.Encoding) int {
+	panic("cannot estimate decode size of parquet MAP type")
+}
+
 func (t *mapType) AssignValue(reflect.Value, Value) error {
 	panic("cannot assign value to a parquet MAP type")
 }
@@ -1924,7 +2020,7 @@ 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) EstimateSize(int) int { return 0 }
 
 func (t *nullType) Compare(Value, Value) int { panic("cannot compare values on parquet NULL type") }
 
@@ -1966,6 +2062,10 @@ func (t *nullType) Decode(dst encoding.Values, _ []byte, _ encoding.Encoding) (e
 	return dst, nil
 }
 
+func (t *nullType) EstimateDecodeSize(_ int, _ []byte, _ encoding.Encoding) int {
+	return 0
+}
+
 func (t *nullType) AssignValue(reflect.Value, Value) error {
 	return nil
 }
@@ -2014,6 +2114,10 @@ func (groupType) Decode(_ encoding.Values, _ []byte, _ encoding.Encoding) (encod
 	panic("cannot decode parquet group")
 }
 
+func (groupType) EstimateDecodeSize(_ int, _ []byte, _ encoding.Encoding) int {
+	panic("cannot estimate decode size of parquet group")
+}
+
 func (groupType) AssignValue(reflect.Value, Value) error {
 	panic("cannot assign value to a parquet group")
 }
@@ -2024,7 +2128,7 @@ func (t groupType) ConvertValue(Value, Type) (Value, error) {
 
 func (groupType) Length() int { return 0 }
 
-func (groupType) EstimateSize(int) int64 { return 0 }
+func (groupType) EstimateSize(int) int { return 0 }
 
 func (groupType) ColumnOrder() *format.ColumnOrder { return nil }
 
diff --git a/writer.go b/writer.go
index 6c4c911..1fbd22e 100644
--- a/writer.go
+++ b/writer.go
@@ -1047,7 +1047,7 @@ func (c *writerColumn) resizeBloomFilter(numValues int64) {
 
 func (c *writerColumn) newColumnBuffer() ColumnBuffer {
 	columnBufferCapacity := sort.Search(math.MaxInt32, func(i int) bool {
-		return c.columnType.EstimateSize(i) >= int64(c.bufferSize)
+		return c.columnType.EstimateSize(i) >= int(c.bufferSize)
 	})
 	column := c.columnType.NewColumnBuffer(int(c.bufferIndex), columnBufferCapacity)
 	switch {


Test patch:
diff --git a/encoding/encoding_test.go b/encoding/encoding_test.go
index 05dd9d3..f330a91 100644
--- a/encoding/encoding_test.go
+++ b/encoding/encoding_test.go
@@ -432,9 +432,13 @@ func testByteArrayEncoding(t *testing.T, e encoding.Encoding) {
 			var err error
 			buffer, err = e.EncodeByteArray(buffer, input, offsets)
 			assertNoError(t, err)
+			estimatedOutputSize := e.EstimateDecodeByteArraySize(buffer)
 			values, _, err = e.DecodeByteArray(values, buffer, offsets)
 			assertNoError(t, err)
 			assertEqualBytes(t, input, values)
+			if len(values) > estimatedOutputSize {
+				t.Errorf("the decode output was larger than the estimate: %d>%d", len(values), estimatedOutputSize)
+			}
 		})
 	}
 }