apacheGH-36850: [Go] Arrow Concatenate fix, ensure allocations are Fr…

…ee'd (apache#36854)

### Rationale for this change

The Concatenate function would capture panic's and return errors, however it wouldn't ensure that any allocations that happened in the `concat` sub function were cleaned up upon doing so. 

This change moves the `recover()` step into the `concat` function and will call `Release()` on the data object in the case of panic or error.  

### What changes are included in this PR?

### Are these changes tested?

A test called `TestConcatPanic` was added that causes the allocator to throw a panic part way during a concatenation, and ensures that the checked allocator still returns 0. 

### Are there any user-facing changes?

* Closes: apache#36850

Authored-by: thorfour <[email protected]>
Signed-off-by: Matt Topol <[email protected]>

go/arrow/array/concat.go

@@ -41,17 +41,6 @@ func Concatenate(arrs []arrow.Array, mem memory.Allocator) (result arrow.Array,
 		return nil, errors.New("array/concat: must pass at least one array")
 	}
 
-	defer func() {
-		if pErr := recover(); pErr != nil {
-			switch e := pErr.(type) {
-			case error:
-				err = fmt.Errorf("arrow/concat: %w", e)
-			default:
-				err = fmt.Errorf("arrow/concat: %v", pErr)
-			}
-		}
-	}()
-
 	// gather Data of inputs
 	data := make([]arrow.ArrayData, len(arrs))
 	for i, ar := range arrs {
@@ -368,8 +357,21 @@ func concatOffsets(buffers []*memory.Buffer, byteWidth int, mem memory.Allocator
 
 // concat is the implementation for actually performing the concatenation of the arrow.ArrayData
 // objects that we can call internally for nested types.
-func concat(data []arrow.ArrayData, mem memory.Allocator) (arrow.ArrayData, error) {
+func concat(data []arrow.ArrayData, mem memory.Allocator) (arr arrow.ArrayData, err error) {
 	out := &Data{refCount: 1, dtype: data[0].DataType(), nulls: 0}
+	defer func() {
+		if pErr := recover(); pErr != nil {
+			switch e := pErr.(type) {
+			case error:
+				err = fmt.Errorf("arrow/concat: %w", e)
+			default:
+				err = fmt.Errorf("arrow/concat: %v", pErr)
+			}
+		}
+		if err != nil {
+			out.Release()
+		}
+	}()
 	for _, d := range data {
 		out.length += d.Len()
 		if out.nulls == UnknownNullCount || d.NullN() == UnknownNullCount {
@@ -445,8 +447,8 @@ func concat(data []arrow.ArrayData, mem memory.Allocator) (arrow.ArrayData, erro
 		if err != nil {
 			return nil, err
 		}
-		out.buffers[2] = concatBuffers(gatherBufferRanges(data, 2, valueRanges), mem)
 		out.buffers[1] = offsetBuffer
+		out.buffers[2] = concatBuffers(gatherBufferRanges(data, 2, valueRanges), mem)
 	case *arrow.ListType:
 		offsetWidth := dt.Layout().Buffers[1].ByteWidth
 		offsetBuffer, valueRanges, err := concatOffsets(gatherFixedBuffers(data, 1, offsetWidth), offsetWidth, mem)

go/arrow/array/concat_test.go

@@ -743,3 +743,43 @@ func TestConcatOverflowRunEndEncoding(t *testing.T) {
 		})
 	}
 }
+
+type panicAllocator struct {
+	n int
+	memory.Allocator
+}
+
+func (p *panicAllocator) Allocate(size int) []byte {
+	if size > p.n {
+		panic("panic allocator")
+	}
+	return p.Allocator.Allocate(size)
+}
+
+func (p *panicAllocator) Reallocate(size int, b []byte) []byte {
+	return p.Allocator.Reallocate(size, b)
+}
+
+func (p *panicAllocator) Free(b []byte) {
+	p.Allocator.Free(b)
+}
+
+func TestConcatPanic(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	allocator := &panicAllocator{
+		n:         400,
+		Allocator: mem,
+	}
+
+	g := gen.NewRandomArrayGenerator(0, memory.DefaultAllocator)
+	ar1 := g.ArrayOf(arrow.STRING, 32, 0)
+	defer ar1.Release()
+	ar2 := g.ArrayOf(arrow.STRING, 32, 0)
+	defer ar2.Release()
+
+	concat, err := array.Concatenate([]arrow.Array{ar1, ar2}, allocator)
+	assert.Error(t, err)
+	assert.Nil(t, concat)
+}