From 4f756596513f69e782e594aff36528dc3b0194e6 Mon Sep 17 00:00:00 2001 From: Richard Artoul Date: Mon, 24 Feb 2020 09:34:42 +0100 Subject: [PATCH] ARROW-7921: [Go] Add Reset method to various components and clean up comments. The reset method allow the data structures to be re-used so they don't have to be allocated over and over again. Closes #6430 from richardartoul/ra/merge-upstream and squashes the following commits: 5a0828187 Add license to test file d76be05ca Add test for data reset d102b1fff Add tests d3e6e6785 cleanup comments c8525aece Add Reset method to int array (#5) 489ca2593 Fix array.setData() to retain before release (#4) 88cd05ff8 Add reset method to Data (#3) 6d1b2775c Add Reset() method to String array (#2) dca230383 Add Reset method to buffer and cleanup comments (#1) Lead-authored-by: Richard Artoul Co-authored-by: Richard Artoul Signed-off-by: Sebastien Binet --- go/arrow/array/array.go | 4 +- go/arrow/array/data.go | 56 ++++++++- go/arrow/array/data_test.go | 51 ++++++++ go/arrow/array/numeric.gen.go | 192 ++++++++++++++++++++++++++--- go/arrow/array/numeric.gen.go.tmpl | 10 ++ go/arrow/array/string.go | 13 +- go/arrow/array/string_test.go | 23 ++++ go/arrow/memory/buffer.go | 23 +++- go/arrow/memory/buffer_test.go | 12 ++ 9 files changed, 358 insertions(+), 26 deletions(-) create mode 100644 go/arrow/array/data_test.go diff --git a/go/arrow/array/array.go b/go/arrow/array/array.go index 9fc4cc0be6db3..9cbaef9ff1236 100644 --- a/go/arrow/array/array.go +++ b/go/arrow/array/array.go @@ -119,11 +119,13 @@ func (a *array) IsValid(i int) bool { } func (a *array) setData(data *Data) { + // Retain before releasing in case a.data is the same as data. + data.Retain() + if a.data != nil { a.data.Release() } - data.Retain() if len(data.buffers) > 0 && data.buffers[0] != nil { a.nullBitmapBytes = data.buffers[0].Bytes() } diff --git a/go/arrow/array/data.go b/go/arrow/array/data.go index e7a69e984b708..264896159baac 100644 --- a/go/arrow/array/data.go +++ b/go/arrow/array/data.go @@ -24,7 +24,7 @@ import ( "github.com/apache/arrow/go/arrow/memory" ) -// A type which represents the memory and metadata for an Arrow array. +// Data represents the memory and metadata of an Arrow array. type Data struct { refCount int64 dtype arrow.DataType @@ -35,6 +35,7 @@ type Data struct { childData []*Data // TODO(sgc): managed by ListArray, StructArray and UnionArray types } +// NewData creates a new Data. func NewData(dtype arrow.DataType, length int, buffers []*memory.Buffer, childData []*Data, nulls, offset int) *Data { for _, b := range buffers { if b != nil { @@ -59,6 +60,42 @@ func NewData(dtype arrow.DataType, length int, buffers []*memory.Buffer, childDa } } +// Reset sets the Data for re-use. +func (d *Data) Reset(dtype arrow.DataType, length int, buffers []*memory.Buffer, childData []*Data, nulls, offset int) { + // Retain new buffers before releasing existing buffers in-case they're the same ones to prevent accidental premature + // release. + for _, b := range buffers { + if b != nil { + b.Retain() + } + } + for _, b := range d.buffers { + if b != nil { + b.Release() + } + } + d.buffers = buffers + + // Retain new children data before releasing existing children data in-case they're the same ones to prevent accidental + // premature release. + for _, d := range childData { + if d != nil { + d.Retain() + } + } + for _, d := range d.childData { + if d != nil { + d.Release() + } + } + d.childData = childData + + d.dtype = dtype + d.length = length + d.nulls = nulls + d.offset = offset +} + // Retain increases the reference count by 1. // Retain may be called simultaneously from multiple goroutines. func (d *Data) Retain() { @@ -85,10 +122,19 @@ func (d *Data) Release() { } } -func (d *Data) DataType() arrow.DataType { return d.dtype } -func (d *Data) NullN() int { return d.nulls } -func (d *Data) Len() int { return d.length } -func (d *Data) Offset() int { return d.offset } +// DataType returns the DataType of the data. +func (d *Data) DataType() arrow.DataType { return d.dtype } + +// NullN returns the number of nulls. +func (d *Data) NullN() int { return d.nulls } + +// Len returns the length. +func (d *Data) Len() int { return d.length } + +// Offset returns the offset. +func (d *Data) Offset() int { return d.offset } + +// Buffers returns the buffers. func (d *Data) Buffers() []*memory.Buffer { return d.buffers } // NewSliceData returns a new slice that shares backing data with the input. diff --git a/go/arrow/array/data_test.go b/go/arrow/array/data_test.go new file mode 100644 index 0000000000000..de87b80afabb2 --- /dev/null +++ b/go/arrow/array/data_test.go @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package array + +import ( + "testing" + + "github.com/apache/arrow/go/arrow" + "github.com/apache/arrow/go/arrow/memory" + "github.com/stretchr/testify/assert" +) + +func TestDataReset(t *testing.T) { + var ( + buffers1 = make([]*memory.Buffer, 0, 3) + buffers2 = make([]*memory.Buffer, 0, 3) + ) + for i := 0; i < cap(buffers1); i++ { + buffers1 = append(buffers1, memory.NewBufferBytes([]byte("some-bytes1"))) + buffers2 = append(buffers2, memory.NewBufferBytes([]byte("some-bytes2"))) + } + + data := NewData(&arrow.StringType{}, 10, buffers1, nil, 0, 0) + data.Reset(&arrow.Int64Type{}, 5, buffers2, nil, 1, 2) + + for i := 0; i < 2; i++ { + assert.Equal(t, buffers2, data.Buffers()) + assert.Equal(t, &arrow.Int64Type{}, data.DataType()) + assert.Equal(t, 1, data.NullN()) + assert.Equal(t, 2, data.Offset()) + assert.Equal(t, 5, data.Len()) + + // Make sure it works when resetting the data with its own buffers (new buffers are retained + // before old ones are released.) + data.Reset(&arrow.Int64Type{}, 5, data.Buffers(), nil, 1, 2) + } +} diff --git a/go/arrow/array/numeric.gen.go b/go/arrow/array/numeric.gen.go index 21c4e4b5e8131..068ab1ec98e79 100644 --- a/go/arrow/array/numeric.gen.go +++ b/go/arrow/array/numeric.gen.go @@ -31,6 +31,7 @@ type Int64 struct { values []int64 } +// NewInt64Data creates a new Int64. func NewInt64Data(data *Data) *Int64 { a := &Int64{} a.refCount = 1 @@ -38,9 +39,18 @@ func NewInt64Data(data *Data) *Int64 { return a } -func (a *Int64) Value(i int) int64 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Int64) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Int64) Value(i int) int64 { return a.values[i] } + +// Values returns the values. func (a *Int64) Int64Values() []int64 { return a.values } +// String returns a string representation of the array. func (a *Int64) String() string { o := new(strings.Builder) o.WriteString("[") @@ -88,6 +98,7 @@ type Uint64 struct { values []uint64 } +// NewUint64Data creates a new Uint64. func NewUint64Data(data *Data) *Uint64 { a := &Uint64{} a.refCount = 1 @@ -95,9 +106,18 @@ func NewUint64Data(data *Data) *Uint64 { return a } -func (a *Uint64) Value(i int) uint64 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Uint64) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Uint64) Value(i int) uint64 { return a.values[i] } + +// Values returns the values. func (a *Uint64) Uint64Values() []uint64 { return a.values } +// String returns a string representation of the array. func (a *Uint64) String() string { o := new(strings.Builder) o.WriteString("[") @@ -145,6 +165,7 @@ type Float64 struct { values []float64 } +// NewFloat64Data creates a new Float64. func NewFloat64Data(data *Data) *Float64 { a := &Float64{} a.refCount = 1 @@ -152,9 +173,18 @@ func NewFloat64Data(data *Data) *Float64 { return a } -func (a *Float64) Value(i int) float64 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Float64) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Float64) Value(i int) float64 { return a.values[i] } + +// Values returns the values. func (a *Float64) Float64Values() []float64 { return a.values } +// String returns a string representation of the array. func (a *Float64) String() string { o := new(strings.Builder) o.WriteString("[") @@ -202,6 +232,7 @@ type Int32 struct { values []int32 } +// NewInt32Data creates a new Int32. func NewInt32Data(data *Data) *Int32 { a := &Int32{} a.refCount = 1 @@ -209,9 +240,18 @@ func NewInt32Data(data *Data) *Int32 { return a } -func (a *Int32) Value(i int) int32 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Int32) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Int32) Value(i int) int32 { return a.values[i] } + +// Values returns the values. func (a *Int32) Int32Values() []int32 { return a.values } +// String returns a string representation of the array. func (a *Int32) String() string { o := new(strings.Builder) o.WriteString("[") @@ -259,6 +299,7 @@ type Uint32 struct { values []uint32 } +// NewUint32Data creates a new Uint32. func NewUint32Data(data *Data) *Uint32 { a := &Uint32{} a.refCount = 1 @@ -266,9 +307,18 @@ func NewUint32Data(data *Data) *Uint32 { return a } -func (a *Uint32) Value(i int) uint32 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Uint32) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Uint32) Value(i int) uint32 { return a.values[i] } + +// Values returns the values. func (a *Uint32) Uint32Values() []uint32 { return a.values } +// String returns a string representation of the array. func (a *Uint32) String() string { o := new(strings.Builder) o.WriteString("[") @@ -316,6 +366,7 @@ type Float32 struct { values []float32 } +// NewFloat32Data creates a new Float32. func NewFloat32Data(data *Data) *Float32 { a := &Float32{} a.refCount = 1 @@ -323,9 +374,18 @@ func NewFloat32Data(data *Data) *Float32 { return a } -func (a *Float32) Value(i int) float32 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Float32) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Float32) Value(i int) float32 { return a.values[i] } + +// Values returns the values. func (a *Float32) Float32Values() []float32 { return a.values } +// String returns a string representation of the array. func (a *Float32) String() string { o := new(strings.Builder) o.WriteString("[") @@ -373,6 +433,7 @@ type Int16 struct { values []int16 } +// NewInt16Data creates a new Int16. func NewInt16Data(data *Data) *Int16 { a := &Int16{} a.refCount = 1 @@ -380,9 +441,18 @@ func NewInt16Data(data *Data) *Int16 { return a } -func (a *Int16) Value(i int) int16 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Int16) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Int16) Value(i int) int16 { return a.values[i] } + +// Values returns the values. func (a *Int16) Int16Values() []int16 { return a.values } +// String returns a string representation of the array. func (a *Int16) String() string { o := new(strings.Builder) o.WriteString("[") @@ -430,6 +500,7 @@ type Uint16 struct { values []uint16 } +// NewUint16Data creates a new Uint16. func NewUint16Data(data *Data) *Uint16 { a := &Uint16{} a.refCount = 1 @@ -437,9 +508,18 @@ func NewUint16Data(data *Data) *Uint16 { return a } -func (a *Uint16) Value(i int) uint16 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Uint16) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Uint16) Value(i int) uint16 { return a.values[i] } + +// Values returns the values. func (a *Uint16) Uint16Values() []uint16 { return a.values } +// String returns a string representation of the array. func (a *Uint16) String() string { o := new(strings.Builder) o.WriteString("[") @@ -487,6 +567,7 @@ type Int8 struct { values []int8 } +// NewInt8Data creates a new Int8. func NewInt8Data(data *Data) *Int8 { a := &Int8{} a.refCount = 1 @@ -494,9 +575,18 @@ func NewInt8Data(data *Data) *Int8 { return a } -func (a *Int8) Value(i int) int8 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Int8) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Int8) Value(i int) int8 { return a.values[i] } + +// Values returns the values. func (a *Int8) Int8Values() []int8 { return a.values } +// String returns a string representation of the array. func (a *Int8) String() string { o := new(strings.Builder) o.WriteString("[") @@ -544,6 +634,7 @@ type Uint8 struct { values []uint8 } +// NewUint8Data creates a new Uint8. func NewUint8Data(data *Data) *Uint8 { a := &Uint8{} a.refCount = 1 @@ -551,9 +642,18 @@ func NewUint8Data(data *Data) *Uint8 { return a } -func (a *Uint8) Value(i int) uint8 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Uint8) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Uint8) Value(i int) uint8 { return a.values[i] } + +// Values returns the values. func (a *Uint8) Uint8Values() []uint8 { return a.values } +// String returns a string representation of the array. func (a *Uint8) String() string { o := new(strings.Builder) o.WriteString("[") @@ -601,6 +701,7 @@ type Timestamp struct { values []arrow.Timestamp } +// NewTimestampData creates a new Timestamp. func NewTimestampData(data *Data) *Timestamp { a := &Timestamp{} a.refCount = 1 @@ -608,9 +709,18 @@ func NewTimestampData(data *Data) *Timestamp { return a } -func (a *Timestamp) Value(i int) arrow.Timestamp { return a.values[i] } +// Reset resets the array for re-use. +func (a *Timestamp) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Timestamp) Value(i int) arrow.Timestamp { return a.values[i] } + +// Values returns the values. func (a *Timestamp) TimestampValues() []arrow.Timestamp { return a.values } +// String returns a string representation of the array. func (a *Timestamp) String() string { o := new(strings.Builder) o.WriteString("[") @@ -658,6 +768,7 @@ type Time32 struct { values []arrow.Time32 } +// NewTime32Data creates a new Time32. func NewTime32Data(data *Data) *Time32 { a := &Time32{} a.refCount = 1 @@ -665,9 +776,18 @@ func NewTime32Data(data *Data) *Time32 { return a } -func (a *Time32) Value(i int) arrow.Time32 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Time32) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Time32) Value(i int) arrow.Time32 { return a.values[i] } + +// Values returns the values. func (a *Time32) Time32Values() []arrow.Time32 { return a.values } +// String returns a string representation of the array. func (a *Time32) String() string { o := new(strings.Builder) o.WriteString("[") @@ -715,6 +835,7 @@ type Time64 struct { values []arrow.Time64 } +// NewTime64Data creates a new Time64. func NewTime64Data(data *Data) *Time64 { a := &Time64{} a.refCount = 1 @@ -722,9 +843,18 @@ func NewTime64Data(data *Data) *Time64 { return a } -func (a *Time64) Value(i int) arrow.Time64 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Time64) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Time64) Value(i int) arrow.Time64 { return a.values[i] } + +// Values returns the values. func (a *Time64) Time64Values() []arrow.Time64 { return a.values } +// String returns a string representation of the array. func (a *Time64) String() string { o := new(strings.Builder) o.WriteString("[") @@ -772,6 +902,7 @@ type Date32 struct { values []arrow.Date32 } +// NewDate32Data creates a new Date32. func NewDate32Data(data *Data) *Date32 { a := &Date32{} a.refCount = 1 @@ -779,9 +910,18 @@ func NewDate32Data(data *Data) *Date32 { return a } -func (a *Date32) Value(i int) arrow.Date32 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Date32) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Date32) Value(i int) arrow.Date32 { return a.values[i] } + +// Values returns the values. func (a *Date32) Date32Values() []arrow.Date32 { return a.values } +// String returns a string representation of the array. func (a *Date32) String() string { o := new(strings.Builder) o.WriteString("[") @@ -829,6 +969,7 @@ type Date64 struct { values []arrow.Date64 } +// NewDate64Data creates a new Date64. func NewDate64Data(data *Data) *Date64 { a := &Date64{} a.refCount = 1 @@ -836,9 +977,18 @@ func NewDate64Data(data *Data) *Date64 { return a } -func (a *Date64) Value(i int) arrow.Date64 { return a.values[i] } +// Reset resets the array for re-use. +func (a *Date64) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Date64) Value(i int) arrow.Date64 { return a.values[i] } + +// Values returns the values. func (a *Date64) Date64Values() []arrow.Date64 { return a.values } +// String returns a string representation of the array. func (a *Date64) String() string { o := new(strings.Builder) o.WriteString("[") @@ -886,6 +1036,7 @@ type Duration struct { values []arrow.Duration } +// NewDurationData creates a new Duration. func NewDurationData(data *Data) *Duration { a := &Duration{} a.refCount = 1 @@ -893,9 +1044,18 @@ func NewDurationData(data *Data) *Duration { return a } -func (a *Duration) Value(i int) arrow.Duration { return a.values[i] } +// Reset resets the array for re-use. +func (a *Duration) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. +func (a *Duration) Value(i int) arrow.Duration { return a.values[i] } + +// Values returns the values. func (a *Duration) DurationValues() []arrow.Duration { return a.values } +// String returns a string representation of the array. func (a *Duration) String() string { o := new(strings.Builder) o.WriteString("[") diff --git a/go/arrow/array/numeric.gen.go.tmpl b/go/arrow/array/numeric.gen.go.tmpl index 1e4a2f210da80..b742823a18d8d 100644 --- a/go/arrow/array/numeric.gen.go.tmpl +++ b/go/arrow/array/numeric.gen.go.tmpl @@ -31,6 +31,7 @@ type {{.Name}} struct { values []{{or .QualifiedType .Type}} } +// New{{.Name}}Data creates a new {{.Name}}. func New{{.Name}}Data(data *Data) *{{.Name}} { a := &{{.Name}}{} a.refCount = 1 @@ -38,9 +39,18 @@ func New{{.Name}}Data(data *Data) *{{.Name}} { return a } +// Reset resets the array for re-use. +func (a *{{.Name}}) Reset(data *Data) { + a.setData(data) +} + +// Value returns the value at the specified index. func (a *{{.Name}}) Value(i int) {{or .QualifiedType .Type}} { return a.values[i] } + +// Values returns the values. func (a *{{.Name}}) {{.Name}}Values() []{{or .QualifiedType .Type}} { return a.values } +// String returns a string representation of the array. func (a *{{.Name}}) String() string { o := new(strings.Builder) o.WriteString("[") diff --git a/go/arrow/array/string.go b/go/arrow/array/string.go index b7e607cfe6f19..42e87d8dfa48c 100644 --- a/go/arrow/array/string.go +++ b/go/arrow/array/string.go @@ -30,7 +30,7 @@ const ( stringArrayMaximumCapacity = math.MaxInt32 ) -// A type which represents an immutable sequence of variable-length UTF-8 strings. +// String represents an immutable sequence of variable-length UTF-8 strings. type String struct { array offsets []int32 @@ -45,11 +45,18 @@ func NewStringData(data *Data) *String { return a } +// Reset resets the String with a different set of Data. +func (a *String) Reset(data *Data) { + a.setData(data) +} + // Value returns the slice at index i. This value should not be mutated. func (a *String) Value(i int) string { i = i + a.array.data.offset return a.values[a.offsets[i]:a.offsets[i+1]] } + +// ValueOffset returns the offset of the value at index i. func (a *String) ValueOffset(i int) int { return int(a.offsets[i]) } func (a *String) String() string { @@ -104,6 +111,7 @@ type StringBuilder struct { builder *BinaryBuilder } +// NewStringBuilder creates a new StringBuilder. func NewStringBuilder(mem memory.Allocator) *StringBuilder { b := &StringBuilder{ builder: NewBinaryBuilder(mem, arrow.BinaryTypes.String), @@ -134,10 +142,12 @@ func (b *StringBuilder) Cap() int { return b.builder.Cap() } // NullN returns the number of null values in the array builder. func (b *StringBuilder) NullN() int { return b.builder.NullN() } +// Append appends a string to the builder. func (b *StringBuilder) Append(v string) { b.builder.Append([]byte(v)) } +// AppendNull appends a null to the builder. func (b *StringBuilder) AppendNull() { b.builder.AppendNull() } @@ -149,6 +159,7 @@ func (b *StringBuilder) AppendValues(v []string, valid []bool) { b.builder.AppendStringValues(v, valid) } +// Value returns the string at index i. func (b *StringBuilder) Value(i int) string { return string(b.builder.Value(i)) } diff --git a/go/arrow/array/string_test.go b/go/arrow/array/string_test.go index a4bdcdd86a71b..896ee9aa07bdf 100644 --- a/go/arrow/array/string_test.go +++ b/go/arrow/array/string_test.go @@ -158,3 +158,26 @@ func TestStringBuilder_Empty(t *testing.T) { assert.Equal(t, want, stringValues(a)) a.Release() } + +// TestStringReset tests the Reset() method on the String type by creating two different Strings and then +// reseting the contents of string2 with the values from string1. +func TestStringReset(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + sb1 := array.NewStringBuilder(mem) + sb2 := array.NewStringBuilder(mem) + defer sb1.Release() + defer sb2.Release() + + sb1.Append("string1") + sb1.AppendNull() + + var ( + string1 = sb1.NewStringArray() + string2 = sb2.NewStringArray() + + string1Data = string1.Data() + ) + string2.Reset(string1Data) + + assert.Equal(t, "string1", string2.Value(0)) +} diff --git a/go/arrow/memory/buffer.go b/go/arrow/memory/buffer.go index 234f5d4337fab..57c0db487d18d 100644 --- a/go/arrow/memory/buffer.go +++ b/go/arrow/memory/buffer.go @@ -22,6 +22,7 @@ import ( "github.com/apache/arrow/go/arrow/internal/debug" ) +// Buffer is a wrapper type for a buffer of bytes. type Buffer struct { refCount int64 buf []byte @@ -35,7 +36,7 @@ func NewBufferBytes(data []byte) *Buffer { return &Buffer{refCount: 0, buf: data, length: len(data)} } -// NewBuffer creates a mutable, resizable buffer with an Allocator for managing memory. +// NewResizableBuffer creates a mutable, resizable buffer with an Allocator for managing memory. func NewResizableBuffer(mem Allocator) *Buffer { return &Buffer{refCount: 1, mutable: true, mem: mem} } @@ -60,15 +61,28 @@ func (b *Buffer) Release() { } } +// Reset resets the buffer for reuse. +func (b *Buffer) Reset(buf []byte) { + b.buf = buf + b.length = len(buf) +} + // Buf returns the slice of memory allocated by the Buffer, which is adjusted by calling Reserve. func (b *Buffer) Buf() []byte { return b.buf } // Bytes returns a slice of size Len, which is adjusted by calling Resize. func (b *Buffer) Bytes() []byte { return b.buf[:b.length] } + +// Mutable returns a bool indicating whether the buffer is mutable or not. func (b *Buffer) Mutable() bool { return b.mutable } -func (b *Buffer) Len() int { return b.length } -func (b *Buffer) Cap() int { return len(b.buf) } +// Len returns the length of the buffer. +func (b *Buffer) Len() int { return b.length } + +// Cap returns the capacity of the buffer. +func (b *Buffer) Cap() int { return len(b.buf) } + +// Reserve reserves the provided amount of capacity for the buffer. func (b *Buffer) Reserve(capacity int) { if capacity > len(b.buf) { newCap := roundUpToMultipleOf64(capacity) @@ -80,10 +94,13 @@ func (b *Buffer) Reserve(capacity int) { } } +// Resize resizes the buffer to the target size. func (b *Buffer) Resize(newSize int) { b.resize(newSize, true) } +// ResizeNoShrink resizes the buffer to the target size, but will not +// shrink it. func (b *Buffer) ResizeNoShrink(newSize int) { b.resize(newSize, false) } diff --git a/go/arrow/memory/buffer_test.go b/go/arrow/memory/buffer_test.go index 918175cb6e269..53b4044b03ee2 100644 --- a/go/arrow/memory/buffer_test.go +++ b/go/arrow/memory/buffer_test.go @@ -43,3 +43,15 @@ func TestNewResizableBuffer(t *testing.T) { assert.Nil(t, buf.Bytes()) assert.Zero(t, buf.Len()) } + +func TestBufferReset(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + buf := memory.NewResizableBuffer(mem) + + newBytes := []byte("some-new-bytes") + buf.Reset(newBytes) + assert.Equal(t, newBytes, buf.Bytes()) + assert.Equal(t, len(newBytes), buf.Len()) +}